From b7dba9e8b752d8a45653508592019ed6850d7e73 Mon Sep 17 00:00:00 2001
From: Martin Trieu <martinkt@google.com>
Date: Tue, 22 Oct 2024 16:09:17 -0700
Subject: [PATCH] address PR comments

---
 .../dataflow/worker/StreamingDataflowWorker.java      | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
index 333f674e94f3..5ee81b2ee362 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
@@ -84,6 +84,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCache;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingRemoteStubFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.ChannelCachingStubFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.IsolationChannel;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactory;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs.WindmillStubFactoryFactoryImpl;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
@@ -603,8 +604,13 @@ private static ChannelCachingStubFactory createStubFactory(
         workerOptions.getGcpCredential(),
         ChannelCache.create(
             serviceAddress ->
-                remoteChannel(
-                    serviceAddress, workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec())));
+                // IsolationChannel will create and manage separate RPC channels to the same
+                // serviceAddress.
+                IsolationChannel.create(
+                    () ->
+                        remoteChannel(
+                            serviceAddress,
+                            workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec()))));
   }
 
   @VisibleForTesting
@@ -860,7 +866,6 @@ ComputationStateCache getComputationStateCache() {
     return computationStateCache;
   }
 
-  @SuppressWarnings("FutureReturnValueIgnored")
   public void start() {
     running.set(true);
     configFetcher.start();