feat: Capture thread state from AsyncLocalStorage store

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

@@ -70,7 +70,7 @@ Stack traces show where each thread is currently executing:
       }
     ]
   },
-  '2': { // Worker thread 
+  '2': { // Worker thread
     frames: [
       {
         function: 'from',
@@ -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,
   threadPoll,
 } from "@sentry-internal/node-native-stacktrace";
+import { AsyncLocalStorage } from "node:async_hooks";
 
-// Register this thread
-registerThread();
+// Create async local storage for state tracking
+const asyncLocalStorage = new AsyncLocalStorage();
+
+// Register this thread with async local storage
+registerThread(asyncLocalStorage);
+
+// Set some state in the async local storage
+asyncLocalStorage.enterWith({ someState: "value" });
 
-// Send heartbeats every 200ms with optional state
+// Send heartbeats every 200ms
 setInterval(() => {
-  threadPoll({
-    endpoint: "/api/current-request",
-    userId: getCurrentUserId(),
-  });
+  threadPoll();
 }, 200);
 ```
 
@@ -160,11 +164,13 @@ setInterval(() => {
 
 ### Functions
 
-#### `registerThread(threadName?: string): void`
+#### `registerThread(asyncLocalStorage?: AsyncLocalStorage, 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.
 
+- `asyncLocalStorage` (optional): AsyncLocalStorage instance for tracking thread
+  state. The current value in the storage will be captured in stack traces.
 - `threadName` (optional): Name for the thread. Defaults to the current thread
   ID.
 
@@ -187,14 +193,10 @@ 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.
 

module.cc

@@ -16,16 +16,16 @@ 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;
+  // Async local storage associated with this thread
+  v8::Global<v8::Value> async_local_storage;
 };
 
 static std::mutex threads_mutex;
@@ -41,21 +41,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,18 +94,76 @@ static void ExecutionInterrupted(Isolate *isolate, void *data) {
     }
   }
 
-  promise->set_value(frames);
+  return frames;
+}
+
+static std::string ToJSONOrEmpty(Isolate *isolate, v8::Local<v8::Value> value) {
+  auto context = isolate->GetCurrentContext();
+  MaybeLocal<String> maybe_json = v8::JSON::Stringify(context, value);
+
+  if (!maybe_json.IsEmpty()) {
+    v8::String::Utf8Value utf8_state(isolate, maybe_json.ToLocalChecked());
+    if (*utf8_state) {
+      return *utf8_state;
+    }
+  }
+  return "";
+}
+
+// Function to fetch the thread state from the async context store
+std::string GetThreadState(Isolate *isolate,
+                           v8::Global<v8::Value> async_local_storage) {
+  // Node.js stores the async local storage in the isolate's
+  // "ContinuationPreservedEmbedderData" map, keyed by the
+  // AsyncLocalStorage instance.
+  // https://github.com/nodejs/node/blob/c6316f9db9869864cea84e5f07585fa08e3e06d2/src/async_context_frame.cc#L37
+  auto data = isolate->GetContinuationPreservedEmbedderData();
+  auto local_async_local_storage = async_local_storage.Get(isolate);
+
+  if (!data.IsEmpty() && data->IsMap()) {
+    auto map = data.As<v8::Map>();
+    auto context = isolate->GetCurrentContext();
+    auto val = map->Get(context, local_async_local_storage);
+
+    if (!val.IsEmpty()) {
+      return ToJSONOrEmpty(isolate, val.ToLocalChecked());
+    }
+  }
+
+  return "";
+}
+
+struct InterruptArgs {
+  std::promise<JsStackTrace> *promise;
+  v8::Global<v8::Value> async_local_storage;
+};
+
+// Function to be called when an isolate's execution is interrupted
+static void ExecutionInterrupted(Isolate *isolate, void *data) {
+  auto args = static_cast<InterruptArgs *>(data);
+
+  v8::Locker locker(isolate);
+  v8::HandleScope scope(isolate);
+
+  auto frames = GetStackFrames(isolate);
+  auto state = GetThreadState(isolate, std::move(args->async_local_storage));
+
+  args->promise->set_value({frames, state});
 }
 
 // Function to capture the stack trace of a single isolate
-JsStackTrace CaptureStackTrace(Isolate *isolate) {
+JsStackTrace CaptureStackTrace(Isolate *isolate,
+                               v8::Global<v8::Value> async_local_storage) {
   std::promise<JsStackTrace> promise;
-  auto future = promise.get_future();
 
   // The v8 isolate must be interrupted to capture the stack trace
   // Execution resumes automatically after ExecutionInterrupted returns
-  isolate->RequestInterrupt(ExecutionInterrupted, &promise);
-  return future.get();
+  isolate->RequestInterrupt(
+      ExecutionInterrupted,
+      new InterruptArgs{&promise,
+                        v8::Global<v8::Value>(isolate, async_local_storage)});
+
+  return promise.get_future().get();
 }
 
 // Function to capture stack traces from all registered threads
@@ -112,21 +175,27 @@ void CaptureStackTraces(const FunctionCallbackInfo<Value> &args) {
 
   {
     std::lock_guard<std::mutex> lock(threads_mutex);
-    for (auto [thread_isolate, thread_info] : threads) {
+    for (auto &[thread_isolate, thread_info] : threads) {
       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_isolate, thread_name](
+              v8::Global<v8::Value> async_local_storage) -> ThreadResult {
+            return ThreadResult{
+                thread_name, CaptureStackTrace(thread_isolate,
+                                               std::move(async_local_storage))};
           },
-          thread_isolate));
+          std::move(thread_info.async_local_storage)));
     }
   }
 
+  v8::Locker locker(capture_from_isolate);
+  v8::HandleScope scope(capture_from_isolate);
+
   Local<Object> output = Object::New(capture_from_isolate);
 
   for (auto &future : futures) {
@@ -137,9 +206,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 +258,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());
@@ -226,7 +296,7 @@ void Cleanup(void *arg) {
 void RegisterThread(const FunctionCallbackInfo<Value> &args) {
   auto isolate = args.GetIsolate();
 
-  if (args.Length() != 1 || !args[0]->IsString()) {
+  if (args.Length() < 1 || !args[0]->IsString()) {
     isolate->ThrowException(Exception::Error(
         String::NewFromUtf8(
             isolate, "registerThread(name) requires a single name argument",
@@ -236,15 +306,20 @@ void RegisterThread(const FunctionCallbackInfo<Value> &args) {
     return;
   }
 
+  v8::Global<v8::Value> async_local_storage;
+  if (args.Length() > 1) {
+    async_local_storage.Reset(isolate, args[1]);
+  }
+
   v8::String::Utf8Value utf8(isolate, args[0]);
   std::string thread_name(*utf8 ? *utf8 : "");
 
   {
     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(),
+                                          std::move(async_local_storage)});
       // Register a cleanup hook to remove this thread when the isolate is
       // destroyed
       node::AddEnvironmentCleanupHook(isolate, Cleanup, isolate);
@@ -280,37 +355,22 @@ 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);
+  bool enable_last_seen = true;
+  if (args.Length() > 0 && args[0]->IsBoolean()) {
+    enable_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 {
+      if (enable_last_seen) {
         thread_info.last_seen = duration_cast<milliseconds>(
             GetUnbiasedMonotonicTime().time_since_epoch());
+      } else {
+        thread_info.last_seen = milliseconds::zero();
       }
     }
   }

package.json

@@ -23,16 +23,17 @@
     "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:prepare": "node ./test/prepare.mjs",
     "test:install": "cross-env ALWAYS_THROW=true yarn install"
   },
   "engines": {
@@ -68,4 +69,4 @@
   "volta": {
     "node": "24.1.0"
   }
-}
+}