valkey: finish the handshake's SELECT before reporting the client connected

src/runtime/valkey_jsc/js_valkey.rs

@@ -719,6 +719,7 @@ impl JSValkeyClient {
                     valkey::TLS::None
                 },
                 database,
+                pending_hello: None,
                 flags: valkey::ConnectionFlags {
                     enable_auto_reconnect: options.enable_auto_reconnect,
                     enable_offline_queue: options.enable_offline_queue,
@@ -822,6 +823,7 @@ impl JSValkeyClient {
                 }),
                 tls,
                 database: client.database,
+                pending_hello: None,
                 flags: valkey::ConnectionFlags {
                     enable_offline_queue: if sub_ctx.is_subscriber {
                         sub_ctx.original_enable_offline_queue

src/runtime/valkey_jsc/valkey.rs

@@ -30,6 +30,8 @@ bun_output::define_scoped_log!(debug, Redis, visible);
 /// Connection flags to track Valkey client state
 pub struct ConnectionFlags {
     pub(crate) is_manually_closed: bool,
+    /// The handshake wrote a `SELECT` behind `HELLO` and its reply is still
+    /// due. The client stays `Connecting` until it arrives.
     pub(crate) is_selecting_db_internal: bool,
     pub(crate) enable_offline_queue: bool,
     pub(crate) enable_auto_reconnect: bool,
@@ -81,7 +83,8 @@ pub enum Status {
     NeverConnected,
     Disconnected,
     Connecting,
-    /// Socket open and HELLO accepted.
+    /// Socket open and handshake accepted: HELLO, then SELECT when the URL
+    /// names a database.
     Connected,
 }
 
@@ -259,6 +262,10 @@ pub struct ValkeyClient {
     pub(crate) password: Box<[u8]>,
     pub(crate) username: Box<[u8]>,
     pub(crate) database: u32,
+    /// The accepted HELLO reply, held while the handshake's SELECT reply is
+    /// due so that `on_valkey_connect` still receives it. While it is held,
+    /// the next reply is the SELECT's.
+    pub(crate) pending_hello: Option<RESPValue>,
     pub(crate) address: Address,
     pub(crate) protocol: Protocol,
 
@@ -665,6 +672,8 @@ impl ValkeyClient {
         // pending activity in `update_poll_ref` and keeps the event loop alive.
         self.read_buffer.clear_and_free();
         self.reply_scanner.reset();
+        self.flags.is_selecting_db_internal = false;
+        self.pending_hello = None;
 
         // A manual close or a failure the client detected itself: no retry.
         if self.flags.is_manually_closed || self.flags.failed {
@@ -702,7 +711,6 @@ impl ValkeyClient {
         );
 
         self.flags.is_reconnecting = true;
-        self.flags.is_selecting_db_internal = false;
 
         self.reject_in_flight_commands(b"Connection closed", RedisError::ConnectionClosed)?;
 
@@ -1003,11 +1011,7 @@ impl ValkeyClient {
             }
             RESPValue::SimpleString(str_) => {
                 if str_.as_ref() == b"OK" {
-                    self.status = Status::Connected;
-                    self.flags.is_reconnecting = false;
-                    self.retry_attempts = 0;
-                    self.on_valkey_connect(value)?;
-                    return Ok(());
+                    return self.hello_accepted(value);
                 }
                 self.fail(
                     b"Authentication failed (unexpected response)",
@@ -1042,11 +1046,7 @@ impl ValkeyClient {
                 }
 
                 // Authentication successful via HELLO
-                self.status = Status::Connected;
-                self.flags.is_reconnecting = false;
-                self.retry_attempts = 0;
-                self.on_valkey_connect(value)?;
-                Ok(())
+                self.hello_accepted(value)
             }
             _ => {
                 self.fail(
@@ -1058,52 +1058,56 @@ impl ValkeyClient {
         }
     }
 
-    /// Handle Valkey protocol response
-    fn handle_response(&mut self, value: &mut RESPValue) -> JsResult<()> {
-        // Special handling for the initial HELLO response
-        if self.status != Status::Connected {
-            self.handle_hello_response(value)?;
-
-            // We've handled the HELLO response without consuming anything from the command queue
+    /// When `authenticate()` wrote a `SELECT` behind `HELLO`, the handshake has
+    /// one more reply to go.
+    fn hello_accepted(&mut self, value: &mut RESPValue) -> JsResult<()> {
+        if self.flags.is_selecting_db_internal {
+            self.pending_hello = Some(core::mem::replace(value, RESPValue::Null));
             return Ok(());
         }
+        self.handshake_complete(value)
+    }
 
-        // Handle initial SELECT response
-        if self.flags.is_selecting_db_internal {
-            self.flags.is_selecting_db_internal = false;
+    fn handle_select_response(
+        &mut self,
+        mut hello: RESPValue,
+        value: &mut RESPValue,
+    ) -> JsResult<()> {
+        self.flags.is_selecting_db_internal = false;
+        match value {
+            RESPValue::Error(err) => self.fail(err, RedisError::ServerError),
+            RESPValue::SimpleString(ok) if ok.as_ref() == b"OK" => {
+                debug!("SELECT {} successful", self.database);
+                self.handshake_complete(&mut hello)
+            }
+            _ => self.fail(
+                b"SELECT command failed with non-OK response",
+                RedisError::InvalidResponse,
+            ),
+        }
+    }
 
-            return match value {
-                RESPValue::Error(err_str) => {
-                    self.fail(err_str, RedisError::InvalidCommand)?;
-                    Ok(())
-                }
-                RESPValue::SimpleString(ok_str) => {
-                    if ok_str.as_ref() != b"OK" {
-                        // SELECT returned something other than "OK"
-                        self.fail(
-                            b"SELECT command failed with non-OK response",
-                            RedisError::InvalidResponse,
-                        )?;
-                        return Ok(());
-                    }
+    /// The one place a connection becomes `Connected`: `onconnect`, the
+    /// `connect()` promise, the queue drain and the auto-flusher all key off it,
+    /// so nothing is written ahead of a handshake reply that is still due.
+    fn handshake_complete(&mut self, hello: &mut RESPValue) -> JsResult<()> {
+        self.status = Status::Connected;
+        self.flags.is_reconnecting = false;
+        self.retry_attempts = 0;
+        self.on_valkey_connect(hello)
+    }
 
-                    // SELECT was successful.
-                    debug!("SELECT {} successful", self.database);
-                    // Connection is now fully ready on the specified database.
-                    // If any commands were queued while waiting for SELECT, try to send them.
-                    self.send_next_command();
-                    Ok(())
-                }
-                _ => {
-                    // Unexpected response type for SELECT
-                    self.fail(
-                        b"Received non-SELECT response while in the SELECT state.",
-                        RedisError::InvalidResponse,
-                    )?;
-                    Ok(())
-                }
-            };
+    /// Handle Valkey protocol response
+    fn handle_response(&mut self, value: &mut RESPValue) -> JsResult<()> {
+        // Replies to the handshake are not paired with anything in the command queue.
+        if self.status != Status::Connected {
+            match self.pending_hello.take() {
+                Some(hello) => self.handle_select_response(hello, value)?,
+                None => self.handle_hello_response(value)?,
+            }
+            return Ok(());
         }
+
         // Check if this is a subscription push message that might not need a promise pair
         let mut should_consume_promise_pair = true;
         let mut pair_maybe: Option<command::PromisePair> = None;
@@ -1269,7 +1273,8 @@ impl ValkeyClient {
             return Ok(());
         }
 
-        // If using a specific database, send SELECT command
+        // If using a specific database, send SELECT right behind HELLO; its
+        // reply completes the handshake (`hello_accepted`).
         if self.database > 0 {
             let mut int_buf = [0u8; 64];
             let db_str = bun_core::fmt::int_as_bytes(&mut int_buf, self.database);
@@ -1298,6 +1303,7 @@ impl ValkeyClient {
         // after a previous connection exhausted retries (#29925).
         self.flags.failed = false;
         self.flags.is_selecting_db_internal = false;
+        self.pending_hello = None;
         if matches!(self.socket, AnySocket::SocketTcp(_)) {
             // if is tcp, we need to start the connection process
             // if is tls, we need to wait for the handshake to complete
@@ -1316,7 +1322,7 @@ impl ValkeyClient {
     /// Test whether we are ready to run "normal" RESP commands, such as
     /// get/set, pub/sub, etc.
     fn connection_ready(&self) -> bool {
-        self.status == Status::Connected && !self.flags.is_selecting_db_internal
+        self.status == Status::Connected
     }
 
     /// Process queued commands in the offline queue

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

@@ -879,7 +879,7 @@ describe("Valkey: Recovering After fail()", () => {
     }
   });
 
-  test("a rejected SELECT after an accepted HELLO fails the connection once and connect() from onclose dials again", async () => {
+  test("a rejected SELECT after an accepted HELLO fails the handshake once and connect() from onclose dials again", async () => {
     // HELLO and SELECT are written together, and both replies come back in one
     // read: connection 1 accepts HELLO and rejects SELECT, connection 2 accepts both.
     const fake = helloServer({
@@ -888,10 +888,25 @@ describe("Valkey: Recovering After fail()", () => {
     const port = await fake.listen();
     const client = new RedisClient(`redis://127.0.0.1:${port}/1`, { autoReconnect: true });
     try {
-      const closes: { message: string; connected: boolean; connections: number }[] = [];
+      const closes: { message: string; connected: boolean; connections: number; connects: number }[] = [];
       let secondConnect: Promise<string> | undefined;
+      // The handshake is HELLO and then SELECT, so nothing may report the
+      // client connected in between: a command sent from there would be
+      // pipelined behind a SELECT the server is about to reject and run
+      // against database 0.
+      let connects = 0;
+      const sentFromOnconnect: Promise<string>[] = [];
+      client.onconnect = () => {
+        connects++;
+        sentFromOnconnect.push(
+          client.ping().then(
+            () => "resolved",
+            (err: Error & { code: string }) => `${err.code}: ${err.message}`,
+          ),
+        );
+      };
       client.onclose = err => {
-        closes.push({ message: err.message, connected: client.connected, connections: fake.connections });
+        closes.push({ message: err.message, connected: client.connected, connections: fake.connections, connects });
         secondConnect ??= client.connect().then(
           () => "connected",
           (err: Error) => `rejected: ${err.message}`,
@@ -902,15 +917,23 @@ describe("Valkey: Recovering After fail()", () => {
         () => "resolved",
         (err: Error & { code: string }) => `${err.code}: ${err.message}`,
       );
-      // connect() settles on the accepted HELLO, before the SELECT reply is
-      // read, so it resolves; the failure that follows is reported by onclose.
-      await client.connect();
-      expect(await queued).toBe("ERR_REDIS_INVALID_COMMAND: ERR DB index is out of range");
+      const firstConnect = await client.connect().then(
+        () => "resolved",
+        (err: Error & { code: string }) => `rejected: ${err.code}`,
+      );
       // The rejection is a failure the client detected, so there is no retry
       // even with autoReconnect on: onclose fires once and the only second
-      // connection is the one dialed from it.
-      expect(closes).toEqual([{ message: "Connection closed", connected: false, connections: 1 }]);
+      // connection is the one dialed from it. onconnect first runs for that one.
+      expect({ firstConnect, queued: await queued, closes }).toEqual({
+        firstConnect: "rejected: ERR_REDIS_CONNECTION_CLOSED",
+        queued: "ERR_REDIS_SERVER_ERROR: ERR DB index is out of range",
+        closes: [{ message: "Connection closed", connected: false, connections: 1, connects: 0 }],
+      });
       expect(await secondConnect).toBe("connected");
+      expect({ connects, sentFromOnconnect: await Promise.all(sentFromOnconnect) }).toEqual({
+        connects: 1,
+        sentFromOnconnect: ["resolved"],
+      });
       expect(await client.ping()).toBe("PONG");
       expect(fake.connections).toBe(2);
     } finally {
@@ -2145,6 +2168,118 @@ describe("Valkey: Recovering After fail()", () => {
   });
 });
 
+// With a database in the URL the handshake is HELLO and then SELECT. The
+// client is connected once both are accepted, not in between.
+describe("Valkey: Handshake SELECT", () => {
+  // Answers HELLO (with a RESP3 map), PING and SET as they arrive, but holds
+  // the reply to SELECT until `answerSelect()` is called. `events` records, in
+  // order, when SELECT was answered and whether a SET got in ahead of that.
+  function selectServer() {
+    const events: string[] = [];
+    const selectArrived = Promise.withResolvers<net.Socket>();
+    let selectAnswered = false;
+    const server = net.createServer(socket => {
+      const state = { buffer: Buffer.alloc(0) };
+      socket.on("data", chunk => {
+        state.buffer = Buffer.concat([state.buffer, chunk]);
+        for (const [name] of readCommands(state)) {
+          switch (name.toUpperCase()) {
+            case "HELLO":
+              socket.write("%2\r\n+server\r\n+redis\r\n+proto\r\n:3\r\n");
+              break;
+            case "SELECT":
+              selectArrived.resolve(socket);
+              break;
+            case "SET":
+              events.push(`SET arrived ${selectAnswered ? "after" : "before"} SELECT was answered`);
+              socket.write("+OK\r\n");
+              break;
+            default:
+              socket.write("+PONG\r\n");
+          }
+        }
+      });
+      socket.on("error", () => {});
+    });
+    return {
+      server,
+      events,
+      selectArrived: selectArrived.promise,
+      answerSelect: async () => {
+        const socket = await selectArrived.promise;
+        selectAnswered = true;
+        events.push("SELECT answered");
+        socket.write("+OK\r\n");
+      },
+      listen: async () => {
+        server.listen(0, "127.0.0.1");
+        await once(server, "listening");
+        return (server.address() as net.AddressInfo).port;
+      },
+    };
+  }
+
+  test("connected, onconnect and connect() wait for the SELECT reply, and so does the first command", async () => {
+    const fake = selectServer();
+    const port = await fake.listen();
+    // No database: this one's handshake is HELLO alone. Its round trips are
+    // what turns the event loop below.
+    const control = new RedisClient(`redis://127.0.0.1:${port}`, { autoReconnect: false });
+    const client = new RedisClient(`redis://127.0.0.1:${port}/2`, { autoReconnect: false });
+    try {
+      await control.connect();
+      let set: Promise<unknown> | undefined;
+      client.onconnect = () => {
+        fake.events.push("onconnect");
+        set = client.set("key", "value");
+      };
+      const connected = client.connect().then(() => fake.events.push("connect() resolved"));
+      await fake.selectArrived;
+      // The stub answered HELLO when it arrived, ahead of SELECT, and now sits
+      // on SELECT. Each round trip on the control connection takes the event
+      // loop through at least two polls, so after these the client has long
+      // read that HELLO reply and done whatever it does on it.
+      for (let i = 0; i < 5; i++) await control.ping();
+      const connectedBeforeSelect = client.connected;
+      await fake.answerSelect();
+      await connected;
+      await set;
+      expect({ connectedBeforeSelect, events: fake.events }).toEqual({
+        connectedBeforeSelect: false,
+        events: ["SELECT answered", "onconnect", "connect() resolved", "SET arrived after SELECT was answered"],
+      });
+    } finally {
+      client.close();
+      control.close();
+      fake.server.close();
+    }
+  });
+
+  test("a SELECT reply that never comes is bounded by connectionTimeout", async () => {
+    const fake = selectServer();
+    const port = await fake.listen();
+    const client = new RedisClient(`redis://127.0.0.1:${port}/2`, { autoReconnect: false, connectionTimeout: 200 });
+    try {
+      const queued = client.get("key").then(
+        () => "resolved",
+        (err: Error & { code: string }) => err.code,
+      );
+      const connect = await client.connect().then(
+        () => "resolved",
+        (err: Error & { code: string }) => `rejected: ${err.code}`,
+      );
+      expect({ connect, connected: client.connected }).toEqual({
+        connect: "rejected: ERR_REDIS_CONNECTION_CLOSED",
+        connected: false,
+      });
+      expect(await queued).toBe("ERR_REDIS_CONNECTION_TIMEOUT");
+    } finally {
+      client.close();
+      fake.server.close();
+    }
+  });
+});
+
 describe("Valkey: Offline Queue", () => {
   // Answers each complete command in the order it arrives: HELLO with `+OK`
   // and everything else with `+PONG`. With `answer: false` nothing is ever