diff --git a/src/runtime/valkey_jsc/js_valkey.rs b/src/runtime/valkey_jsc/js_valkey.rs index 25cbd6ac9004..31966df3138a 100644 --- a/src/runtime/valkey_jsc/js_valkey.rs +++ b/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 diff --git a/src/runtime/valkey_jsc/valkey.rs b/src/runtime/valkey_jsc/valkey.rs index 6b443bdf1d09..58ac6b7e7025 100644 --- a/src/runtime/valkey_jsc/valkey.rs +++ b/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, 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 = 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 diff --git a/test/js/valkey/reliability/connection-failures.test.ts b/test/js/valkey/reliability/connection-failures.test.ts index c9f70240d8a8..6995eb2f19f0 100644 --- a/test/js/valkey/reliability/connection-failures.test.ts +++ b/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 | 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[] = []; + 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(); + 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 | 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