m-trieu commented on code in PR #32774:
URL: https://github.com/apache/beam/pull/32774#discussion_r1801506126
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSender.java:
##########
@@ -131,46 +132,48 @@ private static GetWorkRequest
withRequestBudget(GetWorkRequest request, GetWorkB
return
request.toBuilder().setMaxItems(budget.items()).setMaxBytes(budget.bytes()).build();
}
- @SuppressWarnings("ReturnValueIgnored")
- void startStreams() {
- getWorkStream.get();
- getDataStream.get();
- commitWorkStream.get();
- workCommitter.get().start();
- // *stream.get() is all memoized in a threadsafe manner.
- started.set(true);
- }
-
- void closeAllStreams() {
- // Supplier<Stream>.get() starts the stream which is an expensive
operation as it initiates the
- // streaming RPCs by possibly making calls over the network. Do not close
the streams unless
- // they have already been started.
- if (started.get()) {
- getWorkStream.get().shutdown();
- getDataStream.get().shutdown();
- workCommitter.get().stop();
- commitWorkStream.get().shutdown();
+ synchronized void start() {
+ if (!started.get()) {
+ // Start these 3 streams in parallel since they each may perform
blocking IO.
+ CompletableFuture.allOf(
+ CompletableFuture.runAsync(getWorkStream::start, streamStarter),
+ CompletableFuture.runAsync(getDataStream::start, streamStarter),
+ CompletableFuture.runAsync(commitWorkStream::start,
streamStarter))
+ .join();
+ workCommitter.start();
+ started.set(true);
}
}
@Override
- public void adjustBudget(long itemsDelta, long bytesDelta) {
- getWorkBudget.set(getWorkBudget.get().apply(itemsDelta, bytesDelta));
+ public synchronized void close() {
if (started.get()) {
- getWorkStream.get().adjustBudget(itemsDelta, bytesDelta);
+ getWorkStream.shutdown();
+ getDataStream.shutdown();
+ workCommitter.stop();
+ commitWorkStream.shutdown();
}
}
@Override
- public GetWorkBudget remainingBudget() {
- return started.get() ? getWorkStream.get().remainingBudget() :
getWorkBudget.get();
+ public void setBudget(long items, long bytes) {
+ GetWorkBudget adjustment =
GetWorkBudget.builder().setItems(items).setBytes(bytes).build();
+ getWorkBudget.set(adjustment);
+ if (started.get()) {
+ getWorkStream.adjustBudget(adjustment);
Review Comment:
rebased onto https://github.com/apache/beam/pull/32775
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
