perf: non-blocking sockets

benchmarks/latency-comparison/golang/runners/pgx_runner.go

@@ -31,7 +31,7 @@ func RunPgx(database, sql string, readWrite bool, numOperations, numClients, wai
 	// Connect to Cloud Spanner through PGAdapter.
 	var connString string
 	if useUnixSocket {
-		connString = fmt.Sprintf("host=%s port=%d database=%s", host, port, url.QueryEscape(database))
+		connString = fmt.Sprintf("host=%s port=%d database=%s", host, port, database)
 	} else {
 		connString = fmt.Sprintf("postgres://uid:pwd@%s:%d/%s?sslmode=disable", host, port, url.QueryEscape(database))
 	}

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

@@ -25,6 +25,7 @@
 import com.google.cloud.spanner.Instance;
 import com.google.cloud.spanner.InstanceAdminClient;
 import com.google.cloud.spanner.InstanceNotFoundException;
+import com.google.cloud.spanner.SessionPoolOptions;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
@@ -56,6 +57,7 @@
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.TerminateResponse;
 import com.google.cloud.spanner.pgadapter.wireprotocol.BootstrapMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.SSLMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
@@ -108,7 +110,7 @@ public class ConnectionHandler implements Runnable {
   private static final String CHANNEL_PROVIDER_PROPERTY = "CHANNEL_PROVIDER";
 
   private final ProxyServer server;
-  private Socket socket;
+  protected Socket socket;
   private final Map<String, IntermediatePreparedStatement> statementsMap = new HashMap<>();
   private final Cache<String, Future<DescribeResult>> autoDescribedStatementsCache =
       CacheBuilder.newBuilder()
@@ -120,6 +122,7 @@ public class ConnectionHandler implements Runnable {
   private static final Map<Integer, ConnectionHandler> CONNECTION_HANDLERS =
       new ConcurrentHashMap<>();
   private volatile ConnectionStatus status = ConnectionStatus.UNAUTHENTICATED;
+  private Map<String, String> connectionParameters;
   private Thread thread;
   private final int connectionId;
   private final int secret;
@@ -145,7 +148,7 @@ public class ConnectionHandler implements Runnable {
   private final LinkedList<ParseMessage> skippedAutoDetectParseMessages = new LinkedList<>();
 
   private ExtendedQueryProtocolHandler extendedQueryProtocolHandler;
-  private CopyStatement activeCopyStatement;
+  private volatile CopyStatement activeCopyStatement;
 
   ConnectionHandler(ProxyServer server, Socket socket) {
     this(server, socket, null);
@@ -208,6 +211,9 @@ public void connectToSpanner(String database, @Nullable Credentials credentials)
     if (options.getSessionPoolOptions() != null) {
       connectionOptionsBuilder =
           connectionOptionsBuilder.setSessionPoolOptions(options.getSessionPoolOptions());
+    } else {
+      connectionOptionsBuilder.setSessionPoolOptions(
+          SessionPoolOptions.newBuilder().setTrackStackTraceOfSessionCheckout(false).build());
     }
     if (options.isEnableOpenTelemetryMetrics()) {
       SpannerOptions.enableOpenTelemetryMetrics();
@@ -365,18 +371,29 @@ enum RunConnectionState {
     TERMINATED
   }
 
+  public BootstrapMessage readBootstrapMessage() throws Exception {
+    return BootstrapMessage.create(this);
+  }
+
+  public ControlMessage readControlMessage() throws Exception {
+    return ControlMessage.create(this);
+  }
+
+  protected ConnectionMetadata createConnectionMetadata() throws IOException {
+    return new ConnectionMetadata(this.socket.getInputStream(), this.socket.getOutputStream());
+  }
+
   /**
    * Starts listening for incoming messages on the network socket. Returns RESTART_WITH_SSL if the
    * listening process should be restarted with SSL.
    */
   private RunConnectionState runConnection(boolean ssl) {
     RunConnectionState result = RunConnectionState.TERMINATED;
-    try (ConnectionMetadata connectionMetadata =
-        new ConnectionMetadata(this.socket.getInputStream(), this.socket.getOutputStream())) {
+    try (ConnectionMetadata connectionMetadata = createConnectionMetadata()) {
       this.connectionMetadata = connectionMetadata;
 
       try {
-        this.message = this.server.recordMessage(BootstrapMessage.create(this));
+        this.message = this.server.recordMessage(readBootstrapMessage());
         if (!ssl
             && getServer().getOptions().getSslMode().isSslEnabled()
             && this.message instanceof SSLMessage) {
@@ -435,7 +452,7 @@ && getServer().getOptions().getSslMode().isSslEnabled()
           if (this.spannerConnection != null) {
             this.spannerConnection.close();
           }
-          this.socket.close();
+          closeSocket();
         } catch (SpannerException | IOException e) {
           logger.log(
               Level.WARNING,
@@ -457,6 +474,10 @@ && getServer().getOptions().getSslMode().isSslEnabled()
     return result;
   }
 
+  protected void closeSocket() throws IOException {
+    this.socket.close();
+  }
+
   boolean checkValidConnection(boolean ssl) throws Exception {
     // Allow SSL connections from non-localhost even if the localhost check has not explicitly
     // been disabled.
@@ -546,7 +567,7 @@ void terminate() {
    * @param exception The exception to be related.
    * @throws IOException if there is some issue in the sending of the error messages.
    */
-  void handleError(PGException exception) throws Exception {
+  void handleError(PGException exception) throws IOException {
     logger.log(
         Level.WARNING,
         exception,
@@ -771,22 +792,39 @@ public void clearInvalidMessageCount() {
     this.invalidMessagesCount = 0;
   }
 
+  public boolean supportsPeekNextByte() {
+    return true;
+  }
+
   public ConnectionMetadata getConnectionMetadata() {
     return connectionMetadata;
   }
 
+  protected void setConnectionMetadata(ConnectionMetadata connectionMetadata) {
+    this.connectionMetadata = connectionMetadata;
+  }
+
   public ExtendedQueryProtocolHandler getExtendedQueryProtocolHandler() {
     return extendedQueryProtocolHandler;
   }
 
-  public synchronized ConnectionStatus getStatus() {
+  public ConnectionStatus getStatus() {
     return status;
   }
 
-  public synchronized void setStatus(ConnectionStatus status) {
+  public void setStatus(ConnectionStatus status) {
     this.status = status;
   }
 
+  public void setStatus(ConnectionStatus status, Map<String, String> connectionParameters) {
+    this.status = status;
+    this.connectionParameters = connectionParameters;
+  }
+
+  public Map<String, String> getConnectionParameters() {
+    return this.connectionParameters;
+  }
+
   public WellKnownClient getWellKnownClient() {
     return wellKnownClient;
   }
@@ -894,6 +932,7 @@ boolean isHasDeterminedClientUsingQuery() {
   /** Status of a {@link ConnectionHandler} */
   public enum ConnectionStatus {
     UNAUTHENTICATED,
+    AUTHENTICATING,
     AUTHENTICATED,
     COPY_IN,
     COPY_DONE,

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

@@ -0,0 +1,130 @@
+// Copyright 2024 Google LLC
+//
+// 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 com.google.cloud.spanner.pgadapter;
+
+import com.google.cloud.spanner.pgadapter.metadata.ChannelOutputStream;
+import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
+import com.google.cloud.spanner.pgadapter.metadata.ForwardingInputStream;
+import com.google.cloud.spanner.pgadapter.wireprotocol.BootstrapMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.FlushMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.SyncMessage;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.SocketChannel;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.logging.Logger;
+
+public class NonBlockingConnectionHandler extends ConnectionHandler {
+  private static final int DEFAULT_BUFFER_CAPACITY = 1 << 13;
+
+  private static final Logger logger =
+      Logger.getLogger(NonBlockingConnectionHandler.class.getName());
+
+  private final ByteBuffer headerBuffer = ByteBuffer.allocateDirect(5);
+
+  private ByteBuffer messageBuffer = ByteBuffer.allocateDirect(DEFAULT_BUFFER_CAPACITY);
+
+  private final BlockingQueue<BootstrapMessage> bootstrapMessages = new LinkedBlockingQueue<>();
+
+  private final BlockingQueue<ControlMessage> controlMessages = new LinkedBlockingQueue<>();
+
+  private final ForwardingInputStream forwardingInputStream = new ForwardingInputStream();
+
+  private final SocketChannel channel;
+
+  NonBlockingConnectionHandler(ProxyServer server, SocketChannel channel) {
+    super(server, channel.socket());
+    this.channel = channel;
+    setConnectionMetadata(
+        new ConnectionMetadata(forwardingInputStream, new ChannelOutputStream(channel)));
+  }
+
+  @Override
+  void createSSLSocket() throws IOException {
+    throw new IOException("SSL is not supported for non-blocking connection handlers");
+  }
+
+  ByteBuffer getHeaderBuffer() {
+    this.headerBuffer.rewind();
+    return this.headerBuffer;
+  }
+
+  ByteBuffer getMessageBuffer(int length) {
+    if (this.messageBuffer.capacity() < length) {
+      this.messageBuffer = ByteBuffer.allocateDirect(length);
+    } else {
+      this.messageBuffer.rewind();
+      this.messageBuffer.limit(length);
+    }
+    return this.messageBuffer;
+  }
+
+  void setRawInputStream(InputStream inputStream) {
+    this.forwardingInputStream.setDelegate(inputStream);
+  }
+
+  @Override
+  public boolean supportsPeekNextByte() {
+    return false;
+  }
+
+  @Override
+  protected ConnectionMetadata createConnectionMetadata() {
+    return getConnectionMetadata();
+  }
+
+  @Override
+  protected void closeSocket() throws IOException {
+    try {
+      this.channel.close();
+    } catch (IOException ioException) {
+      ioException.printStackTrace();
+      throw ioException;
+    } catch (Throwable t) {
+      t.printStackTrace();
+    }
+  }
+
+  void addBootstrapMessage(BootstrapMessage bootstrapMessage) {
+    this.bootstrapMessages.add(bootstrapMessage);
+  }
+
+  @Override
+  public BootstrapMessage readBootstrapMessage() throws Exception {
+    return this.bootstrapMessages.take();
+  }
+
+  void addControlMessage(ControlMessage controlMessage) {
+    this.controlMessages.add(controlMessage);
+  }
+
+  @Override
+  public ControlMessage readControlMessage() throws Exception {
+    if (getStatus() == ConnectionStatus.COPY_IN) {
+      while (true) {
+        ControlMessage message = this.controlMessages.take();
+        if (message instanceof FlushMessage || message instanceof SyncMessage) {
+          continue;
+        }
+        return message;
+      }
+    } else {
+      return this.controlMessages.take();
+    }
+  }
+}