Fix aborrt controller

Author: mjameswh

packages/core-bridge/src/helpers/abort_controller.rs

@@ -1,7 +1,6 @@
-use std::{
-    cell::Ref,
-    mem,
-    sync::{Arc, OnceLock},
+use std::sync::{
+    Arc, OnceLock,
+    atomic::{AtomicU8, Ordering, fence},
 };
 
 use neon::{
@@ -10,7 +9,6 @@ use neon::{
     prelude::{Context, JsResult},
     types::{JsFunction, JsObject, JsValue},
 };
-use parking_lot::{Mutex, Once};
 
 use super::{BridgeResult, JsCallback, TryIntoJs, errors::IntoThrow as _};
 
@@ -25,27 +23,61 @@ pub type JsAbortSignal = JsValue;
 /// The JS counterpart object is lazily instantiated when the signal gets converted to JS (through
 /// the `TryIntoJs` trait); this ensures that the Rust side can be created without a JS `Context`.
 pub struct AbortController {
-    inner: AbortControllerInner,
+    inner: Arc<AbortControllerInner>,
     drop_abort_reason: String,
 }
 
 /// An object that models the signal of a JavaScript `AbortController`.
 pub struct AbortSignal {
-    inner: AbortControllerInner,
+    inner: Arc<AbortControllerInner>,
 }
 
-/// The inner state of an `AbortController`, shared between the Rust's controller and its signal.
-type AbortControllerInner = Arc<Mutex<AbortControllerInnerState>>;
+/// The inner state of an `AbortController`, shared between the Rust and JS sides.
+struct AbortControllerInner {
+    // The fact that we require a `Context` to initialize the JS counterpart means that we are running
+    // on the Node's thread, which guarantees that there can't be multiple threads calling into that
+    // function concurrently; that should in theory aleviate the need to use a lock on `js_counterpart`.
+    //
+    // It is however possible for the rust-side controller to get aborted from a non-Node thread
+    // while the JS-side controller is being created on the Node thread, in which case we don't
+    // want the Rust-side thread to get blocked for the JS-side to complete instantiation.
+    //
+    // By modelling the "JS initialization" and "is aborted" states as two distinct independant
+    // structures, we ensure that we're never blocking execution of either thread. This however
+    // means that either step may happen before the other, so we need to be careful not to miss
+    // sending the abort signal. The good news is that nothing bad will happen if we call the JS
+    // abort callback multiple times.
+    js_counterpart: OnceLock<AbortControllerJsCounterpart>,
+    aborted: OnceLock<String>,
+
+    state: AtomicU8,
+}
+
+struct AbortControllerJsCounterpart {
+    controller: Root<JsObject>,
+    abort: JsCallback<(String,), ()>,
+}
+
+const STATE_UNINITIALIZED: u8 = 0;
+const STATE_ARMED: u8 = 1;
+const STATE_ABORTED: u8 = 2;
+const STATE_DISARMED: u8 = 3;
 
 impl AbortController {
-    /// Create a new `AbortController` and `AbortSignal` pair.
+    /// Create a new `AbortController`.
     ///
     /// The `drop_abort_reason` string will be used as the reason for the abort
     /// if the controller is dropped from the Rust side.
     #[must_use]
     pub fn new(drop_abort_reason: String) -> Self {
+        let inner = AbortControllerInner {
+            js_counterpart: OnceLock::new(),
+            aborted: OnceLock::new(),
+            state: AtomicU8::new(STATE_UNINITIALIZED),
+        };
+        let inner = Arc::new(inner);
         Self {
-            inner: Arc::new(Mutex::new(AbortControllerInnerState::Uninitialized)),
+            inner: inner.clone(),
             drop_abort_reason,
         }
     }
@@ -54,7 +86,7 @@ impl AbortController {
     /// be called at any time (i.e. even after the controller has been cancelled),
     /// any number of times (i.e. to pass a same signal to multiple JS functions),
     /// and from any thread.
-    pub fn signal(&self) -> AbortSignal {
+    pub fn get_signal(&self) -> AbortSignal {
         AbortSignal {
             inner: self.inner.clone(),
         }
@@ -65,12 +97,7 @@ impl AbortController {
     /// This method can be called at any time (i.e. even before the controller has been armed,
     /// or after it has been disarmed), and from any thread.
     pub fn abort(&self, reason: impl Into<String>) {
-        let call_abort = self.inner.lock().abort(reason);
-        if let Some(call_abort) = call_abort {
-            call_abort
-                .abort_cb
-                .call_on_js_thread((call_abort.abort_reason,));
-        }
+        self.inner.abort(reason);
     }
 
     /// Disarm the controller, so that it can no longer be aborted.
@@ -80,7 +107,7 @@ impl AbortController {
     /// on the JS side, e.g. when the called JS function has returned, as this will prevent the
     /// overhead of implicit abortion when the controller is dropped.
     pub fn disarm(&self) {
-        self.inner.lock().disarm();
+        self.inner.disarm();
     }
 }
 
@@ -95,220 +122,89 @@ impl Drop for AbortController {
 impl TryIntoJs for AbortSignal {
     type Output = JsAbortSignal;
     fn try_into_js<'cx>(self, cx: &mut impl Context<'cx>) -> JsResult<'cx, JsAbortSignal> {
-        self.inner.get_signal(cx).into_throw(cx)
+        let controller = self.inner.get_or_init_controller(cx).into_throw(cx)?;
+        controller.get(cx, "signal")
     }
 }
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
 impl AbortControllerInner {
     /// Create the JS `AbortController` if it hasn't been created yet.
     /// Returns a reference to the signal object that can be passed to JS.
-    fn get_js_controller<'cx, C: Context<'cx>>(
+    fn get_or_init_controller<'cx, C: Context<'cx>>(
         &self,
         cx: &mut C,
     ) -> BridgeResult<Handle<'cx, JsAbortController>> {
-        // Not using `get_or_init` here because initialization of JS classes can theoretically throw.
-        // But anyway, having a `Context` implies that we are running on the single Node's thread,
-        // so there's no risk of having two threads racing to initialize the JS `AbortController`.
-        if let Some(js_controller) = self.state.get_controller() {
-            return Ok(js_controller.to_inner(cx));
+        if let Some(js_counterpart) = self.js_counterpart.get() {
+            // Already initialized, return the controller
+            return Ok(js_counterpart.controller.to_inner(cx).upcast());
         }
 
-        // Instantiate the JS AbortController counterpart
+        // Not initialized yet, create the JS AbortController
         let global = cx.global_object();
         let abort_controller_class = global.get::<JsFunction, _, _>(cx, "AbortController")?;
         let js_controller = abort_controller_class.construct(cx, [])?;
 
-        // Get the JS `abort` method and make a callback out of it
         let abort_fn = js_controller.get::<JsFunction, _, _>(cx, "abort")?;
         let abort_cb = JsCallback::new(cx, abort_fn, Some(js_controller));
 
-        let js_controller = js_controller.root(cx);
-        let js_controller = match self.js_controller.set(js_controller) {
-            Ok(()) => {
-                // We confirmed instantiation of the JS counterpart; now update the state machine
-                let mut state = self.state.lock();
-                match *state {
-                    AbortControllerInnerState::Uninitialized => {
-                        *state = AbortControllerInnerState::Armed(abort_cb);
-                    }
-                    AbortControllerInnerState::AbortedBeforeArmed(reason) => {
-                        // Controller was aborted before the JS counterpart was initialized, so we must
-                        // immediately and synchronously propagate the abort to the JS controller object.
-
-                        // ... but let's not hold on to the state mutex while calling the abort callback
-                        *state = AbortControllerInnerState::Disarmed;
-                        mem::drop(state);
+        let js_counterpart = AbortControllerJsCounterpart {
+            controller: js_controller.root(cx),
+            abort: abort_cb,
+        };
 
-                        // Call the abort callback
-                        let _ = abort_cb.call(cx, (reason,));
-                    }
-                    AbortControllerInnerState::Armed(_) => {
-                        panic!(
-                            "AbortController: Incoherent state:JS AbortController already initialized"
-                        );
-                    }
-                    _ => {}
+        let controller = match self.js_counterpart.set(js_counterpart) {
+            Ok(()) => {
+                // Ordering: Write to `state` implies previous write to `js_counterpart` is visible to other threads
+                if self.state.fetch_or(STATE_ARMED, Ordering::Release) == STATE_ABORTED {
+                    // Ordering: Previous concurrent write to `aborted` must be visible at this point
+                    fence(Ordering::Acquire);
+
+                    // The controller was aborted before it was armed; immediately call the abort callback
+
+                    // Fire and forget
+                    let _ = self
+                        .js_counterpart
+                        .get()
+                        .unwrap()
+                        .abort
+                        .call(cx, (self.aborted.get().unwrap().clone(),));
                 }
-
-                // If the Rust controller has already been aborted, call the JS abort callback now
-                // if let Some(aborted) = self.aborted.get() {
-                //     // Fire and forget
-                //     let _ = js_counterpart.abort.call_on_js_thread((aborted.clone(),));
-                // }
-                js_controller.clone()
-            }
-            Err(js_controller) => {
-                // That case is totally unexpected, but should that ever happen, it's ok to just
-                // return the existing JS controller; the JS objects we just created have not
-                // yet been exposed anyway, so they will simply get garbage collected.
-                js_controller.clone()
+                js_controller
             }
+            Err(js_counterpart) => js_counterpart.controller.to_inner(cx).upcast(),
         };
 
-        Ok(js_controller.to_inner(cx))
-    }
-
-    /// Get a JS `AbortSignal` object for this `AbortController`. Initializes the
-    /// JS `AbortController` if it hasn't been initialized yet.
-    fn get_signal<'cx, C: Context<'cx>>(
-        &self,
-        cx: &mut C,
-    ) -> BridgeResult<Handle<'cx, JsAbortSignal>> {
-        let abort_controller = self.get_js_controller(cx)?;
-        Ok(abort_controller.get(cx, "signal")?)
+        Ok(controller)
     }
 
     /// Immediately abort the `AbortController`, causing the JS side `signal` to fire.
-    /// The controller is marked as aborted even if the JS counterpart has not yet been initialized;
-    /// in that case, the abort signal will be fired as soon as the JS side controller is initialized.
-    ///
-    /// This method can be called from any thread.
     fn abort(&self, reason: impl Into<String>) {
         let reason = reason.into();
         if self.aborted.set(reason.clone()) == Ok(()) {
             // If we haven't created the JS AbortController yet, there's nothing to abort
-            // VALIDATE: Do we need a memory barrier here to ensure that js_counterpart and aborted are coherent?
-            if let Some(js_counterpart) = self.js_counterpart.get() {
-                // Fire and forget
-                let _ = js_counterpart.abort.call_on_js_thread((reason,));
-            }
-        }
-    }
-}
-
-////////
+            // Ordering: Write to `state` implies previous write to `aborted` is visible to other threads
+            if self.state.fetch_or(STATE_ABORTED, Ordering::Release) == STATE_ARMED {
+                // Ordering: Previous concurrent write to `js_counterpart` must be visible at this point
+                fence(Ordering::Acquire);
 
-struct AbortControllerInnerState2 {
-    //
-    js_controller: Option<Root<JsAbortController>>,
-    abort_cb: Option<JsCallback<(String,), ()>>,
-}
-
-enum AbortControllerInnerState {
-    /// The `AbortController` has just been created; the JS counterpart has not yet been initialized.
-    ///
-    /// Transitions:
-    /// - Calling `abort()` will move to the `AbortedBeforeArmed` state
-    /// - Initializing the JS counterpart will move to the `Armed` state
-    Uninitialized,
-
-    /// The `AbortController` has been aborted before the JS counterpart was initialized.
-    ///
-    /// Transitions:
-    /// - Calling `abort()` is a no-op
-    /// - Initializing the JS counterpart will result in calling the abort
-    ///   callback, and move to the `Disabled` state
-    AbortedBeforeArmed(String),
-
-    /// The JS counterpart has been initialized.
-    ///
-    /// Transitions:
-    /// - Calling `abort()` will result in calling the abort callback, and move to the `Disabled` state
-    /// - Calling `disable()` is a no-op
-    Armed(Root<JsAbortController>, JsCallback<(String,), ()>),
-
-    /// The `AbortController` has been disarmed, though it was actually never armed.
-    ///
-    /// Transitions:
-    /// - Calling `abort()` is a no-op
-    /// - Initializing the JS counterpart will move to the `Disarmed` state
-    UninitializedDisarmed,
-
-    /// The `AbortController` has been disabled, either by being aborted or by calling `disable()`.
-    /// It is still possible to obtain a `signal` from the controller, but it's abort status
-    /// will not change anymore.
-    ///
-    /// Any transitions from this state is a no-op.
-    Disarmed(Root<JsAbortController>),
-}
-
-impl AbortControllerInnerState {
-    fn arm(
-        &mut self,
-        js_controller: Root<JsAbortController>,
-        abort_cb: JsCallback<(String,), ()>,
-    ) -> Option<CallAbortDetails> {
-        match self {
-            Self::Uninitialized => {
-                *self = Self::Armed(js_controller, abort_cb);
-                None
-            }
-            Self::AbortedBeforeArmed(abort_reason) => {
-                let abort_reason = abort_reason.clone();
-                *self = Self::Disarmed(js_controller);
-                Some(CallAbortDetails {
-                    abort_reason,
-                    abort_cb,
-                })
-            }
-            Self::UninitializedDisarmed => {
-                *self = Self::Disarmed(js_controller);
-                None
-            }
-            _ => None,
-        }
-    }
-
-    fn abort(&mut self, reason: impl Into<String>) -> Option<CallAbortDetails> {
-        let current = mem::replace(self, Self::Uninitialized);
-        match current {
-            Self::Uninitialized => {
-                *self = Self::AbortedBeforeArmed(reason.into());
-                None
-            }
-            Self::Armed(js_controller, abort_cb) => {
-                *self = Self::Disarmed(js_controller);
-                Some(CallAbortDetails {
-                    abort_reason: reason.into(),
-                    abort_cb,
-                })
-            }
-            other => {
-                *self = other;
-                None
+                // Fire and forget
+                let _ = self
+                    .js_counterpart
+                    .get()
+                    .unwrap()
+                    .abort
+                    .call_on_js_thread((reason,));
             }
         }
     }
 
-    fn disarm(&mut self) {
-        let current = mem::replace(self, Self::Uninitialized);
-        match current {
-            Self::Uninitialized => {
-                *self = Self::UninitializedDisarmed;
-            }
-            Self::Armed(js_controller, _) => {
-                *self = Self::Disarmed(js_controller);
-            }
-            other => {
-                *self = other;
-            }
-        }
+    fn disarm(&self) {
+        // Ordering: this requires no dependency on any other state
+        self.state.store(STATE_DISARMED, Ordering::Relaxed);
     }
 }
 
-struct CallAbortDetails {
-    abort_reason: String,
-    abort_cb: JsCallback<(String,), ()>,
+#[cfg(feature = "test")]
+mod tests {
+    use super::*;
 }