fix: release pre-acquired stream ids - #965
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe change adds explicit cancellation for pre-acquired stream IDs when request preparation, throttling, channel selection, or write submission fails before in-flight handling. Sequence Diagram(s)sequenceDiagram
participant RequestHandler
participant DriverChannel
participant DefaultWriteCoalescer
participant InFlightHandler
RequestHandler->>DriverChannel: write request
DriverChannel->>DefaultWriteCoalescer: enqueue request
DefaultWriteCoalescer->>InFlightHandler: submit request
DriverChannel-->>RequestHandler: failed future on write failure
RequestHandler->>DriverChannel: cancelPreAcquireId() before submission
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
3c4fed0 to
29080a7
Compare
There was a problem hiding this comment.
Pull request overview
Fixes #947 by releasing pre-acquired stream IDs when requests fail before reaching the in-flight pipeline.
Changes:
- Adds explicit stream-ID reservation cancellation APIs.
- Covers early failures across CQL, graph, continuous, admin, and channel-handler requests.
- Adds regression tests for key failure paths.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
AdminRequestHandler.java |
Releases reservations on pre-write failures. |
ThrottledAdminRequestHandler.java |
Handles throttling failures safely. |
ChannelHandlerRequest.java |
Cancels failed direct-request reservations. |
DriverChannel.java |
Handles rejected writes and exposes cancellation. |
InFlightHandler.java |
Delegates reservation cancellation. |
CqlRequestHandler.java |
Protects CQL request setup. |
CqlPrepareHandler.java |
Protects prepare and reprepare setup. |
DefaultSession.java |
Releases reservations from closed pooled channels. |
GraphRequestHandler.java |
Protects graph request setup. |
ContinuousRequestHandlerBase.java |
Protects continuous request setup and races. |
PoolBehavior.java |
Adds cancellation verification support. |
CqlRequestHandlerTest.java |
Tests decoration failure handling. |
DriverChannelTest.java |
Tests rejected-write cancellation. |
DefaultSessionPoolsTest.java |
Tests closed-channel handling. |
ReprepareOnUpTest.java |
Tests throttled reprepare failures. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
29080a7 to
bff77b2
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 15 out of 15 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java:111
- The new coalescer-exception path is not exercised by this test change: the added test only covers the earlier
closingbranch. Please add a test whoseWriteCoalescer.writeAndFlushthrows and verify both that the returned future fails with that cause and that the pre-acquired ID is cancelled; otherwise the core fallback for synchronous submission failures can regress unnoticed.
try {
return writeCoalescer.writeAndFlush(channel, message);
} catch (Throwable t) {
cancelPreAcquireId();
return channel.newFailedFuture(t);
core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java:321
- There is no test covering this new graph-specific cleanup path. Add a GraphRequestHandlerTest case that makes message or payload conversion throw after pool selection and verifies
cancelPreAcquireId()and no write; the CQL test does not exercise these graph conversion steps.
DriverExecutionProfile executionProfile =
Conversions.resolveExecutionProfile(statement, context);
GraphProtocol graphSubProtocol =
GraphConversions.resolveGraphSubProtocol(statement, graphSupportChecker, context);
Message message =
core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java:246
- This initial prepare cleanup has no regression test for a synchronous setup failure after channel acquisition. Add a CqlPrepareHandlerTest that makes
toPrepareMessageor payload access throw and verifies the pre-acquired ID is cancelled without writing; this reservation-balancing behavior is independent from the CQL execution test.
} finally {
if (!writeSubmitted) {
channel.cancelPreAcquireId();
}
core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.java:85
- This direct channel-handler submission cleanup has no corresponding test. Please exercise a
ChannelHandlerRequestwhose request construction or rawwriteAndFlushthrows and verify the pre-acquired reservation is restored; these protocol-init/heartbeat paths bypassDriverChannel.write, so the DriverChannel test does not cover them.
} finally {
if (!writeSubmitted) {
inFlightHandler.cancelPreAcquireId();
}
core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java:396
- The continuous-request cleanup is untested, including removal of the callback that is added before request conversion. Add a ContinuousCqlRequestHandlerTest case where
getMessageorcreatePayloadfails and assert that the reservation is cancelled, no write occurs, and the callback is not left ininFlightCallbacks.
if (!writeSubmitted) {
if (nodeResponseCallback != null) {
inFlightCallbacks.remove(nodeResponseCallback);
}
channel.cancelPreAcquireId();
nikagra
left a comment
nikagra
left a comment
There was a problem hiding this comment.
Overview
The reservation is made in ChannelSet.next() and consumed only in InFlightHandler.write(); anything that failed between those two points leaked a slot, so getAvailableIds() drifted downward and channels progressively looked saturated. This closes the gap on every early-exit path I can find — the five channel.write(...) call sites, DriverChannel.write() itself, DefaultSession.getChannel() on a closed pooled channel, and the caller-owned reservation in ThrottledAdminRequestHandler.
Verified locally
Built core at bff77b2 on JDK 11 — clean. Ran DriverChannelTest, CqlRequestHandlerTest, DefaultSessionPoolsTest, CqlPrepareHandlerTest, GraphRequestHandlerTest, ContinuousCqlRequestHandlerTest, ReprepareOnUpTest, InFlightHandlerTest, ChannelPoolTest → 127 tests, 0 failures. Surefire enables -ea by default, so StreamIdGenerator's assert available <= maxAvailableIds over-release guard was live for those runs.
I also traced every preAcquireId() producer (ChannelSet.next(), AdminRequestHandler.start(), ChannelHandlerRequest.send()) against every DriverChannel.write() caller and every new cancel site: the 1:1 pairing holds, I found no double-release. In particular ChannelPool.next(routingKey, shardSuggestion) only ever calls one ChannelSet.next() per invocation, and both shouldPreAcquireId=false handlers (ReprepareOnUp, CqlPrepareHandler.prepareOnOtherNode) really do sit on a pool-owned reservation.
What I liked
- The
writeSubmittedflag is the right idiom, not boilerplate to be simplified intocatch (Throwable): setting it beforeaddListenermeans a throw fromaddListenerwon't cancel a reservation the pipeline already owns. Worth keeping as-is. - The
AtomicBooleanCAS inThrottledAdminRequestHandleris genuinely load-bearing —ConcurrencyLimitingRequestThrottler.register()callsonThrottleReady()synchronously, so a throw there cancels once inonThrottleReady's catch and would cancel again instart()'s catch. inFlightCallbacks.remove(nodeResponseCallback)in the continuous handler'sfinally— easy to miss, and a stale callback there would later get aborted spuriously.(result, error)→(preparedId, error)inCqlPrepareHandlerde-shadows theresultfield. Nice drive-by.
Risk
Low: pure accounting, no protocol or wire change, no public API surface (all internal.*), one boolean store on the happy path. The failure mode of getting it wrong is over-release — getInFlight() going negative and a channel accepting more than max_requests_per_connection — and I couldn't construct such a path.
Verdict
Correct as written. The two things I'd want before merge are both documentation: the write() ownership contract and the now-stale preAcquireId() / StreamIdGenerator javadoc, since those comments are the only place the invariant holding all five finally blocks together is written down. Everything else is optional. I filed #980 for the sibling leak on these same paths.
bff77b2 to
ef698f9
Compare
nikagra
left a comment
nikagra
left a comment
There was a problem hiding this comment.
Verdict
Approving. All ten round-1 comments are addressed, including the three documentation items I named as my pre-merge conditions — those comments were the only place the invariant holding the five caller-side finally blocks together was written down, and they now say it.
Re-verified at ef698f9
I re-traced the accounting rather than re-reading the replies:
- Producer/consumer pairing is still 1:1 with the six new cancel sites.
ChannelPool.next()yields exactly oneChannelSet.next()on every branch (L178, L197, L204), so there's no double-acquire feeding a double-release. write()'s no-throw contract now holds by construction, which is what makes the new Javadoc true rather than aspirational: the only throwing statements sit inside the newtry, sowriteSubmittedis alwaystrueoncewrite()returns.- The
catchinwrite()cannot over-release. WithDefaultWriteCoalescer,InFlightHandler.write()runs from the flusher task on the event loop, never synchronously insidewriteAndFlush. WithPassThroughWriteCoalescer, Netty'sinvokeWrite0converts handler exceptions into a failed promise instead of propagating. So nothing can throw afterstreamIds.acquire()consumed the id. The realistic throw isRejectedExecutionExceptionfromeventLoop.execute()— and there the queuedWriteis permanently stranded (runningstaystrue), so the id really is unused and cancelling is correct. holdsExternalReservation.set(false)beforesuper.start()is load-bearing. Inverting those two lines would double-release, sinceAdminRequestHandler.start()'s ownfinallycovers everything from that point on. Worth keeping in that order.- The new continuous test genuinely pins the new branch.
verifyNoWrite()is what rules out thefinallypath at L396 — without it the test would pass either way. - Also checked:
AdminRequestHandlerschedulestimeoutFutureonly inonWriteComplete, so a request sitting in the throttler has no timer that could firesetFinalErrorbehind the reservation, and throttlerclose()drains viaonThrottleFailure, which cancels. And theshould_complete_result_if_first_node_replies_immediatelyrewrite is a literal inlining ofBuilder.withResponse— no semantic change, it just exposes the handle.
On the one thing you declined
Keeping catch (Throwable) is fine. The reasoning refutes the exact snippet I suggested rather than the general shape — rethrowing without cancelling would also have worked — but the contract you chose instead is the better answer anyway: it's documented on the method, it holds by construction, and should_cancel_pre_acquired_id_when_coalesced_write_throws_error pins it. The reservation is released either way, which was my actual concern. No further action.
Non-blocking
Three notes inline. None needs to happen before merge — two are style/consistency, and one belongs to #980, where I've recorded it. Coverage is unchanged from round 1: the four finally paths Copilot flagged are still untested and, as I said then, I'm not blocking on those; the one branch I did ask for is covered.
CI is green across Build, Unit tests and Full verify on JDK 11 and 17 plus every Scylla and Cassandra IT matrix, and #966 landed on 2026-07-20, so the dependency note in the description is resolved.
e5015ce to
624945b
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java (1)
100-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename this test to describe the enqueue retry path.
The name
should_fail_concurrently_enqueued_writes_if_rescheduling_is_rejectedis almost identical toshould_fail_concurrently_enqueued_writes_when_rescheduling_is_rejectedat Line 169, but the two tests cover different code paths. This test rejects the secondeventLoop.execute()call insideFlusher.enqueue. The test at Line 169 rejectseventLoop.schedule()insiderunOnEventLoop.♻️ Suggested rename
`@Test` - public void should_fail_concurrently_enqueued_writes_if_rescheduling_is_rejected() { + public void should_fail_concurrently_enqueued_writes_if_enqueue_retry_is_rejected() {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java` around lines 100 - 101, Rename the test method should_fail_concurrently_enqueued_writes_if_rescheduling_is_rejected to explicitly describe the enqueue retry path and rejection of the second eventLoop.execute() call inside Flusher.enqueue, distinguishing it from the separate runOnEventLoop scheduling test.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java`:
- Around line 100-101: Rename the test method
should_fail_concurrently_enqueued_writes_if_rescheduling_is_rejected to
explicitly describe the enqueue retry path and rejection of the second
eventLoop.execute() call inside Flusher.enqueue, distinguishing it from the
separate runOnEventLoop scheduling test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c935e787-0392-489a-8ce0-89f023a4b08a
📒 Files selected for processing (23)
core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.javacore/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.javacore/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.javacore/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.javacore/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.javacore/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.javacore/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java
🚧 Files skipped from review as they are similar to previous changes (12)
- core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java
- core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.java
- core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java
- core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java
- core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java
- core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java
- core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java
- core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java
- core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java
624945b to
cb3c855
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java`:
- Around line 100-110: Update newHandler and the cleanup tests to use a real
pre-acquired reservation: call and retain channel.preAcquireId() before
handler.start(), configure the mock stateful reservation behavior, and assert
cancelPreAcquireId() clears the reservation rather than only verifying
invocation. Add separate coverage for the shouldPreAcquireId=true path if
handler-owned acquisition is required.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f7e6b1a0-d9fc-47e8-b3e8-a2d40c54eb16
📒 Files selected for processing (24)
core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.javacore/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.javacore/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.javacore/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.javacore/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.javacore/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.javacore/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.javacore/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.javacore/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java
🚧 Files skipped from review as they are similar to previous changes (22)
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.java
- core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.java
- core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java
- core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java
- core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java
- core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java
- core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java
- core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandler.java
- core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java
- core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java
- core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java
- core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java
- core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java
- core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java
- core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java
- core/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.java
- core/src/test/java/com/datastax/oss/driver/internal/core/channel/DriverChannelTest.java
Release stream-id reservations when a channel has been pre-acquired but the request fails before it is submitted to InFlightHandler. Clean up coalesced writes and continuous-request throttler permits when scheduling or request setup is rejected. Cover CQL, prepare, graph, continuous, admin reprepare, direct channel-handler requests, and closed pooled channels. Fixes scylladb#947. Signed-off-by: Dmitry Kropachev <dmitry.kropachev@gmail.com>
cb3c855 to
74327df
Compare
Summary
Fixes #947.
Depends on #966 for latest Scylla IT compatibility.
Tests
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 PATH=/usr/lib/jvm/java-17-openjdk-amd64/bin:$PATH mvn -pl core -Dtest=DefaultWriteCoalescerTest,ChannelHandlerRequestTest,DriverChannelTest,CqlRequestHandlerTest,DefaultSessionPoolsTest,CqlPrepareHandlerTest,GraphRequestHandlerTest,ContinuousCqlRequestHandlerTest,ContinuousGraphRequestHandlerSpeculativeExecutionTest,ReprepareOnUpTest,InFlightHandlerTest test