A gc shared library registers the threads its host calls it on

tslang/lib/AsyncGCThreadsCommon.inc

@@ -2,14 +2,21 @@
 //
 // Shared by TypeScriptRuntime/AsyncGCThreads.cpp (the JIT's DLL) and
 // TypeScriptAsyncRuntime/AsyncGCThreads.cpp (the static library ahead-of-time builds link). In the
-// static library nothing but GC_enable_threads refers to this object, and only a `gc` program
-// calls that, so only a `gc` program pulls it in - and with it the collector's thread API.
+// static library nothing but GC_enable_threads and __tslang_gc_enter refer to this object, and
+// only a `gc` program or library calls those, so only they pull it in - and with it the
+// collector's thread API.
 
 #define GC_THREADS
 #include "gc.h"
 
 #include "TypeScript/AsyncGCThreads.h"
 
+#include <mutex>
+
+#ifndef _WIN32
+#include <pthread.h>
+#endif
+
 namespace
 {
 
@@ -24,8 +31,67 @@ void unregisterThisThread()
     GC_unregister_my_thread();
 }
 
+// A thread that came into a gc shared library through one of its exported functions, from a host
+// that is not a tslang program (an Android app's JNI threads, a C program's): checked once, and
+// registered then if the collector did not know it, until the thread ends.
+struct CallingThread
+{
+    bool checked = false;
+    // registered by __tslang_gc_enter, so unregistered here, at thread exit
+    bool registeredHere = false;
+
+    ~CallingThread()
+    {
+        if (registeredHere)
+        {
+            GC_unregister_my_thread();
+        }
+    }
+};
+
+thread_local CallingThread callingThread;
+
+std::once_flag threadsEnabled;
+
 } // namespace
 
+extern "C" void GC_enable_threads();
+
+// Called first thing by every exported function of a `-mm=gc` shared library (the GC pass injects
+// it). The collector stops and scans only the threads registered with it, and aborts a collection
+// that a thread it does not know started ("Collecting from unknown thread"): a tslang program
+// registers the threads it makes, but a library's host makes its own. GC_init registered the one
+// it ran on; any other is registered on its first call.
+//
+// GC_enable_threads runs once here whatever ran before: GC_register_my_thread refuses every
+// thread until GC_allow_register_threads was called ("Threads explicit registering is not
+// previously enabled"), and the collector may well be initialized already without it - it
+// initializes itself on first use, so a library with no top-level code (no GC_init injected at
+// load) can find it running. Calling it again where GC_init's injection already did is harmless.
+extern "C" void __tslang_gc_enter()
+{
+    auto &thread = callingThread;
+    if (thread.checked)
+    {
+        return;
+    }
+
+    std::call_once(threadsEnabled, [] {
+        if (!GC_is_init_called())
+        {
+            GC_init();
+        }
+
+        GC_enable_threads();
+    });
+
+    thread.checked = true;
+    if (!GC_thread_is_registered())
+    {
+        thread.registeredHere = registerThisThread();
+    }
+}
+
 // Called once from the entry point: the GC pass injects the call beside GC_init, so it happens
 // only in a `gc` build. When linked into the shared TypeScriptRuntime, this symbol is also
 // exported under this name in TypeScriptRuntime.def for the JIT to resolve.

tslang/lib/TypeScript/GCPass.cpp

@@ -130,6 +130,11 @@ class GCPass : public mlir::PassWrapper<GCPass, ModulePass>
             }
         }
 
+        if (tsContext.compileOptions.isDLL)
+        {
+            injectThreadEnter(m);
+        }
+
         LLVM_DEBUG(llvm::dbgs() << "\n!! GCPass: AFTER DUMP: \n" << m << "\n";);
     }
 
@@ -313,6 +318,58 @@ class GCPass : public mlir::PassWrapper<GCPass, ModulePass>
         }
     }
 
+    static bool isExported(LLVM::LLVMFuncOp funcOp)
+    {
+        if (auto passthrough = funcOp->getAttrOfType<mlir::ArrayAttr>("passthrough"))
+        {
+            for (auto attr : passthrough)
+            {
+                if (auto name = dyn_cast<mlir::StringAttr>(attr))
+                {
+                    if (name.getValue() == "export" || name.getValue() == DLL_EXPORT)
+                    {
+                        return true;
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+
+    // A shared library's exported functions are where its host's threads come in, and a host
+    // that is not a tslang program never registered them with the collector: each such function
+    // first registers the thread calling it, once per thread (__tslang_gc_enter, in
+    // AsyncGCThreadsCommon.inc). The global constructors are left alone: GC_init, injected
+    // there, registers the loading thread itself.
+    void injectThreadEnter(mlir::ModuleOp m)
+    {
+        llvm::SmallVector<LLVM::LLVMFuncOp> exported;
+        m.walk([&](LLVM::LLVMFuncOp funcOp) {
+            if (!funcOp.getBody().empty() && isExported(funcOp) && !funcOp.getName().starts_with(MLIR_GCTORS))
+            {
+                exported.push_back(funcOp);
+            }
+        });
+
+        if (exported.empty())
+        {
+            return;
+        }
+
+        PatternRewriter rewriter(m.getContext());
+        TypeHelper th(rewriter.getContext());
+        for (auto funcOp : exported)
+        {
+            LLVMCodeHelper ch(funcOp, rewriter, nullptr, tsContext.compileOptions);
+            auto enterFuncOp = ch.getOrInsertFunction(
+                "__tslang_gc_enter", th.getFunctionType(th.getVoidType(), mlir::ArrayRef<mlir::Type>{}));
+
+            rewriter.setInsertionPointToStart(&*funcOp.getBody().begin());
+            rewriter.create<LLVM::CallOp>(funcOp->getLoc(), enterFuncOp, ValueRange{});
+        }
+    }
+
     void injectInit(LLVM::LLVMFuncOp funcOp)
     {
         PatternRewriter rewriter(funcOp.getContext());

tslang/test/tester/CMakeLists.txt

@@ -2429,6 +2429,30 @@ foreach(android_mode codegen link)
     set_tests_properties(test-compile-android-${android_mode} PROPERTIES SKIP_REGULAR_EXPRESSION "SKIPPED:")
 endforeach()
 
+# A gc shared library called from a non-tslang host's own threads - see foreign-thread-gc.cmake.
+# Windows needs gc.dll (a shared library takes its collector from it); ELF links libgc.a in.
+if (NOT WIN32 OR TSLANG_HAVE_SHARED_GC)
+    add_executable(foreign-thread-gc-host foreign-thread-gc-host.cpp)
+    if (NOT WIN32)
+        target_link_libraries(foreign-thread-gc-host PRIVATE ${CMAKE_DL_LIBS} pthread)
+    endif()
+    if (WIN32)
+        set(gc_foreign_thread_gc_shared "-DGC_SHARED_LIB=${TEST_GC_SHARED_PREFIX}/lib")
+    else()
+        set(gc_foreign_thread_gc_shared "")
+    endif()
+    add_test(NAME test-compile-foreign-thread-gc
+             COMMAND ${CMAKE_COMMAND}
+                     "-DTSLANG=$<TARGET_FILE:tslang>"
+                     "-DHOST=$<TARGET_FILE:foreign-thread-gc-host>"
+                     "-DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/foreign-thread-gc"
+                     "-DGC_LIB=${TEST_GC_LIBDIR}"
+                     ${gc_foreign_thread_gc_shared}
+                     "-DLLVM_LIB=${LLVM_LIBRARY_DIR}"
+                     "-DTSLANG_LIB=${CMAKE_BINARY_DIR}/lib"
+                     -P "${CMAKE_CURRENT_SOURCE_DIR}/foreign-thread-gc.cmake")
+endif()
+
 # One collector per process with the real default library in it - see defaultlib-collector.cmake.
 # Skipped (not failed) when no default library has been built; the release workflows run them after
 # building one. The sources live outside tests/: without the default library they do not compile,

tslang/test/tester/foreign-thread-gc-host.cpp

@@ -0,0 +1,100 @@
+// A host that is not a tslang program - an Android app, a C program - loading a tslang shared
+// library and calling it from threads of its own, which the collector has never heard of. See
+// foreign-thread-gc.cmake.
+//
+//   foreign-thread-gc-host <library> <threads> <calls>
+//
+// Calls work(seed) <calls> times on the loading thread, then on each of <threads> new threads,
+// and checks every result. Exits 0 when all are right.
+
+#include <cstdio>
+#include <cstdlib>
+#include <thread>
+#include <vector>
+
+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+using WorkFn = double (*)(double);
+
+static WorkFn work;
+static int calls;
+
+// what work(seed) in foreign-thread-gc.cmake returns: the sum of seed + i for i in [0, 200)
+static double expected(double seed)
+{
+    return 200 * seed + 199 * 200 / 2;
+}
+
+static void run(long id)
+{
+    for (int i = 0; i < calls; i++)
+    {
+        auto seed = static_cast<double>(id * 100000 + i);
+        auto got = work(seed);
+        if (got != expected(seed))
+        {
+            std::printf("thread %ld, call %d: got %g, expected %g\n", id, i, got, expected(seed));
+            std::fflush(stdout);
+            std::exit(2);
+        }
+    }
+}
+
+int main(int argc, char **argv)
+{
+    if (argc < 4)
+    {
+        std::printf("usage: %s <library> <threads> <calls>\n", argv[0]);
+        return 1;
+    }
+
+    auto threads = std::atoi(argv[2]);
+    calls = std::atoi(argv[3]);
+
+#ifdef _WIN32
+    // the library's own directory first, where tslang put gc.dll beside it
+    auto library = LoadLibraryExA(argv[1], nullptr, LOAD_WITH_ALTERED_SEARCH_PATH);
+    if (!library)
+    {
+        std::printf("LoadLibrary '%s' failed: %lu\n", argv[1], GetLastError());
+        return 1;
+    }
+
+    work = reinterpret_cast<WorkFn>(GetProcAddress(library, "work"));
+#else
+    auto library = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL);
+    if (!library)
+    {
+        std::printf("dlopen: %s\n", dlerror());
+        return 1;
+    }
+
+    work = reinterpret_cast<WorkFn>(dlsym(library, "work"));
+#endif
+    if (!work)
+    {
+        std::printf("'%s' has no 'work'\n", argv[1]);
+        return 1;
+    }
+
+    run(0);
+    std::printf("loading thread: %d calls\n", calls);
+
+    std::vector<std::thread> pool;
+    for (long id = 1; id <= threads; id++)
+    {
+        pool.emplace_back(run, id);
+    }
+
+    for (auto &thread : pool)
+    {
+        thread.join();
+    }
+
+    std::printf("%d other threads: %d calls each\n", threads, calls);
+    return 0;
+}

tslang/test/tester/foreign-thread-gc.cmake

@@ -0,0 +1,89 @@
+# A `-mm=gc` shared library called from threads its host made (foreign-thread-gc-host.cpp).
+#
+# The collector scans, and stops during a collection, only the threads registered with it. A tslang
+# program registers its own: GC_init the one it starts on, the async runtime its pool's. A host
+# that is not a tslang program - an Android app's JNI threads, a C program's - calls in on threads
+# nobody registered, and the first collection one of them started aborted ("Collecting from unknown
+# thread"). An exported function of a gc shared library now registers the thread that calls it
+# (__tslang_gc_enter, AsyncGCThreadsCommon.inc). The library allocates on every call, so a few
+# thousand calls collect many times, on every thread.
+#
+# Also with no top-level code: then nothing ran GC_init at load, and the first call does.
+
+cmake_minimum_required(VERSION 3.17.3)
+
+foreach(var TSLANG HOST WORK_DIR GC_LIB TSLANG_LIB LLVM_LIB)
+    if(NOT DEFINED ${var})
+        message(FATAL_ERROR "${var} is required")
+    endif()
+endforeach()
+
+file(REMOVE_RECURSE "${WORK_DIR}")
+file(MAKE_DIRECTORY "${WORK_DIR}")
+
+set(ENV{GC_LIB_PATH} "")
+set(ENV{GC_SHARED_LIB_PATH} "")
+set(ENV{TSLANG_LIB_PATH} "")
+
+set(work [=[
+class Node {
+    constructor(public value: number, public next: Node | undefined) {}
+}
+
+// a 200-node list per call: a loop of calls collects many times
+export function work(seed: number): number {
+    let head: Node | undefined = undefined;
+    for (let i = 0; i < 200; i++) {
+        head = new Node(seed + i, head);
+    }
+
+    let sum: number = 0;
+    for (let node = head; node !== undefined; node = node.next) {
+        sum += node.value;
+    }
+
+    return sum;
+}
+]=])
+
+file(WRITE "${WORK_DIR}/with-top-level.ts" "${work}\nlet loaded: number = 1;\n")
+file(WRITE "${WORK_DIR}/no-top-level.ts" "${work}")
+
+set(libs "--gc-lib-path=${GC_LIB}" "--tslang-lib-path=${TSLANG_LIB}" "--llvm-lib-path=${LLVM_LIB}")
+if(DEFINED GC_SHARED_LIB)
+    list(APPEND libs "--gc-shared-lib-path=${GC_SHARED_LIB}")
+endif()
+
+if(WIN32)
+    set(prefix "")
+    set(suffix ".dll")
+    set(pic "")
+else()
+    set(prefix "lib")
+    set(suffix ".so")
+    # a shared object's code has to be position independent (R_X86_64_32S otherwise)
+    set(pic "-relocation-model=pic")
+endif()
+
+foreach(name with-top-level no-top-level)
+    set(library "${WORK_DIR}/${prefix}${name}${suffix}")
+    execute_process(COMMAND "${TSLANG}" --emit=dll --opt -mm=gc --no-default-lib ${pic} ${libs} ${name}.ts -o "${library}"
+        WORKING_DIRECTORY "${WORK_DIR}"
+        OUTPUT_VARIABLE out
+        ERROR_VARIABLE err
+        RESULT_VARIABLE status)
+    if(NOT status EQUAL 0)
+        message(FATAL_ERROR "--emit=dll ${name}: exit ${status}\n${out}\n${err}")
+    endif()
+
+    execute_process(COMMAND "${HOST}" "${library}" 8 2000
+        WORKING_DIRECTORY "${WORK_DIR}"
+        OUTPUT_VARIABLE out
+        ERROR_VARIABLE err
+        RESULT_VARIABLE status
+        TIMEOUT 240)
+    if(NOT status EQUAL 0)
+        message(FATAL_ERROR "${name}: the host's threads calling into the library: exit ${status}\n${out}\n${err}")
+    endif()
+    message(STATUS "${name}: ${out}")
+endforeach()