WIP: Fixes for shutting down during async operations

Author: bghgary

Core/AppRuntime/Source/AppRuntime.cpp

@@ -19,6 +19,13 @@ namespace Babylon
 
     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)

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_cancelSource.cancel();
+        });
 
         m_thread.join();
     }
@@ -44,6 +48,8 @@ namespace Babylon
             m_dispatcher.blocking_tick(m_cancelSource);
         }
 
-        m_dispatcher.clear();
+        // There should not be any outstanding work during the shutdown sequence
+        // which should be the only way exit the while loop above.
+        assert(m_dispatcher.empty());
     }
 }

Core/AppRuntime/Source/WorkQueue.h

@@ -19,6 +19,12 @@ namespace Babylon
         template<typename CallableT>
         void Append(CallableT callable)
         {
+            if (m_cancelSource.cancelled())
+            {
+                // There is likely a coding error if this exception is thrown.
+                throw std::runtime_error{"Cannot append to the work queue after it is canceled"};
+            }
+
             // 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)

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);
 
@@ -114,12 +115,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

@@ -271,9 +271,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,41 @@ 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;
+        // 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&, std::function<void()>);
+
+        // 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<std::function<void()>> m_disposingCallbacks{};
     };
 }

Core/JsRuntime/Include/Babylon/JsRuntimeScheduler.h

@@ -1,30 +1,130 @@
 #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.
-     */
+    // 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. The destructor blocks if there are outstanding schedulers not yet 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:
+    //       void MyFunction(const Napi::CallbackInfo& info);
+    //
+    //   private:
+    //       arcana::cancellation_source m_cancellationSource;
+    // 
+    //       // Put this last so that it gets destructed first.
+    //       JsRuntimeScheduler m_runtimeScheduler;
+    //   };
+    //
+    //   void MyClass::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_cancelSource, [thisRef = Napi::Persistent(info.This()), callback = Napi::Persistent(callback)]() {
+    //       {
+    //           callback.Call({});
+    //       });
+    //    }
+    //
+    // **IMPORTANT**:
+    //   1. To prevent continuations from accessing destructed objects, declare the JsRuntimeScheduler at the end of
+    //      the N-API class.  The destructor of the JsRuntimeScheduler will call `Rundown()` which will block until
+    //      all of its schedulers are invoked.  If this is not possible, call `Rundown()` manually in the destructor
+    //      of the N-API class.
+    //   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 GC from collecting
+    //      the N-API object during the asynchronous operation.
     class JsRuntimeScheduler
     {
     public:
         explicit JsRuntimeScheduler(JsRuntime& runtime)
-            : m_runtime{runtime}
+            : m_runtime{&runtime}
+            , m_scheduler{*this}
         {
+            JsRuntime::RegisterDisposing(*m_runtime, [this]()
+            {
+                m_runtime = nullptr;
+            });
         }
 
-        template<typename CallableT>
-        void operator()(CallableT&& callable) const
+        ~JsRuntimeScheduler()
         {
-            m_runtime.Dispatch([callable{std::forward<CallableT>(callable)}](Napi::Env){
-                callable();
-            });
+            Rundown();
+        }
+
+        // Wait until all of the schedulers are invoked.
+        void Rundown()
+        {
+            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;
+        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)
+        {
+            if (m_runtime != nullptr)
+            {
+                m_runtime->Dispatch([callable{std::forward<CallableT>(callable)}](Napi::Env)
+                {
+                    callable();
+                });
+            }
+            else
+            {
+                m_disposingDispatcher(callable);
+            }
+
+            --m_count;
+        }
+
+        JsRuntime* m_runtime;
+        SchedulerImpl m_scheduler;
+        std::atomic<int> m_count{0};
+        arcana::manual_dispatcher<128> m_disposingDispatcher{};
     };
 }

Core/JsRuntime/Source/JsRuntime.cpp

@@ -40,6 +40,20 @@ namespace Babylon
                     .Data();
     }
 
+    void JsRuntime::NotifyDisposing(JsRuntime& runtime)
+    {
+        auto callbacks = std::move(runtime.m_disposingCallbacks);
+        for (const auto& callback : callbacks)
+        {
+            callback();
+        }
+    }
+
+    void JsRuntime::RegisterDisposing(JsRuntime& runtime, std::function<void()> callback)
+    {
+        runtime.m_disposingCallbacks.push_back(std::move(callback));
+    }
+
     void JsRuntime::Dispatch(std::function<void(Napi::Env)> function)
     {
         std::scoped_lock lock{m_mutex};