address PR comments, fix tests

Author: m-trieu

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

@@ -246,7 +246,7 @@ protected final void executeSafely(Runnable runnable) {
    *     serially. If we recently sent a message before we attempt to schedule the health check, the
    *     stream has been restarted/closed, there is a scheduled health check that hasn't completed
    *     or there was a more recent send by the time we enter the synchronized block, we skip the
-   *     attempt to send scheduled the health check.
+   *     attempt to send the health check.
    */
   public final void maybeScheduleHealthCheck(Instant lastSendThreshold) {
     if (debugMetrics.getLastSendTimeMs() < lastSendThreshold.getMillis()
@@ -255,17 +255,17 @@ public final void maybeScheduleHealthCheck(Instant lastSendThreshold) {
       executeSafely(
           () -> {
             synchronized (this) {
-              if (!clientClosed
-                  && debugMetrics.getLastSendTimeMs() < lastSendThreshold.getMillis()) {
-                try {
+              try {
+                if (!clientClosed
+                    && debugMetrics.getLastSendTimeMs() < lastSendThreshold.getMillis()) {
                   sendHealthCheck();
-                } catch (Exception e) {
-                  logger.debug("Received exception sending health check.", e);
                 }
+              } catch (Exception e) {
+                logger.debug("Received exception sending health check.", e);
+              } finally {
+                // Ready to send another health check after we attempt the scheduled health check.
+                isHealthCheckScheduled.set(false);
               }
-
-              // Ready to send another health check after we attempt the scheduled health check.
-              isHealthCheckScheduled.set(false);
             }
           });
     }
@@ -286,11 +286,12 @@ public final void appendSummaryHtml(PrintWriter writer) {
         .ifPresent(
             metrics ->
                 writer.format(
-                    ", %d restarts, last restart reason [ %s ] at [%s], %d errors",
+                    ", %d restarts, last restart reason [ %s ] at [%s], %d errors, isHealthCheckScheduled=[%s]",
                     metrics.restartCount(),
                     metrics.lastRestartReason(),
                     metrics.lastRestartTime().orElse(null),
-                    metrics.errorCount()));
+                    metrics.errorCount(),
+                    isHealthCheckScheduled.get()));
 
     if (summaryMetrics.isClientClosed()) {
       writer.write(", client closed");

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

@@ -117,9 +117,6 @@ public void testMaybeSendHealthCheck_doesNotSendIfLastScheduleLessThanThreshold(
       throw new RuntimeException(e);
     }
 
-    // Sleep a bit to give sendExecutor time to execute the send().
-    Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
-
     // Set a really long reporting threshold.
     Instant reportingThreshold = Instant.now().minus(Duration.standardHours(1));
 
@@ -129,7 +126,8 @@ public void testMaybeSendHealthCheck_doesNotSendIfLastScheduleLessThanThreshold(
     testStream.maybeScheduleHealthCheck(reportingThreshold);
     Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
 
-    callStreamObserver.waitForSend();
+    callStreamObserver.waitForSends(1);
+    // Sleep just to ensure an async health check doesn't show up
     Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
 
     assertThat(testStream.numHealthChecks.get()).isEqualTo(0);
@@ -204,6 +202,7 @@ protected void shutdownInternal() {}
   private static class TestCallStreamObserver extends CallStreamObserver<Integer> {
     private static final Logger LOG = LoggerFactory.getLogger(AbstractWindmillStreamTest.class);
     private final CountDownLatch sendBlocker = new CountDownLatch(1);
+    private final AtomicInteger numSends = new AtomicInteger();
 
     private final boolean waitForSend;
 
@@ -215,12 +214,12 @@ private void unblockSend() {
       sendBlocker.countDown();
     }
 
-    private void waitForSend() {
+    private void waitForSendUnblocked() {
       try {
         int waitedMillis = 0;
         while (!sendBlocker.await(100, TimeUnit.MILLISECONDS)) {
           waitedMillis += 100;
-          LOG.info("Waiting from send for {}ms", waitedMillis);
+          LOG.info("Waiting from send to be unblocked for {}ms", waitedMillis);
         }
       } catch (InterruptedException e) {
         LOG.error("Interrupted waiting for send().");
@@ -230,9 +229,19 @@ private void waitForSend() {
     @Override
     public void onNext(Integer integer) {
       if (waitForSend) {
-        waitForSend();
+        waitForSendUnblocked();
       } else {
-        sendBlocker.countDown();
+        numSends.incrementAndGet();
+      }
+    }
+
+    private void waitForSends(int expectedSends) {
+      int millisWaited = 0;
+      while (numSends.get() < expectedSends) {
+        LOG.info(
+            "Waited {}ms for {} sends, current sends: {}", millisWaited, expectedSends, numSends);
+        Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+        millisWaited += 100;
       }
     }