fix: close proxy websocket upon TCP connection closing (#684)

This fixes a bug where clients using `neondatabase/serverless` would
have their code hang for 60s upon calling `client.end()`. Essentially
the client closes the (proxied) TCP socket, and we need to ensure that
the websocket is closed as well. Otherwise the client keeps blocked
waiting for the websocket to terminate.

Added an integration test with such a case (times out without the
patch). Unfortunately, the test requires bumping node, so I did that as
well, will check to ensure it doesn't impact anything else.

Author: v0idpwn

.tool-versions

@@ -1,4 +1,4 @@
 rust 1.84.1
 elixir 1.18.2-otp-27
-nodejs 18.13.0
-erlang 27.2
+nodejs 19.9.0
+erlang 27.2

VERSION

@@ -1 +1 @@
-2.5.6
+2.5.7

lib/supavisor_web/ws_proxy.ex

@@ -46,6 +46,10 @@ defmodule SupavisorWeb.WsProxy do
     {[{:binary, bin}], state}
   end
 
+  def websocket_info({:tcp_closed, socket}, %{socket: socket}) do
+    {:stop, :normal}
+  end
+
   def websocket_info(msg, state) do
     Logger.error("Undefined websocket_info msg: #{inspect(msg, pretty: true)}")
     {:ok, state}

test/integration/external_test.exs

@@ -25,6 +25,20 @@ defmodule Supavisor.Integration.ExternalTest do
     end
   end
 
+  describe "neondatabase/serverless.js" do
+    @describetag library: "serverless.js", suite: "js"
+
+    @tag runtime: "node", mode: "session"
+    test "Node session", ctx do
+      assert_run(ctx, ~w[neon_serverless/index.js])
+    end
+
+    @tag runtime: "node", mode: "transaction"
+    test "Node transaction", ctx do
+      assert_run(ctx, ~w[neon_serverless/index.js])
+    end
+  end
+
   describe "Postgres.js" do
     @describetag library: "postgres.js", suite: "js"
 
@@ -65,9 +79,10 @@ defmodule Supavisor.Integration.ExternalTest do
 
     env =
       [
+        # {"NODE_DEBUG", "*"},
         {"PGMODE", ctx.mode},
         {"PGDATABASE", ctx.db},
-        {"PGHOST", "localhost"},
+        {"PGHOST", "127.0.0.1"},
         {"PGPORT", to_string(port(ctx.mode))},
         {"PGUSER", ctx.user},
         {"PGPASS", "postgres"},

test/integration/js/neon_serverless/index.js

@@ -0,0 +1,17 @@
+import { t, nt, ot } from '../shared/test.js' // eslint-disable-line
+import { Client, neonConfig } from "@neondatabase/serverless";
+import { WebSocket } from 'ws';
+globalThis.WebSocket = WebSocket;
+
+const connectionString = `postgres://${encodeURIComponent(process.env.PGUSER)}:${encodeURIComponent(process.env.PGPASS)}@${process.env.PGHOST}:${process.env.PGPORT}/${process.env.PGDATABASE}${process.env.PGSSL ? '?ssl=true' : ''}`;
+
+neonConfig.wsProxy = (_host, _port) => `127.0.0.1:4000/v2/`
+neonConfig.useSecureWebSocket = false;
+const client = new Client(connectionString);
+
+t('Connect, run simple query, disconnect', async () => {
+  await client.connect();
+  const result = await client.query("SELECT 1 as one");
+  await client.end();
+  return [1, result.rows[0]["one"]];
+});

test/integration/js/package-lock.json

@@ -8,6 +8,8 @@
     "test:postgres": "node ./postgres/index.js"
   },
   "dependencies": {
-    "postgres": "^3.4.5"
+    "@neondatabase/serverless": "^1.0.1",
+    "postgres": "^3.4.5",
+    "ws": "^8.18.3"
   }
 }

test/integration/js/package.json

@@ -1,4 +1,4 @@
-import { t, nt, ot } from './test.js' // eslint-disable-line
+import { t, nt, ot } from '../shared/test.js' // eslint-disable-line
 import net from 'node:net'
 import fs from 'node:fs'
 import crypto from 'node:crypto'

test/integration/js/postgres/index.js

@@ -83,4 +83,3 @@ function exit() {
 
   if (!success || only || ignored) { process.exit(1) } else { process.exit(0) }
 }
-