Make primary channel failover timeout configurable (apache#39646)

Author: parveensania

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

@@ -34,6 +34,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
@@ -173,6 +174,8 @@ public final class StreamingDataflowWorker {
   private static final String CHANNELZ_PATH = "/channelz";
   private static final String BEAM_FN_API_EXPERIMENT = "beam_fn_api";
   private static final String ELEMENT_METADATA_SUPPORTED_EXPERIMENT = "element_metadata_supported";
+  private static final AtomicLong DIRECTPATH_PRIMARY_NOT_READY_WAIT_MILLIS =
+      new AtomicLong(TimeUnit.SECONDS.toMillis(15));
 
   @SuppressWarnings("unused")
   private static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT =
@@ -867,16 +870,20 @@ private static ChannelCache createChannelCache(
                                   workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
                                   currentFlowControlSettings),
                           MoreCallCredentials.from(
-                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential()))),
+                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential())),
+                          DIRECTPATH_PRIMARY_NOT_READY_WAIT_MILLIS::get),
                   currentFlowControlSettings.getOnReadyThresholdBytes());
             });
 
     configFetcher
         .getGlobalConfigHandle()
         .registerConfigObserver(
-            config ->
-                channelCache.consumeFlowControlSettings(
-                    config.userWorkerJobSettings().getFlowControlSettings()));
+            config -> {
+              DIRECTPATH_PRIMARY_NOT_READY_WAIT_MILLIS.set(
+                  config.userWorkerJobSettings().getDirectpathPrimaryNotReadyWaitMillis());
+              channelCache.consumeFlowControlSettings(
+                  config.userWorkerJobSettings().getFlowControlSettings());
+            });
     return channelCache;
   }
 

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

@@ -45,9 +45,9 @@
  * <p>Routes requests to either primary or fallback channel based on two independent failover modes:
  *
  * <ul>
- *   <li><b>Connection Status Failover:</b> If the primary channel is not ready for 10+ seconds
- *       (e.g., during network issues), routes to fallback channel. Switches back as soon as the
- *       primary channel becomes READY again.
+ *   <li><b>Connection Status Failover:</b> If the primary channel is not ready for the configured
+ *       wait time (e.g., during network issues), routes to fallback channel. Switches back as soon
+ *       as the primary channel becomes READY again.
  *   <li><b>RPC Failover:</b> If primary channel RPCs fail continuously with transient errors
  *       ({@link Status.Code#UNAVAILABLE} or {@link Status.Code#UNKNOWN}), or with {@link
  *       Status.Code#DEADLINE_EXCEEDED} before receiving any response (indicating the connection was
@@ -61,7 +61,6 @@ public final class FailoverChannel extends ManagedChannel {
   private static final AtomicInteger CHANNEL_ID_COUNTER = new AtomicInteger(0);
   // Time to wait before retrying the primary channel after an RPC-based fallback.
   private static final long FALLBACK_COOLING_PERIOD_NANOS = TimeUnit.HOURS.toNanos(1);
-  private static final long PRIMARY_NOT_READY_WAIT_NANOS = TimeUnit.SECONDS.toNanos(10);
   // Minimum duration of continuous RPC failures required before switching to fallback.
   private static final long RPC_FAILURE_THRESHOLD_NANOS = TimeUnit.SECONDS.toNanos(30);
 
@@ -96,10 +95,15 @@ private static final class FailoverState {
 
     private final int channelId;
     private final long rpcFailureThresholdNanos;
+    private final LongSupplier primaryNotReadyWaitMillisSupplier;
 
-    FailoverState(int channelId, long rpcFailureThresholdNanos) {
+    FailoverState(
+        int channelId,
+        long rpcFailureThresholdNanos,
+        LongSupplier primaryNotReadyWaitMillisSupplier) {
       this.channelId = channelId;
       this.rpcFailureThresholdNanos = rpcFailureThresholdNanos;
+      this.primaryNotReadyWaitMillisSupplier = primaryNotReadyWaitMillisSupplier;
     }
 
     /**
@@ -118,14 +122,16 @@ synchronized boolean computeUseFallback(long nowNanos) {
       }
       // Check if primary has been not-ready long enough to switch to fallback.
       // primaryNotReadySinceNanos is set by the state-change callback when primary is not ready.
-      if (!useFallbackDueToRPC
-          && !useFallbackDueToState
-          && primaryNotReadySinceNanos >= 0
-          && nowNanos - primaryNotReadySinceNanos > PRIMARY_NOT_READY_WAIT_NANOS) {
-        useFallbackDueToState = true;
-        LOG.warn(
-            "[channel-{}] Primary connection unavailable. Switching to secondary connection.",
-            channelId);
+      if (!useFallbackDueToRPC && !useFallbackDueToState && primaryNotReadySinceNanos >= 0) {
+        long elapsedPrimaryNotReadyNanos = nowNanos - primaryNotReadySinceNanos;
+        long primaryNotReadyWaitNanos =
+            TimeUnit.MILLISECONDS.toNanos(primaryNotReadyWaitMillisSupplier.getAsLong());
+        if (elapsedPrimaryNotReadyNanos > primaryNotReadyWaitNanos) {
+          useFallbackDueToState = true;
+          LOG.warn(
+              "[channel-{}] Primary connection unavailable. Switching to secondary connection.",
+              channelId);
+        }
       }
       return useFallbackDueToRPC || useFallbackDueToState;
     }
@@ -193,11 +199,13 @@ private FailoverChannel(
       Supplier<ManagedChannel> fallbackSupplier,
       @Nullable CallCredentials fallbackCallCredentials,
       LongSupplier nanoClock,
-      long rpcFailureThresholdNanos) {
+      long rpcFailureThresholdNanos,
+      LongSupplier primaryNotReadyWaitMillisSupplier) {
     this.primary = primary;
     this.fallbackSupplier = Suppliers.memoize(fallbackSupplier::get);
     this.channelId = CHANNEL_ID_COUNTER.getAndIncrement();
-    this.state = new FailoverState(channelId, rpcFailureThresholdNanos);
+    this.state =
+        new FailoverState(channelId, rpcFailureThresholdNanos, primaryNotReadyWaitMillisSupplier);
     this.fallbackCallCredentials = fallbackCallCredentials;
     this.nanoClock = nanoClock;
     // Register callback to monitor primary channel state changes
@@ -207,23 +215,31 @@ private FailoverChannel(
   public static FailoverChannel create(
       ManagedChannel primary,
       Supplier<ManagedChannel> fallbackSupplier,
-      CallCredentials fallbackCallCredentials) {
+      CallCredentials fallbackCallCredentials,
+      LongSupplier primaryNotReadyWaitMillisSupplier) {
     return new FailoverChannel(
         primary,
         fallbackSupplier,
         fallbackCallCredentials,
         System::nanoTime,
-        RPC_FAILURE_THRESHOLD_NANOS);
+        RPC_FAILURE_THRESHOLD_NANOS,
+        primaryNotReadyWaitMillisSupplier);
   }
 
   static FailoverChannel forTest(
       ManagedChannel primary,
       ManagedChannel fallback,
       CallCredentials fallbackCallCredentials,
       LongSupplier nanoClock,
-      long rpcFailureThresholdNanos) {
+      long rpcFailureThresholdNanos,
+      LongSupplier primaryNotReadyWaitMillisSupplier) {
     return new FailoverChannel(
-        primary, () -> fallback, fallbackCallCredentials, nanoClock, rpcFailureThresholdNanos);
+        primary,
+        () -> fallback,
+        fallbackCallCredentials,
+        nanoClock,
+        rpcFailureThresholdNanos,
+        primaryNotReadyWaitMillisSupplier);
   }
 
   /** Returns the fallback channel, creating it from the supplier at most once. */
@@ -399,7 +415,12 @@ private void registerPrimaryStateChangeListener() {
         // never transitions, markPrimaryNotReady() would never be called and state-based
         // failover would not trigger even after the grace period.
         if (currentState == ConnectivityState.READY || currentState == ConnectivityState.IDLE) {
-          state.markPrimaryReady();
+          if (state.markPrimaryReady()) {
+            LOG.info(
+                "[channel-{}] Primary channel observed healthy during state change registration;"
+                    + " switching back from fallback.",
+                channelId);
+          }
         } else {
           // Seed the not-ready timer even if there is no future state transition.
           state.markPrimaryNotReady(nanoClock.getAsLong());
@@ -426,11 +447,12 @@ private void onPrimaryStateChanged() {
     if (newState == ConnectivityState.READY || newState == ConnectivityState.IDLE) {
       if (state.markPrimaryReady()) {
         LOG.info(
-            "[channel-{}] Primary channel recovered; switching back from fallback.", channelId);
+            "[channel-{}] Primary channel observed healthy during state change notification; switching back from fallback.",
+            channelId);
       }
     } else {
       // Primary is not ready; start the grace period timer so computeUseFallback can
-      // switch to fallback once PRIMARY_NOT_READY_WAIT_NANOS elapses.
+      // switch to fallback once the configured wait time elapses.
       state.markPrimaryNotReady(nanoClock.getAsLong());
     }
 

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/FailoverChannelTest.java

@@ -25,6 +25,7 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -87,7 +88,8 @@ private static FailoverChannel createForTest(
         fallback,
         fallbackCallCredentials,
         nanoClock,
-        rpcFailureThresholdNanos != null ? rpcFailureThresholdNanos : 0L);
+        rpcFailureThresholdNanos != null ? rpcFailureThresholdNanos : 0L,
+        () -> TimeUnit.SECONDS.toMillis(10));
   }
 
   /**
@@ -283,11 +285,110 @@ public void testStateFallbackAfterPrimaryNotReady() {
     // Within 10 seconds: grace period not elapsed, routes to primary.
     failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
     verify(mockChannel).newCall(any(), any());
+    // Fallback must not have been used yet.
+    verify(mockFallbackChannel, never()).newCall(any(), any());
 
     // After 10 seconds: routes to fallback.
     time.addAndGet(TimeUnit.SECONDS.toNanos(11));
     failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
     verify(mockFallbackChannel).newCall(any(), any());
+    // Primary must not have been used for this call (still only 1 invocation).
+    verify(mockChannel, times(1)).newCall(any(), any());
+  }
+
+  @Test
+  public void testTimeoutThresholdDecreaseTriggersFallbackEarlier() {
+    ManagedChannel mockChannel = mock(ManagedChannel.class);
+    ManagedChannel mockFallbackChannel = mock(ManagedChannel.class);
+    when(mockChannel.newCall(any(), any())).thenReturn(mock(ClientCall.class));
+    when(mockFallbackChannel.newCall(any(), any())).thenReturn(mock(ClientCall.class));
+    // Simulate primary being TRANSIENT_FAILURE from the start.
+    when(mockChannel.getState(false)).thenReturn(ConnectivityState.TRANSIENT_FAILURE);
+
+    AtomicLong time = new AtomicLong(0);
+    // Start with 10s timeout
+    AtomicLong timeoutThreshold = new AtomicLong(TimeUnit.SECONDS.toMillis(10));
+
+    // Constructor seeds timer at time=0.
+    FailoverChannel failoverChannel =
+        FailoverChannel.forTest(
+            mockChannel, mockFallbackChannel, null, time::get, 0L, timeoutThreshold::get);
+
+    // Call at time=0. elapsed 0 <= 10s —> failover condition false.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel).newCall(any(), any());
+
+    // Advance time by 5 seconds.
+    time.addAndGet(TimeUnit.SECONDS.toNanos(5));
+
+    // Call at time=5s. elapsed 5s <= 10s —> failover condition false.
+    // Primary is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel, times(2)).newCall(any(), any());
+    // Fallback must not have been used yet.
+    verify(mockFallbackChannel, never()).newCall(any(), any());
+
+    // Decrease threshold to 2 seconds.
+    timeoutThreshold.set(TimeUnit.SECONDS.toMillis(2));
+
+    // Call at time=5s. elapsed 5s > 2s —> failover condition true.
+    // Fallback is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockFallbackChannel).newCall(any(), any());
+    // Primary must not have been used for this call (still only 2 invocations).
+    verify(mockChannel, times(2)).newCall(any(), any());
+  }
+
+  @Test
+  public void testTimeoutThresholdIncreaseDelaysFallback() {
+    ManagedChannel mockChannel = mock(ManagedChannel.class);
+    ManagedChannel mockFallbackChannel = mock(ManagedChannel.class);
+    when(mockChannel.newCall(any(), any())).thenReturn(mock(ClientCall.class));
+    when(mockFallbackChannel.newCall(any(), any())).thenReturn(mock(ClientCall.class));
+    // Simulate primary being TRANSIENT_FAILURE from the start.
+    when(mockChannel.getState(false)).thenReturn(ConnectivityState.TRANSIENT_FAILURE);
+
+    AtomicLong time = new AtomicLong(0);
+    // Start with 10s timeout
+    AtomicLong timeoutThreshold = new AtomicLong(TimeUnit.SECONDS.toMillis(10));
+
+    // Constructor seeds timer at time=0.
+    FailoverChannel failoverChannel =
+        FailoverChannel.forTest(
+            mockChannel, mockFallbackChannel, null, time::get, 0L, timeoutThreshold::get);
+
+    // Advance time by 9 seconds.
+    time.addAndGet(TimeUnit.SECONDS.toNanos(9));
+
+    // Call at time=9s. elapsed 9s <= 10s —> failover condition false.
+    // Primary is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel).newCall(any(), any());
+    // Fallback must not have been used yet.
+    verify(mockFallbackChannel, never()).newCall(any(), any());
+
+    // Increase threshold to 20 seconds.
+    timeoutThreshold.set(TimeUnit.SECONDS.toMillis(20));
+
+    // Advance time by 5 seconds (total 14s).
+    time.addAndGet(TimeUnit.SECONDS.toNanos(5));
+
+    // Call at time=14s. elapsed 14s <= 20s —> failover condition false.
+    // Still routes to primary because of increased threshold
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel, times(2)).newCall(any(), any());
+    // Fallback must still not have been used yet.
+    verify(mockFallbackChannel, never()).newCall(any(), any());
+
+    // Advance time by 7s (total 21s).
+    time.addAndGet(TimeUnit.SECONDS.toNanos(7));
+
+    // Call at time=21s. elapsed 21s > 20s —> failover condition true.
+    // Fallback is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockFallbackChannel).newCall(any(), any());
+    // Primary must not have been used for this call (still only 2 invocations).
+    verify(mockChannel, times(2)).newCall(any(), any());
   }
 
   @Test
@@ -307,11 +408,15 @@ public void testStateFallbackWhenPrimaryStartsNonReadyWithoutTransition() {
     // Before grace period, still routes to primary.
     failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
     verify(mockChannel).newCall(any(), any());
+    // Fallback must not have been used yet.
+    verify(mockFallbackChannel, never()).newCall(any(), any());
 
     // After 10 seconds in non-ready state, should route to fallback.
     time.addAndGet(TimeUnit.SECONDS.toNanos(11));
     failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
     verify(mockFallbackChannel).newCall(any(), any());
+    // Primary must not have been used for this call (still only 1 invocation).
+    verify(mockChannel, times(1)).newCall(any(), any());
   }
 
   @Test

runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto

@@ -1018,6 +1018,10 @@ message UserWorkerRunnerV1Settings {
 
   optional int64 max_cached_entry_bytes = 5 [default = -1];
 
+  // Time to wait before switching to fallback connectivity if primary channel is not ready.
+  // Only used if directpath is enabled for the job. Default is 15 seconds.
+  optional int64 directpath_primary_not_ready_wait_millis = 6 [default = 15000];
+
   reserved 1, 2;
 }