diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java b/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java index b16780035f..ed645b6e5e 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/write/Sender.java @@ -393,20 +393,28 @@ private void sendWriteRequest(int destination, short acks, List } else { writeBatchByTable.forEach( (tableId, writeBatches) -> { - if (isLogBatches(writeBatches)) { - sendProduceLogRequestAndHandleResponse( - gateway, - makeProduceLogRequest( - tableId, acks, maxRequestTimeoutMs, writeBatches), - tableId, - writeBatches); - } else { - sendPutKvRequestAndHandleResponse( - gateway, - makePutKvRequest( - tableId, acks, maxRequestTimeoutMs, writeBatches), - tableId, - writeBatches); + try { + if (isLogBatches(writeBatches)) { + sendProduceLogRequestAndHandleResponse( + gateway, + makeProduceLogRequest( + tableId, acks, maxRequestTimeoutMs, writeBatches), + tableId, + writeBatches); + } else { + sendPutKvRequestAndHandleResponse( + gateway, + makePutKvRequest( + tableId, acks, maxRequestTimeoutMs, writeBatches), + tableId, + writeBatches); + } + } catch (Throwable t) { + // A gateway may throw before returning a future, for example when RPC + // encoding runs out of direct memory. No callback is registered in that + // case, so complete the drained batches to release their buffer pages + // and in-flight state. + handleWriteRequestException(t, writeBatches); } }); } diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java index 69d9c669c9..53788d786b 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/write/SenderTest.java @@ -128,6 +128,42 @@ void testSimple() throws Exception { assertThat(future.get()).isNull(); } + @Test + void testSynchronousGatewayErrorReleasesDrainedKvBatch() throws Exception { + sender.destroyResources(); + + Map tableInfos = new HashMap<>(); + tableInfos.put(DATA1_TABLE_PATH, DATA1_TABLE_INFO); + tableInfos.put(DATA1_TABLE_PATH_PK, DATA1_TABLE_INFO_PK); + TestTabletServerGateway failingGateway = + new TestTabletServerGateway(false, Collections.emptySet()) { + @Override + public CompletableFuture putKv(PutKvRequest request) { + throw new OutOfMemoryError("Direct buffer memory"); + } + }; + metadataUpdater = + TestingMetadataUpdater.builder(tableInfos) + .withTabletServerGateway(TestingMetadataUpdater.NODE1.id(), failingGateway) + .build(); + writerMetricGroup = TestingWriterMetricGroup.newInstance(); + sender = setupWithIdempotenceState(); + + TableBucket kvBucket = new TableBucket(DATA1_TABLE_ID_PK, 0); + CompletableFuture future = new CompletableFuture<>(); + appendKvToAccumulator( + kvBucket, + compactedRow(DATA1_ROW_TYPE, new Object[] {1, "a"}), + (tb, leo, e) -> future.complete(e)); + + sender.runOnce(); + + assertThat(future).isCompleted(); + assertThat(future.get()).hasMessageContaining("Direct buffer memory"); + assertThat(sender.numOfInFlightBatches(kvBucket)).isZero(); + assertThat(accumulator.hasIncomplete()).isFalse(); + } + @Test void testRetries() throws Exception { // create a sender with retries = 1.