Improve explicit transaction terminated state handling

An explicit transaction is considered terminated as soon as a failure occurs. Any usage of the terminated transaction and any of its results must be stopped and the transaction must be closed explicitly.

This update aims to ensure that the driver does not send further Bolt messages in regards to the terminated transaction.

In addition, a new `TransactionTerminatedException` class has been introduced. It is a subclass of the previously used `ClientException`, making this exception more specific. The exception will contain a non-null `code` if it is created based on the server's response. It will not have a code if it is generated by the driver. Depending on a failure cause, the result handles may emit other exceptions respectively, matching the existing behaviour.

Author: injectives

driver/src/main/java/org/neo4j/driver/exceptions/TransactionTerminatedException.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed 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 org.neo4j.driver.exceptions;
+
+import java.io.Serial;
+
+/**
+ * Indicates that the transaction has been terminated.
+ * <p>
+ * Any usage of the terminated transaction and any of its results must be stopped and the transaction must be closed
+ * explicitly. Moreover, any error in the transaction result(s) should be considered as a transaction termination and
+ * must be handled in the same way.
+ * <p>
+ * The exception will contain a non-null {@link #code()} if it is created based on the server's response. It will not
+ * have a code if it is generated by the driver.
+ *
+ * @since 5.11
+ */
+public class TransactionTerminatedException extends ClientException {
+    @Serial
+    private static final long serialVersionUID = 7639191706067500206L;
+
+    /**
+     * Creates a new instance.
+     *
+     * @param message the message
+     */
+    public TransactionTerminatedException(String message) {
+        super(message);
+    }
+
+    /**
+     * Creates a new instance.
+     *
+     * @param code    the code
+     * @param message the message
+     */
+    public TransactionTerminatedException(String code, String message) {
+        super(code, message);
+    }
+
+    /**
+     * Creates a new instance.
+     *
+     * @param message the message
+     * @param cause   the cause
+     */
+    public TransactionTerminatedException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

driver/src/main/java/org/neo4j/driver/internal/InternalTransaction.java

@@ -66,6 +66,22 @@ public boolean isOpen() {
         return tx.isOpen();
     }
 
+    /**
+     * <b>THIS IS A PRIVATE API</b>
+     * <p>
+     * Terminates the transaction by sending the Bolt {@code RESET} message and waiting for its response as long as the
+     * transaction has not already been terminated, is not closed or closing.
+     *
+     * @since 5.11
+     * @throws org.neo4j.driver.exceptions.ClientException if the transaction is closed or is closing
+     * @see org.neo4j.driver.exceptions.TransactionTerminatedException
+     */
+    public void terminate() {
+        Futures.blockingGet(
+                tx.terminateAsync(),
+                () -> terminateConnectionOnThreadInterrupt("Thread interrupted while terminating the transaction"));
+    }
+
     private void terminateConnectionOnThreadInterrupt(String reason) {
         tx.connection().terminateAndRelease(reason);
     }

driver/src/main/java/org/neo4j/driver/internal/async/NetworkConnection.java

@@ -22,14 +22,17 @@
 import static org.neo4j.driver.internal.async.connection.ChannelAttributes.poolId;
 import static org.neo4j.driver.internal.async.connection.ChannelAttributes.setTerminationReason;
 import static org.neo4j.driver.internal.util.Futures.asCompletionStage;
+import static org.neo4j.driver.internal.util.LockUtil.executeWithLock;
 
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import java.time.Clock;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import org.neo4j.driver.Logger;
 import org.neo4j.driver.Logging;
 import org.neo4j.driver.internal.BoltServerAddress;
@@ -41,7 +44,12 @@
 import org.neo4j.driver.internal.handlers.ResetResponseHandler;
 import org.neo4j.driver.internal.messaging.BoltProtocol;
 import org.neo4j.driver.internal.messaging.Message;
+import org.neo4j.driver.internal.messaging.request.DiscardAllMessage;
+import org.neo4j.driver.internal.messaging.request.DiscardMessage;
+import org.neo4j.driver.internal.messaging.request.PullAllMessage;
+import org.neo4j.driver.internal.messaging.request.PullMessage;
 import org.neo4j.driver.internal.messaging.request.ResetMessage;
+import org.neo4j.driver.internal.messaging.request.RunWithMetadataMessage;
 import org.neo4j.driver.internal.metrics.ListenerEvent;
 import org.neo4j.driver.internal.metrics.MetricsListener;
 import org.neo4j.driver.internal.spi.Connection;
@@ -53,6 +61,7 @@
  */
 public class NetworkConnection implements Connection {
     private final Logger log;
+    private final Lock lock;
     private final Channel channel;
     private final InboundMessageDispatcher messageDispatcher;
     private final String serverAgent;
@@ -61,12 +70,13 @@ public class NetworkConnection implements Connection {
     private final ExtendedChannelPool channelPool;
     private final CompletableFuture<Void> releaseFuture;
     private final Clock clock;
-
-    private final AtomicReference<Status> status = new AtomicReference<>(Status.OPEN);
     private final MetricsListener metricsListener;
     private final ListenerEvent<?> inUseEvent;
 
     private final Long connectionReadTimeout;
+
+    private Status status = Status.OPEN;
+    private UnmanagedTransaction transaction;
     private ChannelHandler connectionReadTimeoutHandler;
 
     public NetworkConnection(
@@ -76,6 +86,7 @@ public NetworkConnection(
             MetricsListener metricsListener,
             Logging logging) {
         this.log = logging.getLog(getClass());
+        this.lock = new ReentrantLock();
         this.channel = channel;
         this.messageDispatcher = ChannelAttributes.messageDispatcher(channel);
         this.serverAgent = ChannelAttributes.serverAgent(channel);
@@ -93,7 +104,7 @@ public NetworkConnection(
 
     @Override
     public boolean isOpen() {
-        return status.get() == Status.OPEN;
+        return executeWithLock(lock, () -> status == Status.OPEN);
     }
 
     @Override
@@ -110,52 +121,31 @@ public void disableAutoRead() {
         }
     }
 
-    @Override
-    public void flush() {
-        if (verifyOpen(null, null)) {
-            flushInEventLoop();
-        }
-    }
-
     @Override
     public void write(Message message, ResponseHandler handler) {
-        if (verifyOpen(handler, null)) {
+        if (verifyOpen(handler)) {
             writeMessageInEventLoop(message, handler, false);
         }
     }
 
-    @Override
-    public void write(Message message1, ResponseHandler handler1, Message message2, ResponseHandler handler2) {
-        if (verifyOpen(handler1, handler2)) {
-            writeMessagesInEventLoop(message1, handler1, message2, handler2, false);
-        }
-    }
-
     @Override
     public void writeAndFlush(Message message, ResponseHandler handler) {
-        if (verifyOpen(handler, null)) {
+        if (verifyOpen(handler)) {
             writeMessageInEventLoop(message, handler, true);
         }
     }
 
     @Override
-    public void writeAndFlush(Message message1, ResponseHandler handler1, Message message2, ResponseHandler handler2) {
-        if (verifyOpen(handler1, handler2)) {
-            writeMessagesInEventLoop(message1, handler1, message2, handler2, true);
-        }
-    }
-
-    @Override
-    public CompletionStage<Void> reset() {
-        CompletableFuture<Void> result = new CompletableFuture<>();
-        ResetResponseHandler handler = new ResetResponseHandler(messageDispatcher, result);
+    public CompletionStage<Void> reset(Throwable throwable) {
+        var result = new CompletableFuture<Void>();
+        var handler = new ResetResponseHandler(messageDispatcher, result, throwable);
         writeResetMessageIfNeeded(handler, true);
         return result;
     }
 
     @Override
     public CompletionStage<Void> release() {
-        if (status.compareAndSet(Status.OPEN, Status.RELEASED)) {
+        if (executeWithLock(lock, () -> updateStateIfOpen(Status.RELEASED))) {
             ChannelReleasingResetResponseHandler handler = new ChannelReleasingResetResponseHandler(
                     channel, channelPool, messageDispatcher, clock, releaseFuture);
 
@@ -167,7 +157,7 @@ public CompletionStage<Void> release() {
 
     @Override
     public void terminateAndRelease(String reason) {
-        if (status.compareAndSet(Status.OPEN, Status.TERMINATED)) {
+        if (executeWithLock(lock, () -> updateStateIfOpen(Status.TERMINATED))) {
             setTerminationReason(channel, reason);
             asCompletionStage(channel.close())
                     .exceptionally(throwable -> null)
@@ -194,6 +184,25 @@ public BoltProtocol protocol() {
         return protocol;
     }
 
+    @Override
+    public void bindTransaction(UnmanagedTransaction transaction) {
+        executeWithLock(lock, () -> {
+            if (this.transaction != null) {
+                throw new IllegalStateException("transaction is already set");
+            }
+            this.transaction = transaction;
+        });
+    }
+
+    private boolean updateStateIfOpen(Status newStatus) {
+        if (Status.OPEN.equals(status)) {
+            status = newStatus;
+            return true;
+        } else {
+            return false;
+        }
+    }
+
     private void writeResetMessageIfNeeded(ResponseHandler resetHandler, boolean isSessionReset) {
         channel.eventLoop().execute(() -> {
             if (isSessionReset && !isOpen()) {
@@ -208,73 +217,49 @@ private void writeResetMessageIfNeeded(ResponseHandler resetHandler, boolean isS
         });
     }
 
-    private void flushInEventLoop() {
-        channel.eventLoop().execute(() -> {
-            channel.flush();
-            registerConnectionReadTimeout(channel);
-        });
-    }
-
     private void writeMessageInEventLoop(Message message, ResponseHandler handler, boolean flush) {
-        channel.eventLoop().execute(() -> {
-            messageDispatcher.enqueue(handler);
-
-            if (flush) {
-                channel.writeAndFlush(message).addListener(future -> registerConnectionReadTimeout(channel));
-            } else {
-                channel.write(message, channel.voidPromise());
-            }
-        });
-    }
-
-    private void writeMessagesInEventLoop(
-            Message message1, ResponseHandler handler1, Message message2, ResponseHandler handler2, boolean flush) {
-        channel.eventLoop().execute(() -> {
-            messageDispatcher.enqueue(handler1);
-            messageDispatcher.enqueue(handler2);
-
-            channel.write(message1, channel.voidPromise());
-
-            if (flush) {
-                channel.writeAndFlush(message2).addListener(future -> registerConnectionReadTimeout(channel));
-            } else {
-                channel.write(message2, channel.voidPromise());
-            }
-        });
+        channel.eventLoop()
+                .execute(() -> transactionTerminationAwareExecutor(message).accept(causeOfTermination -> {
+                    if (causeOfTermination == null) {
+                        messageDispatcher.enqueue(handler);
+
+                        if (flush) {
+                            channel.writeAndFlush(message)
+                                    .addListener(future -> registerConnectionReadTimeout(channel));
+                        } else {
+                            channel.write(message, channel.voidPromise());
+                        }
+                    } else {
+                        handler.onFailure(causeOfTermination);
+                    }
+                }));
     }
 
     private void setAutoRead(boolean value) {
         channel.config().setAutoRead(value);
     }
 
-    private boolean verifyOpen(ResponseHandler handler1, ResponseHandler handler2) {
-        Status connectionStatus = this.status.get();
-        switch (connectionStatus) {
-            case OPEN:
-                return true;
-            case RELEASED:
+    private boolean verifyOpen(ResponseHandler handler) {
+        var connectionStatus = executeWithLock(lock, () -> status);
+        return switch (connectionStatus) {
+            case OPEN -> true;
+            case RELEASED -> {
                 Exception error =
                         new IllegalStateException("Connection has been released to the pool and can't be used");
-                if (handler1 != null) {
-                    handler1.onFailure(error);
+                if (handler != null) {
+                    handler.onFailure(error);
                 }
-                if (handler2 != null) {
-                    handler2.onFailure(error);
-                }
-                return false;
-            case TERMINATED:
+                yield false;
+            }
+            case TERMINATED -> {
                 Exception terminatedError =
                         new IllegalStateException("Connection has been terminated and can't be used");
-                if (handler1 != null) {
-                    handler1.onFailure(terminatedError);
-                }
-                if (handler2 != null) {
-                    handler2.onFailure(terminatedError);
+                if (handler != null) {
+                    handler.onFailure(terminatedError);
                 }
-                return false;
-            default:
-                throw new IllegalStateException("Unknown status: " + connectionStatus);
-        }
+                yield false;
+            }
+        };
     }
 
     private void registerConnectionReadTimeout(Channel channel) {
@@ -295,6 +280,25 @@ private void registerConnectionReadTimeout(Channel channel) {
         }
     }
 
+    private Consumer<Consumer<Throwable>> transactionTerminationAwareExecutor(Message message) {
+        var result = (Consumer<Consumer<Throwable>>) consumer -> consumer.accept(null);
+        if (isQueryMessage(message)) {
+            var transaction = executeWithLock(lock, () -> this.transaction);
+            if (transaction != null) {
+                result = transaction::executeWithLockedState;
+            }
+        }
+        return result;
+    }
+
+    private boolean isQueryMessage(Message message) {
+        return message instanceof RunWithMetadataMessage
+                || message instanceof PullMessage
+                || message instanceof PullAllMessage
+                || message instanceof DiscardMessage
+                || message instanceof DiscardAllMessage;
+    }
+
     private enum Status {
         OPEN,
         RELEASED,