Add support for Attention token (cancelling running queries).

Client.attention() now is able to cancel ongoing queries by sending an attention token. In attention sent state, incoming frames are discarded until receiving the attention acknowledgement token. Any ongoing query is finishes with an error signal MssqlStatementCancelled.

[closes #215]

Signed-off-by: Mark Paluch <mpaluch@vmware.com>

Author: mp911de

src/main/java/io/r2dbc/mssql/DefaultMssqlResult.java

@@ -115,6 +115,11 @@ public Mono<Integer> getRowsUpdated() {
 
                         sink.next(doneToken.getRowCount());
                     }
+
+                    if (doneToken.isAttentionAck()) {
+                        sink.error(new ExceptionFactory.MssqlStatementCancelled());
+                        return;
+                    }
                 }
 
                 if (message instanceof ErrorToken) {
@@ -191,6 +196,15 @@ private <T> Flux<T> doMap(boolean rows, boolean outparameters, Function<? super
         Flux<T> mapped = messages
             .handle((message, sink) -> {
 
+                if (message instanceof AbstractDoneToken) {
+
+                    AbstractDoneToken doneToken = (AbstractDoneToken) message;
+                    if (doneToken.isAttentionAck()) {
+                        sink.error(new ExceptionFactory.MssqlStatementCancelled());
+                        return;
+                    }
+                }
+
                 if (message.getClass() == ColumnMetadataToken.class) {
 
                     ColumnMetadataToken token = (ColumnMetadataToken) message;

src/main/java/io/r2dbc/mssql/ExceptionFactory.java

@@ -339,6 +339,17 @@ public ErrorDetails getErrorDetails() {
 
     }
 
+    /**
+     * SQL Server-specific {@link R2dbcTransientException} upon statement cancellation due to an attention acknowledgement.
+     */
+    static final class MssqlStatementCancelled extends R2dbcTransientException {
+
+        public MssqlStatementCancelled() {
+            super("Statement cancelled");
+        }
+
+    }
+
     /**
      * SQL Server-specific {@link R2dbcTransientResourceException}.
      */

src/main/java/io/r2dbc/mssql/MssqlConnection.java

@@ -149,6 +149,15 @@ public Mono<Void> beginTransaction(TransactionDefinition transactionDefinition)
         });
     }
 
+    /**
+     * Cancel an ongoing request.
+     *
+     * @return
+     */
+    Mono<Void> cancel() {
+        return this.client.attention();
+    }
+
     @Override
     public Mono<Void> close() {
         return this.client.close();

src/main/java/io/r2dbc/mssql/MssqlSegmentResult.java

@@ -128,6 +128,15 @@ private static Flux<Segment> toSegments(String sql, Codecs codecs, Flux<io.r2dbc
         AtomicReference<MssqlRowMetadata> metadataRef = new AtomicReference<>();
         Flux<Segment> segments = messageStream.handle((message, sink) -> {
 
+            if (message instanceof AbstractDoneToken) {
+
+                AbstractDoneToken doneToken = (AbstractDoneToken) message;
+                if (doneToken.isAttentionAck()) {
+                    sink.error(new ExceptionFactory.MssqlStatementCancelled());
+                    return;
+                }
+            }
+
             if (message.getClass() == ColumnMetadataToken.class) {
 
                 ColumnMetadataToken token = (ColumnMetadataToken) message;

src/main/java/io/r2dbc/mssql/ParametrizedMssqlStatement.java

@@ -48,8 +48,6 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
-import static io.r2dbc.mssql.util.PredicateUtils.or;
-
 /**
  * Parametrized {@link Statement} with parameter markers executed against a Microsoft SQL Server database.
  * <p>
@@ -136,7 +134,7 @@ public Flux<MssqlResult> execute() {
                 Binding binding = this.bindings.bindings.get(0);
                 Flux<Message> exchange = exchange(effectiveFetchSize, useGeneratedKeysClause, sql, binding);
 
-                return exchange.windowUntil(or(DoneInProcToken.class::isInstance)).map(it -> DefaultMssqlResult.toResult(this.parsedQuery.getSql(), this.context, this.codecs, it,
+                return exchange.windowUntil(DoneInProcToken.class::isInstance).map(it -> DefaultMssqlResult.toResult(this.parsedQuery.getSql(), this.context, this.codecs, it,
                     binding.hasOutParameters()));
             }
 
@@ -151,7 +149,7 @@ public Flux<MssqlResult> execute() {
 
                 return exchange.doOnComplete(() -> {
                     tryNextBinding(iterator, sink, cancelled);
-                }).windowUntil(or(DoneInProcToken.class::isInstance)).map(it -> DefaultMssqlResult.toResult(this.parsedQuery.getSql(), this.context, this.codecs, it, binding.hasOutParameters()));
+                }).windowUntil(DoneInProcToken.class::isInstance).map(it -> DefaultMssqlResult.toResult(this.parsedQuery.getSql(), this.context, this.codecs, it, binding.hasOutParameters()));
             }).doOnSubscribe(it -> {
 
                 Binding initial = iterator.next();

src/main/java/io/r2dbc/mssql/RpcQueryMessageFlow.java

@@ -70,7 +70,8 @@ final class RpcQueryMessageFlow {
         DoneInProcToken.class::isInstance,
         IntermediateCount.class::isInstance,
         AbstractInfoToken.class::isInstance,
-        Completion.class::isInstance);
+        Completion.class::isInstance,
+        AbstractDoneToken::isAttentionAck);
 
     private static final Logger logger = Loggers.getLogger(RpcQueryMessageFlow.class);
 
@@ -348,6 +349,14 @@ private static void handleMessage(Client client, int fetchSize, Sinks.Many<Clien
             return;
         }
 
+        if (AbstractDoneToken.isAttentionAck(message)) {
+
+            state.phase = Phase.CLOSED;
+            sink.next(message);
+            onCursorComplete.run();
+            return;
+        }
+
         if (!(message instanceof DoneProcToken)) {
 
             if (emit) {

src/main/java/io/r2dbc/mssql/client/Client.java

@@ -38,6 +38,14 @@
  */
 public interface Client {
 
+    /**
+     * Send an attention token to interrupt an active statement.
+     *
+     * @return
+     * @since 0.9
+     */
+    Mono<Void> attention();
+
     /**
      * Release any resources held by the {@link Client}.
      *

src/main/java/io/r2dbc/mssql/client/ReactorNettyClient.java

@@ -36,7 +36,9 @@
 import io.r2dbc.mssql.message.header.PacketIdProvider;
 import io.r2dbc.mssql.message.tds.ProtocolException;
 import io.r2dbc.mssql.message.tds.Redirect;
+import io.r2dbc.mssql.message.token.AbstractDoneToken;
 import io.r2dbc.mssql.message.token.AbstractInfoToken;
+import io.r2dbc.mssql.message.token.Attention;
 import io.r2dbc.mssql.message.token.EnvChangeToken;
 import io.r2dbc.mssql.message.token.FeatureExtAckToken;
 import io.r2dbc.mssql.message.token.LoginAckToken;
@@ -54,6 +56,7 @@
 import reactor.core.publisher.MonoSink;
 import reactor.core.publisher.SynchronousSink;
 import reactor.netty.Connection;
+import reactor.netty.NettyOutbound;
 import reactor.netty.resources.ConnectionProvider;
 import reactor.netty.tcp.SslProvider;
 import reactor.netty.tcp.TcpClient;
@@ -70,6 +73,7 @@
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -116,6 +120,10 @@ public final class ReactorNettyClient implements Client {
 
     private final AtomicBoolean isClosed = new AtomicBoolean(false);
 
+    private final AtomicLong attentionPropagation = new AtomicLong();
+
+    private final AtomicLong outstandingRequests = new AtomicLong();
+
     private final EmitterProcessor<ClientMessage> requestProcessor = EmitterProcessor.create(false);
 
     private final FluxSink<ClientMessage> requests = this.requestProcessor.sink();
@@ -241,6 +249,31 @@ public void next(Message message) {
                     ReactorNettyClient.this.featureAckChange.accept((FeatureExtAckToken) message);
                 }
 
+                if (AbstractDoneToken.isAttentionAck(message)) {
+
+                    long current;
+                    do {
+                        current = ReactorNettyClient.this.attentionPropagation.get();
+
+                        if (current == 0) {
+                            if (DEBUG_ENABLED) {
+                                logger.debug(ReactorNettyClient.this.context.getMessage("Swallowing attention acknowledged, no pending requests: {}. "), message);
+                            }
+                            return;
+                        }
+
+                    } while (!ReactorNettyClient.this.attentionPropagation.compareAndSet(current, current - 1));
+                }
+
+                long attentionPropagation = ReactorNettyClient.this.attentionPropagation.get();
+
+                if (attentionPropagation > 0 && !AbstractDoneToken.isAttentionAck(message)) {
+                    if (DEBUG_ENABLED) {
+                        logger.debug(ReactorNettyClient.this.context.getMessage("Discard message {}. Draining frames until attention acknowledgement."), message);
+                    }
+                    return;
+                }
+
                 ReactorNettyClient.this.responseProcessor.onNext(message);
             }
         };
@@ -294,11 +327,7 @@ public void onComplete() {
             .concatMap(
                 message -> {
 
-                    if (DEBUG_ENABLED) {
-                        logger.debug(this.context.getMessage("Request: {}"), message);
-                    }
-
-                    Object encoded = message.encode(connection.outbound().alloc(), this.tdsEncoder.getPacketSize());
+                    Object encoded = encodeForSend(message);
 
                     if (encoded instanceof Publisher) {
                         return connection.outbound().sendObject((Publisher) encoded);
@@ -311,6 +340,15 @@ public void onComplete() {
             .subscribe();
     }
 
+    private Object encodeForSend(ClientMessage message) {
+
+        if (DEBUG_ENABLED) {
+            logger.debug(this.context.getMessage("Request: {}"), message);
+        }
+
+        return message.encode(this.connection.outbound().alloc(), this.tdsEncoder.getPacketSize());
+    }
+
     @SuppressWarnings("unchecked")
     private <T> Mono<T> resumeError(Throwable throwable) {
 
@@ -469,6 +507,11 @@ private static SslHandler createSslTunnelHandler(ByteBufAllocator allocator, Ssl
         return new SslHandler(tunnel.getSslProvider().getSslContext().newEngine(allocator));
     }
 
+    @Override
+    public Mono<Void> attention() {
+        return Mono.defer(() -> Mono.fromFuture(send(Mono.just(Attention.create(1, getTransactionDescriptor()))).toFuture()));
+    }
+
     @Override
     public Mono<Void> close() {
 
@@ -566,6 +609,7 @@ public Flux<Message> exchange(Publisher<? extends ClientMessage> requests, Predi
 
             Flux<Message> requestMessages = this.responseProcessor
                 .doOnSubscribe(s -> {
+                    this.outstandingRequests.incrementAndGet();
                     Flux.from(requests).subscribe(t -> {
 
                         if (!isConnected()) {
@@ -598,14 +642,26 @@ public Flux<Message> exchange(Publisher<? extends ClientMessage> requests, Predi
             }
         });
 
-        return handle.doAfterTerminate(this.requestQueue).doOnCancel(() -> {
+        return handle.doAfterTerminate(this.requestQueue).doFinally(it -> this.outstandingRequests.decrementAndGet()).doOnCancel(() -> {
 
             if (!exchangeRequest.isComplete()) {
                 logger.error("Exchange cancelled while exchange is active. This is likely a bug leading to unpredictable outcome.");
             }
         });
     }
 
+    private Mono<Void> send(Publisher<? extends ClientMessage> requests) {
+        return Flux.from(requests).concatMap(message -> {
+            NettyOutbound nettyOutbound = this.connection.outbound().sendObject(encodeForSend(message));
+
+            if (message instanceof Attention && this.outstandingRequests.longValue() != 0) {
+                return Mono.from(nettyOutbound).doOnSuccess(v -> this.attentionPropagation.incrementAndGet());
+            }
+
+            return nettyOutbound;
+        }).then();
+    }
+
     private void handleClose() {
         if (this.isClosed.compareAndSet(false, true)) {
             logger.warn(ReactorNettyClient.this.context.getMessage("Connection has been closed by peer"));

src/main/java/io/r2dbc/mssql/message/token/AbstractDoneToken.java

@@ -65,7 +65,7 @@ public abstract class AbstractDoneToken extends AbstractDataToken implements Res
     /**
      * The DONE message is a server acknowledgement of a client ATTENTION message.
      */
-    static final int DONE_ATTN = 0x10;
+    static final int DONE_ATTN = 0x20;
 
     /**
      * This DONEPROC message is associated with an RPC within a set of batched RPCs. This flag is not set on the last RPC in the RPC batch.
@@ -108,6 +108,22 @@ protected AbstractDoneToken(byte type, int status, int currentCommand, long rowC
         this.rowCount = rowCount;
     }
 
+    /**
+     * Check whether the {@link Message} represents a attention acknowledgement.
+     *
+     * @param message the message to inspect.
+     * @return {@literal true} if the {@link Message} represents a attention acknowledgement.
+     * @since 0.9
+     */
+    public static boolean isAttentionAck(Message message) {
+
+        if (message instanceof AbstractDoneToken) {
+            return ((AbstractDoneToken) message).isAttentionAck();
+        }
+
+        return false;
+    }
+
     /**
      * Check whether the {@link Message} represents a finished done token.
      *
@@ -173,6 +189,13 @@ public int getStatus() {
         return this.status;
     }
 
+    /**
+     * @return {@code true} if this token indicates the response is acknowledging the attention request.
+     */
+    public boolean isAttentionAck() {
+        return (getStatus() & DONE_ATTN) != 0;
+    }
+
     /**
      * @return {@code true} if this token indicates the response is done and has no more rows.
      */
@@ -242,6 +265,7 @@ public String toString() {
         sb.append(", hasCount=").append(hasCount());
         sb.append(", rowCount=").append(getRowCount());
         sb.append(", hasMore=").append(hasMore());
+        sb.append(", attnAck=").append(isAttentionAck());
         sb.append(", currentCommand=").append(getCurrentCommand());
         sb.append(']');
         return sb.toString();