New Lanes finalizer API lanes.finally()

5 files changed, 45 insertions, 2 deletions

diff --git a/src/lanes.cpp b/src/lanes.cpp
index 0ea0900..109fba4 100644
--- a/src/lanes.cpp
+++ b/src/lanes.cpp
@@ -598,6 +598,7 @@ namespace {
             { "set_thread_affinity", LG_set_thread_affinity },
             { "nameof", luaG_nameof },
             { "register", LG_register },
+            { Universe::kFinally, Universe::InitializeFinalizer },
             { "set_singlethreaded", LG_set_singlethreaded },
             { nullptr, nullptr }
         };
diff --git a/src/lanes.lua b/src/lanes.lua
index e75e840..6616667 100644
--- a/src/lanes.lua
+++ b/src/lanes.lua
@@ -799,6 +799,7 @@ local configure = function(settings_)
     lanes.null = core.null
     lanes.require = core.require
     lanes.register = core.register
+    lanes.finally = core.finally
     lanes.set_singlethreaded = core.set_singlethreaded
     lanes.set_thread_affinity = core.set_thread_affinity
     lanes.set_thread_priority = core.set_thread_priority
diff --git a/src/state.cpp b/src/state.cpp
index 5a1d2fb..d9f5499 100644
--- a/src/state.cpp
+++ b/src/state.cpp
@@ -209,7 +209,7 @@ static void copy_one_time_settings(Universe* U_, SourceState L1_, DestState L2_)
 
 // #################################################################################################
 
-static constexpr char const* kOnStateCreate{ "on_state_create" };
+static constexpr char const* kOnStateCreate{ "on_state_create" }; // update lanes.lua if the name changes!
 
 void InitializeOnStateCreate(Universe* U_, lua_State* L_)
 {
diff --git a/src/universe.cpp b/src/universe.cpp
index 5794048..3271e70 100644
--- a/src/universe.cpp
+++ b/src/universe.cpp
@@ -197,6 +197,28 @@ void Universe::initializeAllocatorFunction(lua_State* L_)
 
 // #################################################################################################
 
+// should be called ONLY from the state that created the universe
+int Universe::InitializeFinalizer(lua_State* L_)
+{
+    luaL_argcheck(L_, lua_gettop(L_) <= 1, 1, "too many arguments");                               // L_: f?
+    lua_settop(L_, 1);                                                                             // L_: f|nil
+    luaL_argcheck(L_, lua_isnoneornil(L_, 1) || lua_isfunction(L_, 1), 1, "finalizer should be a function");
+
+    // make sure we are only called from the Master Lua State!
+    kUniverseFullRegKey.pushValue(L_);                                                             // L_: f U
+    if (lua_type_as_enum(L_, -1) != LuaType::USERDATA) {
+        raise_luaL_error(L_, "lanes.%s called from inside a lane", kFinally);
+    }
+    lua_pop(L_, 1);                                                                                // L_: f
+    STACK_GROW(L_, 3);
+    // _R[kFinalizerRegKey] = f
+    kFinalizerRegKey.setValue(L_, [](lua_State* L_) { lua_insert(L_, -2); });                      // L_:
+    // no need to adjust the stack, Lua does this for us
+    return 0;
+}
+
+// #################################################################################################
+
 /*
  * Initialize keeper states
  *
@@ -373,9 +395,22 @@ int universe_gc(lua_State* L_)
 {
     lua_Duration const _shutdown_timeout{ lua_tonumber(L_, lua_upvalueindex(1)) };
     std::string_view const _op_string{ lua_tostringview(L_, lua_upvalueindex(2)) };
-    Universe* const _U{ lua_tofulluserdata<Universe>(L_, 1) };
+    STACK_CHECK_START_ABS(L_, 1);
+    Universe* const _U{ lua_tofulluserdata<Universe>(L_, 1) };                                     // L_: U
     _U->terminateFreeRunningLanes(L_, _shutdown_timeout, which_cancel_op(_op_string));
 
+    // invoke the function installed by lanes.finally()
+    kFinalizerRegKey.pushValue(L_);                                                                // L_: U finalizer|nil
+    if (!lua_isnil(L_, -1)) {
+        lua_pcall(L_, 0, 0, 0);                                                                    // L_: U
+        // discard any error that might have occured
+        lua_settop(L_, 1);
+    } else {
+        lua_pop(L_, 1);                                                                            // L_: U
+    }
+    // in case of error, the message is pushed on the stack
+    STACK_CHECK(L_, 1);
+
     // no need to mutex-protect this as all threads in the universe are gone at that point
     if (_U->timerLinda != nullptr) { // test in case some early internal error prevented Lanes from creating the deep timer
         [[maybe_unused]] int const _prev_ref_count{ _U->timerLinda->refcount.fetch_sub(1, std::memory_order_relaxed) };
@@ -389,6 +424,9 @@ int universe_gc(lua_State* L_)
     // remove the protected allocator, if any
     _U->protectedAllocator.removeFrom(L_);
 
+    // no longer found in the registry
+    kUniverseFullRegKey.setValue(L_, [](lua_State* L_) { lua_pushnil(L_); });
+    kUniverseLightRegKey.setValue(L_, [](lua_State* L_) { lua_pushnil(L_); });
     _U->Universe::~Universe();
 
     return 0;
diff --git a/src/universe.h b/src/universe.h
index f5b31a3..4be6a9a 100644
--- a/src/universe.h
+++ b/src/universe.h
@@ -120,6 +120,8 @@ static constexpr RegistryUniqueKey kUniverseLightRegKey{ 0x48BBE9CEAB0BA04Full }
 class Universe
 {
     public:
+    static constexpr char const* kFinally{ "finally" }; // update lanes.lua if the name changes!
+
 #ifdef PLATFORM_LINUX
     // Linux needs to check, whether it's been run as root
     bool const sudo{ geteuid() == 0 };
@@ -180,6 +182,7 @@ class Universe
 
     void closeKeepers();
     void initializeAllocatorFunction(lua_State* L_);
+    static int InitializeFinalizer(lua_State* L_);
     void initializeKeepers(lua_State* L_);
     void terminateFreeRunningLanes(lua_State* L_, lua_Duration shutdownTimeout_, CancelOp op_);
 };