scwhittle commented on a change in pull request #17004:
URL: https://github.com/apache/beam/pull/17004#discussion_r818942554
##########
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 {
Review comment:
I see that it is recommended to override the array methods for
performance:
https://docs.oracle.com/javase/7/docs/api/java/io/OutputStream.html#write(byte[],%20int,%20int)
##########
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 {
Review comment:
add comment with reason for this
##########
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:
Don't think this needs to be atomic, everything is happening on this
thread
--
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]
