From 3acad289c57d337d84b288abb54902a964a69e73 Mon Sep 17 00:00:00 2001 From: Jan Lukavsky Date: Mon, 20 Apr 2026 12:41:41 +0200 Subject: [PATCH] [beam-core] limit number of inflight requests in ProximaIO.Write --- .../beam/core/ProximaPipelineOptions.java | 10 ++++ .../java/cz/o2/proxima/beam/io/ProximaIO.java | 58 ++++++++++++++----- .../proxima/beam/io/ProximaIOWriteFnTest.java | 4 +- 3 files changed, 56 insertions(+), 16 deletions(-) diff --git a/beam/core/src/main/java/cz/o2/proxima/beam/core/ProximaPipelineOptions.java b/beam/core/src/main/java/cz/o2/proxima/beam/core/ProximaPipelineOptions.java index d82a74cd4..60cd556d7 100644 --- a/beam/core/src/main/java/cz/o2/proxima/beam/core/ProximaPipelineOptions.java +++ b/beam/core/src/main/java/cz/o2/proxima/beam/core/ProximaPipelineOptions.java @@ -60,4 +60,14 @@ public Iterable> getPipelineOptions() { long getProximaIOWriteFinalizeTimeoutMs(); void setProximaIOWriteFinalizeTimeoutMs(long timeout); + + @Default.Integer(1000) + int getProximaIOMaxPendingWrites(); + + void setProximaIOMaxPendingWrites(int value); + + @Default.Integer(50) + int getProximaIOTransactionWriteWeight(); + + void setProximaIOTransactionWriteWeight(int value); } diff --git a/beam/core/src/main/java/cz/o2/proxima/beam/io/ProximaIO.java b/beam/core/src/main/java/cz/o2/proxima/beam/io/ProximaIO.java index 537664446..fa664e456 100644 --- a/beam/core/src/main/java/cz/o2/proxima/beam/io/ProximaIO.java +++ b/beam/core/src/main/java/cz/o2/proxima/beam/io/ProximaIO.java @@ -18,6 +18,7 @@ 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.repository.TransactionMode; import cz.o2.proxima.core.storage.StreamElement; import cz.o2.proxima.core.util.ExceptionUtils; import cz.o2.proxima.core.util.Pair; @@ -75,13 +76,15 @@ private Write(RepositoryFactory repositoryFactory) { @Override public PDone expand(PCollection input) { - long bundleFinalizeTimeoutMs = - input - .getPipeline() - .getOptions() - .as(ProximaPipelineOptions.class) - .getProximaIOWriteFinalizeTimeoutMs(); - input.apply("Write", ParDo.of(new WriteFn(bundleFinalizeTimeoutMs, repositoryFactory))); + ProximaPipelineOptions proximaOpts = + input.getPipeline().getOptions().as(ProximaPipelineOptions.class); + long bundleFinalizeTimeoutMs = proximaOpts.getProximaIOWriteFinalizeTimeoutMs(); + int maxPendingWrites = proximaOpts.getProximaIOMaxPendingWrites(); + int weight = proximaOpts.getProximaIOTransactionWriteWeight(); + input.apply( + "Write", + ParDo.of( + new WriteFn(bundleFinalizeTimeoutMs, maxPendingWrites, weight, repositoryFactory))); return PDone.in(input.getPipeline()); } } @@ -90,15 +93,25 @@ static class WriteFn extends DoFn { private final RepositoryFactory repositoryFactory; private final long bundleFinalizeTimeoutMs; + private final int maxPendingWrites; + private final int transactionalWriteWeight; private transient DirectDataOperator direct; private transient Set>> pendingWrites; private transient AtomicInteger missingResponses; + private transient AtomicInteger inflightWriteWeights; + + WriteFn( + long bundleFinalizeTimeoutMs, + int maxPendingWrites, + int transactionalWriteWeight, + RepositoryFactory repositoryFactory) { - WriteFn(long bundleFinalizeTimeoutMs, RepositoryFactory repositoryFactory) { this.bundleFinalizeTimeoutMs = bundleFinalizeTimeoutMs; this.repositoryFactory = repositoryFactory; + this.maxPendingWrites = maxPendingWrites; + this.transactionalWriteWeight = transactionalWriteWeight; } @VisibleForTesting @@ -116,6 +129,7 @@ public void startBundle() { // we access the collection asynchronously on completion of writes pendingWrites = Collections.synchronizedSet(new HashSet<>()); missingResponses = new AtomicInteger(); + inflightWriteWeights = new AtomicInteger(); } @FinishBundle @@ -160,22 +174,38 @@ public void processElement(@Element StreamElement element) { AtomicReference writeRunnableRef = new AtomicReference<>(); // increment missing responses outside the retry runnable missingResponses.incrementAndGet(); + boolean isTransactional = + element.getAttributeDescriptor().getTransactionMode() != TransactionMode.NONE; + int weight = isTransactional ? transactionalWriteWeight : 1; + synchronized (pendingWrites) { + while (inflightWriteWeights.get() >= maxPendingWrites) { + ExceptionUtils.unchecked(() -> pendingWrites.wait(100)); + } + inflightWriteWeights.addAndGet(weight); + } writeRunnableRef.set( () -> { CompletableFuture> writeResult = new CompletableFuture<>(); writeResult.thenAccept( r -> { - if (Boolean.TRUE.equals(r.getFirst())) { - // remove successfully completed write - missingResponses.decrementAndGet(); - pendingWrites.remove(writeResult); - } else if (r.getSecond() instanceof TransactionRejectedException) { + if (r.getSecond() instanceof TransactionRejectedException) { // restart the writing transaction writeRunnableRef.get().run(); // transaction rejected, restart transaction pendingWrites.remove(writeResult); + } else { + // this is no longer inflight + inflightWriteWeights.addAndGet(-weight); + synchronized (pendingWrites) { + pendingWrites.notify(); + if (Boolean.TRUE.equals(r.getFirst())) { + // remove successfully completed write + missingResponses.decrementAndGet(); + pendingWrites.remove(writeResult); + } + } + // else keep the failed future until finish bundle } - // else keep the failed future until finish bundle }); pendingWrites.add(writeResult); writer.write( diff --git a/beam/core/src/test/java/cz/o2/proxima/beam/io/ProximaIOWriteFnTest.java b/beam/core/src/test/java/cz/o2/proxima/beam/io/ProximaIOWriteFnTest.java index 8058cf848..c7e5450b9 100644 --- a/beam/core/src/test/java/cz/o2/proxima/beam/io/ProximaIOWriteFnTest.java +++ b/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 status = gateway.getAttribute("status"); - private final WriteFn writeFn = new WriteFn(1000L, repository.asFactory()); + private final WriteFn writeFn = new WriteFn(1000L, 100, 50, 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(30000L, repository.asFactory()) { + new WriteFn(30000L, 100, 50, repository.asFactory()) { @Override OnlineAttributeWriter getWriterForElement(StreamElement element) { return mockWriter;