From 5900e56954323659ef56fa140e06493c84243032 Mon Sep 17 00:00:00 2001 From: Arnab Nandy Date: Sun, 19 Jul 2026 21:22:25 +0530 Subject: [PATCH] fix: release pre-acquired stream IDs Signed-off-by: Arnab Nandy --- changelog/README.md | 1 + .../ContinuousRequestHandlerBase.java | 44 +++++++---- .../core/graph/GraphRequestHandler.java | 54 +++++++------ .../internal/core/channel/DriverChannel.java | 17 +++- .../core/channel/InFlightHandler.java | 4 + .../internal/core/cql/CqlPrepareHandler.java | 78 ++++++++++++------- .../internal/core/cql/CqlRequestHandler.java | 55 +++++++------ .../core/channel/DriverChannelTest.java | 12 +++ 8 files changed, 171 insertions(+), 94 deletions(-) diff --git a/changelog/README.md b/changelog/README.md index b90979e2e71..ce15df961f3 100644 --- a/changelog/README.md +++ b/changelog/README.md @@ -23,6 +23,7 @@ under the License. ### 4.19.2 +- [bug] PR 947: Release pre-acquired stream IDs when requests fail before being written - [bug] CASSJAVA-116: Retry or Speculative Execution with RequestIdGenerator throws "Duplicate Key" ### 4.19.1 diff --git a/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java b/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java index 0107ebd1176..7536bcd2b7f 100644 --- a/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java +++ b/core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java @@ -367,23 +367,33 @@ private void sendRequest( abortGlobalRequestOrChosenCallback(AllNodesFailedException.fromErrors(errors)); } } else if (!chosenCallback.isDone()) { - NodeResponseCallback nodeResponseCallback = - new NodeResponseCallback( - statement, - node, - channel, - currentExecutionIndex, - retryCount, - scheduleSpeculativeExecution, - logPrefix); - inFlightCallbacks.add(nodeResponseCallback); - channel - .write( - getMessage(statement), - isTracingEnabled(statement), - createPayload(statement), - nodeResponseCallback) - .addListener(nodeResponseCallback); + boolean writeSubmitted = false; + try { + NodeResponseCallback nodeResponseCallback = + new NodeResponseCallback( + statement, + node, + channel, + currentExecutionIndex, + retryCount, + scheduleSpeculativeExecution, + logPrefix); + inFlightCallbacks.add(nodeResponseCallback); + channel + .write( + getMessage(statement), + isTracingEnabled(statement), + createPayload(statement), + nodeResponseCallback) + .addListener(nodeResponseCallback); + writeSubmitted = true; + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } + } else { + channel.cancelPreAcquireId(); } } diff --git a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java index 015467a66fa..bdbde204f1c 100644 --- a/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java +++ b/core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java @@ -302,29 +302,37 @@ private void sendRequest( NO_SUCCESSFUL_EXECUTION); } } else { - NodeResponseCallback nodeResponseCallback = - new NodeResponseCallback( - statement, - node, - queryPlan, - channel, - currentExecutionIndex, - retryCount, - scheduleNextExecution, - logPrefix); - DriverExecutionProfile executionProfile = - Conversions.resolveExecutionProfile(statement, context); - GraphProtocol graphSubProtocol = - GraphConversions.resolveGraphSubProtocol(statement, graphSupportChecker, context); - Message message = - GraphConversions.createMessageFromGraphStatement( - statement, graphSubProtocol, executionProfile, context, graphBinaryModule); - Map customPayload = - GraphConversions.createCustomPayload( - statement, graphSubProtocol, executionProfile, context, graphBinaryModule); - channel - .write(message, statement.isTracing(), customPayload, nodeResponseCallback) - .addListener(nodeResponseCallback); + boolean writeSubmitted = false; + try { + NodeResponseCallback nodeResponseCallback = + new NodeResponseCallback( + statement, + node, + queryPlan, + channel, + currentExecutionIndex, + retryCount, + scheduleNextExecution, + logPrefix); + DriverExecutionProfile executionProfile = + Conversions.resolveExecutionProfile(statement, context); + GraphProtocol graphSubProtocol = + GraphConversions.resolveGraphSubProtocol(statement, graphSupportChecker, context); + Message message = + GraphConversions.createMessageFromGraphStatement( + statement, graphSubProtocol, executionProfile, context, graphBinaryModule); + Map customPayload = + GraphConversions.createCustomPayload( + statement, graphSubProtocol, executionProfile, context, graphBinaryModule); + channel + .write(message, statement.isTracing(), customPayload, nodeResponseCallback) + .addListener(nodeResponseCallback); + writeSubmitted = true; + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java index d4d1bb600c7..0698ebbc29c 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java @@ -100,10 +100,16 @@ public Future write( Map customPayload, ResponseCallback responseCallback) { if (closing.get()) { + inFlightHandler.cancelPreAcquireId(); return channel.newFailedFuture(new IllegalStateException("Driver channel is closing")); } RequestMessage message = new RequestMessage(request, tracing, customPayload, responseCallback); - return writeCoalescer.writeAndFlush(channel, message); + try { + return writeCoalescer.writeAndFlush(channel, message); + } catch (Throwable t) { + inFlightHandler.cancelPreAcquireId(); + return channel.newFailedFuture(t); + } } /** @@ -220,6 +226,15 @@ public boolean preAcquireId() { return inFlightHandler.preAcquireId(); } + /** + * Cancels the reservation made by a successful {@link #preAcquireId()} call when the + * corresponding request cannot be submitted to {@link #write(Message, boolean, Map, + * ResponseCallback)}. + */ + public void cancelPreAcquireId() { + inFlightHandler.cancelPreAcquireId(); + } + /** * @return the number of requests currently executing on this channel (including {@link * #getOrphanedIds() orphaned ids}). diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java index 90b02f358cd..6834b82f94d 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java @@ -396,6 +396,10 @@ boolean preAcquireId() { return streamIds.preAcquire(); } + void cancelPreAcquireId() { + streamIds.cancelPreAcquire(); + } + int getInFlight() { return streamIds.getMaxAvailableIds() - streamIds.getAvailableIds(); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java index 69e98ca5197..5d64318c735 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java @@ -229,14 +229,22 @@ private void sendRequest(PrepareRequest request, Node node, int retryCount) { if (channel == null) { setFinalError(AllNodesFailedException.fromErrors(this.errors)); } else { - InitialPrepareCallback initialPrepareCallback = - new InitialPrepareCallback(request, node, channel, retryCount); - - Prepare message = toPrepareMessage(request); - - channel - .write(message, false, request.getCustomPayload(), initialPrepareCallback) - .addListener(initialPrepareCallback); + boolean writeSubmitted = false; + try { + InitialPrepareCallback initialPrepareCallback = + new InitialPrepareCallback(request, node, channel, retryCount); + + Prepare message = toPrepareMessage(request); + + channel + .write(message, false, request.getCustomPayload(), initialPrepareCallback) + .addListener(initialPrepareCallback); + writeSubmitted = true; + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } @@ -316,28 +324,38 @@ private CompletionStage prepareOnOtherNode(PrepareRequest request, Node no LOG.trace("[{}] Could not get a channel to reprepare on {}, skipping", logPrefix, node); return CompletableFuture.completedFuture(null); } else { - ThrottledAdminRequestHandler handler = - ThrottledAdminRequestHandler.prepare( - channel, - false, - toPrepareMessage(request), - request.getCustomPayload(), - Conversions.resolveRequestTimeout(request, executionProfile), - throttler, - session.getMetricUpdater(), - logPrefix); - return handler - .start() - .handle( - (result, error) -> { - if (error == null) { - LOG.trace("[{}] Successfully reprepared on {}", logPrefix, node); - } else { - Loggers.warnWithException( - LOG, "[{}] Error while repreparing on {}", node, logPrefix, error); - } - return null; - }); + boolean requestStarted = false; + try { + ThrottledAdminRequestHandler handler = + ThrottledAdminRequestHandler.prepare( + channel, + false, + toPrepareMessage(request), + request.getCustomPayload(), + Conversions.resolveRequestTimeout(request, executionProfile), + throttler, + session.getMetricUpdater(), + logPrefix); + CompletionStage result = + handler + .start() + .handle( + (preparedId, error) -> { + if (error == null) { + LOG.trace("[{}] Successfully reprepared on {}", logPrefix, node); + } else { + Loggers.warnWithException( + LOG, "[{}] Error while repreparing on {}", node, logPrefix, error); + } + return null; + }); + requestStarted = true; + return result; + } finally { + if (!requestStarted) { + channel.cancelPreAcquireId(); + } + } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 29f595e5225..e53aa4594b1 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -409,30 +409,39 @@ private void sendRequest( setFinalError(statement, AllNodesFailedException.fromErrors(this.errors), null, -1); } } else { - Statement finalStatement = statement; - String nodeRequestId = - this.requestIdGenerator - .map((g) -> g.getNodeRequestId(finalStatement, sessionRequestId)) - .orElse(Integer.toString(this.hashCode())); - statement = - this.requestIdGenerator - .map((g) -> g.getDecoratedStatement(finalStatement, nodeRequestId)) - .orElse(finalStatement); + boolean writeSubmitted = false; + try { + Statement finalStatement = statement; + String nodeRequestId = + this.requestIdGenerator + .map((g) -> g.getNodeRequestId(finalStatement, sessionRequestId)) + .orElse(Integer.toString(this.hashCode())); + statement = + this.requestIdGenerator + .map((g) -> g.getDecoratedStatement(finalStatement, nodeRequestId)) + .orElse(finalStatement); - NodeResponseCallback nodeResponseCallback = - new NodeResponseCallback( - statement, - node, - queryPlan, - channel, - currentExecutionIndex, - retryCount, - scheduleNextExecution, - logPrefixJoiner.join(this.sessionName, nodeRequestId, currentExecutionIndex)); - Message message = Conversions.toMessage(statement, executionProfile, context); - channel - .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) - .addListener(nodeResponseCallback); + NodeResponseCallback nodeResponseCallback = + new NodeResponseCallback( + statement, + node, + queryPlan, + channel, + currentExecutionIndex, + retryCount, + scheduleNextExecution, + logPrefixJoiner.join(this.sessionName, nodeRequestId, currentExecutionIndex)); + Message message = Conversions.toMessage(statement, executionProfile, context); + channel + .write( + message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) + .addListener(nodeResponseCallback); + writeSubmitted = true; + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.java index e0660b9609e..f53cbfce53e 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.java @@ -18,6 +18,7 @@ package com.datastax.oss.driver.internal.core.channel; import static com.datastax.oss.driver.Assertions.assertThat; +import static org.mockito.Mockito.verify; import com.datastax.oss.driver.api.core.DefaultProtocolVersion; import com.datastax.oss.driver.api.core.connection.ClosedConnectionException; @@ -143,6 +144,17 @@ public void should_wait_for_coalesced_writes_when_closing_forcefully() { .hasMessageContaining("Channel was force-closed"); } + @Test + public void should_cancel_pre_acquired_id_when_write_is_rejected_before_submission() { + driverChannel.close(); + + Future writeFuture = + driverChannel.write(new Query("test"), false, Frame.NO_PAYLOAD, new MockResponseCallback()); + + assertThat(writeFuture).isFailed(); + verify(streamIds).cancelPreAcquire(); + } + // Simple implementation that holds all the writes, and flushes them when it's explicitly // triggered. private class MockWriteCoalescer implements WriteCoalescer {