dpcollins-google commented on a change in pull request #17004:
URL: https://github.com/apache/beam/pull/17004#discussion_r818944404
##########
File path:
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/GrpcWindmillServer.java
##########
@@ -1462,36 +1463,54 @@ private void issueBatchedRequest(Map<Long,
PendingRequest> requests) {
}
}
- private void issueMultiChunkRequest(final long id, PendingRequest
pendingRequest) {
- Preconditions.checkNotNull(pendingRequest.computation);
- final ByteString serializedCommit =
pendingRequest.request.toByteString();
+ private class ChunkingByteStream extends OutputStream {
+ private final ByteString.Output output =
ByteString.newOutput(COMMIT_STREAM_CHUNK_SIZE);
+ private final Consumer<ByteString> chunkWriter;
- synchronized (this) {
- pending.put(id, pendingRequest);
- for (int i = 0; i < serializedCommit.size(); i +=
COMMIT_STREAM_CHUNK_SIZE) {
- int end = i + COMMIT_STREAM_CHUNK_SIZE;
- ByteString chunk = serializedCommit.substring(i, Math.min(end,
serializedCommit.size()));
-
- StreamingCommitRequestChunk.Builder chunkBuilder =
- StreamingCommitRequestChunk.newBuilder()
- .setRequestId(id)
- .setSerializedWorkItemCommit(chunk)
- .setComputationId(pendingRequest.computation)
- .setShardingKey(pendingRequest.request.getShardingKey());
- int remaining = serializedCommit.size() - end;
- if (remaining > 0) {
- chunkBuilder.setRemainingBytesForWorkItem(remaining);
- }
+ ChunkingByteStream(Consumer<ByteString> chunkWriter) {
+ this.chunkWriter = chunkWriter;
+ }
- StreamingCommitWorkRequest requestChunk =
-
StreamingCommitWorkRequest.newBuilder().addCommitChunk(chunkBuilder).build();
- try {
- send(requestChunk);
- } catch (IllegalStateException e) {
- // Stream was broken, request will be retried when stream is
reopened.
- break;
- }
+ @Override
+ public void close() {
+ flushBytes();
+ }
+
+ @Override
+ public void write(int b) throws IOException {
+ output.write(b);
+ if (output.size() == COMMIT_STREAM_CHUNK_SIZE) {
+ flushBytes();
+ }
+ }
+
+ private void flushBytes() {
+ chunkWriter.accept(output.toByteString());
+ output.reset();
+ }
+ }
+
+ private void issueMultiChunkRequest(final long id, PendingRequest
pendingRequest) {
+ Preconditions.checkNotNull(pendingRequest.computation);
+ AtomicLong remaining = new
AtomicLong(pendingRequest.request.getSerializedSize());
Review comment:
It needs to be atomic because everything captured by a closure needs to
be effectively final... I could alternatively have an explicit anonymous class
if you'd prefer?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
