Merge pull request #1003: [beam-core] fix waiting in finishing bundle in ProximaIO.write

Author: je-ik

beam/core/src/main/java/cz/o2/proxima/beam/io/ProximaIO.java

@@ -124,7 +124,8 @@ public void finishBundle() {
       while (missingResponses.get() > 0) {
         long elapsed = System.currentTimeMillis() - startTime;
         if (elapsed >= bundleFinalizeTimeoutMs) {
-          throw new IllegalStateException("Failed to flush bundle within timeout of 5s");
+          throw new IllegalStateException(
+              "Failed to flush bundle within timeout of " + bundleFinalizeTimeoutMs);
         }
         // clone to avoid ConcurrentModificationException
         final Collection<CompletableFuture<Pair<Boolean, Throwable>>> unfinished;
@@ -135,9 +136,12 @@ public void finishBundle() {
         Optional<Pair<Boolean, Throwable>> failedFuture =
             unfinished.stream()
                 .map(
-                    f ->
-                        ExceptionUtils.uncheckedFactory(
-                            () -> f.get(bundleFinalizeTimeoutMs - elapsed, TimeUnit.MILLISECONDS)))
+                    f -> {
+                      long maxWaitMs =
+                          bundleFinalizeTimeoutMs - System.currentTimeMillis() + startTime;
+                      return ExceptionUtils.uncheckedFactory(
+                          () -> f.get(maxWaitMs, TimeUnit.MILLISECONDS));
+                    })
                 .filter(p -> !p.getFirst())
                 // this will be retried
                 .filter(p -> !(p.getSecond() instanceof TransactionRejectedException))

direct/io-pubsub/module/src/main/java/cz/o2/proxima/direct/io/pubsub/AbstractPubSubWriter.java

@@ -142,6 +142,7 @@ public synchronized void close() {
         }
         executor.shutdown();
         executor.awaitTermination(10, TimeUnit.SECONDS);
+        executor = null;
         publisher.shutdown();
       } catch (Exception ex) {
         log.warn("Failed to shutdown publisher {}", publisher, ex);