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();
     }
@@ -36,14 +40,24 @@ namespace Babylon
 
     void WorkQueue::Run(Napi::Env env)
     {
-        m_env = std::make_optional(env);
         m_dispatcher.set_affinity(std::this_thread::get_id());
 
-        while (!m_cancelSource.cancelled())
+        m_env = std::make_optional(env);
+
+        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());
+
+        // Clear the shutdown queue to make sure the callables are destroyed on this thread.
+        m_shutdownQueue.clear();
+
+        m_env.reset();
     }
 }

Core/AppRuntime/Source/WorkQueue.h

@@ -23,11 +23,29 @@ namespace Babylon
             // copyable callable if necessary.
             if constexpr (std::is_copy_constructible<CallableT>::value)
             {
-                m_dispatcher.queue([this, callable = std::move(callable)]() { Invoke(callable); });
+                if (m_cancellationSource.cancelled())
+                {
+                    m_shutdownQueue.push([callable = std::move(callable)] {});
+                }
+                else
+                {
+                    m_dispatcher.queue([this, callable = std::move(callable)]() {
+                        callable(m_env.value());
+                    });
+                }
             }
             else
             {
-                m_dispatcher.queue([this, callablePtr = std::make_shared<CallableT>(std::move(callable))]() { Invoke(*callablePtr); });
+                if (m_cancellationSource.cancelled())
+                {
+                    m_shutdownQueue.push([callablePtr = std::make_shared<CallableT>(std::move(callable))] {});
+                }
+                else
+                {
+                    m_dispatcher.queue([this, callablePtr = std::make_shared<CallableT>(std::move(callable))]() {
+                        (*callablePtr)(m_env.value());
+                    });
+                }
             }
         }
 
@@ -36,19 +54,18 @@ namespace Babylon
         void Run(Napi::Env);
 
     private:
-        template<typename CallableT>
-        void Invoke(CallableT& callable)
-        {
-            callable(m_env.value());
-        }
-
         std::optional<Napi::Env> m_env{};
 
         std::optional<std::scoped_lock<std::mutex>> m_suspensionLock{};
 
-        arcana::cancellation_source m_cancelSource{};
-        arcana::manual_dispatcher<128> m_dispatcher{};
+        arcana::cancellation_source m_cancellationSource{};
+
+        using DispatcherT = arcana::manual_dispatcher<128>;
+        DispatcherT m_dispatcher{};
+
+        // Put the callables in a separate queue during shutdown to ensure the callables are destroyed on the right thread.
+        arcana::blocking_concurrent_queue<DispatcherT::callback_t> m_shutdownQueue{};
 
-        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{};
     };
 }