feat: Capture thread state by calling global function

Author: timfish

.eslintrc.js

@@ -2,7 +2,7 @@ module.exports = {
   extends: ['@sentry-internal/sdk'],
   env: {
     node: true,
-    es6: true,
+    es2020: true
   },
   parserOptions: {
     sourceType: 'module',

README.md

@@ -107,23 +107,27 @@ Set up automatic detection of blocked event loops:
 
 ### 1. Set up thread heartbeats
 
-Send regular heartbeats with optional state information:
+Send regular heartbeats:
 
 ```ts
 import {
   registerThread,
+  registerThreadStateCallback,
   threadPoll,
 } from "@sentry-internal/node-native-stacktrace";
 
+// Register a function that is called to get the thread state whenever stack traces are captured.
+// This must return a serializable object.
+registerThreadStateCallback(() => {
+  return state.snapshot();
+});
+
 // Register this thread
 registerThread();
 
-// Send heartbeats every 200ms with optional state
+// Send heartbeats every 200ms
 setInterval(() => {
-  threadPoll({
-    endpoint: "/api/current-request",
-    userId: getCurrentUserId(),
-  });
+  threadPoll();
 }, 200);
 ```
 
@@ -162,8 +166,8 @@ setInterval(() => {
 
 #### `registerThread(threadName?: string): void`
 
-Registers the current thread for monitoring. Must be called from each thread you
-want to capture stack traces from.
+Registers the current thread for stack trace capture. Must be called from each
+thread you want to capture stack traces from.
 
 - `threadName` (optional): Name for the thread. Defaults to the current thread
   ID.
@@ -187,18 +191,21 @@ type StackFrame = {
 };
 ```
 
-#### `threadPoll<State>(state?: State, disableLastSeen?: boolean): void`
+#### `threadPoll<State>(disableLastSeen?: boolean): void`
 
-Sends a heartbeat from the current thread with optional state information. The
-state object will be serialized and included as a JavaScript object with the
-corresponding stack trace.
+Sends a heartbeat from the current thread.
 
-- `state` (optional): An object containing state information to include with the
-  stack trace.
 - `disableLastSeen` (optional): If `true`, disables the tracking of the last
   seen time for this thread.
 
 #### `getThreadsLastSeen(): Record<string, number>`
 
 Returns the time in milliseconds since each registered thread called
 `threadPoll()`.
+
+#### `registerThreadStateCallback<State>(callback: () => object | undefined): void`
+
+Registers a callback function that is called whenever stack traces are captured
+from this thread. The return value of the callback is included in the `state`
+property of the corresponding `Thread` object returned by `captureStackTrace()`.
+The callback must return a serializable object.

module.cc

@@ -16,16 +16,14 @@ using namespace v8;
 using namespace node;
 using namespace std::chrono;
 
-static const int kMaxStackFrames = 255;
+static const int kMaxStackFrames = 50;
 
 // Structure to hold information for each thread/isolate
 struct ThreadInfo {
   // Thread name
   std::string thread_name;
   // Last time this thread was seen in milliseconds since epoch
   milliseconds last_seen;
-  // Some JSON serialized state for the thread
-  std::string state;
 };
 
 static std::mutex threads_mutex;
@@ -41,21 +39,26 @@ struct JsStackFrame {
 };
 
 // Type alias for a vector of JsStackFrame
-using JsStackTrace = std::vector<JsStackFrame>;
+using JsStackFrames = std::vector<JsStackFrame>;
+
+struct JsStackTrace {
+  // The frames in the stack trace
+  std::vector<JsStackFrame> frames;
+  // JSON serialized string of the state
+  std::string state;
+};
 
 struct ThreadResult {
   std::string thread_name;
-  std::string state;
-  JsStackTrace stack_frames;
+  JsStackTrace stack_trace;
 };
 
-// Function to be called when an isolate's execution is interrupted
-static void ExecutionInterrupted(Isolate *isolate, void *data) {
-  auto promise = static_cast<std::promise<JsStackTrace> *>(data);
+// Function to get stack frames from a V8 stack trace
+JsStackFrames GetStackFrames(Isolate *isolate) {
   auto stack = StackTrace::CurrentStackTrace(isolate, kMaxStackFrames,
                                              StackTrace::kDetailed);
 
-  JsStackTrace frames;
+  JsStackFrames frames;
   if (!stack.IsEmpty()) {
     for (int i = 0; i < stack->GetFrameCount(); i++) {
       auto frame = stack->GetFrame(isolate, i);
@@ -89,7 +92,47 @@ static void ExecutionInterrupted(Isolate *isolate, void *data) {
     }
   }
 
-  promise->set_value(frames);
+  return frames;
+}
+
+// Function to fetch the thread state from the isolate by calling a global
+// function called __get_thread_state_callback__
+std::string GetThreadState(Isolate *isolate) {
+  auto callback_name =
+      v8::String::NewFromUtf8(isolate, "__get_thread_state_callback__",
+                              v8::NewStringType::kNormal)
+          .ToLocalChecked();
+  auto context = isolate->GetCurrentContext();
+  auto callback =
+      context->Global()->Get(context, callback_name).ToLocalChecked();
+
+  if (callback->IsFunction()) {
+    v8::TryCatch try_catch(isolate);
+
+    auto result = Local<Function>::Cast(callback)->Call(
+        context, Undefined(isolate), 0, {});
+
+    if (!try_catch.HasCaught() && !result.IsEmpty()) {
+      MaybeLocal<String> maybe_json =
+          v8::JSON::Stringify(context, result.ToLocalChecked());
+
+      if (!maybe_json.IsEmpty()) {
+        v8::String::Utf8Value utf8_state(isolate, maybe_json.ToLocalChecked());
+        if (*utf8_state) {
+          return *utf8_state;
+        }
+      }
+    }
+  }
+
+  return "";
+}
+
+// Function to be called when an isolate's execution is interrupted
+static void ExecutionInterrupted(Isolate *isolate, void *data) {
+  auto promise = static_cast<std::promise<JsStackTrace> *>(data);
+
+  promise->set_value({GetStackFrames(isolate), GetThreadState(isolate)});
 }
 
 // Function to capture the stack trace of a single isolate
@@ -116,12 +159,11 @@ void CaptureStackTraces(const FunctionCallbackInfo<Value> &args) {
       if (thread_isolate == capture_from_isolate)
         continue;
       auto thread_name = thread_info.thread_name;
-      auto state = thread_info.state;
 
       futures.emplace_back(std::async(
           std::launch::async,
-          [thread_name, state](Isolate *isolate) -> ThreadResult {
-            return ThreadResult{thread_name, state, CaptureStackTrace(isolate)};
+          [thread_name](Isolate *isolate) -> ThreadResult {
+            return ThreadResult{thread_name, CaptureStackTrace(isolate)};
           },
           thread_isolate));
     }
@@ -137,9 +179,9 @@ void CaptureStackTraces(const FunctionCallbackInfo<Value> &args) {
             .ToLocalChecked();
 
     Local<Array> jsFrames =
-        Array::New(capture_from_isolate, result.stack_frames.size());
-    for (size_t i = 0; i < result.stack_frames.size(); ++i) {
-      const auto &frame = result.stack_frames[i];
+        Array::New(capture_from_isolate, result.stack_trace.frames.size());
+    for (size_t i = 0; i < result.stack_trace.frames.size(); ++i) {
+      const auto &frame = result.stack_trace.frames[i];
       Local<Object> frameObj = Object::New(capture_from_isolate);
       frameObj
           ->Set(current_context,
@@ -189,9 +231,10 @@ void CaptureStackTraces(const FunctionCallbackInfo<Value> &args) {
               jsFrames)
         .Check();
 
-    if (!result.state.empty()) {
+    if (!result.stack_trace.state.empty()) {
       v8::MaybeLocal<v8::String> stateStr = v8::String::NewFromUtf8(
-          capture_from_isolate, result.state.c_str(), NewStringType::kNormal);
+          capture_from_isolate, result.stack_trace.state.c_str(),
+          NewStringType::kNormal);
       if (!stateStr.IsEmpty()) {
         v8::MaybeLocal<v8::Value> maybeStateVal =
             v8::JSON::Parse(current_context, stateStr.ToLocalChecked());
@@ -243,8 +286,7 @@ void RegisterThread(const FunctionCallbackInfo<Value> &args) {
     std::lock_guard<std::mutex> lock(threads_mutex);
     auto found = threads.find(isolate);
     if (found == threads.end()) {
-      threads.emplace(isolate,
-                      ThreadInfo{thread_name, milliseconds::zero(), ""});
+      threads.emplace(isolate, ThreadInfo{thread_name, milliseconds::zero()});
       // Register a cleanup hook to remove this thread when the isolate is
       // destroyed
       node::AddEnvironmentCleanupHook(isolate, Cleanup, isolate);
@@ -280,32 +322,17 @@ steady_clock::time_point GetUnbiasedMonotonicTime() {
 // Function to track a thread and set its state
 void ThreadPoll(const FunctionCallbackInfo<Value> &args) {
   auto isolate = args.GetIsolate();
-  auto context = isolate->GetCurrentContext();
-
-  std::string state_str;
-  if (args.Length() > 0 && args[0]->IsValue()) {
-    MaybeLocal<String> maybe_json = v8::JSON::Stringify(context, args[0]);
-    if (!maybe_json.IsEmpty()) {
-      v8::String::Utf8Value utf8_state(isolate, maybe_json.ToLocalChecked());
-      state_str = *utf8_state ? *utf8_state : "";
-    } else {
-      state_str = "";
-    }
-  } else {
-    state_str = "";
-  }
 
   bool disable_last_seen = false;
-  if (args.Length() > 1 && args[1]->IsBoolean()) {
-    disable_last_seen = args[1]->BooleanValue(isolate);
+  if (args.Length() > 0 && args[0]->IsBoolean()) {
+    disable_last_seen = args[0]->BooleanValue(isolate);
   }
 
   {
     std::lock_guard<std::mutex> lock(threads_mutex);
     auto found = threads.find(isolate);
     if (found != threads.end()) {
       auto &thread_info = found->second;
-      thread_info.state = state_str;
       if (disable_last_seen) {
         thread_info.last_seen = milliseconds::zero();
       } else {

package.json

@@ -23,15 +23,15 @@
     "fix": "yarn fix:eslint && yarn fix:clang",
     "fix:eslint": "eslint . --format stylish --fix",
     "fix:clang": "node scripts/clang-format.mjs --fix",
-    "build": "yarn build:lib && yarn build:bindings:configure && yarn build:bindings",
+    "build": "yarn clean && yarn build:lib && yarn build:bindings:configure && yarn build:bindings",
     "build:lib": "tsc",
     "build:bindings:configure": "node-gyp configure",
     "build:bindings:configure:arm64": "node-gyp configure --arch=arm64 --target_arch=arm64",
     "build:bindings": "node-gyp build && node scripts/copy-target.mjs",
     "build:bindings:arm64": "node-gyp build --arch=arm64 && node scripts/copy-target.mjs",
     "build:dev": "yarn clean && yarn build:bindings:configure && yarn build",
     "build:tarball": "npm pack",
-    "clean": "node-gyp clean && rm -rf lib && rm -rf build",
+    "clean": "node-gyp clean && rm -rf lib && rm -rf build && rm -f *.tgz",
     "test": "yarn test:install && node ./test/prepare.mjs && vitest run --silent=false --disable-console-intercept",
     "test:install": "cross-env ALWAYS_THROW=true yarn install"
   },

src/index.ts

@@ -25,7 +25,7 @@ type StackFrame = {
 
 interface Native {
   registerThread(threadName: string): void;
-  threadPoll(state?: object, disableLastSeen?: boolean): void;
+  threadPoll(disableLastSeen?: boolean): void;
   captureStackTrace<S = unknown>(): Record<string, Thread<S>>;
   getThreadsLastSeen(): Record<string, number>;
 }
@@ -186,15 +186,10 @@ export function registerThread(threadName: string = String(threadId)): void {
 /**
  * Tells the native module that the thread is still running and updates the state.
  *
- * @param state Optional state to pass to the native module.
  * @param disableLastSeen If true, disables the last seen tracking for this thread.
  */
-export function threadPoll(state?: object, disableLastSeen?: boolean): void {
-  if (typeof state === 'object' || disableLastSeen) {
-    native.threadPoll(state, disableLastSeen);
-  } else {
-    native.threadPoll();
-  }
+export function threadPoll(disableLastSeen?: boolean): void {
+  native.threadPoll(!!disableLastSeen);
 }
 
 /**
@@ -212,3 +207,16 @@ export function captureStackTrace<S = unknown>(): Record<string, Thread<S>> {
 export function getThreadsLastSeen(): Record<string, number> {
   return native.getThreadsLastSeen();
 }
+
+type GlobalObjectWithCallback = {
+  __get_thread_state_callback__?: () => object | undefined;
+};
+
+/**
+ * Registers a callback function on the global object that the native module can call to get the current thread state.
+ */
+export function registerThreadStateCallback(callback: () => object | undefined): void {
+  (globalThis as GlobalObjectWithCallback).__get_thread_state_callback__ = callback;
+}
+
+

test/async-storage.mjs

@@ -0,0 +1,27 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { Worker } from 'node:worker_threads';
+import { registerThread, registerThreadStateCallback } from '@sentry-internal/node-native-stacktrace';
+import { longWork } from './long-work.js';
+
+registerThread();
+
+const asyncLocalStorage = new AsyncLocalStorage();
+
+function withTraceId(traceId, fn) {
+  return asyncLocalStorage.run(traceId, fn);
+}
+
+registerThreadStateCallback(() => {
+  const traceId = asyncLocalStorage.getStore();
+  return { traceId };
+});
+
+const watchdog = new Worker('./test/watchdog.js');
+
+for (let i = 0; i < 10; i++) {
+  withTraceId(`trace-${i}`, () => {
+    if(i === 5) {
+      longWork();
+    }
+  });
+}

test/e2e.test.mjs

@@ -90,6 +90,38 @@ describe('e2e Tests', { timeout: 20000 }, () => {
     expect(stacks['2'].frames.length).toEqual(1);
   });
 
+  test('async storage state', { timeout: 20000 }, () => {
+    const testFile = join(__dirname, 'async-storage.mjs');
+    const result = spawnSync('node', [testFile]);
+
+    expect(result.status).toEqual(0);
+
+    const stacks = JSON.parse(result.stdout.toString());
+
+    expect(stacks['0'].frames).toEqual(expect.arrayContaining([
+      {
+        function: 'pbkdf2Sync',
+        filename: expect.any(String),
+        lineno: expect.any(Number),
+        colno: expect.any(Number),
+      },
+      {
+        function: 'longWork',
+        filename: expect.stringMatching(/long-work.js$/),
+        lineno: expect.any(Number),
+        colno: expect.any(Number),
+      },
+      {
+        function: '?',
+        filename: expect.stringMatching(/async-storage.mjs$/),
+        lineno: expect.any(Number),
+        colno: expect.any(Number),
+      },
+    ]));
+
+    expect(stacks['0'].state).toEqual({ traceId: 'trace-5' });
+  });
+
   test('can be disabled', { timeout: 20000 }, () => {
     const testFile = join(__dirname, 'stalled-disabled.js');
     const result = spawnSync('node', [testFile]);