Merge pull request apache#4224

Move BeamFnDataGrpcMultiplexer to java-fn-execution
ifigotin · Jan 17, 2018 · 4da548e · 4da548e
2 parents 8751ea5 + d174ce7
commit 4da548e
Show file tree

Hide file tree

Showing 7 changed files with 87 additions and 49 deletions.
diff --git a/runners/java-fn-execution/build.gradle b/runners/java-fn-execution/build.gradle
@@ -16,7 +16,7 @@
  */
 
 apply from: project(":").file("build_rules.gradle")
-applyJavaNature()
+applyJavaNature(javaVersion: 1.8)
 
 description = "Apache Beam :: Runners :: Java Fn Execution"
 

diff --git a/runners/java-fn-execution/pom.xml b/runners/java-fn-execution/pom.xml
@@ -32,6 +32,20 @@
 
   <packaging>jar</packaging>
 
+  <build>
+    <plugins>
+      <plugin>
+        <!--  Override Beam parent to allow Java8 -->
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-compiler-plugin</artifactId>
+        <configuration>
+          <source>1.8</source>
+          <target>1.8</target>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+
   <dependencies>
     <dependency>
       <groupId>org.apache.beam</groupId>

diff --git a/sdks/java/fn-execution/build.gradle b/sdks/java/fn-execution/build.gradle
@@ -17,7 +17,7 @@
  */
 
 apply from: project(":").file("build_rules.gradle")
-applyJavaNature()
+applyJavaNature(javaVersion: 1.8)
 
 description = "Apache Beam :: SDKs :: Java :: Fn Execution"
 
@@ -37,6 +37,7 @@ dependencies {
   testCompile library.java.hamcrest_core
   testCompile library.java.hamcrest_library
   testCompile library.java.mockito_core
+  testCompile library.java.slf4j_jdk14
 }
 
 task packageTests(type: Jar) {

diff --git a/sdks/java/fn-execution/pom.xml b/sdks/java/fn-execution/pom.xml
@@ -34,6 +34,20 @@
 
   <packaging>jar</packaging>
 
+  <build>
+    <plugins>
+      <plugin>
+        <!--  Override Beam parent to allow Java8 -->
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-compiler-plugin</artifactId>
+        <configuration>
+          <source>1.8</source>
+          <target>1.8</target>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+
   <dependencies>
     <dependency>
       <groupId>org.apache.beam</groupId>

diff --git a/...rness/data/BeamFnDataGrpcMultiplexer.java → ...dk/fn/data/BeamFnDataGrpcMultiplexer.java b/...rness/data/BeamFnDataGrpcMultiplexer.java → ...dk/fn/data/BeamFnDataGrpcMultiplexer.java
@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.fn.harness.data;
+package org.apache.beam.sdk.fn.data;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
@@ -25,16 +25,16 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
-import java.util.function.Function;
+import javax.annotation.Nullable;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements.Data;
 import org.apache.beam.model.pipeline.v1.Endpoints;
-import org.apache.beam.sdk.fn.data.LogicalEndpoint;
+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}.
+ * 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.
@@ -48,27 +48,26 @@
  */
 public class BeamFnDataGrpcMultiplexer {
   private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataGrpcMultiplexer.class);
-  private final Endpoints.ApiServiceDescriptor apiServiceDescriptor;
+  @Nullable private final Endpoints.ApiServiceDescriptor apiServiceDescriptor;
   private final StreamObserver<BeamFnApi.Elements> inboundObserver;
   private final StreamObserver<BeamFnApi.Elements> outboundObserver;
-  @VisibleForTesting
-  final ConcurrentMap<
-          LogicalEndpoint, CompletableFuture<Consumer<BeamFnApi.Elements.Data>>>
+  private final ConcurrentMap<
+            LogicalEndpoint, CompletableFuture<Consumer<BeamFnApi.Elements.Data>>>
       consumers;
 
   public BeamFnDataGrpcMultiplexer(
-      Endpoints.ApiServiceDescriptor apiServiceDescriptor,
-      Function<StreamObserver<BeamFnApi.Elements>,
-               StreamObserver<BeamFnApi.Elements>> outboundObserverFactory) {
+      @Nullable Endpoints.ApiServiceDescriptor apiServiceDescriptor,
+      StreamObserverClientFactory<BeamFnApi.Elements, BeamFnApi.Elements> outboundObserverFactory) {
     this.apiServiceDescriptor = apiServiceDescriptor;
     this.consumers = new ConcurrentHashMap<>();
     this.inboundObserver = new InboundObserver();
-    this.outboundObserver = outboundObserverFactory.apply(inboundObserver);
+    this.outboundObserver = outboundObserverFactory.outboundObserverFor(inboundObserver);
   }
 
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
+        .omitNullValues()
         .add("apiServiceDescriptor", apiServiceDescriptor)
         .add("consumers", consumers)
         .toString();
@@ -82,14 +81,24 @@ public StreamObserver<BeamFnApi.Elements> getOutboundObserver() {
     return outboundObserver;
   }
 
-  public CompletableFuture<Consumer<BeamFnApi.Elements.Data>> futureForKey(
-      LogicalEndpoint key) {
-    return consumers.computeIfAbsent(key, (LogicalEndpoint unused) -> new CompletableFuture<>());
+  private CompletableFuture<Consumer<Data>> receiverFuture(LogicalEndpoint endpoint) {
+    return consumers.computeIfAbsent(
+        endpoint, (LogicalEndpoint unused) -> new CompletableFuture<>());
+  }
+
+  public void registerConsumer(
+      LogicalEndpoint inputLocation, Consumer<BeamFnApi.Elements.Data> dataBytesReceiver) {
+    receiverFuture(inputLocation).complete(dataBytesReceiver);
+  }
+
+  @VisibleForTesting
+  boolean hasConsumer(LogicalEndpoint outputLocation) {
+    return consumers.containsKey(outputLocation);
   }
 
   /**
-   * A multiplexing {@link StreamObserver} that selects the inbound {@link Consumer} to pass
-   * the elements to.
+   * A multiplexing {@link StreamObserver} that selects the inbound {@link Consumer} to
+   * pass the elements to.
    *
    * <p>The inbound observer blocks until the {@link Consumer} is bound allowing for the
    * sending harness to initiate transmitting data without needing for the receiving harness to
@@ -102,7 +111,7 @@ public void onNext(BeamFnApi.Elements value) {
         try {
           LogicalEndpoint key =
               LogicalEndpoint.of(data.getInstructionReference(), data.getTarget());
-          CompletableFuture<Consumer<BeamFnApi.Elements.Data>> consumer = futureForKey(key);
+          CompletableFuture<Consumer<BeamFnApi.Elements.Data>> consumer = receiverFuture(key);
           if (!consumer.isDone()) {
             LOG.debug("Received data for key {} without consumer ready. "
                 + "Waiting for consumer to be registered.", key);
@@ -135,12 +144,17 @@ public void onNext(BeamFnApi.Elements value) {
 
     @Override
     public void onError(Throwable t) {
-      LOG.error("Failed to handle for {}", apiServiceDescriptor, t);
+      LOG.error(
+          "Failed to handle for {}",
+          apiServiceDescriptor == null ? "unknown endpoint" : apiServiceDescriptor,
+          t);
     }
 
     @Override
     public void onCompleted() {
-      LOG.warn("Hanged up for {}.", apiServiceDescriptor);
+      LOG.warn(
+          "Hanged up for {}.",
+          apiServiceDescriptor == null ? "unknown endpoint" : apiServiceDescriptor);
     }
   }
 }
diff --git a/...s/data/BeamFnDataGrpcMultiplexerTest.java → ...n/data/BeamFnDataGrpcMultiplexerTest.java b/...s/data/BeamFnDataGrpcMultiplexerTest.java → ...n/data/BeamFnDataGrpcMultiplexerTest.java
@@ -15,7 +15,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.beam.fn.harness.data;
+package org.apache.beam.sdk.fn.data;
 
 import static org.hamcrest.Matchers.contains;
 import static org.junit.Assert.assertFalse;
@@ -24,15 +24,13 @@
 
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.google.protobuf.ByteString;
-import io.grpc.stub.StreamObserver;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.model.pipeline.v1.Endpoints;
-import org.apache.beam.sdk.fn.data.LogicalEndpoint;
 import org.apache.beam.sdk.fn.test.TestStreams;
 import org.junit.Test;
 
@@ -61,37 +59,33 @@ public class BeamFnDataGrpcMultiplexerTest {
 
   @Test
   public void testOutboundObserver() {
-    Collection<BeamFnApi.Elements> values = new ArrayList<>();
-    BeamFnDataGrpcMultiplexer multiplexer = new BeamFnDataGrpcMultiplexer(
-        DESCRIPTOR,
-        (StreamObserver<BeamFnApi.Elements> inboundObserver)
-            -> TestStreams.withOnNext(values::add).build());
+    final Collection<BeamFnApi.Elements> values = new ArrayList<>();
+    BeamFnDataGrpcMultiplexer multiplexer =
+        new BeamFnDataGrpcMultiplexer(
+            DESCRIPTOR, inboundObserver -> TestStreams.withOnNext(values::add).build());
     multiplexer.getOutboundObserver().onNext(ELEMENTS);
     assertThat(values, contains(ELEMENTS));
   }
 
   @Test
   public void testInboundObserverBlocksTillConsumerConnects() throws Exception {
-    Collection<BeamFnApi.Elements> outboundValues = new ArrayList<>();
-    Collection<BeamFnApi.Elements.Data> inboundValues = new ArrayList<>();
-    BeamFnDataGrpcMultiplexer multiplexer = new BeamFnDataGrpcMultiplexer(
-        DESCRIPTOR,
-        (StreamObserver<BeamFnApi.Elements> inboundObserver)
-            -> TestStreams.withOnNext(outboundValues::add).build());
+    final Collection<BeamFnApi.Elements> outboundValues = new ArrayList<>();
+    final Collection<BeamFnApi.Elements.Data> inboundValues = new ArrayList<>();
+    final BeamFnDataGrpcMultiplexer multiplexer =
+        new BeamFnDataGrpcMultiplexer(
+            DESCRIPTOR, inboundObserver -> TestStreams.withOnNext(outboundValues::add).build());
     ExecutorService executorService = Executors.newCachedThreadPool();
-    executorService.submit(new Runnable() {
-      @Override
-      public void run() {
-        // Purposefully sleep to simulate a delay in a consumer connecting.
-        Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
-        multiplexer.futureForKey(OUTPUT_LOCATION).complete(inboundValues::add);
-      }
-    });
+    executorService.submit(
+        () -> {
+          // Purposefully sleep to simulate a delay in a consumer connecting.
+          Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+          multiplexer.registerConsumer(OUTPUT_LOCATION, inboundValues::add);
+        });
     multiplexer.getInboundObserver().onNext(ELEMENTS);
-    assertTrue(multiplexer.consumers.containsKey(OUTPUT_LOCATION));
+    assertTrue(multiplexer.hasConsumer(OUTPUT_LOCATION));
     // Ensure that when we see a terminal Elements object, we remove the consumer
     multiplexer.getInboundObserver().onNext(TERMINAL_ELEMENTS);
-    assertFalse(multiplexer.consumers.containsKey(OUTPUT_LOCATION));
+    assertFalse(multiplexer.hasConsumer(OUTPUT_LOCATION));
 
     // Assert that normal and terminal Elements are passed to the consumer
     assertThat(inboundValues, contains(ELEMENTS.getData(0), TERMINAL_ELEMENTS.getData(0)));

diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java
@@ -29,6 +29,7 @@
 import org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc;
 import org.apache.beam.model.pipeline.v1.Endpoints;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.fn.data.BeamFnDataGrpcMultiplexer;
 import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver;
 import org.apache.beam.sdk.fn.data.FnDataReceiver;
 import org.apache.beam.sdk.fn.data.LogicalEndpoint;
@@ -87,8 +88,8 @@ public <T> CompletableFuture<Void> receive(
 
     CompletableFuture<Void> readFuture = new CompletableFuture<>();
     BeamFnDataGrpcMultiplexer client = getClientFor(apiServiceDescriptor);
-    client.futureForKey(inputLocation).complete(
-        new BeamFnDataInboundObserver<>(coder, consumer, readFuture));
+    client.registerConsumer(
+        inputLocation, new BeamFnDataInboundObserver<>(coder, consumer, readFuture));
     return readFuture;
   }