StreamingDataflowWorker { public void start() { running.set(true); if (windmillServiceEnabled) { schedulePeriodicGlobalConfigRequests(); } memoryMonitorThread.start(); dispatchThread.start(); commitThread.start(); ExecutionStateSampler.instance().start(); globalWorkerUpdatesTimer = new Timer("GlobalWorkerUpdatesTimer"); globalWorkerUpdatesTimer.schedule( new TimerTask() { @Override public void run() { reportPeriodicWorkerUpdates(); } }, 0, options.getWindmillHarnessUpdateReportingPeriod().getMillis()); refreshWorkTimer = new Timer("RefreshWork"); if (options.getActiveWorkRefreshPeriodMillis() > 0) { refreshWorkTimer.schedule( new TimerTask() { @Override public void run() { refreshActiveWork(); } }, options.getActiveWorkRefreshPeriodMillis(), options.getActiveWorkRefreshPeriodMillis()); } if (windmillServiceEnabled && options.getStuckCommitDurationMillis() > 0) { int periodMillis = Math.max(options.getStuckCommitDurationMillis() / 10, 100); refreshWorkTimer.schedule( new TimerTask() { @Override public void run() { invalidateStuckCommits(); } }, periodMillis, periodMillis); } if (options.getPeriodicStatusPageOutputDirectory() != null) { statusPageTimer = new Timer("DumpStatusPages"); statusPageTimer.schedule( new TimerTask() { @Override public void run() { Collection<Capturable> pages = statusPages.getDebugCapturePages(); if (pages.isEmpty()) { LOG.warn("No captured status pages."); } Long timestamp = Instant.now().getMillis(); for (Capturable page : pages) { PrintWriter writer = null; try { File outputFile = new File( options.getPeriodicStatusPageOutputDirectory(), ("StreamingDataflowWorker" + options.getWorkerId() + "_" + page.pageName() + timestamp.toString()) .replaceAll("/", "_")); writer = new PrintWriter(outputFile, UTF_8.name()); page.captureData(writer); } catch (IOException e) { LOG.warn("Error dumping status page.", e); } finally { if (writer != null) { writer.close(); } } } } }, 60 * 1000, 60 * 1000); } reportHarnessStartup(); } @VisibleForTesting  StreamingDataflowWorker(
      List<MapTask> mapTasks,
      DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
      WorkUnitClient workUnitClient,
      StreamingDataflowWorkerOptions options,
      RunnerApi.@Nullable Pipeline pipeline,
      SdkHarnessRegistry sdkHarnessRegistry,
      boolean publishCounters,
      HotKeyLogger hotKeyLogger); static void main(String[] args); static StreamingDataflowWorker forStreamingFnWorkerHarness(
      List<MapTask> mapTasks,
      WorkUnitClient workUnitClient,
      DataflowWorkerHarnessOptions options,
      RunnerApi.@Nullable Pipeline pipeline,
      SdkHarnessRegistry sdkHarnessRegistry); static StreamingDataflowWorker fromDataflowWorkerHarnessOptions(
      DataflowWorkerHarnessOptions options, SdkHarnessRegistry sdkHarnessRegistry); @VisibleForTesting void setRetryLocallyDelayMs(int retryLocallyDelayMs); @VisibleForTesting void setMaxWorkItemCommitBytes(int maxWorkItemCommitBytes); @VisibleForTesting boolean workExecutorIsEmpty(); void start(); void startStatusPages(); void addWorkerStatusPage(BaseStatusServlet page); void stop(); void waitTillExecutionFinishes(); @VisibleForTesting Iterable<CounterUpdate> buildCounters(); void addFailure(String failureMessage); @VisibleForTesting void reportPeriodicWorkerUpdates(); static final MetricName BIGQUERY_STREAMING_INSERT_THROTTLE_TIME; }
@Test(timeout = 10000) public void testNumberOfWorkerHarnessThreadsIsHonored() throws Exception { int expectedNumberOfThreads = 5; List<ParallelInstruction> instructions = Arrays.asList( makeSourceInstruction(StringUtf8Coder.of()), makeDoFnInstruction(blockingFn, 0, StringUtf8Coder.of()), makeSinkInstruction(StringUtf8Coder.of(), 0)); FakeWindmillServer server = new FakeWindmillServer(errorCollector); StreamingDataflowWorkerOptions options = createTestingPipelineOptions(server); options.setNumberOfWorkerHarnessThreads(expectedNumberOfThreads); StreamingDataflowWorker worker = makeWorker(instructions, options, true ); worker.start(); for (int i = 0; i < expectedNumberOfThreads * 2; ++i) { server.addWorkToOffer(makeInput(i, TimeUnit.MILLISECONDS.toMicros(i))); } BlockingFn.counter.acquire(expectedNumberOfThreads); if (BlockingFn.counter.tryAcquire(500, TimeUnit.MILLISECONDS)) { fail( "Expected number of threads " + expectedNumberOfThreads + " does not match actual " + "number of work items processed concurrently " + BlockingFn.callCounter.get() + "."); } BlockingFn.blocker.countDown(); }
@Test public void testKeyTokenInvalidException() throws Exception { if (streamingEngine) { return; } KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); List<ParallelInstruction> instructions = Arrays.asList( makeSourceInstruction(kvCoder), makeDoFnInstruction(new KeyTokenInvalidFn(), 0, kvCoder), makeSinkInstruction(kvCoder, 1)); FakeWindmillServer server = new FakeWindmillServer(errorCollector); server.addWorkToOffer(makeInput(0, 0, DEFAULT_KEY_STRING, DEFAULT_SHARDING_KEY)); StreamingDataflowWorker worker = makeWorker(instructions, createTestingPipelineOptions(server), true ); worker.start(); server.waitForEmptyWorkQueue(); server.addWorkToOffer(makeInput(1, 0, DEFAULT_KEY_STRING, DEFAULT_SHARDING_KEY)); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1); assertEquals( makeExpectedOutput(1, 0, DEFAULT_KEY_STRING, DEFAULT_SHARDING_KEY, DEFAULT_KEY_STRING) .build(), result.get(1L)); assertEquals(1, result.size()); }
@Test public void testKeyChange() throws Exception { KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); List<ParallelInstruction> instructions = Arrays.asList( makeSourceInstruction(kvCoder), makeDoFnInstruction(new ChangeKeysFn(), 0, kvCoder), makeSinkInstruction(kvCoder, 1)); FakeWindmillServer server = new FakeWindmillServer(errorCollector); for (int i = 0; i < 2; i++) { server.addWorkToOffer( makeInput( i, TimeUnit.MILLISECONDS.toMicros(i), keyStringForIndex(i), DEFAULT_SHARDING_KEY)); server.addWorkToOffer( makeInput( i + 1000, TimeUnit.MILLISECONDS.toMicros(i), keyStringForIndex(i), DEFAULT_SHARDING_KEY + i)); } StreamingDataflowWorker worker = makeWorker(instructions, createTestingPipelineOptions(server), true ); worker.start(); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(4); for (int i = 0; i < 2; i++) { assertTrue(result.containsKey((long) i)); assertEquals( makeExpectedOutput( i, TimeUnit.MILLISECONDS.toMicros(i), keyStringForIndex(i), DEFAULT_SHARDING_KEY, keyStringForIndex(i) + "_data" + i) .build(), result.get((long) i)); assertTrue(result.containsKey((long) i + 1000)); assertEquals( makeExpectedOutput( i + 1000, TimeUnit.MILLISECONDS.toMicros(i), keyStringForIndex(i), DEFAULT_SHARDING_KEY + i, keyStringForIndex(i) + "_data" + (i + 1000)) .build(), result.get((long) i + 1000)); } }
@Test public void testAssignWindows() throws Exception { Duration gapDuration = Duration.standardSeconds(1); CloudObject spec = CloudObject.forClassName("AssignWindowsDoFn"); SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); addString( spec, PropertyNames.SERIALIZED_FN, StringUtils.byteArrayToJsonString( WindowingStrategyTranslation.toMessageProto( WindowingStrategy.of(FixedWindows.of(gapDuration)), sdkComponents) .toByteArray())); ParallelInstruction addWindowsInstruction = new ParallelInstruction() .setSystemName("AssignWindows") .setName("AssignWindows") .setOriginalName("AssignWindowsOriginal") .setParDo( new ParDoInstruction() .setInput(new InstructionInput().setProducerInstructionIndex(0).setOutputNum(0)) .setNumOutputs(1) .setUserFn(spec)) .setOutputs( Arrays.asList( new InstructionOutput() .setOriginalName(DEFAULT_OUTPUT_ORIGINAL_NAME) .setSystemName(DEFAULT_OUTPUT_SYSTEM_NAME) .setName("output") .setCodec( CloudObjects.asCloudObject( WindowedValue.getFullCoder( StringUtf8Coder.of(), IntervalWindow.getCoder()), null)))); List<ParallelInstruction> instructions = Arrays.asList( makeSourceInstruction(StringUtf8Coder.of()), addWindowsInstruction, makeSinkInstruction(StringUtf8Coder.of(), 1)); FakeWindmillServer server = new FakeWindmillServer(errorCollector); int timestamp1 = 0; int timestamp2 = 1000000; server.addWorkToOffer(makeInput(timestamp1, timestamp1)); server.addWorkToOffer(makeInput(timestamp2, timestamp2)); StreamingDataflowWorker worker = makeWorker(instructions, createTestingPipelineOptions(server), false ); worker.start(); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(2); assertThat( result.get((long) timestamp1), equalTo( setMessagesMetadata( PaneInfo.NO_FIRING, intervalWindowBytes(WINDOW_AT_ZERO), makeExpectedOutput(timestamp1, timestamp1)) .build())); assertThat( result.get((long) timestamp2), equalTo( setMessagesMetadata( PaneInfo.NO_FIRING, intervalWindowBytes(WINDOW_AT_ONE_SECOND), makeExpectedOutput(timestamp2, timestamp2)) .build())); }
@Test public void testUnboundedSourcesDrain() throws Exception { List<Integer> finalizeTracker = Lists.newArrayList(); TestCountingSource.setFinalizeTracker(finalizeTracker); FakeWindmillServer server = new FakeWindmillServer(errorCollector); StreamingDataflowWorker worker = makeWorker( makeUnboundedSourcePipeline(), createTestingPipelineOptions(server), true ); worker.start(); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 2" + " cache_token: 3" + " }" + "}", null)); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1); Windmill.WorkItemCommitRequest commit = result.get(2L); UnsignedLong finalizeId = UnsignedLong.fromLongBits(commit.getSourceStateUpdates().getFinalizeIds(0)); assertThat( commit, equalTo( setMessagesMetadata( PaneInfo.NO_FIRING, CoderUtils.encodeToByteArray( CollectionCoder.of(GlobalWindow.Coder.INSTANCE), Arrays.asList(GlobalWindow.INSTANCE)), parseCommitRequest( "key: \"0000000000000001\" " + "sharding_key: 1 " + "work_token: 2 " + "cache_token: 3 " + "source_backlog_bytes: 7 " + "output_messages {" + " destination_stream_id: \"out\"" + " bundles {" + " key: \"0000000000000001\"" + " messages {" + " timestamp: 0" + " data: \"0:0\"" + " }" + " messages_ids: \"\"" + " }" + "} " + "source_state_updates {" + " state: \"\000\"" + " finalize_ids: " + finalizeId + "} " + "source_watermark: 1000")) .build())); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 3" + " cache_token: 3" + " source_state {" + " only_finalize: true" + " finalize_ids: " + finalizeId + " }" + " }" + "}", null)); result = server.waitForAndGetCommits(1); commit = result.get(3L); assertThat( commit, equalTo( parseCommitRequest( "key: \"0000000000000001\" " + "sharding_key: 1 " + "work_token: 3 " + "cache_token: 3 " + "source_state_updates {" + " only_finalize: true" + "} ") .build())); assertThat(finalizeTracker, contains(0)); }
@Test public void testLimitOnOutputBundleSize() throws Exception { List<Integer> finalizeTracker = Lists.newArrayList(); TestCountingSource.setFinalizeTracker(finalizeTracker); final int numMessagesInCustomSourceShard = 100000; final int inflatedSizePerMessage = 10000; FakeWindmillServer server = new FakeWindmillServer(errorCollector); StreamingDataflowWorker worker = makeWorker( makeUnboundedSourcePipeline( numMessagesInCustomSourceShard, new InflateDoFn(inflatedSizePerMessage)), createTestingPipelineOptions(server), false ); worker.start(); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 1" + " cache_token: 1" + " }" + "}", null)); Matcher<Integer> isWithinBundleSizeLimits = both(greaterThan(StreamingDataflowWorker.MAX_SINK_BYTES * 9 / 10)) .and(lessThan(StreamingDataflowWorker.MAX_SINK_BYTES * 11 / 10)); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1); Windmill.WorkItemCommitRequest commit = result.get(1L); assertThat(commit.getSerializedSize(), isWithinBundleSizeLimits); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 2" + " cache_token: 1" + " }" + "}", null)); result = server.waitForAndGetCommits(1); commit = result.get(2L); assertThat(commit.getSerializedSize(), isWithinBundleSizeLimits); }
@Test public void testLimitOnOutputBundleSizeWithMultipleSinks() throws Exception { List<Integer> finalizeTracker = Lists.newArrayList(); TestCountingSource.setFinalizeTracker(finalizeTracker); final int numMessagesInCustomSourceShard = 100000; final int inflatedSizePerMessage = 10000; List<ParallelInstruction> instructions = new ArrayList<>(); instructions.addAll( makeUnboundedSourcePipeline( numMessagesInCustomSourceShard, new InflateDoFn(inflatedSizePerMessage))); instructions.add( makeSinkInstruction( DEFAULT_DESTINATION_STREAM_ID + "-1", StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE)); instructions.add( makeSinkInstruction( DEFAULT_DESTINATION_STREAM_ID + "-2", StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE)); FakeWindmillServer server = new FakeWindmillServer(errorCollector); StreamingDataflowWorker worker = makeWorker(instructions, createTestingPipelineOptions(server), true ); worker.start(); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 1" + " cache_token: 1" + " }" + "}", null)); Matcher<Integer> isWithinBundleSizeLimits = both(greaterThan(StreamingDataflowWorker.MAX_SINK_BYTES * 9 / 10)) .and(lessThan(StreamingDataflowWorker.MAX_SINK_BYTES * 11 / 10)); Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1); Windmill.WorkItemCommitRequest commit = result.get(1L); assertThat(commit.getSerializedSize(), isWithinBundleSizeLimits); server.addWorkToOffer( buildInput( "work {" + " computation_id: \"computation\"" + " input_data_watermark: 0" + " work {" + " key: \"0000000000000001\"" + " sharding_key: 1" + " work_token: 2" + " cache_token: 1" + " }" + "}", null)); result = server.waitForAndGetCommits(1); commit = result.get(2L); assertThat(commit.getSerializedSize(), isWithinBundleSizeLimits); }