From 74327dfacd16231b8451d27a70d0312a48aef17a Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Fri, 31 Jul 2026 14:28:35 -0400 Subject: [PATCH] fix: release pre-acquired stream ids 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 #947. Signed-off-by: Dmitry Kropachev --- .../ContinuousRequestHandlerBase.java | 78 ++++-- .../core/graph/GraphRequestHandler.java | 54 ++-- .../adminrequest/AdminRequestHandler.java | 25 +- .../ThrottledAdminRequestHandler.java | 44 +++- .../core/channel/ChannelHandlerRequest.java | 29 ++- .../core/channel/DefaultWriteCoalescer.java | 153 ++++++++++- .../internal/core/channel/DriverChannel.java | 70 ++++- .../core/channel/InFlightHandler.java | 10 + .../core/channel/StreamIdGenerator.java | 4 +- .../internal/core/cql/CqlPrepareHandler.java | 52 ++-- .../internal/core/cql/CqlRequestHandler.java | 55 ++-- .../internal/core/session/DefaultSession.java | 1 + .../ContinuousCqlRequestHandlerTest.java | 31 +++ ...equestHandlerSpeculativeExecutionTest.java | 54 ++++ .../core/graph/GraphRequestHandlerTest.java | 33 +++ .../ThrottledAdminRequestHandlerTest.java | 127 +++++++++ .../channel/ChannelHandlerRequestTest.java | 134 ++++++++++ .../channel/DefaultWriteCoalescerTest.java | 240 ++++++++++++++++++ .../core/channel/DriverChannelTest.java | 89 +++++++ .../core/cql/CqlPrepareHandlerTest.java | 23 ++ .../core/cql/CqlRequestHandlerTest.java | 43 +++- .../internal/core/cql/PoolBehavior.java | 8 + .../core/session/DefaultSessionPoolsTest.java | 32 +++ .../core/session/ReprepareOnUpTest.java | 56 ++++ 24 files changed, 1325 insertions(+), 120 deletions(-) create mode 100644 core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java create mode 100644 core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.java create mode 100644 core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java 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..c17ccd50ff8 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 @@ -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 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 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 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)); } 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..cf6875dd362 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); + Future writeFuture = + channel.write(message, statement.isTracing(), customPayload, nodeResponseCallback); + writeSubmitted = true; + writeFuture.addListener(nodeResponseCallback); + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java index 5078428c21a..16ad89f6126 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/AdminRequestHandler.java @@ -131,11 +131,34 @@ public CompletionStage 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 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. + * + *

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 future) { if (future.isSuccess()) { LOG.debug("[{}] Successfully wrote {}, waiting for response", logPrefix, this); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java index 40ab21b759a..55b8e821249 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandler.java +++ b/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 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 customPayload, Duration timeout, @@ -118,7 +120,7 @@ protected ThrottledAdminRequestHandler( Class 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 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); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.java index 3ba3d70eb8d..538f28aed90 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequest.java +++ b/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(); + } + } + } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.java index 232fa83be44..9722f66a0fa 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescer.java @@ -23,13 +23,17 @@ import io.netty.channel.Channel; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelPromise; +import io.netty.channel.DefaultChannelPromise; import io.netty.channel.EventLoop; +import io.netty.util.concurrent.AbstractEventExecutor; +import io.netty.util.concurrent.Future; import java.util.HashSet; import java.util.Queue; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import net.jcip.annotations.ThreadSafe; @@ -59,19 +63,17 @@ public DefaultWriteCoalescer(DriverContext context) { @Override public ChannelFuture writeAndFlush(Channel channel, Object message) { - ChannelPromise writePromise = channel.newPromise(); + Flusher flusher = flushers.computeIfAbsent(channel.eventLoop(), Flusher::new); + ChannelPromise writePromise = + new DefaultChannelPromise(channel, flusher.listenerNotificationExecutor); Write write = new Write(channel, message, writePromise); - enqueue(write, channel.eventLoop()); - return writePromise; - } - - private void enqueue(Write write, EventLoop eventLoop) { - Flusher flusher = flushers.computeIfAbsent(eventLoop, Flusher::new); flusher.enqueue(write); + return writePromise; } private class Flusher { private final EventLoop eventLoop; + private final RejectionSafeEventExecutor listenerNotificationExecutor; // These variables are accessed both from client threads and the event loop private final Queue writes = new ConcurrentLinkedQueue<>(); @@ -82,13 +84,55 @@ private class Flusher { private Flusher(EventLoop eventLoop) { this.eventLoop = eventLoop; + this.listenerNotificationExecutor = new RejectionSafeEventExecutor(eventLoop); } private void enqueue(Write write) { boolean added = writes.offer(write); assert added; // always true (see MpscLinkedAtomicQueue implementation) if (running.compareAndSet(false, true)) { - eventLoop.execute(this::runOnEventLoop); + try { + eventLoop.execute(this::runOnEventLoop); + } catch (Throwable t) { + // execute() rejected the task, so this write can never reach the pipeline. Remove it + // before making the flusher schedulable again; otherwise a concurrent retry could submit + // a request whose caller has already cancelled its pre-acquired stream id. + boolean removed = writes.remove(write); + assert removed; + running.set(false); + + // Other writers might have enqueued while running was true. Make sure they are not left + // behind with no task scheduled. If the event loop rejects again, fail those queued + // writes while allowing subsequent enqueues to schedule a new run. + if (!writes.isEmpty() && running.compareAndSet(false, true)) { + try { + eventLoop.execute(this::runOnEventLoop); + } catch (Throwable retryFailure) { + failPendingWrites(retryFailure); + if (retryFailure != t) { + t.addSuppressed(retryFailure); + } + } + } + throw t; + } + } + } + + private void failPendingWrites(Throwable failure) { + while (true) { + Write write; + while ((write = writes.poll()) != null) { + write.fail(failure); + } + + // Allow concurrent enqueues to schedule a new run. If a write was added while we still + // owned the running flag and nobody claimed it after we released it, reclaim ownership and + // fail that write too. + running.set(false); + if (writes.isEmpty() || !running.compareAndSet(false, true)) { + return; + } } } @@ -99,7 +143,11 @@ private void runOnEventLoop() { while ((write = writes.poll()) != null) { Channel channel = write.channel; channels.add(channel); - channel.write(write.message, write.writePromise); + try { + channel.write(write.message, write.writePromise); + } catch (Throwable t) { + write.fail(t); + } } for (Channel channel : channels) { @@ -124,12 +172,88 @@ private void runOnEventLoop() { // on. If not, we need to do it ourselves. boolean shouldRestartMyself = running.compareAndSet(false, true); - if (shouldRestartMyself && !eventLoop.isShuttingDown()) { - eventLoop.schedule(this::runOnEventLoop, rescheduleIntervalNanos, TimeUnit.NANOSECONDS); + if (shouldRestartMyself) { + if (eventLoop.isShuttingDown()) { + failPendingWrites(new RejectedExecutionException("Event loop is shutting down")); + } else { + try { + eventLoop.schedule(this::runOnEventLoop, rescheduleIntervalNanos, TimeUnit.NANOSECONDS); + } catch (Throwable t) { + failPendingWrites(t); + } + } } } } + /** + * Runs promise listeners on their channel's event loop in the normal case, but falls back to the + * completing thread when that event loop rejects the notification task during shutdown. + */ + private static class RejectionSafeEventExecutor extends AbstractEventExecutor { + private final EventLoop eventLoop; + + private RejectionSafeEventExecutor(EventLoop eventLoop) { + super(eventLoop.parent()); + this.eventLoop = eventLoop; + } + + @Override + public boolean inEventLoop() { + return eventLoop.inEventLoop(); + } + + @Override + public boolean inEventLoop(Thread thread) { + return eventLoop.inEventLoop(thread); + } + + @Override + public void execute(Runnable command) { + try { + eventLoop.execute(command); + } catch (RejectedExecutionException e) { + command.run(); + } + } + + @Override + public Future shutdownGracefully(long quietPeriod, long timeout, TimeUnit unit) { + return eventLoop.shutdownGracefully(quietPeriod, timeout, unit); + } + + @Override + public Future terminationFuture() { + return eventLoop.terminationFuture(); + } + + @Override + @Deprecated + public void shutdown() { + eventLoop.shutdownGracefully(); + } + + @Override + public boolean isShuttingDown() { + return eventLoop.isShuttingDown(); + } + + @Override + public boolean isShutdown() { + return eventLoop.isShutdown(); + } + + @Override + public boolean isTerminated() { + return eventLoop.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + return eventLoop.awaitTermination(timeout, unit); + } + } + private static class Write { private final Channel channel; private final Object message; @@ -140,5 +264,12 @@ private Write(Channel channel, Object message, ChannelPromise writePromise) { this.message = message; this.writePromise = writePromise; } + + private void fail(Throwable failure) { + if (message instanceof DriverChannel.RequestMessage) { + ((DriverChannel.RequestMessage) message).cancelPreAcquireId(); + } + writePromise.tryFailure(failure); + } } } 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..1423d0c103b 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 @@ -44,12 +44,14 @@ import io.netty.channel.EventLoop; import io.netty.util.AttributeKey; import io.netty.util.concurrent.Future; +import io.netty.util.concurrent.ImmediateEventExecutor; import io.netty.util.concurrent.Promise; import java.net.SocketAddress; import java.nio.ByteBuffer; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import net.jcip.annotations.ThreadSafe; /** @@ -91,6 +93,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 +107,27 @@ public Future write( Map customPayload, ResponseCallback responseCallback) { if (closing.get()) { - return channel.newFailedFuture(new IllegalStateException("Driver channel is closing")); + cancelPreAcquireId(); + return ImmediateEventExecutor.INSTANCE.newFailedFuture( + new IllegalStateException("Driver channel is closing")); + } + RequestMessage message = + new RequestMessage(request, tracing, customPayload, responseCallback, inFlightHandler); + try { + Future writeFuture = writeCoalescer.writeAndFlush(channel, message); + writeFuture.addListener( + future -> { + if (!future.isSuccess()) { + message.cancelPreAcquireId(); + } + }); + return writeFuture; + } catch (Throwable t) { + // This method must not throw after taking ownership because callers only cancel failures + // that happen before they invoke it. + message.cancelPreAcquireId(); + return ImmediateEventExecutor.INSTANCE.newFailedFuture(t); } - RequestMessage message = new RequestMessage(request, tracing, customPayload, responseCallback); - return writeCoalescer.writeAndFlush(channel, message); } /** @@ -193,7 +217,9 @@ public int getAvailableIds() { * *

There must be exactly one invocation of this method before each call to {@link * #write(Message, boolean, Map, ResponseCallback)}. If this method returns true, the client - * must proceed with the write. If it returns false, it must not proceed. + * must 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. * *

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 +246,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}). @@ -323,20 +357,48 @@ public String toString() { // This is essentially a stripped-down Frame. We can't materialize the frame before writing, // because we need the stream id, which is assigned from within the event loop. static class RequestMessage { + private static final AtomicIntegerFieldUpdater OWNS_PRE_ACQUIRE_ID = + AtomicIntegerFieldUpdater.newUpdater(RequestMessage.class, "ownsPreAcquireId"); + final Message request; final boolean tracing; final Map customPayload; final ResponseCallback responseCallback; + private final InFlightHandler preAcquireIdOwner; + + @SuppressWarnings("unused") // accessed through OWNS_PRE_ACQUIRE_ID + private volatile int ownsPreAcquireId; RequestMessage( Message message, boolean tracing, Map customPayload, ResponseCallback responseCallback) { + this(message, tracing, customPayload, responseCallback, null); + } + + RequestMessage( + Message message, + boolean tracing, + Map customPayload, + ResponseCallback responseCallback, + InFlightHandler preAcquireIdOwner) { this.request = message; this.tracing = tracing; this.customPayload = customPayload; this.responseCallback = responseCallback; + this.preAcquireIdOwner = preAcquireIdOwner; + this.ownsPreAcquireId = preAcquireIdOwner == null ? 0 : 1; + } + + boolean markPreAcquireIdAsSubmitted() { + return preAcquireIdOwner == null || OWNS_PRE_ACQUIRE_ID.compareAndSet(this, 1, 0); + } + + void cancelPreAcquireId() { + if (OWNS_PRE_ACQUIRE_ID.compareAndSet(this, 1, 0)) { + preAcquireIdOwner.cancelPreAcquireId(); + } } } 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..68450a936fd 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 @@ -117,6 +117,12 @@ public void write(ChannelHandlerContext ctx, Object in, ChannelPromise promise) } private void write(ChannelHandlerContext ctx, RequestMessage message, ChannelPromise promise) { + // From this point on, this handler owns the pre-acquired stream id and is responsible for + // releasing it on every failure path. + if (!message.markPreAcquireIdAsSubmitted()) { + promise.tryFailure(new IllegalStateException("Stream id reservation was already cancelled")); + return; + } if (closingGracefully) { promise.setFailure(new IllegalStateException("Channel is closing")); streamIds.cancelPreAcquire(); @@ -396,6 +402,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/channel/StreamIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.java index 3384bc57c94..5384a93e98d 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/StreamIdGenerator.java +++ b/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. * - *

{@link #preAcquire()} / {@link #getAvailableIds()} follow atomic semantics. See {@link - * DriverChannel#preAcquireId()} for more explanations. + *

{@link #preAcquire()}, {@link #cancelPreAcquire()}, and {@link #getAvailableIds()} follow + * atomic semantics. See {@link DriverChannel#preAcquireId()} for more explanations. * *

Other methods are not synchronized, they are only called by {@link InFlightHandler} on the I/O * thread. 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..b05fec2cb7e 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); + + Future writeFuture = + channel.write(message, false, request.getCustomPayload(), initialPrepareCallback); + writeSubmitted = true; + writeFuture.addListener(initialPrepareCallback); + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } @@ -316,20 +324,26 @@ 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); + ThrottledAdminRequestHandler handler; + try { + handler = + ThrottledAdminRequestHandler.prepare( + channel, + false, + toPrepareMessage(request), + request.getCustomPayload(), + Conversions.resolveRequestTimeout(request, executionProfile), + throttler, + session.getMetricUpdater(), + logPrefix); + } catch (Throwable t) { + channel.cancelPreAcquireId(); + throw t; + } return handler .start() .handle( - (result, error) -> { + (preparedId, error) -> { if (error == null) { LOG.trace("[{}] Successfully reprepared on {}", logPrefix, node); } else { 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..ce4b40d6d29 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); + Future writeFuture = + channel.write( + message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback); + writeSubmitted = true; + writeFuture.addListener(nodeResponseCallback); + } finally { + if (!writeSubmitted) { + channel.cancelPreAcquireId(); + } + } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java b/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java index c9fee86f2c1..ffd68c00271 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/session/DefaultSession.java @@ -269,6 +269,7 @@ public DriverChannel getChannel( return null; } else if (channel.closeFuture().isDone()) { LOG.trace("[{}] Pool returned closed connection to {}, skipping", logPrefix, node); + channel.cancelPreAcquireId(); return null; } else { return channel; diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java index a816183e9ee..e24579e7fe6 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousCqlRequestHandlerTest.java @@ -31,6 +31,7 @@ import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.matches; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -49,7 +50,9 @@ import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; import com.datastax.oss.driver.api.core.cql.ExecutionInfo; import com.datastax.oss.driver.api.core.cql.Row; +import com.datastax.oss.driver.api.core.cql.SimpleStatement; import com.datastax.oss.driver.api.core.servererrors.BootstrappingException; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; import com.datastax.oss.driver.api.core.tracker.RequestTracker; import com.datastax.oss.driver.internal.core.ProtocolFeature; import com.datastax.oss.driver.internal.core.cql.PoolBehavior; @@ -211,6 +214,34 @@ public void should_throw_if_protocol_version_does_not_support_continuous_paging( } } + @Test + public void should_unwind_execution_if_request_setup_fails_before_write() { + RuntimeException failure = new RuntimeException("mock failure"); + SimpleStatement statement = Mockito.spy(SimpleStatement.newInstance("mock query")); + doThrow(failure).when(statement).getCustomPayload(); + RequestThrottler throttler = mock(RequestThrottler.class); + RequestHandlerTestHarness.Builder builder = + continuousHarnessBuilder().withProtocolVersion(DSE_V2); + PoolBehavior node1Behavior = builder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = builder.build()) { + when(harness.getContext().getRequestThrottler()).thenReturn(throttler); + ContinuousCqlRequestHandler handler = + new ContinuousCqlRequestHandler( + statement, harness.getSession(), harness.getContext(), "test"); + CompletionStage resultSetFuture = handler.handle(); + + assertThatThrownBy(() -> handler.onThrottleReady(false)).isSameAs(failure); + + assertThatStage(resultSetFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(handler.getActiveExecutionsCount()).isZero(); + assertThat(handler.getInFlightCallbackCount()).isZero(); + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + verify(throttler).signalError(handler, failure); + } + } + @Test @UseDataProvider(value = "allDseProtocolVersions", location = DseTestDataProviders.class) public void should_time_out_if_first_page_takes_too_long(DseProtocolVersion version) diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java index c67be162181..f9ef1ebf6f0 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/ContinuousGraphRequestHandlerSpeculativeExecutionTest.java @@ -26,6 +26,7 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -260,6 +261,59 @@ public void should_not_start_execution_if_result_complete( } } + @Test + @UseDataProvider(location = DseTestDataProviders.class, value = "idempotentGraphConfig") + public void should_cancel_pre_acquired_id_if_result_completes_after_channel_acquired( + boolean defaultIdempotence, GraphStatement statement) throws Exception { + GraphRequestHandlerTestHarness.Builder harnessBuilder = + GraphRequestHandlerTestHarness.builder().withDefaultIdempotence(defaultIdempotence); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + PoolBehavior node2Behavior = harnessBuilder.customBehavior(node2); + + try (GraphRequestHandlerTestHarness harness = harnessBuilder.build()) { + SpeculativeExecutionPolicy speculativeExecutionPolicy = + harness.getContext().getSpeculativeExecutionPolicy(DriverExecutionProfile.DEFAULT_NAME); + long firstExecutionDelay = 100L; + when(speculativeExecutionPolicy.nextExecution( + any(Node.class), eq(null), eq(statement), eq(1))) + .thenReturn(firstExecutionDelay); + + GraphBinaryModule module = createGraphBinaryModule(harness.getContext()); + CompletionStage resultSetFuture = + new ContinuousGraphRequestHandler( + statement, + harness.getSession(), + harness.getContext(), + "test", + module, + graphSupportChecker) + .handle(); + node1Behavior.verifyWrite(); + node1Behavior.setWriteSuccess(); + + CapturedTimeout speculativeExecution = harness.nextScheduledTimeout(); + assertThat(speculativeExecution.getDelay(TimeUnit.MILLISECONDS)) + .isEqualTo(firstExecutionDelay); + + // Simulate the initial execution winning after the speculative execution reserves an ID, + // but before it submits its write. + doAnswer( + invocation -> { + node1Behavior.setResponseSuccess( + defaultDseFrameOf(singleGraphRow(GraphProtocol.GRAPH_BINARY_1_0, module))); + return node2Behavior.getChannel(); + }) + .when(harness.getSession()) + .getChannel(eq(node2), anyString()); + + speculativeExecution.task().run(speculativeExecution); + + assertThatStage(resultSetFuture).isSuccess(); + node2Behavior.verifyNoWrite(); + node2Behavior.verifyPreAcquireCancelled(); + } + } + @Test @UseDataProvider(location = DseTestDataProviders.class, value = "idempotentGraphConfig") public void should_fail_if_no_nodes(boolean defaultIdempotence, GraphStatement statement) { diff --git a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java index 9f325003610..852991b9607 100644 --- a/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandlerTest.java @@ -27,10 +27,12 @@ import static com.datastax.oss.driver.api.core.type.codec.TypeCodecs.BIGINT; import static com.datastax.oss.driver.api.core.type.codec.TypeCodecs.TEXT; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.matches; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -577,6 +579,37 @@ public void should_honor_statement_consistency_level() { assertThat(q.options.consistency).isEqualTo(DefaultConsistencyLevel.THREE.getProtocolCode()); } + @Test + public void should_cancel_pre_acquired_id_if_graph_payload_conversion_fails() { + RuntimeException failure = new RuntimeException("mock failure"); + ScriptGraphStatement graphStatement = + Mockito.spy(ScriptGraphStatement.newInstance("mock query")); + doThrow(failure).when(graphStatement).getCustomPayload(); + + GraphRequestHandlerTestHarness.Builder builder = GraphRequestHandlerTestHarness.builder(); + PoolBehavior nodeBehavior = builder.customBehavior(node); + try (GraphRequestHandlerTestHarness harness = builder.build()) { + GraphSupportChecker graphSupportChecker = mock(GraphSupportChecker.class); + when(graphSupportChecker.inferGraphProtocol(any(), any(), any())) + .thenReturn(GRAPH_BINARY_1_0); + GraphBinaryModule module = createGraphBinaryModule(harness.getContext()); + + assertThatThrownBy( + () -> + new GraphRequestHandler( + graphStatement, + harness.getSession(), + harness.getContext(), + "test", + module, + graphSupportChecker)) + .isSameAs(failure); + + nodeBehavior.verifyNoWrite(); + nodeBehavior.verifyPreAcquireCancelled(); + } + } + @DataProvider public static Object[][] dseVersionsWithDefaultGraphProtocol() { // Default GraphSON sub protocol version differs based on DSE version, so test with a version diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java new file mode 100644 index 00000000000..bd0cefcc14a --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/adminrequest/ThrottledAdminRequestHandlerTest.java @@ -0,0 +1,127 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.datastax.oss.driver.internal.core.adminrequest; + +import static com.datastax.oss.driver.Assertions.assertThatStage; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; +import com.datastax.oss.driver.internal.core.channel.DriverChannel; +import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; +import com.datastax.oss.protocol.internal.Frame; +import com.datastax.oss.protocol.internal.request.Query; +import java.time.Duration; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +public class ThrottledAdminRequestHandlerTest { + + @Mock private DriverChannel channel; + @Mock private RequestThrottler throttler; + @Mock private SessionMetricUpdater metricUpdater; + private final AtomicInteger availableIds = new AtomicInteger(1); + + @Before + public void setup() { + MockitoAnnotations.initMocks(this); + when(channel.preAcquireId()).thenAnswer(invocation -> availableIds.compareAndSet(1, 0)); + doAnswer( + invocation -> { + if (!availableIds.compareAndSet(0, 1)) { + throw new AssertionError("No caller-owned reservation to cancel"); + } + return null; + }) + .when(channel) + .cancelPreAcquireId(); + } + + @Test + public void should_release_permit_and_reservation_when_metric_update_throws() { + RuntimeException failure = new RuntimeException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(); + handler.start(); + doThrow(failure) + .when(metricUpdater) + .updateTimer( + eq(DefaultSessionMetric.THROTTLING_DELAY), + isNull(), + anyLong(), + eq(TimeUnit.NANOSECONDS)); + + assertThatThrownBy(() -> handler.onThrottleReady(true)).isSameAs(failure); + + assertThat(availableIds.get()).isEqualTo(1); + verify(throttler).signalError(handler, failure); + verify(channel, never()).write(any(), anyBoolean(), anyMap(), any()); + assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); + } + + @Test + public void should_release_permit_when_synchronous_write_throws() { + RuntimeException failure = new RuntimeException("mock failure"); + ThrottledAdminRequestHandler handler = newHandler(); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleReady(false); + return null; + }) + .when(throttler) + .register(handler); + doThrow(failure).when(channel).write(any(), anyBoolean(), anyMap(), eq(handler)); + + assertThatThrownBy(handler::start).isSameAs(failure); + + assertThat(availableIds.get()).isEqualTo(1); + verify(throttler).signalError(handler, failure); + assertThatStage(handler.result).isFailed(error -> assertThat(error).isSameAs(failure)); + } + + private ThrottledAdminRequestHandler newHandler() { + assertThat(channel.preAcquireId()).isTrue(); + assertThat(availableIds.get()).isZero(); + return ThrottledAdminRequestHandler.query( + channel, + false, + new Query("mock query"), + Frame.NO_PAYLOAD, + Duration.ZERO, + throttler, + metricUpdater, + "test", + "mock query"); + } +} diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.java new file mode 100644 index 00000000000..b1d1e4e5f59 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelHandlerRequestTest.java @@ -0,0 +1,134 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.datastax.oss.driver.internal.core.channel; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.datastax.oss.protocol.internal.Message; +import com.datastax.oss.protocol.internal.request.Query; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPipeline; +import io.netty.channel.ChannelPromise; +import io.netty.channel.DefaultChannelPromise; +import io.netty.channel.EventLoop; +import io.netty.util.concurrent.ImmediateEventExecutor; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +public class ChannelHandlerRequestTest { + + @Mock private ChannelHandlerContext context; + @Mock private Channel channel; + @Mock private ChannelPipeline pipeline; + @Mock private EventLoop eventLoop; + @Mock private InFlightHandler inFlightHandler; + + @Before + public void setup() { + MockitoAnnotations.initMocks(this); + when(context.channel()).thenReturn(channel); + when(context.pipeline()).thenReturn(pipeline); + when(pipeline.get(InFlightHandler.class)).thenReturn(inFlightHandler); + when(channel.eventLoop()).thenReturn(eventLoop); + when(eventLoop.inEventLoop()).thenReturn(true); + when(inFlightHandler.preAcquireId()).thenReturn(true); + } + + @Test + public void should_cancel_pre_acquired_id_when_request_construction_fails() { + RuntimeException failure = new RuntimeException("mock failure"); + TestRequest request = new TestRequest(failure); + + assertThatThrownBy(request::send).isSameAs(failure); + + verify(inFlightHandler).cancelPreAcquireId(); + verify(channel, never()).writeAndFlush(any()); + } + + @Test + public void should_cancel_pre_acquired_id_when_raw_write_throws() { + RuntimeException failure = new RuntimeException("mock failure"); + TestRequest request = new TestRequest(null); + doThrow(failure).when(channel).writeAndFlush(any()); + + assertThatThrownBy(request::send).isSameAs(failure); + + verify(inFlightHandler).cancelPreAcquireId(); + } + + @Test + public void should_cancel_pre_acquired_id_when_raw_write_fails_asynchronously() { + RuntimeException failure = new RuntimeException("mock failure"); + TestRequest request = new TestRequest(null); + request.expectFailure = true; + ChannelPromise writePromise = + new DefaultChannelPromise(channel, ImmediateEventExecutor.INSTANCE); + when(channel.writeAndFlush(any())).thenReturn(writePromise); + + request.send(); + writePromise.setFailure(failure); + + verify(inFlightHandler).cancelPreAcquireId(); + assertThat(request.failureCause).isSameAs(failure); + } + + private class TestRequest extends ChannelHandlerRequest { + + private final RuntimeException requestFailure; + private boolean expectFailure; + private Throwable failureCause; + + private TestRequest(RuntimeException requestFailure) { + super(context, 1000); + this.requestFailure = requestFailure; + } + + @Override + String describe() { + return "test request"; + } + + @Override + Message getRequest() { + if (requestFailure != null) { + throw requestFailure; + } + return new Query("mock query"); + } + + @Override + void onResponse(Message response) {} + + @Override + void fail(String message, Throwable cause) { + if (!expectFailure) { + throw new AssertionError("Unexpected failure callback", cause); + } + failureCause = cause; + } + } +} diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java new file mode 100644 index 00000000000..8a1e9721258 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/DefaultWriteCoalescerTest.java @@ -0,0 +1,240 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.datastax.oss.driver.internal.core.channel; + +import static com.datastax.oss.driver.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.DefaultProtocolVersion; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.config.DriverConfig; +import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; +import com.datastax.oss.driver.api.core.context.DriverContext; +import com.datastax.oss.protocol.internal.Frame; +import com.datastax.oss.protocol.internal.request.Query; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.channel.ChannelPromise; +import io.netty.channel.EventLoop; +import java.time.Duration; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +public class DefaultWriteCoalescerTest { + + @Mock private DriverContext context; + @Mock private DriverConfig config; + @Mock private DriverExecutionProfile executionProfile; + @Mock private Channel channel; + @Mock private EventLoop eventLoop; + + @Before + public void setup() { + MockitoAnnotations.initMocks(this); + when(context.getConfig()).thenReturn(config); + when(config.getDefaultProfile()).thenReturn(executionProfile); + when(executionProfile.getDuration(DefaultDriverOption.COALESCER_INTERVAL)) + .thenReturn(Duration.ZERO); + when(channel.eventLoop()).thenReturn(eventLoop); + when(eventLoop.inEventLoop()).thenReturn(true); + } + + @Test + public void should_remove_rejected_write_and_schedule_concurrently_enqueued_writes() { + DefaultWriteCoalescer coalescer = new DefaultWriteCoalescer(context); + Object rejectedMessage = new Object(); + Object queuedMessage = new Object(); + AtomicReference queuedFuture = new AtomicReference<>(); + RejectedExecutionException failure = new RejectedExecutionException("mock failure"); + + doAnswer( + invocation -> { + // Simulate another caller enqueueing while the first scheduling attempt owns the + // running flag. + queuedFuture.set(coalescer.writeAndFlush(channel, queuedMessage)); + throw failure; + }) + .doAnswer( + invocation -> { + invocation.getArgument(0, Runnable.class).run(); + return null; + }) + .when(eventLoop) + .execute(any(Runnable.class)); + + assertThatThrownBy(() -> coalescer.writeAndFlush(channel, rejectedMessage)).isSameAs(failure); + + verify(channel, never()).write(eq(rejectedMessage), any(ChannelPromise.class)); + verify(channel).write(queuedMessage, (ChannelPromise) queuedFuture.get()); + verify(channel).flush(); + } + + @Test + public void should_fail_concurrently_enqueued_writes_if_rescheduling_is_rejected() { + DefaultWriteCoalescer coalescer = new DefaultWriteCoalescer(context); + Object rejectedMessage = new Object(); + StreamIdGenerator streamIds = new StreamIdGenerator(1); + assertThat(streamIds.preAcquire()).isTrue(); + DriverChannel.RequestMessage queuedMessage = newRequestMessage(streamIds); + AtomicReference queuedFuture = new AtomicReference<>(); + RejectedExecutionException firstFailure = new RejectedExecutionException("first mock failure"); + RejectedExecutionException retryFailure = new RejectedExecutionException("retry mock failure"); + + doAnswer( + invocation -> { + // Simulate another caller enqueueing while the first scheduling attempt owns the + // running flag. + queuedFuture.set(coalescer.writeAndFlush(channel, queuedMessage)); + throw firstFailure; + }) + .doThrow(retryFailure) + .when(eventLoop) + .execute(any(Runnable.class)); + + assertThatThrownBy(() -> coalescer.writeAndFlush(channel, rejectedMessage)) + .isSameAs(firstFailure) + .hasSuppressedException(retryFailure); + + verify(channel, never()).write(any(), any(ChannelPromise.class)); + assertThat(queuedFuture.get()).isFailed(error -> assertThat(error).isSameAs(retryFailure)); + assertThat(streamIds.getAvailableIds()).isEqualTo(1); + + // Listener notification must not depend on the same event loop that rejected the write task. + when(eventLoop.inEventLoop()).thenReturn(false); + AtomicReference listenerFailure = new AtomicReference<>(); + queuedFuture.get().addListener(future -> listenerFailure.set(future.cause())); + assertThat(listenerFailure.get()).isSameAs(retryFailure); + } + + @Test + public void should_fail_concurrently_enqueued_writes_when_event_loop_is_shutting_down() { + DefaultWriteCoalescer coalescer = new DefaultWriteCoalescer(context); + StreamIdGenerator streamIds = new StreamIdGenerator(1); + assertThat(streamIds.preAcquire()).isTrue(); + DriverChannel.RequestMessage queuedMessage = newRequestMessage(streamIds); + AtomicReference queuedFuture = new AtomicReference<>(); + + doAnswer( + invocation -> { + invocation.getArgument(0, Runnable.class).run(); + return null; + }) + .when(eventLoop) + .execute(any(Runnable.class)); + doAnswer( + invocation -> { + queuedFuture.set(coalescer.writeAndFlush(channel, queuedMessage)); + return null; + }) + .when(channel) + .flush(); + when(eventLoop.isShuttingDown()).thenReturn(true); + + coalescer.writeAndFlush(channel, new Object()); + + assertThat(queuedFuture.get()) + .isFailed(error -> assertThat(error).isInstanceOf(RejectedExecutionException.class)); + assertThat(streamIds.getAvailableIds()).isEqualTo(1); + } + + @Test + public void should_fail_concurrently_enqueued_writes_when_rescheduling_is_rejected() { + DefaultWriteCoalescer coalescer = new DefaultWriteCoalescer(context); + StreamIdGenerator streamIds = new StreamIdGenerator(1); + assertThat(streamIds.preAcquire()).isTrue(); + DriverChannel.RequestMessage queuedMessage = newRequestMessage(streamIds); + AtomicReference queuedFuture = new AtomicReference<>(); + RejectedExecutionException failure = new RejectedExecutionException("mock failure"); + + doAnswer( + invocation -> { + invocation.getArgument(0, Runnable.class).run(); + return null; + }) + .when(eventLoop) + .execute(any(Runnable.class)); + doAnswer( + invocation -> { + queuedFuture.set(coalescer.writeAndFlush(channel, queuedMessage)); + return null; + }) + .when(channel) + .flush(); + when(eventLoop.schedule(any(Runnable.class), anyLong(), eq(TimeUnit.NANOSECONDS))) + .thenThrow(failure); + + coalescer.writeAndFlush(channel, new Object()); + + assertThat(queuedFuture.get()).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(streamIds.getAvailableIds()).isEqualTo(1); + } + + @Test + public void should_release_pre_acquired_id_when_channel_write_throws() { + DefaultWriteCoalescer coalescer = new DefaultWriteCoalescer(context); + StreamIdGenerator streamIds = new StreamIdGenerator(1); + assertThat(streamIds.preAcquire()).isTrue(); + DriverChannel.RequestMessage message = newRequestMessage(streamIds); + RuntimeException failure = new RuntimeException("mock failure"); + + doAnswer( + invocation -> { + invocation.getArgument(0, Runnable.class).run(); + return null; + }) + .when(eventLoop) + .execute(any(Runnable.class)); + doThrow(failure).when(channel).write(eq(message), any(ChannelPromise.class)); + + ChannelFuture writeFuture = coalescer.writeAndFlush(channel, message); + + assertThat(writeFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + assertThat(streamIds.getAvailableIds()).isEqualTo(1); + } + + private DriverChannel.RequestMessage newRequestMessage(StreamIdGenerator streamIds) { + InFlightHandler inFlightHandler = + new InFlightHandler( + DefaultProtocolVersion.V3, + streamIds, + Integer.MAX_VALUE, + 0, + org.mockito.Mockito.mock(ChannelPromise.class), + null, + "test"); + return new DriverChannel.RequestMessage( + new Query("mock query"), + false, + Frame.NO_PAYLOAD, + org.mockito.Mockito.mock(ResponseCallback.class), + inFlightHandler); + } +} 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..b9977dca7e8 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,8 @@ package com.datastax.oss.driver.internal.core.channel; import static com.datastax.oss.driver.Assertions.assertThat; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.datastax.oss.driver.api.core.DefaultProtocolVersion; import com.datastax.oss.driver.api.core.connection.ClosedConnectionException; @@ -143,6 +145,87 @@ 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() { + // Given + when(streamIds.preAcquire()).thenReturn(true); + assertThat(driverChannel.preAcquireId()).isTrue(); + driverChannel.close(); + + // When + Future writeFuture = + driverChannel.write(new Query("test"), false, Frame.NO_PAYLOAD, new MockResponseCallback()); + + // Then + assertThat(writeFuture).isFailed(); + verify(streamIds).cancelPreAcquire(); + } + + @Test + public void should_cancel_pre_acquired_id_when_coalesced_write_throws() { + // Given + RuntimeException failure = new RuntimeException("mock failure"); + driverChannel = + new DriverChannel( + new EmbeddedEndPoint(), + channel, + (channel, message) -> { + throw failure; + }, + DefaultProtocolVersion.V3); + when(streamIds.preAcquire()).thenReturn(true); + assertThat(driverChannel.preAcquireId()).isTrue(); + + // When + Future writeFuture = + driverChannel.write(new Query("test"), false, Frame.NO_PAYLOAD, new MockResponseCallback()); + + // Then + assertThat(writeFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(streamIds).cancelPreAcquire(); + } + + @Test + public void should_cancel_pre_acquired_id_when_coalesced_write_throws_error() { + // Given + AssertionError failure = new AssertionError("mock failure"); + driverChannel = + new DriverChannel( + new EmbeddedEndPoint(), + channel, + (channel, message) -> { + throw failure; + }, + DefaultProtocolVersion.V3); + when(streamIds.preAcquire()).thenReturn(true); + assertThat(driverChannel.preAcquireId()).isTrue(); + + // When + Future writeFuture = + driverChannel.write(new Query("test"), false, Frame.NO_PAYLOAD, new MockResponseCallback()); + + // Then + assertThat(writeFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(streamIds).cancelPreAcquire(); + } + + @Test + public void should_cancel_pre_acquired_id_when_coalesced_write_fails_asynchronously() { + // Given + RuntimeException failure = new RuntimeException("mock failure"); + when(streamIds.preAcquire()).thenReturn(true); + assertThat(driverChannel.preAcquireId()).isTrue(); + + // When + Future writeFuture = + driverChannel.write(new Query("test"), false, Frame.NO_PAYLOAD, new MockResponseCallback()); + writeCoalescer.failWrites(failure); + + // Then + assertThat(writeFuture).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(streamIds).cancelPreAcquire(); + } + // Simple implementation that holds all the writes, and flushes them when it's explicitly // triggered. private class MockWriteCoalescer implements WriteCoalescer { @@ -161,5 +244,11 @@ void triggerFlush() { channel.writeAndFlush(entry.getKey(), entry.getValue()); } } + + void failWrites(Throwable failure) { + for (Map.Entry entry : messages) { + entry.getValue().tryFailure(failure); + } + } } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java index 1924ef5a9af..24280113697 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlPrepareHandlerTest.java @@ -20,10 +20,13 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static com.datastax.oss.driver.Assertions.assertThatStage; import static com.datastax.oss.driver.internal.core.cql.CqlRequestHandlerTestBase.defaultFrameOf; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -356,6 +359,26 @@ public void should_propagate_custom_payload_on_all_nodes() { } } + @Test + public void should_cancel_pre_acquired_id_if_initial_prepare_payload_access_fails() { + RuntimeException failure = new RuntimeException("mock failure"); + DefaultPrepareRequest prepareRequest = spy(new DefaultPrepareRequest("mock query")); + doThrow(failure).when(prepareRequest).getCustomPayload(); + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + assertThatThrownBy( + () -> + new CqlPrepareHandler( + prepareRequest, harness.getSession(), harness.getContext(), "test")) + .isSameAs(failure); + + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + } + } + private static Message simplePrepared() { RowsMetadata variablesMetadata = new RowsMetadata( diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java index c1a2765eef0..f9068d137f2 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandlerTest.java @@ -19,6 +19,9 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static com.datastax.oss.driver.Assertions.assertThatStage; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -37,6 +40,7 @@ import com.datastax.oss.driver.api.core.cql.Row; import com.datastax.oss.driver.api.core.cql.Statement; import com.datastax.oss.driver.api.core.metadata.Node; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.internal.core.session.RepreparePayload; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; import com.datastax.oss.protocol.internal.request.Prepare; @@ -60,10 +64,12 @@ public class CqlRequestHandlerTest extends CqlRequestHandlerTestBase { @Test public void should_complete_result_if_first_node_replies_immediately() { - try (RequestHandlerTestHarness harness = - RequestHandlerTestHarness.builder() - .withResponse(node1, defaultFrameOf(singleRow())) - .build()) { + RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + node1Behavior.setWriteSuccess(); + node1Behavior.setResponseSuccess(defaultFrameOf(singleRow())); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { CompletionStage resultSetFuture = new CqlRequestHandler( @@ -89,6 +95,7 @@ public void should_complete_result_if_first_node_replies_immediately() { assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(0); assertThat(executionInfo.getWarnings()).isEmpty(); }); + node1Behavior.verifyPreAcquireNotCancelled(); } } @@ -149,6 +156,34 @@ public void should_fail_if_nodes_unavailable() { } } + @Test + public void should_cancel_pre_acquired_id_if_request_decoration_fails_before_write() { + RequestIdGenerator requestIdGenerator = mock(RequestIdGenerator.class); + RuntimeException failure = new RuntimeException("mock failure"); + when(requestIdGenerator.getSessionRequestId()).thenReturn("session"); + when(requestIdGenerator.getNodeRequestId(any(), eq("session"))).thenReturn("node"); + when(requestIdGenerator.getDecoratedStatement(any(), eq("node"))).thenThrow(failure); + + RequestHandlerTestHarness.Builder harnessBuilder = + RequestHandlerTestHarness.builder().withRequestIdGenerator(requestIdGenerator); + PoolBehavior node1Behavior = harnessBuilder.customBehavior(node1); + + try (RequestHandlerTestHarness harness = harnessBuilder.build()) { + // This only verifies stream-id cleanup; the other failure-path leaks are tracked in #980. + assertThatThrownBy( + () -> + new CqlRequestHandler( + UNDEFINED_IDEMPOTENCE_STATEMENT, + harness.getSession(), + harness.getContext(), + "test")) + .isSameAs(failure); + + node1Behavior.verifyNoWrite(); + node1Behavior.verifyPreAcquireCancelled(); + } + } + @Test public void should_time_out_if_first_node_takes_too_long_to_respond() throws Exception { RequestHandlerTestHarness.Builder harnessBuilder = RequestHandlerTestHarness.builder(); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java index 9b018f17531..9f95f2d3c18 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/PoolBehavior.java @@ -87,6 +87,14 @@ public void verifyNoWrite() { .write(any(Message.class), anyBoolean(), anyMap(), any(ResponseCallback.class)); } + public void verifyPreAcquireCancelled() { + verify(channel).cancelPreAcquireId(); + } + + public void verifyPreAcquireNotCancelled() { + verify(channel, never()).cancelPreAcquireId(); + } + public void setWriteSuccess() { writePromise.setSuccess(null); } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java index 58d1783038d..7eeb5bb7129 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/session/DefaultSessionPoolsTest.java @@ -45,6 +45,7 @@ import com.datastax.oss.driver.api.core.session.Session; import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy; import com.datastax.oss.driver.api.core.tracker.RequestTracker; +import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.context.EventBus; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; import com.datastax.oss.driver.internal.core.context.NettyOptions; @@ -62,6 +63,7 @@ import com.datastax.oss.driver.internal.core.pool.ChannelPoolFactory; import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap; +import io.netty.channel.ChannelFuture; import io.netty.channel.DefaultEventLoopGroup; import io.netty.util.concurrent.DefaultPromise; import io.netty.util.concurrent.GlobalEventExecutor; @@ -892,6 +894,36 @@ public void should_set_keyspace_on_pool_if_recreated_while_switching_keyspace() verify(pool2, VERIFY_TIMEOUT).setKeyspace(newKeyspace); } + @Test + public void should_cancel_pre_acquired_id_if_pool_returns_closed_channel() { + ChannelPool pool1 = mockPool(node1); + ChannelPool pool2 = mockPool(node2); + ChannelPool pool3 = mockPool(node3); + MockChannelPoolFactoryHelper factoryHelper = + MockChannelPoolFactoryHelper.builder(channelPoolFactory) + .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1) + .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2) + .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3) + .build(); + + CompletionStage initFuture = newSession(); + factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL); + factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL); + factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL); + assertThatStage(initFuture).isSuccess(); + DefaultSession session = + (DefaultSession) CompletableFutures.getCompleted(initFuture.toCompletableFuture()); + DriverChannel channel = mock(DriverChannel.class); + ChannelFuture closeFuture = mock(ChannelFuture.class); + when(closeFuture.isDone()).thenReturn(true); + when(channel.closeFuture()).thenReturn(closeFuture); + when(pool1.next(null, null)).thenReturn(channel); + + assertThat(session.getChannel(node1, "test")).isNull(); + + verify(channel).cancelPreAcquireId(); + } + private ChannelPool mockPool(Node node) { ChannelPool pool = mock(ChannelPool.class); when(pool.getNode()).thenReturn(node); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java index 555ed2e8806..966451bf9c2 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/session/ReprepareOnUpTest.java @@ -19,12 +19,20 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static com.datastax.oss.driver.Assertions.assertThatStage; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.datastax.oss.driver.api.core.DefaultProtocolVersion; +import com.datastax.oss.driver.api.core.RequestThrottlingException; import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.config.DriverConfig; import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; +import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler; +import com.datastax.oss.driver.api.core.session.throttling.Throttled; import com.datastax.oss.driver.internal.core.adminrequest.AdminResult; import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; @@ -70,6 +78,7 @@ public class ReprepareOnUpTest { @Mock private TopologyMonitor topologyMonitor; @Mock private MetricsFactory metricsFactory; @Mock private SessionMetricUpdater metricUpdater; + @Mock private RequestThrottler throttler; private Runnable whenPrepared; private CompletionStage done; @@ -90,6 +99,7 @@ public void setup() { when(context.getMetricsFactory()).thenReturn(metricsFactory); when(metricsFactory.getSessionUpdater()).thenReturn(metricUpdater); + when(context.getRequestThrottler()).thenReturn(throttler); done = new CompletableFuture<>(); whenPrepared = () -> ((CompletableFuture) done).complete(null); @@ -337,6 +347,52 @@ public void should_limit_number_of_statements_reprepared_in_parallel() { assertThatStage(done).isSuccess(v -> assertThat(reprepareOnUp.queries).isEmpty()); } + @Test + public void should_cancel_pre_acquired_id_if_reprepare_query_fails_before_write() { + RuntimeException failure = new RuntimeException("mock failure"); + doThrow(failure).when(throttler).register(any()); + ReprepareOnUp reprepareOnUp = + new ReprepareOnUp( + "test", + pool, + ImmediateEventExecutor.INSTANCE, + getMockPayloads('a'), + context, + whenPrepared); + + assertThatThrownBy( + () -> reprepareOnUp.queryAsync(new Query("mock query"), Collections.emptyMap(), "mock")) + .isSameAs(failure); + + verify(channel).cancelPreAcquireId(); + } + + @Test + public void should_cancel_pre_acquired_id_if_reprepare_query_is_rejected_by_throttler() { + RequestThrottlingException failure = new RequestThrottlingException("mock failure"); + doAnswer( + invocation -> { + invocation.getArgument(0, Throttled.class).onThrottleFailure(failure); + return null; + }) + .when(throttler) + .register(any()); + ReprepareOnUp reprepareOnUp = + new ReprepareOnUp( + "test", + pool, + ImmediateEventExecutor.INSTANCE, + getMockPayloads('a'), + context, + whenPrepared); + + CompletionStage result = + reprepareOnUp.queryAsync(new Query("mock query"), Collections.emptyMap(), "mock"); + + assertThatStage(result).isFailed(error -> assertThat(error).isSameAs(failure)); + verify(channel).cancelPreAcquireId(); + } + private Map getMockPayloads(char... values) { ImmutableMap.Builder builder = ImmutableMap.builder(); for (char value : values) {