chore: refactor ProxyServer to use channels instead of sockets

.github/workflows/integration.yaml

@@ -75,7 +75,7 @@ jobs:
         uses: actions/setup-java@v3
         with:
           distribution: zulu
-          java-version: 8
+          java-version: 17
       - run: java -version
       - name: Setup Go
         uses: actions/setup-go@v3
@@ -100,6 +100,8 @@ jobs:
           psql -h /pg -U postgres -c "CREATE DATABASE pgadapter"
       - name: Run unit tests
         run: mvn test -B -Ptest-all
+      - name: Run unit tests for native image
+        run: mvn test -B -Pnative-image
       - name: Setup GCloud
         uses: google-github-actions/setup-gcloud@v0
         with:

.github/workflows/native-image.yaml

@@ -0,0 +1,65 @@
+on:
+  # This allows manual activation of this action for testing.
+  workflow_dispatch:
+  pull_request:
+  schedule:
+    - cron: '0 2 * * 1,2,3,4,5'
+name: native-image
+env:
+  GOOGLE_CLOUD_PROJECT: "span-cloud-testing"
+  GOOGLE_CLOUD_INSTANCE: "pgadapter-testing"
+  GOOGLE_CLOUD_DATABASE: "testdb_integration"
+  GOOGLE_CLOUD_ENDPOINT: "spanner.googleapis.com"
+jobs:
+  check-env:
+    outputs:
+      has-key: ${{ steps.project-id.outputs.defined }}
+    runs-on: ubuntu-latest
+    steps:
+      - id: project-id
+        env:
+          GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }}
+        if: "${{ env.GCP_PROJECT_ID != '' }}"
+        run: echo "::set-output name=defined::true"
+
+  build-native-image:
+    needs: [check-env]
+    if: needs.check-env.outputs.has-key == 'true'
+    timeout-minutes: 60
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: graalvm/setup-graalvm@v1
+        with:
+          version: 'latest'
+          java-version: '17'
+          components: 'native-image'
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Setup GCloud
+        uses: google-github-actions/setup-gcloud@v0
+        with:
+          project_id: ${{ secrets.GCP_PROJECT_ID }}
+          service_account_key: ${{ secrets.JSON_SERVICE_ACCOUNT_CREDENTIALS }}
+          export_default_credentials: true
+      - name: Set up swap space
+        if: runner.os == 'Linux'
+        uses: pierotofy/[email protected]
+        with:
+          swap-size-gb: 12
+      - name: Build and run PGAdapter native image
+        run: |
+          mvn package -Passembly -Pnative-image -DskipTests
+          cd "target/pgadapter"
+          native-image -J-Xmx14g -H:IncludeResources=".*metadata.*json$" -jar pgadapter.jar --no-fallback
+          ./pgadapter -p ${{env.GOOGLE_CLOUD_PROJECT}} -i ${{env.GOOGLE_CLOUD_INSTANCE}} &
+      - name: Run integration tests
+        run: |
+          mvn verify \
+                -Dclirr.skip=true \
+                -DskipITs=false \
+                -DPG_ADAPTER_PROJECT="${{env.GOOGLE_CLOUD_PROJECT}}" \
+                -DPG_ADAPTER_INSTANCE="${{env.GOOGLE_CLOUD_INSTANCE}}" \
+                -DPG_ADAPTER_DATABASE="${{env.GOOGLE_CLOUD_DATABASE}}" \
+                -DPG_ADAPTER_ADDRESS="localhost" \
+                -DPG_ADAPTER_SOCKET_DIR="/tmp" \
+                -DPG_ADAPTER_PORT="5432"

pom.xml

@@ -284,6 +284,51 @@
         </plugins>
       </build>
     </profile>
+    <profile>
+    <id>native-image</id>
+      <dependencies>
+        <!--
+          Change the scope of the junixsocket dependencies to test, as we don't need them in the
+          final build, but we do need them for some tests.
+        -->
+        <dependency>
+          <groupId>com.kohlschutter.junixsocket</groupId>
+          <artifactId>junixsocket-core</artifactId>
+          <version>${junixsocket.version}</version>
+          <type>pom</type>
+          <scope>test</scope>
+        </dependency>
+        <dependency>
+          <groupId>com.kohlschutter.junixsocket</groupId>
+          <artifactId>junixsocket-common</artifactId>
+          <version>${junixsocket.version}</version>
+          <scope>test</scope>
+        </dependency>
+      </dependencies>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.apache.maven.plugins</groupId>
+            <artifactId>maven-compiler-plugin</artifactId>
+            <configuration>
+              <excludes combine.self="override">
+                <exclude>com/google/cloud/spanner/pgadapter/Java8ServerSocketFactory.java</exclude>
+              </excludes>
+              <source>17</source>
+              <target>17</target>
+              <release>17</release>
+            </configuration>
+          </plugin>
+          <plugin>
+            <groupId>org.codehaus.mojo</groupId>
+            <artifactId>animal-sniffer-maven-plugin</artifactId>
+            <configuration>
+              <skip>true</skip>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
   </profiles>
 
   <build>
@@ -321,6 +366,15 @@
           </execution>
         </executions>
       </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-compiler-plugin</artifactId>
+        <configuration>
+          <excludes>
+            <exclude>com/google/cloud/spanner/pgadapter/Java17ServerSocketFactory.java</exclude>
+          </excludes>
+        </configuration>
+      </plugin>
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>

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

@@ -53,7 +53,10 @@
 import java.io.DataOutputStream;
 import java.io.EOFException;
 import java.io.IOException;
-import java.net.Socket;
+import java.net.InetSocketAddress;
+import java.nio.channels.ByteChannel;
+import java.nio.channels.Channels;
+import java.nio.channels.SocketChannel;
 import java.security.SecureRandom;
 import java.text.MessageFormat;
 import java.util.HashMap;
@@ -83,7 +86,15 @@ public class ConnectionHandler extends Thread {
   private static final String CHANNEL_PROVIDER_PROPERTY = "CHANNEL_PROVIDER";
 
   private final ProxyServer server;
-  private final Socket socket;
+  private final ByteChannel socketChannel;
+
+  /**
+   * Remote address of the client that is connected. This is null for Unix domain socket
+   * connections.
+   */
+  @Nullable private final InetSocketAddress remoteAddress;
+
+  private final String remoteClient;
   private final Map<String, IntermediatePreparedStatement> statementsMap = new HashMap<>();
   private final Map<String, IntermediatePortalStatement> portalsMap = new HashMap<>();
   private static final Map<Integer, IntermediateStatement> activeStatementsMap =
@@ -103,18 +114,30 @@ public class ConnectionHandler extends Thread {
   private WellKnownClient wellKnownClient;
   private ExtendedQueryProtocolHandler extendedQueryProtocolHandler;
 
-  ConnectionHandler(ProxyServer server, Socket socket) {
+  ConnectionHandler(ProxyServer server, ByteChannel byteChannel) throws IOException {
     super("ConnectionHandler-" + CONNECTION_HANDLER_ID_GENERATOR.incrementAndGet());
     this.server = server;
-    this.socket = socket;
+    this.socketChannel = byteChannel;
+    // byteChannel is an instance of SocketChannel in case of a TCP connection.
+    // Unix domain socket connections use other implementations. The exact implementation depends on
+    // whether the connection uses the third-party AFUNIXSocketFactory or the Java 16+ Unix domain
+    // socket implementation.
+    if (byteChannel instanceof SocketChannel
+        && ((SocketChannel) byteChannel).getRemoteAddress() instanceof InetSocketAddress) {
+      SocketChannel socketChannel = (SocketChannel) byteChannel;
+      this.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
+      this.remoteClient = remoteAddress.getAddress().getHostAddress();
+    } else {
+      this.remoteAddress = null;
+      this.remoteClient = "(local)";
+    }
     this.secret = new SecureRandom().nextInt();
     setDaemon(true);
     logger.log(
         Level.INFO,
         () ->
             String.format(
-                "Connection handler with ID %s created for client %s",
-                getName(), socket.getInetAddress().getHostAddress()));
+                "Connection handler with ID %s created for client %s", getName(), remoteClient));
   }
 
   @InternalApi
@@ -217,15 +240,16 @@ public void run() {
         Level.INFO,
         () ->
             String.format(
-                "Connection handler with ID %s starting for client %s",
-                getName(), socket.getInetAddress().getHostAddress()));
+                "Connection handler with ID %s starting for client %s", getName(), remoteClient));
 
     try (ConnectionMetadata connectionMetadata =
-        new ConnectionMetadata(this.socket.getInputStream(), this.socket.getOutputStream())) {
+        new ConnectionMetadata(
+            Channels.newInputStream(socketChannel), Channels.newOutputStream(socketChannel))) {
       this.connectionMetadata = connectionMetadata;
       if (!server.getOptions().disableLocalhostCheck()
-          && !this.socket.getInetAddress().isAnyLocalAddress()
-          && !this.socket.getInetAddress().isLoopbackAddress()) {
+          && this.remoteAddress != null
+          && !this.remoteAddress.getAddress().isAnyLocalAddress()
+          && !this.remoteAddress.getAddress().isLoopbackAddress()) {
         handleError(
             PGException.newBuilder()
                 .setMessage("This proxy may only be accessed from localhost.")
@@ -268,15 +292,15 @@ public void run() {
           () ->
               String.format(
                   "Exception on connection handler with ID %s for client %s: %s",
-                  getName(), socket.getInetAddress().getHostAddress(), e));
+                  getName(), remoteClient, e));
     } finally {
       logger.log(
           Level.INFO, () -> String.format("Closing connection handler with ID %s", getName()));
       try {
         if (this.spannerConnection != null) {
           this.spannerConnection.close();
         }
-        this.socket.close();
+        this.socketChannel.close();
       } catch (SpannerException | IOException e) {
         logger.log(
             Level.WARNING,
@@ -332,9 +356,7 @@ public void handleTerminate() {
   void terminate() {
     handleTerminate();
     try {
-      if (!socket.isClosed()) {
-        socket.close();
-      }
+      socketChannel.close();
     } catch (IOException exception) {
       logger.log(
           Level.WARNING,

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

@@ -0,0 +1,79 @@
+// Copyright 2022 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 static java.net.StandardSocketOptions.SO_RCVBUF;
+
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import java.io.File;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.StandardProtocolFamily;
+import java.net.UnixDomainSocketAddress;
+import java.nio.channels.ServerSocketChannel;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+// If your IDE is complaining about this file not being compatible with the language level
+// of the project, then that is an indication that the IDE did not pick up the exclusion of this
+// file that is defined in the pom.xml.
+// This is a known issue in IntelliJ: https://youtrack.jetbrains.com/issue/IDEA-87868
+//
+// Follow these steps to make IntelliJ ignore this compilation error:
+// 1. Go to Settings.
+// 2. Select Build, Execution, Deployment| Compiler | Excludes
+// 3. Add this file to the list of excludes.
+// See also https://www.jetbrains.com/help/idea/specifying-compilation-settings.html#5a737cfc
+
+class Java17ServerSocketFactory implements ServerSocketFactory {
+
+  public Java17ServerSocketFactory() {}
+
+  public ServerSocketChannel createTcpServerSocketChannel(OptionsMetadata options)
+      throws IOException {
+    ServerSocketChannel channel = ServerSocketChannel.open();
+    InetSocketAddress address = new InetSocketAddress(options.getProxyPort());
+    channel.configureBlocking(true);
+    channel.bind(address, options.getMaxBacklog());
+
+    return channel;
+  }
+
+  public ServerSocketChannel creatUnixDomainSocketChannel(OptionsMetadata options, int localPort)
+      throws IOException {
+    File socketFile = new File(options.getSocketFile(localPort));
+    Path path = Path.of(socketFile.toURI());
+    if (socketFile.getParentFile() != null && !socketFile.getParentFile().exists()) {
+      socketFile.mkdirs();
+    }
+
+    // deleteOnExit() does not work when PGAdapter is built as a native image. It also does not get
+    // deleted if the JVM crashes. We therefore need to try to delete the file at startup to ensure
+    // we don't get a lot of 'address already in use' errors.
+    try {
+      Files.deleteIfExists(path);
+    } catch (IOException ignore) {
+      // Ignore and let the Unix domain socket subsystem throw an 'Address in use' error.
+    }
+
+    UnixDomainSocketAddress address = UnixDomainSocketAddress.of(path);
+    ServerSocketChannel domainSocketChannel = ServerSocketChannel.open(StandardProtocolFamily.UNIX);
+    domainSocketChannel.configureBlocking(true);
+    domainSocketChannel.setOption(SO_RCVBUF, 65536);
+    domainSocketChannel.bind(address, options.getMaxBacklog());
+
+    return domainSocketChannel;
+  }
+}