[Dataflow Streaming] Speed up non-portable worker precommit tests

runners/google-cloud-dataflow-java/worker/build.gradle

@@ -156,6 +156,8 @@ applyJavaNature(
 // Configure the worker root project
 
 tasks.withType(Test).configureEach {
+  maxParallelForks 4
+
   // WindmillStateTestUtils.assertNoReference walks every object reachable from the Windmill
   // state cache (Guava Cache / ConcurrentHashMap) to ensure no per-work-item WindmillStateReader
   // leaks into the global cache. It uses reflection (Field.setAccessible) into JDK internals

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java

@@ -532,7 +532,7 @@ public Instant startTime() {
 
   public void waitForEmptyWorkQueue() {
     while (!workToOffer.isEmpty()) {
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+      Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
     }
   }
 
@@ -542,17 +542,17 @@ public Map<Long, WorkItemCommitRequest> waitForAndGetCommitsWithTimeout(
     Instant waitStart = Instant.now();
     while (commitsReceived.size() < commitsRequested + numCommits
         && Instant.now().isBefore(waitStart.plus(timeout))) {
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+      Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
     }
     commitsRequested += numCommits;
     return commitsReceived;
   }
 
   public Map<Long, WorkItemCommitRequest> waitForAndGetCommits(int numCommits) {
     LOG.debug("waitForAndGetCommitsRequest: {}", numCommits);
-    int maxTries = 100;
+    int maxTries = 1000;
     while (maxTries-- > 0 && commitsReceived.size() < commitsRequested + numCommits) {
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+      Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
     }
 
     assertFalse(
@@ -588,9 +588,9 @@ public void clearMultiKeyCommitsReceived() {
   public ConcurrentHashMap<Long, Consumer<Windmill.CommitStatus>> waitForDroppedCommits(
       int droppedCommits) {
     LOG.debug("waitForDroppedCommits: {}", droppedCommits);
-    int maxTries = 100;
+    int maxTries = 1000;
     while (maxTries-- > 0 && droppedStreamingCommits.size() < droppedCommits) {
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+      Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
     }
     assertEquals(droppedCommits, droppedStreamingCommits.size());
     return droppedStreamingCommits;

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerApplianceTest.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Unit tests for {@link StreamingDataflowWorker} executing in Appliance mode ({@code
+ * streamingEngine = false}).
+ */
+@RunWith(JUnit4.class)
+public class StreamingDataflowWorkerApplianceTest extends StreamingDataflowWorkerTest {
+  public StreamingDataflowWorkerApplianceTest() {
+    super(false);
+  }
+}

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerStreamingEngineTest.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Unit tests for {@link StreamingDataflowWorker} executing in Streaming Engine mode ({@code
+ * streamingEngine = true}).
+ */
+@RunWith(JUnit4.class)
+public class StreamingDataflowWorkerStreamingEngineTest extends StreamingDataflowWorkerTest {
+  public StreamingDataflowWorkerStreamingEngineTest() {
+    super(true);
+  }
+}

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

@@ -217,17 +217,15 @@
 import org.junit.rules.Timeout;
 import org.junit.runner.Description;
 import org.junit.runner.RunWith;
-import org.junit.runners.Parameterized;
 import org.junit.runners.model.Statement;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /** Unit tests for {@link StreamingDataflowWorker}. */
-@RunWith(Parameterized.class)
 @SuppressWarnings({"deprecation", "unused"})
-public class StreamingDataflowWorkerTest {
+public abstract class StreamingDataflowWorkerTest {
 
   private static final Logger LOG = LoggerFactory.getLogger(StreamingDataflowWorkerTest.class);
   private static final IntervalWindow DEFAULT_WINDOW =
@@ -292,7 +290,7 @@ private static Function<GetDataRequest, GetDataResponse> emptyDataResponderWithF
     };
   }
 
-  private final boolean streamingEngine;
+  protected final boolean streamingEngine;
   private final Supplier<Long> idGenerator =
       new Supplier<Long>() {
         private final AtomicLong idGenerator = new AtomicLong(1L);
@@ -324,15 +322,10 @@ public Long get() {
           errorCollector, computationId -> computationStateCache.get(computationId));
   private StreamingCounters streamingCounters;
 
-  public StreamingDataflowWorkerTest(Boolean streamingEngine) {
+  protected StreamingDataflowWorkerTest(Boolean streamingEngine) {
     this.streamingEngine = streamingEngine;
   }
 
-  @Parameterized.Parameters(name = "{index}: [streamingEngine={0}]")
-  public static Iterable<Object[]> data() {
-    return Arrays.asList(new Object[][] {{false}, {true}});
-  }
-
   private static CounterUpdate getCounter(Iterable<CounterUpdate> counters, String name) {
     for (CounterUpdate counter : counters) {
       if (counter.getNameAndKind().getName().equals(name)) {
@@ -2404,9 +2397,9 @@ public void testExceptions() throws Exception {
     server.waitForEmptyWorkQueue();
 
     // Wait until the worker has given up.
-    int maxTries = 100;
+    int maxTries = 1000;
     while (maxTries-- > 0 && !worker.workExecutorIsEmpty()) {
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+      Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
     }
     assertTrue(worker.workExecutorIsEmpty());
 
@@ -5972,7 +5965,7 @@ static class WorkDoFn extends DoFn<KV<String, String>, KV<String, String>> {
     @ProcessElement
     public void processElement(ProcessContext c, @StateId("state") ValueState<String> state) {
       try {
-        Thread.sleep(1000);
+        Thread.sleep(10);
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
       }

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/status/WorkerStatusPagesTest.java

@@ -95,7 +95,8 @@ public void testUnknownHandler() throws Exception {
   }
 
   private String getPage(String requestURL) throws Exception {
-    String request = String.format("GET %s HTTP/1.1\nhost: localhost\n\n", requestURL);
+    String request =
+        String.format("GET %s HTTP/1.1\r\nhost: localhost\r\nConnection: close\r\n\r\n", requestURL);
     return connector.getResponses(request);
   }
 }

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingGlobalConfigHandleImplTest.java

@@ -250,7 +250,6 @@ public void registerConfigObserver_shouldNotCallCallbackForIfConfigRemainsSame()
     // call setter again with same config
     globalConfigHandle.setConfig(configToSet.get());
     assertTrue(latch.await(10, TimeUnit.SECONDS));
-    Thread.sleep(TimeUnit.SECONDS.toMillis(10));
     assertEquals(1, callbackCount.get());
   }
 

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java

@@ -316,7 +316,7 @@ public void testRecordTotalTimeMaxActiveThreadsUsed() throws Exception {
     stop.countDown();
     while (executor.activeCount() != 0) {
       // Waiting for all threads to be ended.
-      Thread.sleep(200);
+      Thread.sleep(10);
     }
     // Max pool size was reached so the allThreadsActiveTime() was updated.
     assertThat(executor.allThreadsActiveTime(), greaterThan(0L));
@@ -357,7 +357,7 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsIncrease
     stop.countDown();
     while (executor.activeCount() != 0) {
       // Waiting for all threads to be ended.
-      Thread.sleep(200);
+      Thread.sleep(10);
     }
     // Max pool size was updated during execution but allThreadsActiveTime() was still recorded
     // for the thread which reached the old max pool size.
@@ -393,7 +393,7 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsReduced(
     processStop1.countDown();
     while (executor.activeCount() != 0) {
       // Waiting for all threads to be ended.
-      Thread.sleep(200);
+      Thread.sleep(10);
     }
 
     // Reduce max pool size to 1
@@ -409,7 +409,7 @@ public void testRecordTotalTimeMaxActiveThreadsUsedWhenMaximumPoolSizeIsReduced(
 
     while (executor.activeCount() != 0) {
       // Waiting for all threads to be ended.
-      Thread.sleep(200);
+      Thread.sleep(10);
     }
 
     // allThreadsActiveTime() should be recorded

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

@@ -213,7 +213,8 @@ private void waitForHealthChecks(int expectedHealthChecks) {
             waitedMillis,
             numHealthChecks.get(),
             expectedHealthChecks);
-        Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+        Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
+        waitedMillis += 10;
       }
     }
 
@@ -248,8 +249,8 @@ private void waitForSends(int expectedSends) {
       while (numSends.get() < expectedSends) {
         LOG.info(
             "Waited {}ms for {} sends, current sends: {}", millisWaited, expectedSends, numSends);
-        Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
-        millisWaited += 100;
+        Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS);
+        millisWaited += 10;
       }
     }
 

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

@@ -22,6 +22,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 import org.joda.time.DateTime;
 import org.joda.time.Instant;
@@ -33,12 +34,14 @@
 public class StreamDebugMetricsTest {
 
   @Test
-  public void testSummaryMetrics_noRestarts() throws InterruptedException {
-    StreamDebugMetrics streamDebugMetrics = StreamDebugMetrics.forTesting(Instant::now);
+  public void testSummaryMetrics_noRestarts() {
+    AtomicLong nowMs = new AtomicLong(1000L);
+    StreamDebugMetrics streamDebugMetrics =
+        StreamDebugMetrics.forTesting(() -> new Instant(nowMs.get()));
     streamDebugMetrics.recordStart();
     streamDebugMetrics.recordSend();
     streamDebugMetrics.recordResponse();
-    Thread.sleep(1000);
+    nowMs.addAndGet(100);
     StreamDebugMetrics.Snapshot metricsSnapshot = streamDebugMetrics.getSummaryMetrics();
     assertFalse(metricsSnapshot.shutdownTime().isPresent());
     assertTrue(metricsSnapshot.sleepLeft() <= 0);

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateInternalsTest.java

@@ -1656,7 +1656,6 @@ false, key(NAMESPACE, tag), STATE_FAMILY, VarIntCoder.of()))
             new Iterator<Map.Entry<ByteString, Iterable<Integer>>>() {
               final int targetEntries = 1_000_000; // return 1 million entries, which is 10 GBs
               final byte[] entryKey = new byte[10_000]; // each key is 10KB
-              final Random rand = new Random();
               int returnedEntries = 0;
 
               @Override
@@ -1667,7 +1666,10 @@ public boolean hasNext() {
               @Override
               public Map.Entry<ByteString, Iterable<Integer>> next() {
                 returnedEntries++;
-                rand.nextBytes(entryKey);
+                entryKey[0] = (byte) (returnedEntries >> 24);
+                entryKey[1] = (byte) (returnedEntries >> 16);
+                entryKey[2] = (byte) (returnedEntries >> 8);
+                entryKey[3] = (byte) returnedEntries;
                 return multimapEntry(entryKey, 1);
               }
             },
@@ -1696,7 +1698,6 @@ true, key(NAMESPACE, tag), STATE_FAMILY, VarIntCoder.of()))
             new Iterator<Map.Entry<ByteString, Iterable<Integer>>>() {
               final int targetEntries = 1_000_000; // return 1 million entries, which is 10 GBs
               final byte[] entryKey = new byte[10_000]; // each key is 10KB
-              final Random rand = new Random();
               int returnedEntries = 0;
 
               @Override
@@ -1707,7 +1708,10 @@ public boolean hasNext() {
               @Override
               public Map.Entry<ByteString, Iterable<Integer>> next() {
                 returnedEntries++;
-                rand.nextBytes(entryKey);
+                entryKey[0] = (byte) (returnedEntries >> 24);
+                entryKey[1] = (byte) (returnedEntries >> 16);
+                entryKey[2] = (byte) (returnedEntries >> 8);
+                entryKey[3] = (byte) returnedEntries;
                 return multimapEntry(entryKey);
               }
             },
@@ -1744,7 +1748,6 @@ false, key(NAMESPACE, tag), STATE_FAMILY, ByteArrayCoder.of()))
             new Iterator<byte[]>() {
               final int targetValues = 1_000_000; // return 1 million values, which is 10 GBs
               final byte[] value = new byte[10_000]; // each value is 10KB
-              final Random rand = new Random();
               int returnedValues = 0;
 
               @Override
@@ -1755,7 +1758,10 @@ public boolean hasNext() {
               @Override
               public byte[] next() {
                 returnedValues++;
-                rand.nextBytes(value);
+                value[0] = (byte) (returnedValues >> 24);
+                value[1] = (byte) (returnedValues >> 16);
+                value[2] = (byte) (returnedValues >> 8);
+                value[3] = (byte) returnedValues;
                 return value;
               }
             };