prodriguezdefino commented on code in PR #40606:
URL: https://github.com/apache/beam/pull/40606#discussion_r4300031495
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java:
##########
@@ -177,19 +197,23 @@ public void awaitCompletion() throws Exception {
// collection.
Iterator<Elements.Data> dataIterator;
Iterator<Elements.Timers> timersIterator;
+ ElementWithCallback wrapper;
{
- BeamFnApi.Elements elements = queue.take();
+ wrapper = queue.take();
+ BeamFnApi.Elements elements = wrapper.elements;
dataIterator = createDiscardingIterator(elements.getDataList());
timersIterator = createDiscardingIterator(elements.getTimersList());
}
// The SDK is no longer blocked on receiving more data from the Runner.
consumingReceivedData.set(true);
- if (multiplexElements(dataIterator, timersIterator)) {
+ boolean isDone = multiplexElements(dataIterator, timersIterator);
+ wrapper.callback.run();
Review Comment:
Now `awaitCompletion` keeps the callback of the element it took from the
queue and runs it after cancelling the queue, so a failed bundle releases the
permit of that element as well as the queued ones.
`BeamFnDataGrpcMultiplexerTest.testPermitsReleasedWhenBundleFails` sends three
messages to a bundle whose consumer throws and checks for three `request(1)`
calls on a mocked stream.
--
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]
