Add a FailoverChannel wrapper on top of IsolationChannel to maintain a set of primary and failover channel. (#37840)

Author: parveensania

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

@@ -87,8 +87,10 @@
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcDispatcherClient;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillServer;
 import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.GrpcWindmillStreamFactory;
+import org.apache.beam.runners.dataflow.worker.windmill.client.grpc.auth.VendoredCredentialsAdapter;
 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.FailoverChannel;
 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;
@@ -113,6 +115,7 @@
 import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.util.construction.CoderTranslation;
 import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.auth.MoreCallCredentials;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats;
@@ -382,7 +385,8 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
       MemoryMonitor memoryMonitor,
       GrpcDispatcherClient dispatcherClient) {
     WeightedSemaphore<Commit> maxCommitByteSemaphore = Commits.maxCommitByteSemaphore();
-    ChannelCache channelCache = createChannelCache(options, checkNotNull(configFetcher));
+    ChannelCache channelCache =
+        createChannelCache(options, checkNotNull(configFetcher), dispatcherClient);
     @SuppressWarnings("methodref.receiver.bound")
     FanOutStreamingEngineWorkerHarness fanOutStreamingEngineWorkerHarness =
         FanOutStreamingEngineWorkerHarness.create(
@@ -785,20 +789,32 @@ private static void validateWorkerOptions(DataflowWorkerHarnessOptions options)
   }
 
   private static ChannelCache createChannelCache(
-      DataflowWorkerHarnessOptions workerOptions, ComputationConfig.Fetcher configFetcher) {
+      DataflowWorkerHarnessOptions workerOptions,
+      ComputationConfig.Fetcher configFetcher,
+      GrpcDispatcherClient dispatcherClient) {
     ChannelCache channelCache =
         ChannelCache.create(
             (currentFlowControlSettings, serviceAddress) -> {
-              // IsolationChannel will create and manage separate RPC channels to the same
-              // serviceAddress.
+              // IsolationChannel wrapping FailoverChannel so that each active RPC gets its own
+              // FailoverChannel instance. The fallback channel is created lazily, at most once,
+              // only if failover is actually needed.
               return IsolationChannel.create(
                   () ->
-                      remoteChannel(
-                          serviceAddress,
-                          workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
-                          currentFlowControlSettings),
+                      FailoverChannel.create(
+                          remoteChannel(
+                              serviceAddress,
+                              workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
+                              currentFlowControlSettings),
+                          () ->
+                              remoteChannel(
+                                  dispatcherClient.getDispatcherEndpoints().iterator().next(),
+                                  workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
+                                  currentFlowControlSettings),
+                          MoreCallCredentials.from(
+                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential()))),
                   currentFlowControlSettings.getOnReadyThresholdBytes());
             });
+
     configFetcher
         .getGlobalConfigHandle()
         .registerConfigObserver(

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

@@ -410,15 +410,18 @@ private GlobalDataStreamSender getOrCreateGlobalDataSteam(
   }
 
   private WindmillStreamSender createAndStartWindmillStreamSender(Endpoint endpoint) {
+    GetWorkRequest.Builder getWorkRequestBuilder =
+        GetWorkRequest.newBuilder()
+            .setClientId(jobHeader.getClientId())
+            .setJobId(jobHeader.getJobId())
+            .setProjectId(jobHeader.getProjectId())
+            .setWorkerId(jobHeader.getWorkerId());
+    endpoint.workerToken().ifPresent(getWorkRequestBuilder::setBackendWorkerToken);
+
     WindmillStreamSender windmillStreamSender =
         WindmillStreamSender.create(
             WindmillConnection.from(endpoint, this::createWindmillStub),
-            GetWorkRequest.newBuilder()
-                .setClientId(jobHeader.getClientId())
-                .setJobId(jobHeader.getJobId())
-                .setProjectId(jobHeader.getProjectId())
-                .setWorkerId(jobHeader.getWorkerId())
-                .build(),
+            getWorkRequestBuilder.build(),
             GetWorkBudget.noBudget(),
             streamFactory,
             workItemScheduler,

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java

@@ -108,7 +108,7 @@ public CloudWindmillServiceV1Alpha1Stub getWindmillServiceStub() {
         : randomlySelectNextStub(windmillServiceStubs));
   }
 
-  ImmutableSet<HostAndPort> getDispatcherEndpoints() {
+  public ImmutableSet<HostAndPort> getDispatcherEndpoints() {
     return dispatcherStubs.get().dispatcherEndpoints();
   }