[proxima-beam-core] configurable finish bundle timeout for ProximaIO

Author: je-ik

beam/core/src/main/java/cz/o2/proxima/beam/core/ProximaPipelineOptions.java

@@ -55,4 +55,9 @@ public Iterable<Class<? extends PipelineOptions>> getPipelineOptions() {
   boolean getEnforceStableInputForRemoteConsumer();
 
   void setEnforceStableInputForRemoteConsumer(boolean enforce);
+
+  @Default.Long(30000L)
+  long getProximaIOWriteFinalizeTimeoutMs();
+
+  void setProximaIOWriteFinalizeTimeoutMs(long timeout);
 }

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

@@ -15,6 +15,7 @@
  */
 package cz.o2.proxima.beam.io;
 
+import cz.o2.proxima.beam.core.ProximaPipelineOptions;
 import cz.o2.proxima.core.annotations.Experimental;
 import cz.o2.proxima.core.repository.RepositoryFactory;
 import cz.o2.proxima.core.storage.StreamElement;
@@ -74,21 +75,29 @@ private Write(RepositoryFactory repositoryFactory) {
 
     @Override
     public PDone expand(PCollection<StreamElement> input) {
-      input.apply("Write", ParDo.of(new WriteFn(repositoryFactory)));
+      long bundleFinalizeTimeoutMs =
+          input
+              .getPipeline()
+              .getOptions()
+              .as(ProximaPipelineOptions.class)
+              .getProximaIOWriteFinalizeTimeoutMs();
+      input.apply("Write", ParDo.of(new WriteFn(bundleFinalizeTimeoutMs, repositoryFactory)));
       return PDone.in(input.getPipeline());
     }
   }
 
   static class WriteFn extends DoFn<StreamElement, Void> {
 
     private final RepositoryFactory repositoryFactory;
+    private final long bundleFinalizeTimeoutMs;
 
     private transient DirectDataOperator direct;
 
     private transient Set<CompletableFuture<Pair<Boolean, Throwable>>> pendingWrites;
     private transient AtomicInteger missingResponses;
 
-    WriteFn(RepositoryFactory repositoryFactory) {
+    WriteFn(long bundleFinalizeTimeoutMs, RepositoryFactory repositoryFactory) {
+      this.bundleFinalizeTimeoutMs = bundleFinalizeTimeoutMs;
       this.repositoryFactory = repositoryFactory;
     }
 
@@ -113,7 +122,7 @@ public void startBundle() {
     public void finishBundle() {
       long startTime = System.currentTimeMillis();
       while (missingResponses.get() > 0) {
-        if (System.currentTimeMillis() - startTime > 5000) {
+        if (System.currentTimeMillis() - startTime > bundleFinalizeTimeoutMs) {
           throw new IllegalStateException("Failed to flush bundle within timeout of 5s");
         }
         // clone to avoid ConcurrentModificationException
@@ -124,7 +133,10 @@ public void finishBundle() {
         }
         Optional<Pair<Boolean, Throwable>> failedFuture =
             unfinished.stream()
-                .map(f -> ExceptionUtils.uncheckedFactory(() -> f.get(5, TimeUnit.SECONDS)))
+                .map(
+                    f ->
+                        ExceptionUtils.uncheckedFactory(
+                            () -> f.get(bundleFinalizeTimeoutMs, TimeUnit.MILLISECONDS)))
                 .filter(p -> !p.getFirst())
                 // this will be retried
                 .filter(p -> !(p.getSecond() instanceof TransactionRejectedException))

beam/core/src/test/java/cz/o2/proxima/beam/io/ProximaIOWriteFnTest.java

@@ -49,7 +49,7 @@ public class ProximaIOWriteFnTest {
       Repository.ofTest(ConfigFactory.load("test-reference.conf").resolve());
   private final EntityDescriptor gateway = repository.getEntity("gateway");
   private final AttributeDescriptor<byte[]> status = gateway.getAttribute("status");
-  private final WriteFn writeFn = new WriteFn(repository.asFactory());
+  private final WriteFn writeFn = new WriteFn(1000L, repository.asFactory());
   private RandomAccessReader reader;
 
   @Before
@@ -88,7 +88,7 @@ public void testTransactionRejection() {
     AtomicInteger written = new AtomicInteger();
     OnlineAttributeWriter mockWriter = createSerializableWriter(fails, written);
     WriteFn modifiedWriteFn =
-        new WriteFn(repository.asFactory()) {
+        new WriteFn(30000L, repository.asFactory()) {
           @Override
           OnlineAttributeWriter getWriterForElement(StreamElement element) {
             return mockWriter;