Make primary channel failover timeout configurable

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.Consumer;
 import java.util.function.Function;
@@ -172,6 +173,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_NANOS =
+      new AtomicLong(TimeUnit.SECONDS.toNanos(15));
 
   @SuppressWarnings("unused")
   private static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT =
@@ -847,16 +850,20 @@ private static ChannelCache createChannelCache(
                                   workerOptions.getWindmillServiceRpcChannelAliveTimeoutSec(),
                                   currentFlowControlSettings),
                           MoreCallCredentials.from(
-                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential()))),
+                              new VendoredCredentialsAdapter(workerOptions.getGcpCredential())),
+                          DIRECTPATH_PRIMARY_NOT_READY_WAIT_NANOS::get),
                   currentFlowControlSettings.getOnReadyThresholdBytes());
             });
 
     configFetcher
         .getGlobalConfigHandle()
         .registerConfigObserver(
-            config ->
-                channelCache.consumeFlowControlSettings(
-                    config.userWorkerJobSettings().getFlowControlSettings()));
+            config -> {
+              DIRECTPATH_PRIMARY_NOT_READY_WAIT_NANOS.set(
+                  config.userWorkerJobSettings().getDirectpathPrimaryNotReadyWaitNanos());
+              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 primaryNotReadyWaitNanosSupplier;
 
-    FailoverState(int channelId, long rpcFailureThresholdNanos) {
+    FailoverState(
+        int channelId,
+        long rpcFailureThresholdNanos,
+        LongSupplier primaryNotReadyWaitNanosSupplier) {
       this.channelId = channelId;
       this.rpcFailureThresholdNanos = rpcFailureThresholdNanos;
+      this.primaryNotReadyWaitNanosSupplier = primaryNotReadyWaitNanosSupplier;
     }
 
     /**
@@ -121,7 +125,7 @@ synchronized boolean computeUseFallback(long nowNanos) {
       if (!useFallbackDueToRPC
           && !useFallbackDueToState
           && primaryNotReadySinceNanos >= 0
-          && nowNanos - primaryNotReadySinceNanos > PRIMARY_NOT_READY_WAIT_NANOS) {
+          && nowNanos - primaryNotReadySinceNanos > primaryNotReadyWaitNanosSupplier.getAsLong()) {
         useFallbackDueToState = true;
         LOG.warn(
             "[channel-{}] Primary connection unavailable. Switching to secondary connection.",
@@ -193,11 +197,13 @@ private FailoverChannel(
       Supplier<ManagedChannel> fallbackSupplier,
       @Nullable CallCredentials fallbackCallCredentials,
       LongSupplier nanoClock,
-      long rpcFailureThresholdNanos) {
+      long rpcFailureThresholdNanos,
+      LongSupplier primaryNotReadyWaitNanosSupplier) {
     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, primaryNotReadyWaitNanosSupplier);
     this.fallbackCallCredentials = fallbackCallCredentials;
     this.nanoClock = nanoClock;
     // Register callback to monitor primary channel state changes
@@ -207,23 +213,31 @@ private FailoverChannel(
   public static FailoverChannel create(
       ManagedChannel primary,
       Supplier<ManagedChannel> fallbackSupplier,
-      CallCredentials fallbackCallCredentials) {
+      CallCredentials fallbackCallCredentials,
+      LongSupplier primaryNotReadyWaitNanosSupplier) {
     return new FailoverChannel(
         primary,
         fallbackSupplier,
         fallbackCallCredentials,
         System::nanoTime,
-        RPC_FAILURE_THRESHOLD_NANOS);
+        RPC_FAILURE_THRESHOLD_NANOS,
+        primaryNotReadyWaitNanosSupplier);
   }
 
   static FailoverChannel forTest(
       ManagedChannel primary,
       ManagedChannel fallback,
       CallCredentials fallbackCallCredentials,
       LongSupplier nanoClock,
-      long rpcFailureThresholdNanos) {
+      long rpcFailureThresholdNanos,
+      LongSupplier primaryNotReadyWaitNanosSupplier) {
     return new FailoverChannel(
-        primary, () -> fallback, fallbackCallCredentials, nanoClock, rpcFailureThresholdNanos);
+        primary,
+        () -> fallback,
+        fallbackCallCredentials,
+        nanoClock,
+        rpcFailureThresholdNanos,
+        primaryNotReadyWaitNanosSupplier);
   }
 
   /** Returns the fallback channel, creating it from the supplier at most once. */
@@ -399,7 +413,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 +445,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 registration; 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

@@ -87,7 +87,8 @@ private static FailoverChannel createForTest(
         fallback,
         fallbackCallCredentials,
         nanoClock,
-        rpcFailureThresholdNanos != null ? rpcFailureThresholdNanos : 0L);
+        rpcFailureThresholdNanos != null ? rpcFailureThresholdNanos : 0L,
+        () -> TimeUnit.SECONDS.toNanos(10));
   }
 
   /**
@@ -290,6 +291,91 @@ public void testStateFallbackAfterPrimaryNotReady() {
     verify(mockFallbackChannel).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.toNanos(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 -> 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 -> false.
+    // Primary is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel, org.mockito.Mockito.times(2)).newCall(any(), any());
+
+    // Decrease threshold to 2 seconds.
+    timeoutThreshold.set(TimeUnit.SECONDS.toNanos(2));
+
+    // Call at time=5s. elapsed 5s > 2s -> true.
+    // Fallback is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockFallbackChannel).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.toNanos(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 -> false.
+    // Primary is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel).newCall(any(), any());
+
+    // Increase threshold to 20 seconds.
+    timeoutThreshold.set(TimeUnit.SECONDS.toNanos(20));
+
+    // Advance time by 5 seconds (total 14s).
+    time.addAndGet(TimeUnit.SECONDS.toNanos(5));
+
+    // Call at time=14s. elapsed 14s <= 20s -> false.
+    // Still routes to primary because of increased threshold
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockChannel, org.mockito.Mockito.times(2)).newCall(any(), any());
+
+    // Advance time by 7s (total 21s).
+    time.addAndGet(TimeUnit.SECONDS.toNanos(7));
+
+    // Call at time=21s. elapsed 21s > 20s -> true.
+    // Fallback is used.
+    failoverChannel.newCall(methodDescriptor, CallOptions.DEFAULT);
+    verify(mockFallbackChannel).newCall(any(), any());
+  }
+
   @Test
   public void testStateFallbackWhenPrimaryStartsNonReadyWithoutTransition() {
     // Primary starts in a non-ready state and stays there. Even without a state transition,

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 is not ready.
+  // Only used if direcpath is enabled for the job. Default is 15 seconds.
+  optional int64 directpath_primary_not_ready_wait_nanos = 6 [default = 15000000000];
+
   reserved 1, 2;
 }