From 6b46f3750cd773369e47c56ba8d2dce2c5a9ff64 Mon Sep 17 00:00:00 2001
From: Delio Vicini <vicini@google.com>
Date: Tue, 3 Sep 2024 16:34:08 +0200
Subject: [PATCH] Ensure that Logger and FileResolver are initialized for
 external threads

---
 src/core/tests/test_thread.py | 52 +++++++++++++++++++++++++++++++++--
 src/core/thread.cpp           |  9 ++++++
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/src/core/tests/test_thread.py b/src/core/tests/test_thread.py
index 62487a73b..bc082ba04 100644
--- a/src/core/tests/test_thread.py
+++ b/src/core/tests/test_thread.py
@@ -2,12 +2,58 @@
 import mitsuba as mi
 from threading import Thread
 
+
 def test01_use_scoped_thread_environment(variant_scalar_rgb):
-    def use_scoped_set_thred_env(env):
+    def use_scoped_set_thread_env(env):
         with mi.ScopedSetThreadEnvironment(environment):
-            mi.Log(mi.LogLevel.Info, "Log from a thread environment.")
+            mi.Log(mi.LogLevel.Info, 'Log from a thread environment.')
 
     environment = mi.ThreadEnvironment()
-    thread = Thread(target = use_scoped_set_thred_env, args = (environment, ))
+    thread = Thread(target=use_scoped_set_thread_env, args=(environment, ))
+    thread.start()
+    thread.join()
+
+
+def test02_log_from_new_thread(variant_scalar_rgb, tmp_path):
+
+    # We use a StreamAppender to capture the output.
+    log_path = tmp_path / 'log.txt'
+    appender = mi.StreamAppender(str(log_path))
+
+    log_str = 'Log from a thread environment.'
+
+    def print_to_log():
+        logger = mi.Thread.thread().logger()
+        assert logger is not None
+        logger.add_appender(appender)
+        mi.set_log_level(mi.LogLevel.Info)
+        mi.Log(mi.LogLevel.Info, log_str)
+
+    thread = Thread(target=print_to_log)
     thread.start()
     thread.join()
+
+    log = appender.read_log()
+    assert 'print_to_log' in log
+    assert log_str in log
+
+
+def test03_access_file_resolver_from_new_thread(variant_scalar_rgb):
+    def access_fresolver():
+        fs = mi.Thread.thread().file_resolver()
+        assert fs is not None
+        fs.resolve('./some_path')
+        n_paths = len(fs)
+        fs.prepend('./some_folder')
+        fs.prepend('./some_folder2')
+        assert len(fs) == n_paths + 2
+
+    fs = mi.Thread.thread().file_resolver()
+    n_paths = len(fs)
+
+    thread = Thread(target=access_fresolver)
+    thread.start()
+    thread.join()
+
+    # The original file resolver remains unchanged.
+    assert len(fs) == n_paths
diff --git a/src/core/thread.cpp b/src/core/thread.cpp
index 53747dc1f..09e3c83bf 100644
--- a/src/core/thread.cpp
+++ b/src/core/thread.cpp
@@ -23,6 +23,8 @@
 NAMESPACE_BEGIN(mitsuba)
 
 static size_t global_thread_count = 0;
+
+static ref<Thread> main_thread = nullptr;
 static thread_local ref<Thread> self = nullptr;
 static std::atomic<uint32_t> thread_ctr { 0 };
 #if defined(__linux__) || defined(__APPLE__)
@@ -476,6 +478,11 @@ bool Thread::register_external_thread(const std::string &prefix) {
     self->d->running = true;
     self->d->external_thread = true;
 
+    // An external thread will re-use the main thread's Logger (thread safe) 
+    // and create a new FileResolver (since the FileResolver is not thread safe).
+    self->d->logger = main_thread->d->logger;
+    self->d->fresolver = new FileResolver();
+
     const std::string &thread_name = self->name();
     #if defined(__linux__)
         pthread_setname_np(pthread_self(), thread_name.c_str());
@@ -520,6 +527,7 @@ void Thread::static_initialization() {
     self = new MainThread();
     self->d->running = true;
     self->d->fresolver = new FileResolver();
+    main_thread = self;
 }
 
 void Thread::static_shutdown() {
@@ -529,6 +537,7 @@ void Thread::static_shutdown() {
 
     thread()->d->running = false;
     self = nullptr;
+    main_thread = nullptr;
     #if defined(__linux__) || defined(__APPLE__)
         pthread_key_delete(this_thread_id);
     #endif