valkey: close the socket on every fail() and mark the client disconnected before onclose runs

src/runtime/valkey_jsc/js_valkey.rs

@@ -1906,15 +1906,12 @@
     fn fail_handshake(
         this: &JSValkeyClient,
         _vm: &VirtualMachine,
         err_value: JSValue,
     ) -> JsTerminatedResult<()> {
         let _exit = this.vm().enter_event_loop_scope();
-        this.client_mut().flags.is_manually_closed = true;
-        let this_br = BackRef::new(this);
-        let _close = scopeguard::guard(this_br, |p| p.client_mut().close());
         narrow_terminated(
             this.client_mut()
                 .fail_with_js_value(&this.global_object, err_value),
         )
     }
 
@@ -1937,10 +1934,10 @@
         let _guard = this.ref_scope();
         // Ensure the socket pointer is updated.
         this.client_mut().socket = Socket::SocketTcp(uws::SocketTCP::detached());
-        let _defer = scopeguard::guard(BackRef::new(this), |p| {
-            p.client_mut().status = valkey::Status::Disconnected;
-            p.update_poll_ref();
-        });
+        // Before `on_close()`: it runs `onclose` and settles the connect()
+        // promise, and a connect() called from either must see Disconnected.
+        this.client_mut().status = valkey::Status::Disconnected;
+        let _defer = scopeguard::guard(BackRef::new(this), |p| p.update_poll_ref());
 
         let _ = this.client_mut().on_close(); // TODO: properly propagate exception upwards
     }
@@ -1962,10 +1959,8 @@
         // Ensure the socket pointer is updated.
         this.client_mut().socket = Socket::SocketTcp(uws::SocketTCP::detached());
         let _guard = this.ref_scope();
-        let _defer = scopeguard::guard(BackRef::new(this), |p| {
-            p.client_mut().status = valkey::Status::Disconnected;
-            p.update_poll_ref();
-        });
+        this.client_mut().status = valkey::Status::Disconnected;
+        let _defer = scopeguard::guard(BackRef::new(this), |p| p.update_poll_ref());
 
         narrow_terminated(this.client_mut().on_close())
     }

src/runtime/valkey_jsc/valkey.rs

@@ -33,12 +33,12 @@
     pub(crate) is_selecting_db_internal: bool,
     pub(crate) enable_offline_queue: bool,
     pub(crate) enable_auto_reconnect: bool,
-    /// Sticky until the next accepted HELLO, so it overlaps `Connecting`
-    /// (`reconnect()` reads it there) and `failed` (`update_poll_ref` reads it
-    /// there); that is why it is not a `Status` variant.
+    /// Set from the close that schedules a retry until the next accepted HELLO
+    /// or `fail()`, so it overlaps `Disconnected` and `Connecting`; that is why
+    /// it is not a `Status` variant.
     pub(crate) is_reconnecting: bool,
-    /// Sticky until `on_open`/`connect()`, and orthogonal to `Status`: `fail()`
-    /// while `Connected` leaves the socket open and `status` unchanged.
+    /// Sticky until `on_open`/`connect()`; the socket is closed when it is set,
+    /// so it overlaps `Disconnected`.
     pub(crate) failed: bool,
     pub(crate) enable_auto_pipelining: bool,
     pub(crate) finalized: bool,
@@ -606,18 +606,17 @@
         if self.flags.failed {
             return Ok(());
         }
         self.flags.failed = true;
+        self.flags.is_reconnecting = false;
         let val = Self::reject_all_pending_commands(
             &mut self.in_flight,
             &mut self.queue,
             global_this,
             jsvalue,
         );
 
-        if !self.connection_ready() {
-            self.flags.is_manually_closed = true;
-            self.close();
-        }
+        self.flags.is_manually_closed = true;
+        self.close();
         val
     }
 
@@ -652,6 +651,10 @@
     pub fn on_close(&mut self) -> JsTerminated<()> {
         self.unregister_auto_flusher();
         self.write_buffer.clear_and_free();
+        // A partial reply can never complete now; left in place it counts as
+        // pending activity in `update_poll_ref` and keeps the event loop alive.
+        self.read_buffer.clear_and_free();
+        self.reply_scanner.reset();
 
         // If manually closing, don't attempt to reconnect
         if self.flags.is_manually_closed {
@@ -743,6 +746,10 @@
             data.len(),
             bstr::BStr::new(data)
         );
+        // Handling a reply can close this socket and, from `onclose` or a
+        // rejection handler, dial the next one; the remaining replies came from
+        // the closed connection and must not reach the new one.
+        let socket = *self.socket.socket();
         // Path 1: Buffer already has data, append and process from buffer
         if !self.read_buffer.remaining().is_empty() {
             self.read_buffer
@@ -810,7 +817,7 @@
                 let mut value_to_handle = value; // Use temp var for defer
                 self.handle_response(&mut value_to_handle)?;
 
-                if self.status == Status::Disconnected || self.flags.failed {
+                if *self.socket.socket() != socket {
                     return Ok(());
                 }
                 self.send_next_command();
@@ -869,8 +876,7 @@
             let mut value_to_handle = value; // Use temp var for defer
             self.handle_response(&mut value_to_handle)?;
 
-            // Check connection status after handling
-            if self.status == Status::Disconnected || self.flags.failed {
+            if *self.socket.socket() != socket {
                 return Ok(());
             }
 

test/js/valkey/reliability/connection-failures.test.ts

@@ -1,5 +1,6 @@
 import { RedisClient } from "bun";
 import { describe, expect, mock, test } from "bun:test";
+import { bunEnv, bunExe } from "harness";
 import net from "net";
 import { DEFAULT_REDIS_OPTIONS, DEFAULT_REDIS_URL, delay, isEnabled } from "../test-utils";
 
@@ -441,3 +442,193 @@ describe("Valkey: Auto-Reconnect In-Flight Commands", () => {
     }
   });
 });
+
+describe("Valkey: Recovering After fail()", () => {
+  // Answers the chunk carrying HELLO with `+OK` and the one carrying PING with
+  // `+PONG` unless `replies` says otherwise for that connection.
+  function helloServer(replies: Partial<Record<"HELLO" | "PING", (connection: number) => string>> = {}) {
+    let connections = 0;
+    const server = net.createServer(socket => {
+      connections += 1;
+      const connection = connections;
+      socket.on("data", chunk => {
+        const text = chunk.toString("latin1");
+        if (text.includes("HELLO")) socket.write(replies.HELLO?.(connection) ?? "+OK\r\n");
+        if (text.includes("PING")) socket.write(replies.PING?.(connection) ?? "+PONG\r\n");
+      });
+      socket.on("error", () => {});
+    });
+    return {
+      server,
+      get connections() {
+        return connections;
+      },
+      listen: () =>
+        new Promise<number>(resolve =>
+          server.listen(0, "127.0.0.1", () => resolve((server.address() as net.AddressInfo).port)),
+        ),
+    };
+  }
+
+  // Calls connect() from the first onclose and reports how that attempt ended.
+  function connectFromOnclose(client: RedisClient): Promise<string> {
+    const { promise, resolve } = Promise.withResolvers<string>();
+    client.onclose = () => {
+      client.onclose = () => {};
+      resolve(
+        client.connect().then(
+          () => "connected",
+          (err: Error) => `rejected: ${err.message}`,
+        ),
+      );
+    };
+    return promise;
+  }
+
+  test("a failure while connected closes the socket, fires onclose, and connect() reconnects", async () => {
+    // 0x01 is not a RESP type byte, so the first connection fails after the
+    // handshake, on the same path as an idle timeout or any other protocol error.
+    const fake = helloServer({ PING: connection => (connection === 1 ? "\x01\r\n" : "+PONG\r\n") });
+    const port = await fake.listen();
+    const closed = Promise.withResolvers<Error>();
+    const client = new RedisClient(`redis://127.0.0.1:${port}`, { autoReconnect: false });
+    try {
+      client.onclose = err => closed.resolve(err);
+      await client.connect();
+      expect(client.connected).toBe(true);
+      await expect(client.ping()).rejects.toMatchObject({ code: "ERR_REDIS_INVALID_RESPONSE_TYPE" });
+      expect(await closed.promise).toBeInstanceOf(Error);
+      expect(client.connected).toBe(false);
+      await client.connect();
+      expect(await client.ping()).toBe("PONG");
+      expect(fake.connections).toBe(2);
+    } finally {
+      client.close();
+      fake.server.close();
+    }
+  });
+
+  test("a connect() issued from onclose after a refused connection rejects instead of hanging", async () => {
+    // Nothing listens on the port a just-closed listener used.
+    const fake = helloServer();
+    const port = await fake.listen();
+    await new Promise(resolve => fake.server.close(resolve));
+    const client = new RedisClient(`redis://127.0.0.1:${port}`, { autoReconnect: false });
+    try {
+      const secondConnect = connectFromOnclose(client);
+      await expect(client.connect()).rejects.toMatchObject({ code: "ERR_REDIS_CONNECTION_CLOSED" });
+      expect(await secondConnect).toBe("rejected: Connection closed");
+    } finally {
+      client.close();
+    }
+  });
+
+  test("a connect() issued from onclose is not fed the replies left over from the failed connection", async () => {
+    // With a database in the URL, HELLO and SELECT are written together, so a
+    // server that rejects HELLO delivers both error replies in one read.
+    const fake = helloServer({
+      HELLO: connection =>
+        connection === 1 ? "-WRONGPASS invalid password\r\n-NOAUTH Authentication required.\r\n" : "+OK\r\n+OK\r\n",
+    });
+    const port = await fake.listen();
+    const client = new RedisClient(`redis://127.0.0.1:${port}/1`, { autoReconnect: false });
+    try {
+      const secondConnect = connectFromOnclose(client);
+      await expect(client.connect()).rejects.toMatchObject({ code: "ERR_REDIS_CONNECTION_CLOSED" });
+      expect(await secondConnect).toBe("connected");
+      expect(await client.ping()).toBe("PONG");
+      expect(fake.connections).toBe(2);
+    } finally {
+      client.close();
+      fake.server.close();
+    }
+  });
+
+  test("a connect() issued from onclose after a failed TLS handshake gets to dial again", async () => {
+    let handshakes = 0;
+    const server = net.createServer(socket => {
+      // The first bytes are the ClientHello; dropping the connection there
+      // fails the client's handshake.
+      socket.once("data", () => {
+        handshakes += 1;
+        socket.destroy();
+      });
+      socket.on("error", () => {});
+    });
+    await new Promise<void>(resolve => server.listen(0, "127.0.0.1", resolve));
+    const { port } = server.address() as net.AddressInfo;
+    const client = new RedisClient(`rediss://127.0.0.1:${port}`, { autoReconnect: false });
+    try {
+      const secondConnect = connectFromOnclose(client);
+      await expect(client.connect()).rejects.toMatchObject({ code: "ERR_REDIS_CONNECTION_CLOSED" });
+      expect({ secondConnect: await secondConnect, handshakes }).toEqual({
+        secondConnect: "rejected: Connection closed",
+        handshakes: 2,
+      });
+    } finally {
+      client.close();
+      server.close();
+    }
+  });
+
+  test("close() discards a half-received reply instead of letting it keep the process alive", async () => {
+    // Announces a 4 byte bulk string and stops halfway through it.
+    const fake = helloServer({ PING: () => "$4\r\nPO" });
+    const port = await fake.listen();
+    try {
+      await using proc = Bun.spawn({
+        cmd: [
+          bunExe(),
+          "-e",
+          `
+          const client = new Bun.RedisClient("redis://127.0.0.1:${port}", { autoReconnect: false });
+          await client.connect();
+          const ping = client.ping().catch(err => console.log("ping rejected", err.code));
+          while (client.bufferedAmount === 0) await Bun.sleep(1);
+          console.log("buffered before close", client.bufferedAmount);
+          client.close();
+          await ping;
+          console.log("buffered after close", client.bufferedAmount);
+          `,
+        ],
+        env: bunEnv,
+        stdout: "pipe",
+        stderr: "pipe",
+      });
+      const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+      expect({ stdout, stderr, exitCode }).toEqual({
+        stdout: "buffered before close 6\nping rejected ERR_REDIS_CONNECTION_CLOSED\nbuffered after close 0\n",
+        stderr: "",
+        exitCode: 0,
+      });
+    } finally {
+      fake.server.close();
+    }
+  });
+
+  test("the process exits once auto-reconnect gives up", async () => {
+    const fake = helloServer();
+    const port = await fake.listen();
+    await new Promise(resolve => fake.server.close(resolve));
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `
+        const client = new Bun.RedisClient("redis://127.0.0.1:${port}", { autoReconnect: true, maxRetries: 1 });
+        client.onclose = err => console.log("onclose", err.code);
+        await client.connect().catch(err => console.log("connect rejected", err.code));
+        `,
+      ],
+      env: bunEnv,
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    expect({ stdout, stderr, exitCode }).toEqual({
+      stdout: "onclose ERR_REDIS_CONNECTION_CLOSED\nconnect rejected ERR_REDIS_CONNECTION_CLOSED\n",
+      stderr: "",
+      exitCode: 0,
+    });
+  });
+});