use direct executor to deflake tests

apache · Nov 26, 2024 · 48a048e · 48a048e
1 parent a06454a
commit 48a048e
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 110 deletions.
diff --git a/...he/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java b/...he/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java
@@ -68,6 +68,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Streams;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -125,7 +126,8 @@ private FanOutStreamingEngineWorkerHarness(
       GetWorkBudgetDistributor getWorkBudgetDistributor,
       GrpcDispatcherClient dispatcherClient,
       Function<WindmillStream.CommitWorkStream, WorkCommitter> workCommitterFactory,
-      ThrottlingGetDataMetricTracker getDataMetricTracker) {
+      ThrottlingGetDataMetricTracker getDataMetricTracker,
+      ExecutorService workerMetadataConsumer) {
     this.jobHeader = jobHeader;
     this.getDataMetricTracker = getDataMetricTracker;
     this.started = false;
@@ -138,9 +140,7 @@ private FanOutStreamingEngineWorkerHarness(
     this.windmillStreamManager =
         Executors.newCachedThreadPool(
             new ThreadFactoryBuilder().setNameFormat(STREAM_MANAGER_THREAD_NAME).build());
-    this.workerMetadataConsumer =
-        Executors.newSingleThreadExecutor(
-            new ThreadFactoryBuilder().setNameFormat(WORKER_METADATA_CONSUMER_THREAD_NAME).build());
+    this.workerMetadataConsumer = workerMetadataConsumer;
     this.getWorkBudgetDistributor = getWorkBudgetDistributor;
     this.totalGetWorkBudget = totalGetWorkBudget;
     this.activeMetadataVersion = Long.MIN_VALUE;
@@ -171,7 +171,11 @@ public static FanOutStreamingEngineWorkerHarness create(
         getWorkBudgetDistributor,
         dispatcherClient,
         workCommitterFactory,
-        getDataMetricTracker);
+        getDataMetricTracker,
+        Executors.newSingleThreadExecutor(
+            new ThreadFactoryBuilder()
+                .setNameFormat(WORKER_METADATA_CONSUMER_THREAD_NAME)
+                .build()));
   }
 
   @VisibleForTesting
@@ -195,7 +199,10 @@ static FanOutStreamingEngineWorkerHarness forTesting(
             getWorkBudgetDistributor,
             dispatcherClient,
             workCommitterFactory,
-            getDataMetricTracker);
+            getDataMetricTracker,
+            // Run the workerMetadataConsumer on the direct calling thread to make testing more
+            // deterministic.
+            MoreExecutors.newDirectExecutorService());
     fanOutStreamingEngineWorkProvider.start();
     return fanOutStreamingEngineWorkProvider;
   }

diff --git a/...ker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java b/...ker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java
@@ -88,17 +88,23 @@ private static Optional<WindmillServiceAddress> parseDirectEndpoint(
             .map(address -> AuthenticatedGcpServiceAddress.create(authenticatingService, address))
             .map(WindmillServiceAddress::create);
 
-    return directEndpointIpV6Address.isPresent()
-        ? directEndpointIpV6Address
-        : tryParseEndpointIntoHostAndPort(endpointProto.getDirectEndpoint())
-            .map(WindmillServiceAddress::create);
+    Optional<WindmillServiceAddress> windmillServiceAddress =
+        directEndpointIpV6Address.isPresent()
+            ? directEndpointIpV6Address
+            : tryParseEndpointIntoHostAndPort(endpointProto.getDirectEndpoint())
+                .map(WindmillServiceAddress::create);
+
+    if (!windmillServiceAddress.isPresent()) {
+      LOG.warn("Endpoint {} could not be parsed into a WindmillServiceAddress.", endpointProto);
+    }
+
+    return windmillServiceAddress;
   }
 
   private static Optional<HostAndPort> tryParseEndpointIntoHostAndPort(String directEndpoint) {
     try {
       return Optional.of(HostAndPort.fromString(directEndpoint));
     } catch (IllegalArgumentException e) {
-      LOG.warn("{} cannot be parsed into a gcpServiceAddress", directEndpoint);
       return Optional.empty();
     }
   }
@@ -113,19 +119,12 @@ private static Optional<HostAndPort> tryParseDirectEndpointIntoIpV6Address(
     try {
       directEndpointAddress = Inet6Address.getByName(endpointProto.getDirectEndpoint());
     } catch (UnknownHostException e) {
-      LOG.warn(
-          "Error occurred trying to parse direct_endpoint={} into IPv6 address. Exception={}",
-          endpointProto.getDirectEndpoint(),
-          e.toString());
       return Optional.empty();
     }
 
     // Inet6Address.getByAddress returns either an IPv4 or an IPv6 address depending on the format
     // of the direct_endpoint string.
     if (!(directEndpointAddress instanceof Inet6Address)) {
-      LOG.warn(
-          "{} is not an IPv6 address. Direct endpoints are expected to be in IPv6 format.",
-          endpointProto.getDirectEndpoint());
       return Optional.empty();
     }
 

diff --git a/...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java b/...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDispatcherClient.java
@@ -150,6 +150,8 @@ public CloudWindmillMetadataServiceV1Alpha1Stub getWindmillMetadataServiceStubBl
       }
     }
 
+    LOG.info("Windmill Service endpoint initialized after {} seconds.", secondsWaited);
+
     ImmutableList<CloudWindmillMetadataServiceV1Alpha1Stub> windmillMetadataServiceStubs =
         dispatcherStubs.get().windmillMetadataServiceStubs();
 
@@ -190,7 +192,7 @@ public void onJobConfig(StreamingGlobalConfig config) {
 
   public synchronized void consumeWindmillDispatcherEndpoints(
       ImmutableSet<HostAndPort> dispatcherEndpoints) {
-    consumeWindmillDispatcherEndpoints(dispatcherEndpoints, /*forceRecreateStubs=*/ false);
+    consumeWindmillDispatcherEndpoints(dispatcherEndpoints, /* forceRecreateStubs= */ false);
   }
 
   private synchronized void consumeWindmillDispatcherEndpoints(

diff --git a/...java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java b/...java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/ChannelCache.java
@@ -18,6 +18,7 @@
 package org.apache.beam.runners.dataflow.worker.windmill.client.grpc.stubs;
 
 import java.io.PrintWriter;
+import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -31,6 +32,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.LoadingCache;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalListener;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalListeners;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.MoreExecutors;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -50,14 +52,11 @@ public final class ChannelCache implements StatusDataProvider {
 
   private ChannelCache(
       Function<WindmillServiceAddress, ManagedChannel> channelFactory,
-      RemovalListener<WindmillServiceAddress, ManagedChannel> onChannelRemoved) {
+      RemovalListener<WindmillServiceAddress, ManagedChannel> onChannelRemoved,
+      Executor channelCloser) {
     this.channelCache =
         CacheBuilder.newBuilder()
-            .removalListener(
-                RemovalListeners.asynchronous(
-                    onChannelRemoved,
-                    Executors.newCachedThreadPool(
-                        new ThreadFactoryBuilder().setNameFormat("GrpcChannelCloser").build())))
+            .removalListener(RemovalListeners.asynchronous(onChannelRemoved, channelCloser))
             .build(
                 new CacheLoader<WindmillServiceAddress, ManagedChannel>() {
                   @Override
@@ -72,11 +71,13 @@ public static ChannelCache create(
     return new ChannelCache(
         channelFactory,
         // Shutdown the channels as they get removed from the cache, so they do not leak.
-        notification -> shutdownChannel(notification.getValue()));
+        notification -> shutdownChannel(notification.getValue()),
+        Executors.newCachedThreadPool(
+            new ThreadFactoryBuilder().setNameFormat("GrpcChannelCloser").build()));
   }
 
   @VisibleForTesting
-  static ChannelCache forTesting(
+  public static ChannelCache forTesting(
       Function<WindmillServiceAddress, ManagedChannel> channelFactory, Runnable onChannelShutdown) {
     return new ChannelCache(
         channelFactory,
@@ -85,7 +86,9 @@ static ChannelCache forTesting(
         notification -> {
           shutdownChannel(notification.getValue());
           onChannelShutdown.run();
-        });
+        },
+        // Run the removal on the calling thread for better determinism in tests.
+        MoreExecutors.directExecutor());
   }
 
   private static void shutdownChannel(ManagedChannel channel) {
@@ -108,6 +111,7 @@ public void remove(WindmillServiceAddress windmillServiceAddress) {
 
   public void clear() {
     channelCache.invalidateAll();
+    channelCache.cleanUp();
   }
 
   /**