Fixes for shutting down during async operations

Core/AppRuntime/Source/AppRuntime.cpp

@@ -12,13 +12,21 @@ namespace Babylon
         : m_workQueue{std::make_unique<WorkQueue>([this] { RunPlatformTier(); })}
         , m_unhandledExceptionHandler{unhandledExceptionHandler}
     {
-        Dispatch([this](Napi::Env env) {
+        Dispatch([this](Napi::Env env)
+        {
             JsRuntime::CreateForJavaScript(env, [this](auto func) { Dispatch(std::move(func)); });
         });
     }
 
     AppRuntime::~AppRuntime()
     {
+        // Notify the JsRuntime on the JavaScript thread that the JavaScript
+        // runtime shutdown sequence has begun. The JsRuntimeScheduler will
+        // use this signal to gracefully cancel asynchronous operations.
+        Dispatch([](Napi::Env env)
+        {
+            JsRuntime::NotifyDisposing(JsRuntime::GetFromJavaScript(env));
+        });
     }
 
     void AppRuntime::Run(Napi::Env env)
@@ -38,8 +46,10 @@ namespace Babylon
 
     void AppRuntime::Dispatch(Dispatchable<void(Napi::Env)> func)
     {
-        m_workQueue->Append([this, func{std::move(func)}](Napi::Env env) mutable {
-            Execute([this, env, func{std::move(func)}]() mutable {
+        m_workQueue->Append([this, func{std::move(func)}](Napi::Env env) mutable
+        {
+            Execute([this, env, func{std::move(func)}]() mutable
+            {
                 try
                 {
                     func(env);

Core/AppRuntime/Source/WorkQueue.cpp

@@ -14,8 +14,12 @@ namespace Babylon
             Resume();
         }
 
-        m_cancelSource.cancel();
-        m_dispatcher.cancelled();
+        // Dispatch a cancel to signal the Run function to gracefully end.
+        // It must be dispatched and not canceled directly to ensure that
+        // existing work is executed and executed in the correct order.
+        m_dispatcher([this]() {
+            m_cancellationSource.cancel();
+        });
 
         m_thread.join();
     }
@@ -39,11 +43,15 @@ namespace Babylon
         m_env = std::make_optional(env);
         m_dispatcher.set_affinity(std::this_thread::get_id());
 
-        while (!m_cancelSource.cancelled())
+        while (!m_cancellationSource.cancelled())
         {
-            m_dispatcher.blocking_tick(m_cancelSource);
+            m_dispatcher.blocking_tick(m_cancellationSource);
         }
 
-        m_dispatcher.clear();
+        // Drain the queue to complete work dispatched after cancellation.
+        m_dispatcher.tick(arcana::cancellation::none());
+
+        // There should no longer be any outstanding work once the queue is drained.
+        assert(m_dispatcher.empty());
     }
 }

Core/AppRuntime/Source/WorkQueue.h

@@ -19,6 +19,12 @@ namespace Babylon
         template<typename CallableT>
         void Append(CallableT callable)
         {
+            if (m_cancellationSource.cancelled())
+            {
+                // There is likely a coding error if this exception is thrown.
+                throw std::runtime_error{"Cannot append to the work queue while shutting down"};
+            }
+
             // Manual dispatcher queueing requires a copyable CallableT, we use a shared pointer trick to make a
             // copyable callable if necessary.
             if constexpr (std::is_copy_constructible<CallableT>::value)
@@ -46,9 +52,9 @@ namespace Babylon
 
         std::optional<std::scoped_lock<std::mutex>> m_suspensionLock{};
 
-        arcana::cancellation_source m_cancelSource{};
+        arcana::cancellation_source m_cancellationSource{};
         arcana::manual_dispatcher<128> m_dispatcher{};
 
-        std::thread m_thread;
+        std::thread m_thread{};
     };
 }

Core/Graphics/InternalInclude/Babylon/Graphics/DeviceContext.h

@@ -6,6 +6,7 @@
 
 #include <napi/env.h>
 
+#include <bx/allocator.h>
 #include <bgfx/bgfx.h>
 #include <bgfx/platform.h>
 
@@ -94,7 +95,7 @@ namespace Babylon::Graphics
 
         Update GetUpdate(const char* updateName);
 
-        void RequestScreenShot(std::function<void(std::vector<uint8_t>)> callback);
+        arcana::task<std::vector<uint8_t>, std::exception_ptr> RequestScreenShotAsync();
 
         arcana::task<void, std::exception_ptr> ReadTextureAsync(bgfx::TextureHandle handle, gsl::span<uint8_t> data, uint8_t mipLevel = 0);
 
@@ -115,12 +116,16 @@ namespace Babylon::Graphics
         void RemoveTexture(bgfx::TextureHandle handle);
         TextureInfo GetTextureInfo(bgfx::TextureHandle handle);
 
+        bx::AllocatorI& Allocator() { return m_allocator; }
+
     private:
         friend UpdateToken;
 
         DeviceImpl& m_graphicsImpl;
 
         std::unordered_map<uint16_t, TextureInfo> m_textureHandleToInfo{};
         std::mutex m_textureHandleToInfoMutex{};
+
+        bx::DefaultAllocator m_allocator{};
     };
 }

Core/Graphics/Source/DeviceContext.cpp

@@ -56,9 +56,9 @@ namespace Babylon::Graphics
         return {m_graphicsImpl.GetSafeTimespanGuarantor(updateName), *this};
     }
 
-    void DeviceContext::RequestScreenShot(std::function<void(std::vector<uint8_t>)> callback)
+    arcana::task<std::vector<uint8_t>, std::exception_ptr> DeviceContext::RequestScreenShotAsync()
     {
-        return m_graphicsImpl.RequestScreenShot(std::move(callback));
+        return m_graphicsImpl.RequestScreenShotAsync();
     }
 
     arcana::task<void, std::exception_ptr> DeviceContext::ReadTextureAsync(bgfx::TextureHandle handle, gsl::span<uint8_t> data, uint8_t mipLevel)

Core/Graphics/Source/DeviceImpl.cpp

@@ -272,9 +272,16 @@ namespace Babylon::Graphics
         m_bgfxCallback.SetDiagnosticOutput(std::move(diagnosticOutput));
     }
 
-    void DeviceImpl::RequestScreenShot(std::function<void(std::vector<uint8_t>)> callback)
+    arcana::task<std::vector<uint8_t>, std::exception_ptr> DeviceImpl::RequestScreenShotAsync()
     {
-        m_screenShotCallbacks.push(std::move(callback));
+        arcana::task_completion_source<std::vector<uint8_t>, std::exception_ptr> taskCompletionSource{};
+
+        m_screenShotCallbacks.push([taskCompletionSource](std::vector<uint8_t> bytes) mutable
+        {
+            taskCompletionSource.complete(std::move(bytes));
+        });
+
+        return taskCompletionSource.as_task();
     }
 
     arcana::task<void, std::exception_ptr> DeviceImpl::ReadTextureAsync(bgfx::TextureHandle handle, gsl::span<uint8_t> data, uint8_t mipLevel)

Core/Graphics/Source/DeviceImpl.h

@@ -66,7 +66,7 @@ namespace Babylon::Graphics
 
         Update GetUpdate(const char* updateName);
 
-        void RequestScreenShot(std::function<void(std::vector<uint8_t>)> callback);
+        arcana::task<std::vector<uint8_t>, std::exception_ptr> RequestScreenShotAsync();
 
         arcana::task<void, std::exception_ptr> ReadTextureAsync(bgfx::TextureHandle handle, gsl::span<uint8_t> data, uint8_t mipLevel);
 

Core/JsRuntime/Include/Babylon/JsRuntime.h

@@ -4,6 +4,7 @@
 
 #include <functional>
 #include <mutex>
+#include <vector>
 
 namespace Babylon
 {
@@ -22,31 +23,49 @@ namespace Babylon
             }
         };
 
-        struct InternalState;
-        friend struct InternalState;
+        JsRuntime(const JsRuntime&) = delete;
+        JsRuntime& operator=(const JsRuntime&) = delete;
 
         // Any JavaScript errors that occur will bubble up as a Napi::Error C++ exception.
         // JsRuntime expects the provided dispatch function to handle this exception,
         // such as with a try/catch and logging the exception message.
         using DispatchFunctionT = std::function<void(std::function<void(Napi::Env)>)>;
 
+        // Creates the JsRuntime object owned by the JavaScript environment.
         // Note: It is the contract of JsRuntime that its dispatch function must be usable
         // at the moment of construction. JsRuntime cannot be built with dispatch function
         // that captures a reference to a not-yet-completed object that will be completed
         // later -- an instance of an inheriting type, for example. The dispatch function
         // must be safely callable as soon as it is passed to the JsRuntime constructor.
         static JsRuntime& CreateForJavaScript(Napi::Env, DispatchFunctionT);
+
+        // Gets the JsRuntime from the given N-API environment.
         static JsRuntime& GetFromJavaScript(Napi::Env);
-        void Dispatch(std::function<void(Napi::Env)>);
 
-    protected:
-        JsRuntime(const JsRuntime&) = delete;
-        JsRuntime& operator=(const JsRuntime&) = delete;
+        struct IDisposingCallback
+        {
+            virtual void Disposing() = 0;
+        };
+
+        // Notifies the JsRuntime that the JavaScript environment will begin shutting down.
+        // Calling this function will signal callbacks registered with RegisterDisposing.
+        static void NotifyDisposing(JsRuntime&);
+
+        // Registers a callback for when the JavaScript environment will begin shutting down.
+        static void RegisterDisposing(JsRuntime&, IDisposingCallback*);
+
+        // Unregisters a callback for when the JavaScript environment will begin shutting down.
+        static void UnregisterDisposing(JsRuntime&, IDisposingCallback*);
+
+        // Dispatches work onto the JavaScript thread and provides access to the N-API
+        // environment.
+        void Dispatch(std::function<void(Napi::Env)>);
 
     private:
         JsRuntime(Napi::Env, DispatchFunctionT);
 
-        DispatchFunctionT m_dispatchFunction{};
         std::mutex m_mutex{};
+        DispatchFunctionT m_dispatchFunction{};
+        std::vector<IDisposingCallback*> m_disposingCallbacks{};
     };
 }

Core/JsRuntime/Include/Babylon/JsRuntimeScheduler.h

@@ -1,30 +1,146 @@
 #pragma once
 
 #include "JsRuntime.h"
+#include <arcana/threading/dispatcher.h>
+#include <atomic>
+#include <cassert>
 
 namespace Babylon
 {
-    /**
-     * Scheduler that invokes continuations via JsRuntime::Dispatch.
-     * Intended to be consumed by arcana.cpp tasks.
-     */
-    class JsRuntimeScheduler
+    // This class encapsulates a coding pattern for invoking continuations on the JavaScript thread while properly
+    // handling garbage collection and shutdown scenarios.  This class provides and manages the schedulers intended
+    // for a N-API object to use with arcana tasks.  It is different than the typical scheduler as this class itself
+    // is not a scheduler directly, but instead hands out scheduler via its `Get()` function.  It provides special
+    // handling for when the JsRuntime begins shutting down, i.e., when JsRuntime::NotifyDisposing is called.
+    //   1. Calling `Rundown` blocks execution until all outstanding schedulers are invoked on the JavaScript thread.
+    //   2. Once the JsRuntime begins shutting down, all schedulers will reroute its dispatch calls from the
+    //      JsRuntime to a separate dispatcher owned by the JsRuntimeScheduler itself.  This class will then be able
+    //      to pump this dispatcher in its destructor to prevent deadlocks.
+    //
+    // The typical pattern for an arcana task will look something like this:
+    //   class MyClass
+    //   {
+    //   public:
+    //       ~MyClass()
+    //       {
+    //           m_cancellationSource.cancel();
+    //
+    //           // Wait for asynchronous operations to complete.
+    //           m_runtimeScheduler.Rundown();
+    //       }
+    // 
+    //       void MyFunction(const Napi::CallbackInfo& info)
+    //       {
+    //           const auto callback{info[0].As<Napi::Function>()};
+    //           
+    //           arcana::make_task(arcana::threadpool_scheduler, m_cancellationSource, []() {
+    //               // do some asynchronous work
+    //           }).then(m_runtimeScheduler.Get(), m_cancellationSource, [thisRef = Napi::Persistent(info.This()), callback = Napi::Persistent(callback)]() {
+    //               callback.Call({});
+    //           });
+    //       }
+    //
+    //   private:
+    //       arcana::cancellation_source m_cancellationSource;
+    //       JsRuntimeScheduler m_runtimeScheduler;
+    //   };
+    //
+    // **IMPORTANT**:
+    //   1. To prevent continuations from accessing freed memory, the destructor of the N-API class is expected to call
+    //      `Rundown()` which blocks execution until all of its schedulers are invoked. Failing to do so will result in
+    //      an assert if there are outstanding schedulers not yet invoked.
+    //   2. The last continuation that accesses members of the N-API object, including the cancellation associated with
+    //      the continuation, must capture a persistent reference to the N-API object itself to prevent the GC from
+    //      collecting the N-API object during the asynchronous operation. Failing to do so will result in a hang
+    //      when `Rundown()` is called in the N-API class destructor.
+    class JsRuntimeScheduler : public JsRuntime::IDisposingCallback
     {
     public:
         explicit JsRuntimeScheduler(JsRuntime& runtime)
-            : m_runtime{runtime}
+            : m_runtime{&runtime}
+            , m_scheduler{*this}
         {
+            std::scoped_lock lock{m_mutex};
+            JsRuntime::RegisterDisposing(*m_runtime, this);
         }
 
-        template<typename CallableT>
-        void operator()(CallableT&& callable) const
+        ~JsRuntimeScheduler()
+        {
+            assert(m_count == 0 && "Schedulers for the JavaScript thread are not yet invoked");
+
+            std::scoped_lock lock{m_mutex};
+            if (m_runtime != nullptr)
+            {
+                JsRuntime::UnregisterDisposing(*m_runtime, this);
+            }
+        }
+
+        // Wait until all of the schedulers are invoked.
+        void Rundown()
         {
-            m_runtime.Dispatch([callable{std::forward<CallableT>(callable)}](Napi::Env) {
-                callable();
-            });
+            while (m_count > 0)
+            {
+                m_disposingDispatcher.blocking_tick(arcana::cancellation::none());
+            }
+        }
+
+        // Get a scheduler to invoke continuations on the JavaScript thread.
+        const auto& Get()
+        {
+            ++m_count;
+            return m_scheduler;
         }
 
     private:
-        JsRuntime& m_runtime;
+        void Disposing() override
+        {
+            std::scoped_lock lock{m_mutex};
+            m_runtime = nullptr;
+        }
+
+        class SchedulerImpl
+        {
+        public:
+            explicit SchedulerImpl(JsRuntimeScheduler& parent) : m_parent{parent}
+            {
+            }
+
+            template<typename CallableT>
+            void operator()(CallableT&& callable) const
+            {
+                m_parent.Dispatch(callable);
+            }
+
+        private:
+            JsRuntimeScheduler& m_parent;
+        };
+
+        template<typename CallableT>
+        void Dispatch(CallableT&& callable)
+        {
+            std::scoped_lock lock{m_mutex};
+
+            if (m_runtime != nullptr)
+            {
+                m_runtime->Dispatch([callable{std::forward<CallableT>(callable)}](Napi::Env)
+                {
+                    callable();
+                });
+            }
+            else
+            {
+                m_disposingDispatcher(callable);
+            }
+
+            --m_count;
+        }
+
+        mutable std::mutex m_mutex{};
+        JsRuntime* m_runtime{};
+        std::function<void()> m_disposingCallback{};
+
+        SchedulerImpl m_scheduler;
+        std::atomic<int> m_count{0};
+        arcana::manual_dispatcher<128> m_disposingDispatcher{};
     };
 }