fix: buffer Close messages in extended query protocol

Buffer Close ('C') messages in ExtendedQueryProtocolHandler until a
Flush or Sync message is received, matching the behavior of other
extended query protocol messages.

Previously, Close messages were executed eagerly upon receipt. In
pipelined batches (such as Parse + Close + Sync or Execute + Close + Sync),
this caused CloseComplete ('3') to be returned to the client before
responses to earlier messages in the batch. This broke connection poolers
(such as PgBouncer in transaction pooling mode) and client drivers that
track in-flight requests in FIFO order.

Statements and portals are unregistered from ConnectionHandler when
buffering so that subsequent messages in the same pipeline can reuse the
name. Statement closure and CloseComplete response emission are deferred
to flush. If an earlier message in the pipeline fails, buffered messages
are aborted in reverse (LIFO) order, restoring any statement or portal
whose Close was skipped.

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java

@@ -715,6 +715,15 @@ public void closePortal(String portalName) throws Exception {
     portal.close();
   }
 
+  /**
+   * Unregisters a portal from the connection map without closing it. Returns the unregistered
+   * portal, or null if no portal was registered with this name. Unlike {@link
+   * #closePortal(String)}, this method does not throw an exception if the portal is not found.
+   */
+  public IntermediatePortalStatement unregisterPortal(String portalName) {
+    return this.portalsMap.remove(portalName);
+  }
+
   public boolean hasPortal(String portalName) {
     return this.portalsMap.containsKey(portalName);
   }
@@ -848,6 +857,16 @@ public void closeStatement(String statementName) {
     this.statementsMap.remove(statementName);
   }
 
+  /**
+   * Unregisters a prepared statement from the connection map without closing it. Returns the
+   * unregistered statement, or null if no statement was registered with this name. Unlike {@link
+   * #closeStatement(String)}, this method does not throw an exception if the statement is not
+   * found.
+   */
+  public IntermediatePreparedStatement unregisterStatement(String statementName) {
+    return this.statementsMap.remove(statementName);
+  }
+
   public void closeAllStatements() {
     Set<String> names = new HashSet<>(this.statementsMap.keySet());
     for (String statementName : names) {

src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java

@@ -48,7 +48,7 @@ public class ExtendedQueryProtocolHandler {
   private static final Logger logger =
       Logger.getLogger(ExtendedQueryProtocolHandler.class.getName());
 
-  private final LinkedList<AbstractQueryProtocolMessage> messages = new LinkedList<>();
+  private final List<AbstractQueryProtocolMessage> messages = new LinkedList<>();
   private final ConnectionHandler connectionHandler;
   private final BackendConnection backendConnection;
 
@@ -210,7 +210,11 @@ private void flushMessages(boolean includeReadyResponse) throws Exception {
             Logging.format(
                 "Flushing message", Action.Finished, () -> String.format("Message: %s", message)));
         if (message.isReturnedErrorResponse()) {
-          for (int j = i + 1; j < messages.size(); j++) {
+          // Abort remaining messages in reverse (LIFO) order to correctly unwind state mutations in
+          // the opposite order that they were registered/buffered (e.g. if a pipeline contains a
+          // Close followed by a Parse/Bind reusing the same statement or portal name, or both
+          // creation and closure of a statement in the same aborted pipeline).
+          for (int j = messages.size() - 1; j > i; j--) {
             messages.get(j).abort();
           }
           break;

src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CloseMessage.java

@@ -17,14 +17,17 @@
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.error.PGException;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
+import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.wireoutput.CloseCompleteResponse;
 import java.text.MessageFormat;
 import javax.annotation.Nullable;
 
 /** Close the designated statement. */
 @InternalApi
-public class CloseMessage extends ControlMessage {
+public class CloseMessage extends AbstractQueryProtocolMessage {
 
   protected static final char IDENTIFIER = 'C';
 
@@ -51,19 +54,53 @@ public CloseMessage(ConnectionHandler connection) throws Exception {
     this.statement = statement;
   }
 
-  /** Close the statement server-side and clean up by deleting their metadata locally. */
   @Override
-  protected void sendPayload() throws Exception {
+  void buffer(BackendConnection backendConnection) throws Exception {
+    // Unregister the statement or portal immediately from the connection map so that subsequent
+    // messages in the same pipeline (e.g. a Parse or Bind reusing the same name) will see the name
+    // available, or attempts to use this closed statement will fail.
+    // Resource cleanup (statement.close()) is deferred until flush() so that:
+    // 1. If an earlier message in the pipeline fails, abort() can restore the statement.
+    // 2. Preceding messages in the pipeline (e.g. Execute) can finish streaming rows before
+    // closure.
     if (this.statement != null) {
       if (this.type == PreparedType.Portal) {
-        this.statement.close(); // Only portals need to be closed server side, since PS is not bound
-        this.connection.closePortal(this.name);
+        this.connection.unregisterPortal(this.name);
       } else {
-        this.connection.closeStatement(this.name);
+        this.connection.unregisterStatement(this.name);
       }
     }
-    CloseCompleteResponse.send(this.outputStream);
-    this.outputStream.flush();
+  }
+
+  @Override
+  public void flush() throws Exception {
+    try {
+      if (this.statement != null) {
+        this.statement.close();
+      }
+      CloseCompleteResponse.send(this.outputStream);
+    } catch (Exception exception) {
+      handleError(exception);
+    }
+  }
+
+  @Override
+  public void abort() {
+    // Only restore the statement or portal if it did not fail during creation (e.g. an
+    // InvalidStatement created by a failed Parse or Bind message must not be resurrected).
+    if (this.statement != null && !this.statement.hasException()) {
+      if (this.type == PreparedType.Portal) {
+        this.connection.registerPortal(this.name, (IntermediatePortalStatement) this.statement);
+      } else {
+        this.connection.registerStatement(
+            this.name, (IntermediatePreparedStatement) this.statement);
+      }
+    }
+  }
+
+  @Override
+  public String getSql() {
+    return this.statement == null ? "" : this.statement.getSql();
   }
 
   @Override