From bdaa1ffef440941df9c62e8cb6fb1fef23b7f080 Mon Sep 17 00:00:00 2001
From: Morten Piibeleht <morten.piibeleht@juliahub.com>
Date: Thu, 12 Oct 2023 13:39:33 +1300
Subject: [PATCH] feat: add way to register post-__init__ callbacks (#66)

(cherry picked from commit 7ef6c567f9a7aec33be1d7a704694735ef48ee4c)
---
 src/DataSets.jl  | 50 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl |  9 ++++++++-
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/DataSets.jl b/src/DataSets.jl
index 5214879..46bbe83 100644
--- a/src/DataSets.jl
+++ b/src/DataSets.jl
@@ -173,6 +173,56 @@ function __init__()
                     =# project=proj exception=(exc,catch_backtrace())
             end
         end
+        # Call any post-__init__() callbacks that were registered before __init__() was called,
+        # or had chance to finish.
+        lock(_PROJECT_INIT_LOCK) do
+            _PROJECT_INITIALIZED[] = true
+            for f in _PROJECT_INIT_CALLBACKS
+                _invoke_init_cb(f)
+            end
+            # No need to keep the callbacks around, and maybe the GC can free some memory.
+            empty!(_PROJECT_INIT_CALLBACKS)
+        end
+    end
+end
+
+# The register_post_init_callback() can be used to add a callback that will get called
+# when DataSets.__init__() has run. Note: if f() throws an error, it does not cause a crash.
+#
+# This is useful for sysimages where the module is already be loaded (in Base.loaded_modules),
+# but __init__() has not been called yet. In particular, this means that other packages' __init__
+# functions can be sure that when they call initialization code that affects DataSets (in particular,
+# DataSets.PROJECT), then that code runs after __init__() has run.
+#
+# In the non-sysimage case, DataSets.__init__() would normally have already been called when
+# once register_post_init_callback() becomes available, and so in those situations, the callback
+# gets called immediately. However, in a system image, DataSets may have to queue up (FIFO) the
+# callback functions and wait until DataSets.__init__() has finished.
+#
+# Since the __init__() functions in sysimages can run in parallel, we use a lock just in case,
+# to make sure that two parallel calls would succeed.
+const _PROJECT_INIT_LOCK = ReentrantLock()
+const _PROJECT_INITIALIZED = Ref{Bool}(false)
+const _PROJECT_INIT_CALLBACKS = Base.Callable[]
+function register_post_init_callback(f::Base.Callable)
+    invoke = lock(_PROJECT_INIT_LOCK) do
+        if _PROJECT_INITIALIZED[]
+            return true
+        end
+        push!(_PROJECT_INIT_CALLBACKS, f)
+        return false
+    end
+    # We'll invoke outside of the lock, so that a long-running f() call
+    # wouldn't block other calls to register_post_init_callback()
+    invoke && _invoke_init_cb(f)
+    return nothing
+end
+
+function _invoke_init_cb(f::Base.Callable)
+    try
+        Base.invokelatest(f)
+    catch e
+        @error "Failed to run init callback: $f" exception = (e, catch_backtrace())
     end
 end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index b5c7b1b..4e34625 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,7 +7,14 @@ using ResourceContexts
 
 using DataSets: FileSystemRoot
 
-#-------------------------------------------------------------------------------
+@testset "register_post_init_callback" begin
+    init_was_called = Ref(false)
+    DataSets.register_post_init_callback() do
+        init_was_called[] = true
+    end
+    @test init_was_called[]
+end
+
 @testset "DataSet config" begin
     proj = DataSets.load_project("Data.toml")