/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.fn.data;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import com.google.common.util.concurrent.SettableFuture;
import io.grpc.stub.StreamObserver;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.sdk.fn.stream.StreamObserverFactory.StreamObserverClientFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A gRPC multiplexer for a specific {@link Endpoints.ApiServiceDescriptor}.
 *
 * <p>Multiplexes data for inbound consumers based upon their individual {@link
 * org.apache.beam.model.fnexecution.v1.BeamFnApi.Target}s.
 *
 * <p>Multiplexing inbound and outbound streams is as thread safe as the consumers of those streams.
 * For inbound streams, this is as thread safe as the inbound observers. For outbound streams, this
 * is as thread safe as the underlying stream observer.
 *
 * <p>TODO: Add support for multiplexing over multiple outbound observers by stickying the output
 * location with a specific outbound observer.
 */
public class BeamFnDataGrpcMultiplexer {
  private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataGrpcMultiplexer.class);
  private final Endpoints.ApiServiceDescriptor apiServiceDescriptor;
  private final StreamObserver<BeamFnApi.Elements> inboundObserver;
  private final StreamObserver<BeamFnApi.Elements> outboundObserver;
  private final ConcurrentMap<
            LogicalEndpoint, SettableFuture<DataBytesReceiver>>
      consumers;

  public BeamFnDataGrpcMultiplexer(
      Endpoints.ApiServiceDescriptor apiServiceDescriptor,
      StreamObserverClientFactory<BeamFnApi.Elements, BeamFnApi.Elements> outboundObserverFactory) {
    this.apiServiceDescriptor = apiServiceDescriptor;
    this.consumers = new ConcurrentHashMap<>();
    this.inboundObserver = new InboundObserver();
    this.outboundObserver = outboundObserverFactory.outboundObserverFor(inboundObserver);
  }

  @Override
  public String toString() {
    return MoreObjects.toStringHelper(this)
        .add("apiServiceDescriptor", apiServiceDescriptor)
        .add("consumers", consumers)
        .toString();
  }

  public StreamObserver<BeamFnApi.Elements> getInboundObserver() {
    return inboundObserver;
  }

  public StreamObserver<BeamFnApi.Elements> getOutboundObserver() {
    return outboundObserver;
  }

  private SettableFuture<DataBytesReceiver> receiverFuture(LogicalEndpoint endpoint) {
    SettableFuture<DataBytesReceiver> future = consumers.get(endpoint);
    if (future == null) {
      future = SettableFuture.create();
      SettableFuture<DataBytesReceiver> present =
          consumers.putIfAbsent(endpoint, future);
      if (present != null) {
        // This call lost the race; use the existing future
        future = present;
      }
    }
    return future;
  }

  public void registerReceiver(LogicalEndpoint inputLocation, DataBytesReceiver dataBytesReceiver) {
    receiverFuture(inputLocation).set(dataBytesReceiver);
  }

  @VisibleForTesting
  boolean hasReceiver(LogicalEndpoint outputLocation) {
    return consumers.containsKey(outputLocation);
  }

  /**
   * A multiplexing {@link StreamObserver} that selects the inbound {@link DataBytesReceiver} to
   * pass the elements to.
   *
   * <p>The inbound observer blocks until the {@link DataBytesReceiver} is bound allowing for the
   * sending harness to initiate transmitting data without needing for the receiving harness to
   * signal that it is ready to consume that data.
   */
  private final class InboundObserver implements StreamObserver<BeamFnApi.Elements> {
    @Override
    public void onNext(BeamFnApi.Elements value) {
      for (BeamFnApi.Elements.Data data : value.getDataList()) {
        try {
          LogicalEndpoint key =
              LogicalEndpoint.of(data.getInstructionReference(), data.getTarget());
          SettableFuture<DataBytesReceiver> consumer = receiverFuture(key);
          if (!consumer.isDone()) {
            LOG.debug("Received data for key {} without consumer ready. "
                + "Waiting for consumer to be registered.", key);
          }
          consumer.get().receive(data);
          if (data.getData().isEmpty()) {
            consumers.remove(key);
          }
        /*
         * TODO: On failure we should fail any bundles that were impacted eagerly
         * instead of relying on the Runner harness to do all the failure handling.
         */
        } catch (ExecutionException | InterruptedException e) {
          LOG.error(
              "Client interrupted during handling of data for instruction {} and target {}",
              data.getInstructionReference(),
              data.getTarget(),
              e);
          outboundObserver.onError(e);
        } catch (RuntimeException e) {
          LOG.error(
              "Client failed to handle data for instruction {} and target {}",
              data.getInstructionReference(),
              data.getTarget(),
              e);
          outboundObserver.onError(e);
        }
      }
    }

    @Override
    public void onError(Throwable t) {
      LOG.error("Failed to handle for {}", apiServiceDescriptor, t);
    }

    @Override
    public void onCompleted() {
      LOG.warn("Hanged up for {}.", apiServiceDescriptor);
    }
  }
}

@@ -1167,16 +1167,50 @@
 eption;%0A
+import javax.annotation.Nullable;%0A
 import o
@@ -2223,24 +2223,34 @@
 r.class);%0A  
+@Nullable 
 private fina
@@ -2596,16 +2596,26 @@
 (%0A      
+@Nullable 
 Endpoint
@@ -6186,16 +6186,27 @@
 G.error(
+%0A          
 %22Failed 
@@ -6223,16 +6223,78 @@
 for %7B%7D%22,
+%0A          apiServiceDescriptor == null ? %22unknown endpoint%22 :
  apiServ
@@ -6307,16 +6307,26 @@
 criptor,
+%0A         
  t);%0A   
@@ -6390,16 +6390,27 @@
 OG.warn(
+%0A          
 %22Hanged 
@@ -6421,16 +6421,78 @@
 or %7B%7D.%22,
+%0A          apiServiceDescriptor == null ? %22unknown endpoint%22 :
  apiServ

Mark BeamFnDataGrpcMultiplexer ApiServiceDescriptor as Nullable