fix: release pre-acquired stream ids

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

@@ -282,10 +282,12 @@ public void onThrottleFailure(@NonNull RequestThrottlingException error) {
     abortGlobalRequestOrChosenCallback(error);
   }
 
-  private void abortGlobalRequestOrChosenCallback(@NonNull Throwable error) {
-    if (!chosenCallback.completeExceptionally(error)) {
+  private boolean abortGlobalRequestOrChosenCallback(@NonNull Throwable error) {
+    boolean completedChosenCallback = chosenCallback.completeExceptionally(error);
+    if (!completedChosenCallback) {
       chosenCallback.thenAccept(callback -> callback.abort(error, false));
     }
+    return completedChosenCallback;
   }
 
   public CompletionStage<ResultSetT> handle() {
@@ -367,23 +369,51 @@ 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;
+      Throwable terminalPreWriteFailure = null;
+      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);
+      } catch (Throwable t) {
+        if (!writeSubmitted && activeExecutionsCount.decrementAndGet() == 0) {
+          if (abortGlobalRequestOrChosenCallback(t)) {
+            terminalPreWriteFailure = t;
+          }
+        }
+        throw t;
+      } finally {
+        if (!writeSubmitted) {
+          if (nodeResponseCallback != null) {
+            inFlightCallbacks.remove(nodeResponseCallback);
+          }
+          try {
+            channel.cancelPreAcquireId();
+          } finally {
+            if (terminalPreWriteFailure != null) {
+              throttler.signalError(this, terminalPreWriteFailure);
+            }
+          }
+        }
+      }
+    } else {
+      channel.cancelPreAcquireId();
     }
   }
 
@@ -488,6 +518,16 @@ CompletableFuture<ResultSetT> getPendingResult() {
     }
   }
 
+  @VisibleForTesting
+  int getActiveExecutionsCount() {
+    return activeExecutionsCount.get();
+  }
+
+  @VisibleForTesting
+  int getInFlightCallbackCount() {
+    return inFlightCallbacks.size();
+  }
+
   private void recordError(@NonNull Node node, @NonNull Throwable error) {
     errors.add(new AbstractMap.SimpleEntry<>(node, error));
   }

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,55 @@ 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);
+      }
+      holdsExternalReservation.set(false);
+      super.start();
+    } catch (Throwable t) {
+      cancelExternalReservation();
+      setFinalError(t);
+      throw t;
     }
-    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,31 @@ 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;
+      DriverChannel.RequestMessage message = null;
+      try {
+        message =
+            new DriverChannel.RequestMessage(
+                getRequest(), false, Frame.NO_PAYLOAD, this, inFlightHandler);
+        ChannelFuture writeFuture = channel.writeAndFlush(message);
+        DriverChannel.RequestMessage submittedMessage = message;
+        writeFuture.addListener(
+            future -> {
+              if (!future.isSuccess()) {
+                submittedMessage.cancelPreAcquireId();
+              }
+              writeListener(future);
+            });
+        writeSubmitted = true;
+      } finally {
+        if (!writeSubmitted) {
+          if (message == null) {
+            inFlightHandler.cancelPreAcquireId();
+          } else {
+            message.cancelPreAcquireId();
+          }
+        }
+      }
     }
   }