4.x: Pre-acquired stream id can leak when request fails before InFlightHandler write

[dkropachev]

Problem

ChannelPool.next() can return a channel after ChannelSet.next() already called DriverChannel.preAcquireId(). CqlRequestHandler.sendRequest() then still does work before the write actually reaches InFlightHandler.

If that work throws synchronously, or if DriverChannel.write() returns a failed future before submitting a RequestMessage to the pipeline, the pre-acquired stream id is not released. The request may fail or retry, but the channel accounting can remain out of sync.

Proof

• Pool selection pre-acquires before returning the channel:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelSet.java

  Lines 92 to 115 in 3cafb2a

  	/** @return null if the set is empty or all are full */
  	DriverChannel next() {
  	DriverChannel[] snapshot = this.channels;
  	switch (snapshot.length) {
  	case 0:
  	return null;
  	case 1:
  	DriverChannel onlyChannel = snapshot[0];
  	return onlyChannel.preAcquireId() ? onlyChannel : null;
  	default:
  	for (int i = 0; i < MAX_ITERATIONS; i++) {
  	DriverChannel best = null;
  	int bestScore = 0;
  	for (DriverChannel channel : snapshot) {
  	int score = channel.getAvailableIds();
  	if (score > bestScore) {
  	bestScore = score;
  	best = channel;
  	}
  	}
  	if (best == null) {
  	return null;
  	} else if (best.preAcquireId()) {
  	return best;

• DefaultSession.getChannel() returns that pre-acquired channel to request execution:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java

  Lines 261 to 274 in 3cafb2a

  	ChannelPool pool = poolManager.getPools().get(node);
  	if (pool == null) {
  	LOG.trace("[{}] No pool to {}, skipping", logPrefix, node);
  	return null;
  	} else {
  	DriverChannel channel = pool.next(routingKey, shardSuggestion);
  	if (channel == null) {
  	LOG.trace("[{}] Pool returned no channel for {}, skipping", logPrefix, node);
  	return null;
  	} else if (channel.closeFuture().isDone()) {
  	LOG.trace("[{}] Pool returned closed connection to {}, skipping", logPrefix, node);
  	return null;
  	} else {
  	return channel;

• sendRequest() still decorates the statement, converts it to a protocol message, reads custom payload, and calls channel.write() after channel acquisition:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

  Lines 384 to 435 in 3cafb2a

  	|| (channel =
  	session.getChannel(
  	node,
  	handlerLogPrefix,
  	getRoutingToken(statement),
  	getShardFromTabletMap(statement, node, getRoutingToken(statement))))
  	== null) {
  	while (!result.isDone() && (node = queryPlan.poll()) != null) {
  	channel =
  	session.getChannel(
  	node,
  	handlerLogPrefix,
  	getRoutingToken(statement),
  	getShardFromTabletMap(statement, node, getRoutingToken(statement)));
  	if (channel != null) {
  	break;
  	} else {
  	recordError(node, new NodeUnavailableException(node));
  	}
  	}
  	}
  	if (channel == null) {
  	// We've reached the end of the query plan without finding any node to write to
  	if (!result.isDone() && activeExecutionsCount.decrementAndGet() == 0) {
  	// We're the last execution so fail the result
  	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);
  	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);

• preAcquireId() contract says a successful caller must proceed with exactly one write:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java

  Lines 194 to 202 in 3cafb2a

  	* <p>There must be <b>exactly one</b> invocation of this method before each call to {@link
  	* #write(Message, boolean, Map, ResponseCallback)}. If this method returns true, the client
  	* <b>must</b> proceed with the write. If it returns false, it <b>must not</b> proceed.
  	*
  	* <p>This method is used together with {@link #getAvailableIds()} to track how many requests are
  	* currently executing on the channel, and avoid submitting a request that would result in a
  	* {@link BusyConnectionException}. The two methods follow atomic semantics: {@link
  	* #getAvailableIds()} returns the exact count of clients that have called {@link #preAcquireId()}
  	* and not yet released their stream id at this point in time.

• DriverChannel.write() can return a failed future immediately when the channel is closing, before a RequestMessage reaches InFlightHandler:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannel.java

  Lines 97 to 106 in 3cafb2a

  	public Future<Void> write(
  	Message request,
  	boolean tracing,
  	Map<String, ByteBuffer> customPayload,
  	ResponseCallback responseCallback) {
  	if (closing.get()) {
  	return channel.newFailedFuture(new IllegalStateException("Driver channel is closing"));
  	}
  	RequestMessage message = new RequestMessage(request, tracing, customPayload, responseCallback);
  	return writeCoalescer.writeAndFlush(channel, message);

• InFlightHandler only cancels the pre-acquire for failures it sees while handling the RequestMessage:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java

  Lines 119 to 133 in 3cafb2a

  	private void write(ChannelHandlerContext ctx, RequestMessage message, ChannelPromise promise) {
  	if (closingGracefully) {
  	promise.setFailure(new IllegalStateException("Channel is closing"));
  	streamIds.cancelPreAcquire();
  	return;
  	}
  	int streamId = streamIds.acquire();
  	if (streamId < 0) {
  	// Should not happen with the preAcquire mechanism, but handle gracefully
  	promise.setFailure(
  	new BusyConnectionException(
  	String.format(
  	"Couldn't acquire a stream id from InFlightHandler on %s", ctx.channel())));
  	streamIds.cancelPreAcquire();
  	return;

• CQL write-future failure handling records the error and tries the next node, but does not release the pre-acquired id:

  java-driver/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

  Lines 680 to 704 in 3cafb2a

  	public void operationComplete(Future<java.lang.Void> future) throws Exception {
  	if (!future.isSuccess()) {
  	Throwable error = future.cause();
  	if (error instanceof EncoderException
  	&& error.getCause() instanceof FrameTooLongException) {
  	trackNodeError(node, error.getCause(), NANOTIME_NOT_MEASURED_YET);
  	setFinalError(statement, error.getCause(), node, execution);
  	} else {
  	LOG.trace(
  	"[{}] Failed to send request on {}, trying next node (cause: {})",
  	logPrefix,
  	channel,
  	error);
  	recordError(node, error);
  	trackNodeError(node, error, NANOTIME_NOT_MEASURED_YET);
  	((DefaultNode) node)
  	.getMetricUpdater()
  	.incrementCounter(DefaultNodeMetric.UNSENT_REQUESTS, executionProfile.getName());
  	sendRequest(
  	statement,
  	null,
  	queryPlan,
  	execution,
  	retryCount,
  	scheduleNextExecution); // try next node

Impact

Repeated early failures can make getAvailableIds() drift from real stream-id usage. Over time this can make a channel or pool look busier than it is, and eventually look saturated.

Expected

Every successful preAcquireId() should either reach InFlightHandler.write() or be explicitly released/cancelled on all early failure paths.

Possible fixes

• Move all throwable request-building work before acquiring a channel.
• Or add a release/cancel path when a pre-acquired channel fails before InFlightHandler receives the RequestMessage.

The same pattern may also need checking in graph and continuous request handlers.