@@ -76,6 +76,7 @@ public class BulkIngester<Context> implements AutoCloseable {
76
76
private final FnCondition addCondition = new FnCondition (lock , this ::canAddOperation );
77
77
private final FnCondition sendRequestCondition = new FnCondition (lock , this ::canSendRequest );
78
78
private final FnCondition closeCondition = new FnCondition (lock , this ::closedAndFlushed );
79
+ private AtomicInteger listenerInProgressCount = new AtomicInteger ();
79
80
80
81
private static class RequestExecution <Context > {
81
82
public final long id ;
@@ -235,7 +236,7 @@ private boolean canAddOperation() {
235
236
}
236
237
237
238
private boolean closedAndFlushed () {
238
- return isClosed && operations .isEmpty () && requestsInFlightCount == 0 ;
239
+ return isClosed && operations .isEmpty () && requestsInFlightCount == 0 && listenerInProgressCount . get () == 0 ;
239
240
}
240
241
241
242
//----- Ingester logic
@@ -314,14 +315,32 @@ public void flush() {
314
315
if (resp != null ) {
315
316
// Success
316
317
if (listener != null ) {
317
- scheduler .submit (() -> listener .afterBulk (exec .id , exec .request ,
318
- exec .contexts , resp ));
318
+ listenerInProgressCount .incrementAndGet ();
319
+ scheduler .submit (() -> {
320
+ try {
321
+ listener .afterBulk (exec .id , exec .request , exec .contexts , resp );
322
+ }
323
+ finally {
324
+ if (listenerInProgressCount .decrementAndGet () == 0 ){
325
+ closeCondition .signalIfReady ();
326
+ }
327
+ }
328
+ });
319
329
}
320
330
} else {
321
331
// Failure
322
332
if (listener != null ) {
323
- scheduler .submit (() -> listener .afterBulk (exec .id , exec .request ,
324
- exec .contexts , thr ));
333
+ listenerInProgressCount .incrementAndGet ();
334
+ scheduler .submit (() -> {
335
+ try {
336
+ listener .afterBulk (exec .id , exec .request , exec .contexts , thr );
337
+ }
338
+ finally {
339
+ if (listenerInProgressCount .decrementAndGet () == 0 ){
340
+ closeCondition .signalIfReady ();
341
+ }
342
+ }
343
+ });
325
344
}
326
345
}
327
346
0 commit comments