[BEAM-11227] Trying ChannelBuilder's executor option

This is suggestion by Sanjay from gRPC team grpc/grpc-java#8174 (comment)
suztomo · May 13, 2021 · d19e324 · d19e324
1 parent a942baf
commit d19e324
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/...est/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java b/...est/java/org/apache/beam/runners/dataflow/worker/fn/logging/BeamFnLoggingServiceTest.java
@@ -31,6 +31,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.model.fnexecution.v1.BeamFnLoggingGrpc;
@@ -125,6 +126,15 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
     LOG.info("Starting testMultipleClientsFailingIsHandledGracefullyByServer");
     Collection<Callable<Void>> tasks = new ArrayList<>();
     ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
+    ExecutorService channelExecutor =
+        Executors.newCachedThreadPool(
+            new ThreadFactory() {
+              @Override
+              public Thread newThread(Runnable r) {
+                return new Thread(r, "BeamFnLoggingServiceTest-channel-executor");
+              }
+            });
+
     try (BeamFnLoggingService service =
         new BeamFnLoggingService(
             findOpenPort(),
@@ -139,6 +149,7 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
               CountDownLatch waitForTermination = new CountDownLatch(1);
               ManagedChannel channel =
                   InProcessChannelBuilder.forName(service.getApiServiceDescriptor().getUrl())
+                      .executor(channelExecutor)
                       .build();
               StreamObserver<BeamFnApi.LogEntry.List> outboundObserver =
                   BeamFnLoggingGrpc.newStub(channel)
@@ -175,6 +186,7 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
       LOG.info("executorService terminated");
     } finally {
       LOG.info("Finished testMultipleClientsFailingIsHandledGracefullyByServer");
+      channelExecutor.shutdown();
     }
   }
 

diff --git a/...ion/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java b/...ion/src/test/java/org/apache/beam/runners/fnexecution/logging/GrpcLoggingServiceTest.java
@@ -29,6 +29,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -106,6 +107,16 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
     ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
     GrpcLoggingService service =
         GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
+
+    ExecutorService channelExecutor =
+        Executors.newCachedThreadPool(
+            new ThreadFactory() {
+              @Override
+              public Thread newThread(Runnable r) {
+                return new Thread(r, "GrpcLoggingServiceTest-channel-executor");
+              }
+            });
+
     try (GrpcFnServer<GrpcLoggingService> server =
         GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
 
@@ -116,7 +127,8 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
             () -> {
               CountDownLatch waitForTermination = new CountDownLatch(1);
               String url = server.getApiServiceDescriptor().getUrl();
-              ManagedChannel channel = InProcessChannelBuilder.forName(url).build();
+              ManagedChannel channel =
+                  InProcessChannelBuilder.forName(url).executor(channelExecutor).build();
               StreamObserver<LogEntry.List> outboundObserver =
                   BeamFnLoggingGrpc.newStub(channel)
                       .logging(
@@ -141,6 +153,7 @@ public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Excep
       LOG.info("executorService terminated");
     } finally {
       LOG.info("Finishing testMultipleClientsFailingIsHandledGracefullyByServer");
+      channelExecutor.shutdown();
     }
   }