fix: release pre-acquired stream ids

core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java

@@ -367,23 +367,37 @@ 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;
+      NodeResponseCallback nodeResponseCallback = null;
+      try {
+        nodeResponseCallback =
+            new NodeResponseCallback(
+                statement,
+                node,
+                channel,
+                currentExecutionIndex,
+                retryCount,
+                scheduleSpeculativeExecution,
+                logPrefix);
+        inFlightCallbacks.add(nodeResponseCallback);
+        Future<java.lang.Void> writeFuture =
+            channel.write(
+                getMessage(statement),
+                isTracingEnabled(statement),
+                createPayload(statement),
+                nodeResponseCallback);
+        writeSubmitted = true;
+        writeFuture.addListener(nodeResponseCallback);
+      } finally {
+        if (!writeSubmitted) {
+          if (nodeResponseCallback != null) {
+            inFlightCallbacks.remove(nodeResponseCallback);
+          }
+          channel.cancelPreAcquireId();
+        }
+      }
+    } else {
+      channel.cancelPreAcquireId();
     }
   }
 

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<String, ByteBuffer> 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<String, ByteBuffer> customPayload =
+            GraphConversions.createCustomPayload(
+                statement, graphSubProtocol, executionProfile, context, graphBinaryModule);
+        Future<java.lang.Void> writeFuture =
+            channel.write(message, statement.isTracing(), customPayload, nodeResponseCallback);
+        writeSubmitted = true;
+        writeFuture.addListener(nodeResponseCallback);
+      } finally {
+        if (!writeSubmitted) {
+          channel.cancelPreAcquireId();
+        }
+      }
     }
   }
 

core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java

@@ -131,11 +131,34 @@ public CompletionStage<ResultT> start() {
               String.format(
                   "%s has reached its maximum number of simultaneous requests", channel)));
     } else {
-      channel.write(message, false, customPayload, this).addListener(this::onWriteComplete);
+      boolean writeSubmitted = false;
+      try {
+        Future<Void> writeFuture = channel.write(message, false, customPayload, this);
+        writeSubmitted = true;
+        writeFuture.addListener(this::onWriteComplete);
+      } finally {
+        if (!writeSubmitted) {
+          channel.cancelPreAcquireId();
+        }
+      }
     }
     return result;
   }
 
+  /**
+   * Cancels a stream id reservation supplied by the caller.
+   *
+   * <p>This is only valid when {@code shouldPreAcquireId} is {@code false}; otherwise this handler
+   * does not own a reservation before {@link #start()}.
+   */
+  protected final void cancelCallerOwnedPreAcquireId() {
+    if (shouldPreAcquireId) {
+      throw new IllegalStateException(
+          "Cannot cancel a caller-owned reservation when this handler pre-acquires its own id");
+    }
+    channel.cancelPreAcquireId();
+  }
+
   private void onWriteComplete(Future<? super Void> future) {
     if (future.isSuccess()) {
       LOG.debug("[{}] Successfully wrote {}, waiting for response", logPrefix, this);

core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java

@@ -38,6 +38,7 @@
 import java.util.Map;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import net.jcip.annotations.ThreadSafe;
 
 @ThreadSafe
@@ -104,10 +105,11 @@ public static ThrottledAdminRequestHandler<ByteBuffer> prepare(
   private final long startTimeNanos;
   private final RequestThrottler throttler;
   private final SessionMetricUpdater metricUpdater;
+  private final AtomicBoolean holdsExternalReservation;
 
   protected ThrottledAdminRequestHandler(
       DriverChannel channel,
-      boolean preAcquireId,
+      boolean shouldPreAcquireId,
       Message message,
       Map<String, ByteBuffer> customPayload,
       Duration timeout,
@@ -118,7 +120,7 @@ protected ThrottledAdminRequestHandler(
       Class<? extends Result> expectedResponseType) {
     super(
         channel,
-        preAcquireId,
+        shouldPreAcquireId,
         message,
         customPayload,
         timeout,
@@ -128,33 +130,54 @@ protected ThrottledAdminRequestHandler(
     this.startTimeNanos = System.nanoTime();
     this.throttler = throttler;
     this.metricUpdater = metricUpdater;
+    this.holdsExternalReservation = new AtomicBoolean(!shouldPreAcquireId);
   }
 
   @Override
   public CompletionStage<ResultT> start() {
     // Don't write request yet, wait for green light from throttler
-    throttler.register(this);
+    try {
+      throttler.register(this);
+    } catch (Throwable t) {
+      cancelExternalReservation();
+      throw t;
+    }
     return result;
   }
 
   @Override
   public void onThrottleReady(boolean wasDelayed) {
-    if (wasDelayed) {
-      metricUpdater.updateTimer(
-          DefaultSessionMetric.THROTTLING_DELAY,
-          null,
-          System.nanoTime() - startTimeNanos,
-          TimeUnit.NANOSECONDS);
+    try {
+      if (wasDelayed) {
+        metricUpdater.updateTimer(
+            DefaultSessionMetric.THROTTLING_DELAY,
+            null,
+            System.nanoTime() - startTimeNanos,
+            TimeUnit.NANOSECONDS);
+      }
+    } catch (Throwable t) {
+      cancelExternalReservation();
+      throw t;
     }
+    holdsExternalReservation.set(false);
     super.start();
   }
 
   @Override
   public void onThrottleFailure(@NonNull RequestThrottlingException error) {
+    cancelExternalReservation();
     metricUpdater.incrementCounter(DefaultSessionMetric.THROTTLING_ERRORS, null);
     setFinalError(error);
   }
 
+  private void cancelExternalReservation() {
+    // register() can invoke onThrottleReady() synchronously. If that callback throws, both
+    // onThrottleReady() and start() catch the same failure, so cancellation must be idempotent.
+    if (holdsExternalReservation.compareAndSet(true, false)) {
+      cancelCallerOwnedPreAcquireId();
+    }
+  }
+
   @Override
   protected boolean setFinalResult(ResultT result) {
     boolean wasSet = super.setFinalResult(result);

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

@@ -72,10 +72,18 @@ void send() {
               String.format(
                   "%s has reached its maximum number of simultaneous requests", channel)));
     } else {
-      DriverChannel.RequestMessage message =
-          new DriverChannel.RequestMessage(getRequest(), false, Frame.NO_PAYLOAD, this);
-      ChannelFuture writeFuture = channel.writeAndFlush(message);
-      writeFuture.addListener(this::writeListener);
+      boolean writeSubmitted = false;
+      try {
+        DriverChannel.RequestMessage message =
+            new DriverChannel.RequestMessage(getRequest(), false, Frame.NO_PAYLOAD, this);
+        ChannelFuture writeFuture = channel.writeAndFlush(message);
+        writeSubmitted = true;
+        writeFuture.addListener(this::writeListener);
+      } finally {
+        if (!writeSubmitted) {
+          inFlightHandler.cancelPreAcquireId();
+        }
+      }
     }
   }
 

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

@@ -91,6 +91,11 @@ public class DriverChannel {
   }
 
   /**
+   * Once this method is entered, it takes ownership of the caller's pre-acquired stream id: it
+   * either submits the request to {@link InFlightHandler} or cancels the reservation and returns a
+   * failed future. Callers must invoke {@link #cancelPreAcquireId()} themselves only if they fail
+   * before reaching this method.
+   *
    * @return a future that succeeds when the request frame was successfully written on the channel.
    *     Beyond that, the caller will be notified through the {@code responseCallback}.
    */
@@ -100,10 +105,19 @@ public Future<Void> write(
       Map<String, ByteBuffer> customPayload,
       ResponseCallback responseCallback) {
     if (closing.get()) {
+      cancelPreAcquireId();
       return channel.newFailedFuture(new IllegalStateException("Driver channel is closing"));
     }
-    RequestMessage message = new RequestMessage(request, tracing, customPayload, responseCallback);
-    return writeCoalescer.writeAndFlush(channel, message);
+    try {
+      RequestMessage message =
+          new RequestMessage(request, tracing, customPayload, responseCallback);
+      return writeCoalescer.writeAndFlush(channel, message);
+    } catch (Throwable t) {
+      // This method must not throw after taking ownership: callers cancel the reservation
+      // themselves when a failure escapes before write() returns.
+      cancelPreAcquireId();
+      return channel.newFailedFuture(t);
+    }
   }
 
   /**
@@ -193,7 +207,9 @@ public int getAvailableIds() {
    *
    * <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.
+   * <b>must</b> proceed with the write or call {@link #cancelPreAcquireId()} if it fails before
+   * reaching the write. If it returns false, it must neither proceed with the write nor cancel the
+   * reservation.
    *
    * <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
@@ -220,6 +236,14 @@ public boolean preAcquireId() {
     return inFlightHandler.preAcquireId();
   }
 
+  /**
+   * Cancels the reservation made by a successful {@link #preAcquireId()} call when its matching
+   * 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}).

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();
   }

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

@@ -24,8 +24,8 @@
 /**
  * Manages the set of identifiers used to distinguish multiplexed requests on a channel.
  *
- * <p>{@link #preAcquire()} / {@link #getAvailableIds()} follow atomic semantics. See {@link
- * DriverChannel#preAcquireId()} for more explanations.
+ * <p>{@link #preAcquire()}, {@link #cancelPreAcquire()}, and {@link #getAvailableIds()} follow
+ * atomic semantics. See {@link DriverChannel#preAcquireId()} for more explanations.
  *
  * <p>Other methods are not synchronized, they are only called by {@link InFlightHandler} on the I/O
  * thread.