fix(fetch): release body resolver realm roots on completion

moli-renderer-v8/src/network_host/body_source.rs

@@ -5,7 +5,7 @@ use crate::context_bootstrap::{
 };
 use crate::protocol_types::SubresourceResponseBody;
 use crate::types::NetworkBodySourceId;
-use crate::util::{RealmObjectHandle, get_private_value, set_private_value};
+use crate::util::{PendingPromiseResolver, get_private_value, set_private_value};
 use crate::web_api_interfaces;
 use crate::worker::get_worker_state;
 use std::collections::{HashMap, HashSet};
@@ -245,7 +245,7 @@ enum PendingBodyMaterializationKind {
 }
 
 pub(crate) struct PendingBodyMaterialization {
-    resolver: RealmObjectHandle,
+    resolver: PendingPromiseResolver,
     kind: PendingBodyMaterializationKind,
 }
 
@@ -905,13 +905,18 @@ fn pending_body_active_branch_count(
 }
 
 fn cancel_pending_network_body_source_in_maps(
+    scope: &mut v8::PinScope<'_, '_>,
     sources: &mut HashMap<NetworkBodySourceId, PendingNetworkBodySourceState>,
     clones: &mut HashMap<NetworkBodySourceId, Vec<NetworkBodySourceId>>,
     id: NetworkBodySourceId,
 ) -> Option<NetworkBodySourceId> {
     let root_id = pending_body_clone_root(clones, id);
     let has_descendants = !pending_body_clone_descendants(clones, id).is_empty();
-    sources.remove(&id);
+    if let Some(state) = sources.remove(&id) {
+        for materialization in state.materializations {
+            let _ = materialization.resolver.into_local(scope);
+        }
+    }
     if !has_descendants {
         clones.remove(&id);
         remove_pending_body_clone_link(clones, id);
@@ -1553,7 +1558,7 @@ fn inspect_pending_body_source_for_materialization<'s>(
             sources.remove(&id);
         } else {
             state.materializations.push(PendingBodyMaterialization {
-                resolver: RealmObjectHandle::new(scope, resolver.into()),
+                resolver: PendingPromiseResolver::new(scope, resolver),
                 kind,
             });
         }
@@ -1797,11 +1802,9 @@ fn resolve_pending_body_materializations(
     let mut bytes: Vec<u8> = take_pending_network_body_source_bytes(scope, id).unwrap_or_default();
     let materialization_count = materializations.len();
     for (index, materialization) in materializations.into_iter().enumerate() {
-        let Some(resolver) = materialization.resolver.to_local(scope) else {
+        let Some(resolver) = materialization.resolver.into_local(scope) else {
             continue;
         };
-        // The private handle is populated only from a PromiseResolver.
-        let resolver = unsafe { v8::Local::<v8::PromiseResolver>::cast_unchecked(resolver) };
         // Most body sources have exactly one pending materialization. Move the
         // completed buffer into that consumer and clone only when fan-out is
         // actually needed.
@@ -1839,11 +1842,9 @@ fn reject_pending_body_materializations_with_reason<'s>(
     reason: v8::Local<'s, v8::Value>,
 ) {
     for materialization in materializations {
-        let Some(resolver) = materialization.resolver.to_local(scope) else {
+        let Some(resolver) = materialization.resolver.into_local(scope) else {
             continue;
         };
-        // The private handle is populated only from a PromiseResolver.
-        let resolver = unsafe { v8::Local::<v8::PromiseResolver>::cast_unchecked(resolver) };
         let _ = resolver.reject(scope, reason);
     }
 }
@@ -2046,6 +2047,7 @@ fn pending_body_source_cancel_callback<'s>(
     if let Some(id) = registry_body_source_id(scope, args.this()) {
         if let Some(host) = context_host_mut(scope) {
             if let Some(root_id) = cancel_pending_network_body_source_in_maps(
+                scope,
                 &mut host.pending_network_body_sources,
                 &mut host.pending_network_body_clones,
                 id,
@@ -2056,6 +2058,7 @@ fn pending_body_source_cancel_callback<'s>(
             let mut worker_state = worker_state.borrow_mut();
             let worker_state = &mut *worker_state;
             let _ = cancel_pending_network_body_source_in_maps(
+                scope,
                 &mut worker_state.pending_network_body_sources,
                 &mut worker_state.pending_network_body_clones,
                 id,

moli-renderer-v8/src/script_vm/tests/extracted/realms_and_teardown.rs

@@ -167,6 +167,203 @@ fn retained_response_body_keeps_bytes_until_the_last_v8_reference() {
     }
 }
 
+enum RetiredBodyCompletion<'a> {
+    Bytes(&'a [u8]),
+    Error,
+    Cancel,
+}
+
+fn assert_body_result_is_collectible_in_retained_realm(
+    method: &str,
+    completion: RetiredBodyCompletion<'_>,
+    expected_state: Option<v8::PromiseState>,
+) {
+    let (mut vm, body_source_id) = new_vm_with_pending_response_for_teardown_test();
+    vm.eval(&format!(
+        "globalThis.bodyPromise = response.{method}();\n\
+         bodyPromise.catch(() => {{}});\n\
+         globalThis.keepRealmAlive = () => document.URL;"
+    ))
+    .unwrap();
+    if matches!(completion, RetiredBodyCompletion::Cancel) {
+        vm.with_default_context_scope_and_checkpoint_for_test(|scope, _| {
+            let global = scope.get_current_context().global(scope);
+            let response = global
+                .get(scope, crate::util::v8str(scope, "response").into())
+                .unwrap();
+            let response = v8::Local::<v8::Object>::try_from(response).unwrap();
+            let source = crate::util::get_private_value(scope, response, "__lmNetworkBodySource")
+                .expect("the pending Response has a native body source");
+            let _ = global.set(
+                scope,
+                crate::util::v8str(scope, "cancelBody").into(),
+                source,
+            );
+            Ok(())
+        })
+        .unwrap();
+    }
+    let (callback, promise) = vm
+        .with_default_context_scope_and_checkpoint_for_test(|scope, _| {
+            let global = scope.get_current_context().global(scope);
+            let callback = global
+                .get(scope, crate::util::v8str(scope, "keepRealmAlive").into())
+                .unwrap();
+            let callback = v8::Local::<v8::Function>::try_from(callback).unwrap();
+            let promise = global
+                .get(scope, crate::util::v8str(scope, "bodyPromise").into())
+                .unwrap();
+            let promise = v8::Local::<v8::Promise>::try_from(promise).unwrap();
+            Ok((
+                v8::Global::new(scope, callback),
+                v8::Weak::new(scope, promise),
+            ))
+        })
+        .unwrap();
+    vm.eval("globalThis.response = null; globalThis.bodyPromise = null")
+        .unwrap();
+    let weak_host = vm.context_host_weak_for_test();
+    let isolate = vm.renderer_document_isolate.clone();
+    drop(vm);
+    for _ in 0..2 {
+        isolate.with_renderer_document_isolate_mut(|isolate| isolate.low_memory_notification());
+    }
+    assert!(weak_host.upgrade().is_some());
+
+    let result = isolate.with_renderer_document_isolate_mut(|isolate| {
+        let scope = std::pin::pin!(v8::HandleScope::new(isolate));
+        let scope = &mut scope.init();
+        let caller_context = v8::Context::new(scope, Default::default());
+        let scope = &mut v8::ContextScope::new(scope, caller_context);
+        let callback = v8::Local::new(scope, &callback);
+        let context = callback.get_creation_context(scope).unwrap();
+        let scope = &mut v8::ContextScope::new(scope, context);
+        let promise = promise
+            .to_local(scope)
+            .expect("a pending body resolver must survive retirement and GC");
+        assert_eq!(promise.state(), v8::PromiseState::Pending);
+        match completion {
+            RetiredBodyCompletion::Bytes(bytes) => {
+                crate::network_host::enqueue_pending_network_body_chunk(
+                    scope,
+                    body_source_id,
+                    bytes.to_vec(),
+                );
+                crate::network_host::close_pending_network_body_stream(scope, body_source_id);
+            }
+            RetiredBodyCompletion::Error => {
+                crate::network_host::error_pending_network_body_stream(
+                    scope,
+                    body_source_id,
+                    "body failed".to_owned(),
+                );
+            }
+            RetiredBodyCompletion::Cancel => {
+                // Exercise the native cancellation protocol directly.
+                let global = context.global(scope);
+                let source = global
+                    .get(scope, crate::util::v8str(scope, "cancelBody").into())
+                    .unwrap();
+                let source = v8::Local::<v8::Object>::try_from(source).unwrap();
+                let cancel = source
+                    .get(scope, crate::util::v8str(scope, "cancel").into())
+                    .unwrap();
+                let cancel = v8::Local::<v8::Function>::try_from(cancel).unwrap();
+                crate::script_execution::call_function(scope, cancel, source.into(), &[])
+                    .expect("the retained native body source should be cancellable");
+                let _ = global.set(
+                    scope,
+                    crate::util::v8str(scope, "cancelBody").into(),
+                    v8::null(scope).into(),
+                );
+            }
+        }
+        context
+            .get_microtask_queue()
+            .unwrap()
+            .perform_checkpoint(scope);
+        expected_state.map(|state| {
+            assert_eq!(promise.state(), state);
+            let result = v8::Local::<v8::Object>::try_from(promise.result(scope)).unwrap();
+            v8::Weak::new(scope, result)
+        })
+    });
+    for _ in 0..2 {
+        isolate.with_renderer_document_isolate_mut(|isolate| isolate.low_memory_notification());
+    }
+    assert!(
+        weak_host.upgrade().is_some(),
+        "the unrelated callback must still retain the retired Document"
+    );
+    isolate.with_renderer_document_isolate_mut(|isolate| {
+        let scope = std::pin::pin!(v8::HandleScope::new(isolate));
+        let scope = &mut scope.init();
+        if let Some(result) = result {
+            assert!(
+                result.to_local(scope).is_none(),
+                "a completed body result must be collectible before its retired realm"
+            );
+        }
+        assert!(
+            promise.to_local(scope).is_none(),
+            "a completed or cancelled body resolver must no longer be rooted by its realm"
+        );
+    });
+    isolate.with_renderer_document_isolate_mut(|_| drop(callback));
+    for _ in 0..2 {
+        isolate.with_renderer_document_isolate_mut(|isolate| isolate.low_memory_notification());
+    }
+    assert!(
+        weak_host.upgrade().is_none(),
+        "releasing the unrelated callback must release the retired Document"
+    );
+}
+
+#[test]
+fn resolved_json_body_result_is_collectible_in_retained_realm() {
+    assert_body_result_is_collectible_in_retained_realm(
+        "json",
+        RetiredBodyCompletion::Bytes(br#"{"ok":true}"#),
+        Some(v8::PromiseState::Fulfilled),
+    );
+}
+
+#[test]
+fn resolved_array_buffer_body_result_is_collectible_in_retained_realm() {
+    assert_body_result_is_collectible_in_retained_realm(
+        "arrayBuffer",
+        RetiredBodyCompletion::Bytes(b"body bytes"),
+        Some(v8::PromiseState::Fulfilled),
+    );
+}
+
+#[test]
+fn invalid_json_body_error_is_collectible_in_retained_realm() {
+    assert_body_result_is_collectible_in_retained_realm(
+        "json",
+        RetiredBodyCompletion::Bytes(b"invalid json"),
+        Some(v8::PromiseState::Rejected),
+    );
+}
+
+#[test]
+fn network_body_error_is_collectible_in_retained_realm() {
+    assert_body_result_is_collectible_in_retained_realm(
+        "json",
+        RetiredBodyCompletion::Error,
+        Some(v8::PromiseState::Rejected),
+    );
+}
+
+#[test]
+fn cancelled_body_promise_is_collectible_in_retained_realm() {
+    assert_body_result_is_collectible_in_retained_realm(
+        "json",
+        RetiredBodyCompletion::Cancel,
+        None,
+    );
+}
+
 #[test]
 fn retained_response_body_keeps_error_reason_until_the_last_v8_reference() {
     let (mut vm, body_source_id) = new_vm_with_pending_response_for_teardown_test();

moli-renderer-v8/src/util.rs

@@ -149,6 +149,90 @@ impl RealmObjectHandle {
     }
 }
 
+/// A uniquely owned pending resolver needs a realm root only until its native
+/// operation ends. Identity caches use `RealmObjectHandle` instead, since their
+/// roots intentionally last as long as the realm.
+pub(crate) enum PendingPromiseResolver {
+    Strong(v8::Global<v8::PromiseResolver>),
+    Traced(v8::Weak<v8::PromiseResolver>),
+}
+
+impl PendingPromiseResolver {
+    const ROOTS: &str = "__moliPendingRealmResolvers";
+
+    pub(crate) fn new(
+        scope: &mut v8::PinScope<'_, '_>,
+        resolver: v8::Local<'_, v8::PromiseResolver>,
+    ) -> Self {
+        if page_context_is_detached_document(scope.get_current_context()) {
+            Self::traced(scope, resolver)
+        } else {
+            Self::Strong(v8::Global::new(scope, resolver))
+        }
+    }
+
+    fn traced(
+        scope: &mut v8::PinScope<'_, '_>,
+        resolver: v8::Local<'_, v8::PromiseResolver>,
+    ) -> Self {
+        let context = resolver
+            .get_creation_context(scope)
+            .unwrap_or_else(|| scope.get_current_context());
+        // Create the collection in its owning realm even when main-document
+        // teardown visits a resolver belonging to a shared child world.
+        let scope = &mut v8::ContextScope::new(scope, context);
+        let anchor = context.get_extras_binding_object(scope);
+        let roots = get_private_value(scope, anchor, Self::ROOTS)
+            .and_then(|value| v8::Local::<v8::Set>::try_from(value).ok())
+            .unwrap_or_else(|| {
+                let roots = v8::Set::new(scope);
+                set_private_value(scope, anchor, Self::ROOTS, roots.into());
+                roots
+            });
+        let _ = roots.add(scope, resolver.into());
+        Self::Traced(v8::Weak::new(scope, resolver))
+    }
+
+    pub(crate) fn to_local<'s>(
+        &self,
+        scope: &mut v8::PinScope<'s, '_>,
+    ) -> Option<v8::Local<'s, v8::PromiseResolver>> {
+        match self {
+            Self::Strong(resolver) => Some(v8::Local::new(scope, resolver)),
+            Self::Traced(resolver) => resolver.to_local(scope),
+        }
+    }
+
+    pub(crate) fn retain_in_realm(&mut self, scope: &mut v8::PinScope<'_, '_>) {
+        if let Self::Strong(resolver) = self {
+            let resolver = v8::Local::new(scope, &*resolver);
+            *self = Self::traced(scope, resolver);
+        }
+    }
+
+    /// Release the realm edge before settling or cancelling the operation.
+    /// The returned Local keeps the resolver alive for the rest of this scope.
+    pub(crate) fn into_local<'s>(
+        self,
+        scope: &mut v8::PinScope<'s, '_>,
+    ) -> Option<v8::Local<'s, v8::PromiseResolver>> {
+        let resolver = self.to_local(scope)?;
+        if matches!(self, Self::Traced(_)) {
+            let context = resolver
+                .get_creation_context(scope)
+                .unwrap_or_else(|| scope.get_current_context());
+            let scope = &mut v8::ContextScope::new(scope, context);
+            let anchor = context.get_extras_binding_object(scope);
+            if let Some(roots) = get_private_value(scope, anchor, Self::ROOTS)
+                .and_then(|value| v8::Local::<v8::Set>::try_from(value).ok())
+            {
+                let _ = roots.delete(scope, resolver.into());
+            }
+        }
+        Some(resolver)
+    }
+}
+
 /// Value edges owned by native state become ordinary V8 edges when a realm
 /// retires. A private holder supports primitives as well as objects and keeps
 /// cross-realm values attached to the realm that owns the native state.