[JSC] bun:ffi: a threadsafe JSFFICallback's entry thunk stays callable after its cell and VM are gone

Source/JavaScriptCore/ffi/BunFFI.h

@@ -50,6 +50,8 @@ JS_EXPORT_PRIVATE RefPtr<Signature> signatureFromJS(JSGlobalObject*, JSValue des
 class ThreadsafeInvocation;
 JS_EXPORT_PRIVATE JSFFICallback* createCallback(JSGlobalObject*, Ref<Signature>&&, JSObject* callable, bool threadsafe, void* embedderContext);
 JS_EXPORT_PRIVATE void runThreadsafeInvocation(ThreadsafeInvocation&);
+// For an invocation the embedder's dispatch accepted but will not run (owning thread; takes the API lock itself).
+JS_EXPORT_PRIVATE void retireThreadsafeInvocation(ThreadsafeInvocation&);
 
 } // namespace FFI
 

Source/JavaScriptCore/ffi/FFICallbackThunk.cpp

@@ -46,6 +46,7 @@
 #include "LinkBuffer.h"
 #include "MarkedVector.h"
 #include "Options.h"
+#include <wtf/RawPointer.h>
 #include <wtf/MathExtras.h>
 
 namespace JSC {
@@ -274,9 +275,15 @@ MacroAssemblerCodeRef<JITThunkPtrTag> generateCallbackThunk(VM&, JSFFICallback&
     }
 
     jit.addPtr(CCallHelpers::TrustedImm32(slotsOffsetFromFP), GPRInfo::callFrameRegister, GPRInfo::argumentGPR1);
-    jit.move(CCallHelpers::TrustedImmPtr(&callback), GPRInfo::argumentGPR0);
-    auto dispatchOperation = callback.isThreadsafe() ? tagCFunction<OperationPtrTag>(ffiCallbackDispatchThreadsafe) : tagCFunction<OperationPtrTag>(ffiCallbackDispatch);
-    jit.move(CCallHelpers::TrustedImmPtr(dispatchOperation), thunkScratchGPR);
+    // A threadsafe thunk may run when the cell (and its VM) no longer exist: it is handed the handle,
+    // which outlives them, and never the cell.
+    if (auto* handle = callback.threadsafeHandle()) {
+        jit.move(CCallHelpers::TrustedImmPtr(handle), GPRInfo::argumentGPR0);
+        jit.move(CCallHelpers::TrustedImmPtr(tagCFunction<OperationPtrTag>(ffiCallbackDispatchThreadsafe)), thunkScratchGPR);
+    } else {
+        jit.move(CCallHelpers::TrustedImmPtr(&callback), GPRInfo::argumentGPR0);
+        jit.move(CCallHelpers::TrustedImmPtr(tagCFunction<OperationPtrTag>(ffiCallbackDispatch)), thunkScratchGPR);
+    }
     jit.call(thunkScratchGPR, OperationPtrTag);
 
     const CCallHelpers::Address returnSlot(GPRInfo::callFrameRegister, slotsOffsetFromFP + static_cast<int>(argumentCount * slotSize));
@@ -353,17 +360,18 @@ class CallbackEntryScope {
 
 } // namespace FFI
 
-JSC_DEFINE_JIT_OPERATION(ffiCallbackDispatchThreadsafe, EncodedJSValue, (JSFFICallback* callback, uint64_t* slots))
+// Any thread, any time — possibly after the callback's cell, global object and VM are gone. Touches only
+// the handle.
+JSC_DEFINE_JIT_OPERATION(ffiCallbackDispatchThreadsafe, EncodedJSValue, (FFI::ThreadsafeCallbackHandle* handle, uint64_t* slots))
 {
-    ASSERT(callback->isThreadsafe());
-    FFI::Signature& signature = callback->signature();
-    const unsigned argumentCount = signature.argumentCount();
+    const unsigned argumentCount = handle->signature().argumentCount();
     auto dispatch = FFI::FFIContext::threadsafeDispatch();
     RELEASE_ASSERT(dispatch);
-    if (callback->tryBeginThreadsafeInvocation()) [[likely]] {
-        auto invocation = FFI::ThreadsafeInvocation::create(callback, callback->embedderContext(), std::span<const uint64_t>(slots, argumentCount));
+    if (handle->tryBeginInvocation()) [[likely]] {
+        auto invocation = FFI::ThreadsafeInvocation::create(*handle, std::span<const uint64_t>(slots, argumentCount));
         dispatch(invocation.get());
-    }
+    } else if (Options::verboseFFI()) [[unlikely]]
+        dataLogLn("FFI: dropped a call to closed threadsafe callback ", RawPointer(handle), " (", handle->signature().toString(), "); returning zero");
     slots[argumentCount] = 0;
     return { encodedJSUndefined(), nullptr };
 }
@@ -440,21 +448,31 @@ JSC_DEFINE_JIT_OPERATION(ffiCallbackDispatch, EncodedJSValue, (JSFFICallback* ca
 
 namespace FFI {
 
+void retireThreadsafeInvocation(ThreadsafeInvocation& invocation)
+{
+    ThreadsafeCallbackHandle& handle = invocation.handle();
+    if (handle.endInvocation() && handle.callback()) {
+        JSLockHolder locker(handle.callback()->vm());
+        handle.callback()->unroot();
+    }
+}
+
+// The owning thread, routed here by the embedder while the callback's context still accepts tasks.
 void runThreadsafeInvocation(ThreadsafeInvocation& invocation)
 {
-    JSFFICallback* callback = invocation.callback();
-    struct RetireInvocation {
-        JSFFICallback* callback;
-        ~RetireInvocation()
-        {
-            if (callback->endThreadsafeInvocation())
-                callback->unroot();
-        }
-    } retire { callback };
+    JSFFICallback* callback = invocation.handle().callback();
+    if (!callback) {
+        retireThreadsafeInvocation(invocation);
+        return;
+    }
 
     JSGlobalObject* globalObject = callback->globalObject();
     VM& vm = globalObject->vm();
     JSLockHolder locker(vm);
+    struct RetireInvocation {
+        ThreadsafeInvocation& invocation;
+        ~RetireInvocation() { retireThreadsafeInvocation(invocation); }
+    } retire { invocation };
     auto scope = DECLARE_THROW_SCOPE(vm);
 
     Signature& signature = callback->signature();
@@ -491,6 +509,10 @@ void runThreadsafeInvocation(ThreadsafeInvocation&)
 {
     RELEASE_ASSERT_NOT_REACHED(); // unreachable: no threadsafe callback exists to have queued this
 }
+void retireThreadsafeInvocation(ThreadsafeInvocation&)
+{
+    RELEASE_ASSERT_NOT_REACHED();
+}
 } // namespace FFI
 
 #endif // FFI_CALLBACK_THUNK_SUPPORTED

Source/JavaScriptCore/ffi/FFICallbackThunk.h

@@ -45,6 +45,7 @@ class JSFFICallback;
 class VM;
 
 namespace FFI {
+class ThreadsafeCallbackHandle;
 
 #if FFI_CALLBACK_THUNK_SUPPORTED
 MacroAssemblerCodeRef<JITThunkPtrTag> generateCallbackThunk(VM&, JSFFICallback&);
@@ -54,7 +55,7 @@ MacroAssemblerCodeRef<JITThunkPtrTag> generateCallbackThunk(VM&, JSFFICallback&)
 
 #if FFI_CALLBACK_THUNK_SUPPORTED
 JSC_DECLARE_JIT_OPERATION(ffiCallbackDispatch, EncodedJSValue, (JSFFICallback*, uint64_t*));
-JSC_DECLARE_JIT_OPERATION(ffiCallbackDispatchThreadsafe, EncodedJSValue, (JSFFICallback*, uint64_t*));
+JSC_DECLARE_JIT_OPERATION(ffiCallbackDispatchThreadsafe, EncodedJSValue, (FFI::ThreadsafeCallbackHandle*, uint64_t*));
 #endif // FFI_CALLBACK_THUNK_SUPPORTED
 
 } // namespace JSC

Source/JavaScriptCore/ffi/FFIContext.cpp

@@ -38,6 +38,7 @@ namespace JSC { namespace FFI {
 
 WTF_MAKE_TZONE_ALLOCATED_IMPL(FFIContext);
 
+WTF_MAKE_TZONE_ALLOCATED_IMPL(ThreadsafeCallbackHandle);
 WTF_MAKE_TZONE_ALLOCATED_IMPL(ThreadsafeInvocation);
 
 FFIContext::ThreadsafeDispatchFunction FFIContext::s_threadsafeDispatch { nullptr };

Source/JavaScriptCore/ffi/FFIContext.h

@@ -29,9 +29,12 @@
 
 #if USE(BUN_JSC_ADDITIONS)
 
+#include "FFISignature.h"
 #include "HeapObserver.h"
 #include "JSExportMacros.h"
+#include "MacroAssemblerCodeRef.h"
 #include "WriteBarrier.h"
+#include <atomic>
 #include <span>
 #include <wtf/ForbidHeapAllocation.h>
 #include <wtf/MallocSpan.h>
@@ -50,29 +53,91 @@ class VM;
 
 namespace JSC { namespace FFI {
 
+// What the entry thunk of a `threadsafe` JSFFICallback targets in place of the cell. Native code may
+// call that thunk from any thread and at any time — including after the callback's global object and VM
+// are gone (a terminated worker), which the embedder cannot always sequence against `close()`. So
+// nothing the foreign thread touches lives in the GC cell: the thunk code, the signature, the
+// closed/pending state and the embedder's routing token live here, and once an entrypoint has been handed
+// out the handle (and so the thunk) is never freed: a call arriving after close() or after the cell and
+// VM are gone counts itself out and returns zero rather than jumping into freed code. Only the owning
+// thread reads `callback()`.
+class ThreadsafeCallbackHandle final : public ThreadSafeRefCounted<ThreadsafeCallbackHandle> {
+    WTF_MAKE_TZONE_ALLOCATED(ThreadsafeCallbackHandle);
+public:
+    static Ref<ThreadsafeCallbackHandle> create(JSFFICallback* callback, Ref<Signature>&& signature, void* embedderContext)
+    {
+        return adoptRef(*new ThreadsafeCallbackHandle(callback, WTF::move(signature), embedderContext));
+    }
+
+    Signature& signature() const { return m_signature.get(); }
+    void* embedderContext() const { return m_embedderContext; }
+    JSFFICallback* callback() const { return m_callback; } // owning thread; null once the cell is destroyed
+
+    void setEntryCode(MacroAssemblerCodeRef<JITThunkPtrTag>&& code) { m_entryCode = WTF::move(code); }
+    const MacroAssemblerCodeRef<JITThunkPtrTag>& entryCode() const { return m_entryCode; }
+
+    static constexpr unsigned closedBit = 0x80000000u;
+    static constexpr unsigned countMask = 0x7fffffffu;
+    // Any thread: count an invocation in unless closed.
+    bool tryBeginInvocation()
+    {
+        unsigned state = m_state.load(std::memory_order_acquire);
+        do {
+            if (state & closedBit)
+                return false;
+        } while (!m_state.compare_exchange_weak(state, state + 1, std::memory_order_acq_rel, std::memory_order_acquire));
+        return true;
+    }
+    // Owning thread: whether this was the last pending invocation of a closed callback.
+    bool endInvocation()
+    {
+        unsigned state = m_state.fetch_sub(1, std::memory_order_acq_rel) - 1;
+        return (state & closedBit) && !(state & countMask);
+    }
+    // Owning thread: whether invocations are still pending.
+    bool markClosedAndReportPending() { return m_state.fetch_or(closedBit, std::memory_order_acq_rel) & countMask; }
+    void cellDestroyed()
+    {
+        m_callback = nullptr;
+        m_state.fetch_or(closedBit, std::memory_order_acq_rel);
+    }
+
+private:
+    ThreadsafeCallbackHandle(JSFFICallback* callback, Ref<Signature>&& signature, void* embedderContext)
+        : m_signature(WTF::move(signature))
+        , m_embedderContext(embedderContext)
+        , m_callback(callback)
+    {
+    }
+
+    const Ref<Signature> m_signature;
+    void* const m_embedderContext;
+    JSFFICallback* m_callback;
+    MacroAssemblerCodeRef<JITThunkPtrTag> m_entryCode;
+    std::atomic<unsigned> m_state { 0 }; // closedBit | pending-invocation count
+};
+
 class ThreadsafeInvocation final : public ThreadSafeRefCounted<ThreadsafeInvocation> {
     WTF_MAKE_TZONE_ALLOCATED(ThreadsafeInvocation);
 public:
-    static Ref<ThreadsafeInvocation> create(JSFFICallback* callback, void* embedderContext, std::span<const uint64_t> slots)
+    static Ref<ThreadsafeInvocation> create(ThreadsafeCallbackHandle& handle, std::span<const uint64_t> slots)
     {
-        return adoptRef(*new ThreadsafeInvocation(callback, embedderContext, slots));
+        return adoptRef(*new ThreadsafeInvocation(handle, slots));
     }
 
-    JSFFICallback* callback() const { return m_callback; }
-    void* embedderContext() const { return m_embedderContext; }
+    ThreadsafeCallbackHandle& handle() const { return m_handle.get(); }
+    void* embedderContext() const { return m_handle->embedderContext(); }
     std::span<uint64_t> slots() { return m_slots.mutableSpan(); }
     std::span<const uint64_t> slots() const { return m_slots.span(); }
 
 private:
-    ThreadsafeInvocation(JSFFICallback* callback, void* embedderContext, std::span<const uint64_t> slots)
-        : m_callback(callback)
-        , m_embedderContext(embedderContext)
+    ThreadsafeInvocation(ThreadsafeCallbackHandle& handle, std::span<const uint64_t> slots)
+        : m_handle(handle)
         , m_slots(slots)
     {
     }
 
-    JSFFICallback* m_callback;
-    void* m_embedderContext;
+    const Ref<ThreadsafeCallbackHandle> m_handle;
     Vector<uint64_t, 8> m_slots;
 };
 

Source/JavaScriptCore/ffi/JSFFICallback.cpp

@@ -60,7 +60,18 @@ JSFFICallback::JSFFICallback(VM& vm, Structure* structure, Ref<FFI::Signature>&&
 {
 }
 
-JSFFICallback::~JSFFICallback() = default;
+JSFFICallback::~JSFFICallback()
+{
+    if (!m_threadsafe)
+        return;
+    m_threadsafe->cellDestroyed();
+    // Once an entrypoint has been handed out, native code may call it at any later time -- there is no
+    // happens-before between a foreign thread leaving the thunk and this thread freeing it, close()d or
+    // not -- so the handle, and with it the thunk, stays alive for the rest of the process. Late calls
+    // count themselves out and return zero.
+    if (m_threadsafe->entryCode())
+        m_threadsafe->ref();
+}
 
 void JSFFICallback::destroy(JSCell* cell)
 {
@@ -106,15 +117,22 @@ JSFFICallback* JSFFICallback::create(VM& vm, JSGlobalObject* globalObject, Struc
         return nullptr;
     }
     JSFFICallback* callback = new (NotNull, allocateCell<JSFFICallback>(vm)) JSFFICallback(vm, structure, WTF::move(signature));
-    callback->m_threadsafe = threadsafe;
-    callback->m_embedderContext = embedderContext;
+    if (threadsafe)
+        callback->m_threadsafe = FFI::ThreadsafeCallbackHandle::create(callback, callback->m_signature.copyRef(), embedderContext);
     callback->finishCreation(vm, callable);
     RETURN_IF_EXCEPTION(scope, nullptr);
-    if (!callback->m_entryCode) [[unlikely]] {
+    if (!callback->entryCode()) [[unlikely]] {
         throwOutOfMemoryError(globalObject, scope);
         return nullptr;
     }
     globalObject->ffiContext().addLiveCallback(vm, *globalObject, callback);
+    // A threadsafe callback's handle reads the cell through a raw pointer that only ~JSFFICallback clears,
+    // and cells are swept lazily: the cell must not be able to die (and sit unswept, still reachable from
+    // a delivered invocation) while its VM lives on -- as it could if it were rooted only through a global
+    // object that is itself collected. Root it VM-wide until it is closed and drained; it then dies only
+    // by close() or with the VM, whose last sweep runs the destructor before anything else can run.
+    if (threadsafe)
+        vm.heap.protect(callback);
     return callback;
 #endif
 }
@@ -127,22 +145,31 @@ void JSFFICallback::finishCreation(VM& vm, JSObject* callable)
     m_callable.set(vm, this, callable);
 
 #if FFI_CALLBACK_THUNK_SUPPORTED
-    m_entryCode = FFI::generateCallbackThunk(vm, *this);
-    if (!m_entryCode) [[unlikely]]
+    auto entryCode = FFI::generateCallbackThunk(vm, *this);
+    if (!entryCode) [[unlikely]]
         return; // create() throws the OutOfMemoryError.
+    if (m_threadsafe)
+        m_threadsafe->setEntryCode(WTF::move(entryCode));
+    else
+        m_entryCode = WTF::move(entryCode);
 
     if (Options::verboseFFI()) [[unlikely]]
         dataLogLn("FFI: created JSFFICallback ", RawPointer(this), " signature ", m_signature->toString(), " entrypoint ", RawPointer(nativeEntrypoint()));
 
     constexpr unsigned attributes = static_cast<unsigned>(PropertyAttribute::ReadOnly | PropertyAttribute::DontEnum | PropertyAttribute::DontDelete);
     putDirect(vm, Identifier::fromString(vm, "ptr"_s), FFI::pointerToJSValue(globalObject(), reinterpret_cast<uint64_t>(nativeEntrypoint())), attributes);
-    putDirect(vm, Identifier::fromString(vm, "threadsafe"_s), jsBoolean(m_threadsafe), attributes);
+    putDirect(vm, Identifier::fromString(vm, "threadsafe"_s), jsBoolean(isThreadsafe()), attributes);
 #endif
 }
 
+const MacroAssemblerCodeRef<JITThunkPtrTag>& JSFFICallback::entryCode() const
+{
+    return m_threadsafe ? m_threadsafe->entryCode() : m_entryCode;
+}
+
 void* JSFFICallback::nativeEntrypoint() const
 {
-    return untagCFunctionPtr<void*, JITThunkPtrTag>(m_entryCode.code().taggedPtr());
+    return untagCFunctionPtr<void*, JITThunkPtrTag>(entryCode().code().taggedPtr());
 }
 
 const char* JSFFICallback::setReturnCString(const CString& string)
@@ -159,7 +186,8 @@ void JSFFICallback::close()
         return;
     m_closed = true;
     VM& vm = this->vm();
-    if (!markClosedAndReportPending())
+    // A threadsafe callback with invocations still queued stays rooted until the last of them has run.
+    if (!m_threadsafe || !m_threadsafe->markClosedAndReportPending())
         unroot();
     constexpr unsigned attributes = static_cast<unsigned>(PropertyAttribute::ReadOnly | PropertyAttribute::DontEnum | PropertyAttribute::DontDelete);
     putDirect(vm, Identifier::fromString(vm, "ptr"_s), jsNull(), attributes);
@@ -169,6 +197,8 @@ void JSFFICallback::unroot()
 {
     if (auto* globalObject = this->globalObject())
         globalObject->ffiContext().removeLiveCallback(*globalObject, this);
+    if (m_threadsafe)
+        vm().heap.unprotect(this);
 }
 
 template<typename Visitor>