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

src/runtime/valkey_jsc/js_valkey.rs

@@ -1142,6 +1142,37 @@ impl JSValkeyClient {
         }
     }
 
+    /// Runs `ValkeyClient::on_close()` for a dial that failed before there was
+    /// a socket, so the connect() promise, `onclose`, the retry policy and the
+    /// poll ref are handled as for a dial that failed asynchronously. Deferred
+    /// because `onclose` may call connect(), and a dial that fails the same way
+    /// from in there would otherwise re-enter `on_close()` on the same stack.
+    ///
+    /// Until the task runs the client is `Connecting`, as it would be with a
+    /// dial in flight: JS that runs in between (timers due in the same tick,
+    /// or the caller of connect() itself) then gets the cached promise from
+    /// connect() instead of a second dial, and a disconnect() marks the close
+    /// as manual for the task to honour. `update_poll_ref` keeps the wrapper
+    /// and the event loop alive for it like a dial would.
+    fn close_without_socket_next_tick(&self) {
+        self.client_mut().status = valkey::Status::Connecting;
+        self.update_poll_ref();
+        self.enqueue_deferred_close(DeferredClose::WithoutSocket);
+    }
+
+    fn enqueue_deferred_close(&self, what: DeferredClose) {
+        // Released by the task, whether it runs or the VM tears down first.
+        self.ref_();
+        let task = jsc::Task::from_boxed(Box::new(ValkeyDeferredClose {
+            ctx: self.as_ctx_ptr(),
+            what,
+        }));
+        // SAFETY: VM-owned event loop pointer; uniquely accessed on the JS thread.
+        unsafe {
+            (*self.vm().event_loop()).enqueue_task(task);
+        }
+    }
+
     pub(crate) fn on_reconnect_timer(&self) -> JsResult<()> {
         debug!("Reconnect timer fired, attempting to reconnect");
 
@@ -1175,15 +1206,14 @@ impl JSValkeyClient {
         });
 
         if let Err(err) = self.connect() {
-            self.poll_ref.with_mut(|r| r.disable());
-            return self.fail_with_js_value(
-                self.global_object
-                    .err(
-                        jsc::ErrorCode::SOCKET_CLOSED_BEFORE_CONNECTION,
-                        format_args!("{} reconnecting", err.name()),
-                    )
-                    .to_js(),
+            debug!(
+                "reconnect failed before a socket was opened: {}",
+                err.name()
             );
+            // Same outcome as a dial that fails asynchronously: another retry,
+            // or fail() and a settled connect() promise once retries are used up.
+            self.close_without_socket_next_tick();
+            return Ok(());
         }
 
         // Reset the socket timeout
@@ -1221,7 +1251,6 @@ impl JSValkeyClient {
                     let error = global_object.take_exception(err);
                     let client = self.client_mut();
                     client.flags.connection_promise_returns_client = false;
-                    client.flags.is_manually_closed = true;
                     let rejected = match Js::connection_promise_get_cached(this_value) {
                         Some(promise) => {
                             Js::connection_promise_set_cached(
@@ -1234,10 +1263,16 @@ impl JSValkeyClient {
                         }
                         None => Ok(()),
                     };
-                    let failed =
-                        rejected.and_then(|()| client.fail_with_js_value(&global_object, error));
-                    let closed = self.client_mut().close();
-                    return failed.and(closed);
+                    // `fail_with_js_value` closes the socket itself; a second close here would
+                    // hit whatever a connect() from `onclose` just opened.
+                    return match rejected {
+                        Ok(()) => client.fail_with_js_value(&global_object, error),
+                        Err(_) => {
+                            client.flags.is_manually_closed = true;
+                            let closed = client.close(uws::CloseCode::Failure);
+                            rejected.and(closed)
+                        }
+                    };
                 }
             };
             Js::hello_set_cached(this_value, &global_object, hello_value);
@@ -1343,11 +1378,16 @@ impl JSValkeyClient {
     // Callback for when Valkey client needs to reconnect
     pub(crate) fn on_valkey_reconnect(&self) {
         // SAFETY: adopts connect()'s socket keep-alive ref for the just-closed
-        // socket. Reached only from `ValkeyClient::on_close()`'s reconnect
-        // branch, which never calls `on_valkey_close()`, so this scope is the
-        // sole releaser. The caller holds its own scoped ref, so count > 0.
+        // socket (or the one `ValkeyDeferredClose::run` took in its place).
+        // Reached only from `ValkeyClient::on_close()`'s reconnect branch,
+        // which never calls `on_valkey_close()`, so this scope is the sole
+        // releaser. The caller holds its own scoped ref, so count > 0.
         let _socket_ref = unsafe { ScopedRef::adopt(self.as_ctx_ptr()) };
 
+        // This timer was bounding the attempt that just ended; left armed it
+        // fires during the retry delay, and `fail()` then has no socket to
+        // close and nothing settles connect(). `reconnect()` arms a new one.
+        self.timer.disarm(self);
         self.reconnect_timer
             .arm(self, self.client.get().get_reconnect_delay());
     }
@@ -1356,8 +1396,9 @@ impl JSValkeyClient {
     pub(crate) fn on_valkey_close(&self) -> JsResult<()> {
         let global_object = self.global_object;
 
-        // SAFETY: adopts connect()'s socket keep-alive ref; the caller holds
-        // its own scoped ref so count stays > 0 until this drops.
+        // SAFETY: adopts connect()'s socket keep-alive ref (or the one
+        // `ValkeyDeferredClose::run` took in its place); the caller holds its
+        // own scoped ref so count stays > 0 until this drops.
         let _socket_ref = unsafe { ScopedRef::adopt(self.as_ctx_ptr()) };
         let _defer = scopeguard::guard(BackRef::new(self), |p| p.update_poll_ref());
 
@@ -1398,32 +1439,13 @@ impl JSValkeyClient {
         self.client_mut().fail(message, err)
     }
 
-    pub(crate) fn fail_with_js_value(&self, value: JSValue) -> JsResult<()> {
-        let Some(this_value) = self.this_value.get().try_get() else {
-            return Ok(());
-        };
-        let global_object = self.global_object;
-        if let Some(on_close) = Js::onclose_get_cached(this_value) {
-            let _exit = self.vm().enter_event_loop_scope();
-            on_close.call(&global_object, this_value, &[value])?;
-        }
-        Ok(())
-    }
-
     fn close_socket_next_tick(&self) {
         if self.client.get().socket.is_closed() {
             return;
         }
 
-        self.ref_();
         // socket close can potentially call JS so we need to enqueue the deinit
-        let task = jsc::Task::from_boxed(Box::new(ValkeyDeferredClose {
-            ctx: self.as_ctx_ptr(),
-        }));
-        // SAFETY: VM-owned event loop pointer; uniquely accessed on the JS thread.
-        unsafe {
-            (*self.vm().event_loop()).enqueue_task(task);
-        }
+        self.enqueue_deferred_close(DeferredClose::Socket);
     }
 
     pub fn finalize(self: Box<Self>) {
@@ -1457,12 +1479,6 @@ impl JSValkeyClient {
 
         let _guard = self.ref_scope();
 
-        // Socket keep-alive ref, released by on_valkey_close/on_valkey_reconnect.
-        // Taken before the TLS-context check so the `tls_ctx_failed` branch's
-        // `on_valkey_close()` has a ref to consume instead of over-releasing.
-        // Forgotten on success (the socket adopts it).
-        let socket_ref = self.ref_scope();
-
         let is_tls = self.client.get().tls != valkey::TLS::None;
         let vm = self.client.get().vm.as_mut();
         let loop_ = vm.uws_loop();
@@ -1499,11 +1515,7 @@ impl JSValkeyClient {
                 b"Failed to create TLS context",
                 protocol::RedisError::ConnectionClosed,
             )?;
-            // `on_valkey_close()` consumes the socket ref; hand it over so it
-            // isn't released twice.
-            socket_ref.forget();
-            self.client_mut().on_valkey_close()?;
-            self.client_mut().status = valkey::Status::Disconnected;
+            self.close_without_socket_next_tick();
             return Ok(());
         }
         let ssl_ctx: Option<*mut uws::SslCtx> = match &self.client.get().tls {
@@ -1531,6 +1543,9 @@ impl JSValkeyClient {
         // `owner_ptr` opaquely (no overlapping write).
         let owner_ptr: *mut JSValkeyClient = std::ptr::from_ref::<JSValkeyClient>(self).cast_mut();
         let client_ptr: *mut valkey::ValkeyClient = self.client.as_ptr();
+        // Socket keep-alive ref, released by on_valkey_close/on_valkey_reconnect.
+        // Forgotten once there is a socket to own it.
+        let socket_ref = self.ref_scope();
         // SAFETY: `client_ptr` is live; `group` is the lazy-initialised per-VM
         // `SocketGroup` (stable for the VM's lifetime). `ssl_ctx` is a +1-ref
         // BoringSSL `SSL_CTX*` (or None) forwarded opaquely to usockets.
@@ -1844,12 +1859,8 @@ impl<const SSL: bool> SocketHandler<SSL> {
         err_value: JSValue,
     ) -> JsResult<()> {
         let _exit = this.vm().enter_event_loop_scope();
-        this.client_mut().flags.is_manually_closed = true;
-        let failed = this
-            .client_mut()
-            .fail_with_js_value(&this.global_object, err_value);
-        let closed = this.client_mut().close();
-        failed.and(closed)
+        this.client_mut()
+            .fail_with_js_value(&this.global_object, err_value)
     }
 
     pub(crate) const ON_HANDSHAKE: Option<
@@ -1866,10 +1877,10 @@ impl<const SSL: bool> SocketHandler<SSL> {
         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());
 
         this.client_mut().on_close()
     }
@@ -1891,10 +1902,8 @@ impl<const SSL: bool> SocketHandler<SSL> {
         // 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());
 
         this.client_mut().on_close()
     }
@@ -2010,27 +2019,67 @@ impl Options {
     }
 }
 
+#[derive(Clone, Copy)]
+enum DeferredClose {
+    /// Close the socket the finalized wrapper left behind.
+    Socket,
+    /// Run the close path for a dial that never produced a socket
+    /// (`close_without_socket_next_tick`).
+    WithoutSocket,
+}
+
 pub(crate) struct ValkeyDeferredClose {
-    ctx: *const JSValkeyClient,
+    ctx: *mut JSValkeyClient,
+    what: DeferredClose,
 }
 
 impl ValkeyDeferredClose {
     #[allow(clippy::boxed_local, reason = "reclaim point for the boxed task")]
     pub(crate) fn run(self: Box<Self>) {
-        let ctx = self.ctx;
-        // SAFETY: single-threaded; intrusive ref taken before enqueue guarantees liveness.
-        unsafe {
-            crate::dispatch::fold((*ctx).client_mut().close());
-            JSValkeyClient::deref(ctx.cast_mut());
+        // SAFETY: adopts the ref `enqueue_deferred_close` took, which kept the
+        // client alive until now; released when this scope ends.
+        let _enqueue_ref = unsafe { ScopedRef::adopt(self.ctx) };
+        // SAFETY: live per the ref above; tasks run on the JS thread.
+        let this = unsafe { &*self.ctx };
+        match self.what {
+            DeferredClose::Socket => {
+                crate::dispatch::fold(this.client_mut().close(uws::CloseCode::FastShutdown))
+            }
+            DeferredClose::WithoutSocket => {
+                // Holding Connecting (see `close_without_socket_next_tick`) is
+                // what keeps a dial from starting in between; if one did, its
+                // own callbacks own the close path now, so only drop our ref.
+                if !this.client.get().socket.is_closed() {
+                    return;
+                }
+                // `on_close()` ends in `on_valkey_close`/`on_valkey_reconnect`,
+                // which release the ref the socket would have held.
+                this.ref_();
+                this.client_mut().status = valkey::Status::Disconnected;
+                let closed = this.client_mut().on_close();
+                this.update_poll_ref();
+                crate::dispatch::fold(closed);
+            }
         }
     }
 }
 
 impl bun_event_loop::Taskable for ValkeyDeferredClose {
     const TAG: bun_event_loop::TaskTag = bun_event_loop::task_tag::ValkeyDeferredClose;
-    /// The deferred close is script-free bookkeeping; do it.
     unsafe fn release_unrun(this: *mut Self) {
         // SAFETY: fn contract — boxed at the enqueue site.
-        unsafe { bun_core::heap::take(this) }.run();
+        let task = unsafe { bun_core::heap::take(this) };
+        match task.what {
+            // Script-free bookkeeping; do it.
+            DeferredClose::Socket => task.run(),
+            // The VM is going away: `on_close()` would run `onclose`, so only
+            // give back what `close_without_socket_next_tick` took.
+            DeferredClose::WithoutSocket => {
+                // SAFETY: as in `run`.
+                let _enqueue_ref = unsafe { ScopedRef::adopt(task.ctx) };
+                // SAFETY: live per the ref above.
+                unsafe { &*task.ctx }.poll_ref.with_mut(|r| r.disable());
+            }
+        }
     }
 }