Support capturing all threads

Author: timfish

README.md

@@ -1,18 +1,19 @@
 # `cross-thread-stack-trace`
 
-Native Node module to capture stack traces across threads. This allows capturing
-main thread stack traces from a worker thread, even if the main thread event
-loop is blocked.
+Native Node module to capture stack traces from all registered threads.
 
-Main thread:
+This allows capturing main and worker thread stack traces from another watchdog
+thread, even if the event loops are blocked.
+
+In the main or worker threads:
 
 ```ts
-const { setMainIsolate } = require("cross-thread-stack-trace");
+const { registerThread } = require("cross-thread-stack-trace");
 
-setMainIsolate();
+registerThread();
 ```
 
-Worker thread:
+Watchdog thread:
 
 ```ts
 const { captureStackTrace } = require("cross-thread-stack-trace");
@@ -30,80 +31,58 @@ npm i && npm test
 Results in:
 
 ```js
-[
-    {
-        function: "from",
-        filename: "node:buffer",
-        lineno: 298,
-        colno: 28,
-    },
+{
+  main: [
     {
-        function: "pbkdf2Sync",
-        filename: "node:internal/crypto/pbkdf2",
-        lineno: 78,
-        colno: 17,
+      function: 'from',
+      filename: 'node:buffer',
+      lineno: 298,
+      colno: 28
     },
     {
-        function: "longWork",
-        filename:
-            "/Users/tim/Documents/Repositories/cross-thread-stack-trace/test.js",
-        lineno: 15,
-        colno: 29,
+      function: 'pbkdf2Sync',
+      filename: 'node:internal/crypto/pbkdf2',
+      lineno: 78,
+      colno: 17
     },
     {
-        function: "?",
-        filename:
-            "/Users/tim/Documents/Repositories/cross-thread-stack-trace/test.js",
-        lineno: 19,
-        colno: 1,
+      function: 'longWork',
+      filename: '/Users/tim/Documents/Repositories/cross-thread-stack-trace/test/test.js',
+      lineno: 20,
+      colno: 29
     },
     {
-        function: "?",
-        filename: "node:internal/modules/cjs/loader",
-        lineno: 1730,
-        colno: 14,
-    },
-    {
-        function: "?",
-        filename: "node:internal/modules/cjs/loader",
-        lineno: 1895,
-        colno: 10,
-    },
+      function: '?',
+      filename: '/Users/tim/Documents/Repositories/cross-thread-stack-trace/test/test.js',
+      lineno: 24,
+      colno: 1
+    }
+  ],
+  'worker-2': [
     {
-        function: "?",
-        filename: "node:internal/modules/cjs/loader",
-        lineno: 1465,
-        colno: 32,
+      function: 'from',
+      filename: 'node:buffer',
+      lineno: 298,
+      colno: 28
     },
     {
-        function: "?",
-        filename: "node:internal/modules/cjs/loader",
-        lineno: 1282,
-        colno: 12,
+      function: 'pbkdf2Sync',
+      filename: 'node:internal/crypto/pbkdf2',
+      lineno: 78,
+      colno: 17
     },
     {
-        function: "traceSync",
-        filename: "node:diagnostics_channel",
-        lineno: 322,
-        colno: 14,
+      function: 'longWork',
+      filename: '/Users/tim/Documents/Repositories/cross-thread-stack-trace/test/worker.js',
+      lineno: 10,
+      colno: 29
     },
     {
-        function: "wrapModuleLoad",
-        filename: "node:internal/modules/cjs/loader",
-        lineno: 235,
-        colno: 24,
-    },
-    {
-        function: "executeUserEntryPoint",
-        filename: "node:internal/modules/run_main",
-        lineno: 170,
-        colno: 5,
-    },
-    {
-        function: "?",
-        filename: "node:internal/main/run_main_module",
-        lineno: 36,
-        colno: 49,
-    },
-];
+      function: '?',
+      filename: '/Users/tim/Documents/Repositories/cross-thread-stack-trace/test/worker.js',
+      lineno: 14,
+      colno: 1
+    }
+  ]
+}
 ```

index.d.ts

@@ -1,5 +1,5 @@
 
-export declare function setMainIsolate(): void;
+export declare function registerThread(): void;
 
 export type StackFrame = {
     function: string;
@@ -8,4 +8,8 @@ export type StackFrame = {
     colno: number;
 };
 
-export declare function captureStackTrace(): StackFrame[];
+export type Trace = {
+    main: StackFrame[];
+} & Record<string, StackFrame[]>;
+
+export declare function captureStackTrace(excludeWorkers: boolean): Trace;

index.js

@@ -1,6 +1,9 @@
+const { isMainThread, threadId } = require('node:worker_threads');
 const native = require('./build/release/cross-thread-stack-trace.node');
 
-exports.setMainIsolate = native.setMainIsolate;
-exports.captureStackTrace = function () {
-    return JSON.parse(native.captureStackTrace());
+exports.registerThread = function () {
+    native.registerThread(isMainThread ? -1 : threadId );
+};
+exports.captureStackTrace = function (excludeWorkers) {
+    return JSON.parse(native.captureStackTrace(excludeWorkers));
 };

module.cc

@@ -1,49 +1,37 @@
 #include <node.h>
 #include <mutex>
+#include <map>
+#include <sstream>
+#include <future>
 
 using namespace v8;
 using namespace node;
 
-static v8::Isolate *main_thread_isolate;
-
-static std::mutex interrupt_mutex;
-static std::condition_variable interrupt_cv;
-static bool interrupt_done = false;
-
 static const int kMaxStackFrames = 255;
-static const int kMaxStackJsonSize = 10240;
+
+static std::unordered_map<v8::Isolate *, int> threads = {};
 
 static void ExecutionInterrupted(Isolate *isolate, void *data)
 {
-    char *buffer = static_cast<char *>(data);
-
-    v8::RegisterState state;
-    v8::SampleInfo info;
-    void *samples[kMaxStackFrames];
+    auto promise = static_cast<std::promise<std::string> *>(data);
 
-    uint32_t pos = 0;
+    Local<StackTrace> stack = StackTrace::CurrentStackTrace(isolate, kMaxStackFrames, StackTrace::kDetailed);
 
-    // Initialise the register state
-    state.pc = nullptr;
-    state.fp = &state;
-    state.sp = &state;
-
-    isolate->GetStackSample(state, samples, kMaxStackFrames, &info);
-
-    Local<StackTrace> stack = StackTrace::CurrentStackTrace(isolate, 255, StackTrace::kDetailed);
     if (stack.IsEmpty())
     {
-        snprintf(buffer, kMaxStackJsonSize, "[]");
+        promise->set_value("[]");
         return;
     }
 
-    pos += snprintf(&buffer[pos], kMaxStackJsonSize, "[");
-    int count = stack->GetFrameCount();
+    std::ostringstream out;
 
-    for (int i = 0; i < count; i++)
+    out << "[";
+    auto count = stack->GetFrameCount();
+
+    for (auto i = 0; i < count; i++)
     {
-        Local<StackFrame> frame = stack->GetFrame(isolate, i);
-        Local<String> fn_name = frame->GetFunctionName();
+        auto frame = stack->GetFrame(isolate, i);
+        auto fn_name = frame->GetFunctionName();
 
         if (frame->IsEval())
         {
@@ -60,74 +48,113 @@ static void ExecutionInterrupted(Isolate *isolate, void *data)
 
         String::Utf8Value function_name(isolate, fn_name);
         String::Utf8Value script_name(isolate, frame->GetScriptName());
-        const int line_number = frame->GetLineNumber();
-        const int column = frame->GetColumn();
+        auto line_number = frame->GetLineNumber();
+        auto column = frame->GetColumn();
 
-        pos += snprintf(&buffer[pos], kMaxStackJsonSize,
-                        "{\"function\":\"%s\",\"filename\":\"%s\",\"lineno\":%d,\"colno\":%d}",
-                        *function_name,
-                        *script_name,
-                        line_number,
-                        column);
+        out << "{\"function\":\"" << *function_name
+            << "\",\"filename\":\"" << *script_name
+            << "\",\"lineno\":" << line_number
+            << ",\"colno\":" << column << "}";
 
         if (i < count - 1)
         {
-            pos += snprintf(&buffer[pos], kMaxStackJsonSize, ",");
+            out << ",";
         }
     }
 
-    pos += snprintf(&buffer[pos], kMaxStackJsonSize, "]");
+    out << "]";
 
-    {
-        std::lock_guard<std::mutex> lock(interrupt_mutex);
-        interrupt_done = true;
-    }
-    interrupt_cv.notify_one();
+    promise->set_value(out.str());
 }
 
-void CaptureStackTrace(const FunctionCallbackInfo<Value> &args)
+std::string CaptureStackTrace(Isolate *isolate)
 {
-    char buffer[kMaxStackJsonSize] = {0};
+    std::promise<std::string> promise;
+    auto future = promise.get_future();
+
+    isolate->RequestInterrupt(ExecutionInterrupted, &promise);
+    return future.get();
+}
 
-    if (auto isolate = main_thread_isolate)
+void CaptureStackTraces(const FunctionCallbackInfo<Value> &args)
+{
+    bool exclude_workers = args.Length() == 1 && args[0]->IsBoolean() && args[0].As<Boolean>()->Value();
+    auto capture_from_isolate = args.GetIsolate();
+
+    std::vector<std::future<std::string>> futures;
+
+    for (auto &thread : threads)
     {
-        // Reset the interrupt_done flag
+        auto thread_isolate = thread.first;
+        if (thread_isolate != capture_from_isolate)
         {
-            std::lock_guard<std::mutex> lock(interrupt_mutex);
-            interrupt_done = false;
+            int thread_id = thread.second;
+
+            if (exclude_workers && thread_id != -1)
+            {
+                continue;
+            }
+
+            auto thread_name = thread_id == -1 ? "main" : "worker-" + std::to_string(thread_id);
+
+            futures.emplace_back(std::async(std::launch::async, [thread_name](Isolate *isolate)
+                                            { return "\"" + thread_name + "\":" + CaptureStackTrace(isolate); }, thread_isolate));
         }
+    }
 
-        isolate->RequestInterrupt(ExecutionInterrupted, buffer);
+    std::ostringstream out;
 
-        // Wait for the interrupt to complete
-        std::unique_lock<std::mutex> lock(interrupt_mutex);
-        interrupt_cv.wait(lock, []
-                          { return interrupt_done; });
+    auto count = futures.size();
+    out << "{";
+    for (auto &future : futures)
+    {
+        out << future.get();
+        if (--count > 0)
+        {
+            out << ",";
+        }
     }
+    out << "}";
 
-    Local<String> result = String::NewFromUtf8(args.GetIsolate(), buffer, NewStringType::kNormal).ToLocalChecked();
-    args.GetReturnValue().Set(result);
+    args.GetReturnValue().Set(String::NewFromUtf8(capture_from_isolate, out.str().c_str(), NewStringType::kNormal).ToLocalChecked());
 }
 
-void SetMainIsolate(const FunctionCallbackInfo<Value> &args)
+void Cleanup(void *arg)
 {
-    main_thread_isolate = args.GetIsolate();
+    auto isolate = static_cast<Isolate *>(arg);
+    threads.erase(isolate);
+}
+
+void RegisterThread(const FunctionCallbackInfo<Value> &args)
+{
+    auto isolate = args.GetIsolate();
+
+    if (args.Length() != 1 || !args[0]->IsNumber())
+    {
+        isolate->ThrowException(Exception::Error(String::NewFromUtf8(isolate, "registerThread() requires a single threadId argument", NewStringType::kInternalized).ToLocalChecked()));
+        return;
+    }
+
+    int thread_id = args[0].As<Number>()->Value();
+
+    threads.emplace(isolate, thread_id);
+    node::AddEnvironmentCleanupHook(isolate, Cleanup, isolate);
 }
 
 extern "C" NODE_MODULE_EXPORT void
 NODE_MODULE_INITIALIZER(Local<Object> exports,
                         Local<Value> module,
                         Local<Context> context)
 {
-    Isolate *isolate = context->GetIsolate();
+    auto isolate = context->GetIsolate();
 
     exports->Set(context,
                  String::NewFromUtf8(isolate, "captureStackTrace", NewStringType::kInternalized).ToLocalChecked(),
-                 FunctionTemplate::New(isolate, CaptureStackTrace)->GetFunction(context).ToLocalChecked())
+                 FunctionTemplate::New(isolate, CaptureStackTraces)->GetFunction(context).ToLocalChecked())
         .Check();
 
     exports->Set(context,
-                 String::NewFromUtf8(isolate, "setMainIsolate", NewStringType::kInternalized).ToLocalChecked(),
-                 FunctionTemplate::New(isolate, SetMainIsolate)->GetFunction(context).ToLocalChecked())
+                 String::NewFromUtf8(isolate, "registerThread", NewStringType::kInternalized).ToLocalChecked(),
+                 FunctionTemplate::New(isolate, RegisterThread)->GetFunction(context).ToLocalChecked())
         .Check();
 }