How to use AsyncLocalStorage

[phoenix741]

Hello everyone,

I have a question for you, i develop a software that use napi-rs to create a node_modules that contains core functionality and node for the rest.

I want to use AsyncLocalStorage for logging (adding information into log).

I have write a rust logger to intercept rust log and transmit this log to nodejs :

use log::{info, LevelFilter};
use log::{Level, Metadata, Record};

use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction};
use napi::{Error, JsFunction, Result};
use woodstock::config::Context;

use crate::config::context::{JsBackupContext, LogLevel};

#[napi(object)]
pub struct JsBackupLog {
  pub level: LogLevel,
  pub context: String,
  pub message: String,
}

#[napi(object)]
pub struct JsBackupLogMessage {
  pub progress: Option<JsBackupLog>,
  pub error: Option<String>,
  pub complete: bool,
}

struct JavascriptLog {
  tsfn: ThreadsafeFunction<JsBackupLogMessage, ErrorStrategy::Fatal>,
}

impl JavascriptLog {
  pub fn new(tsfn: ThreadsafeFunction<JsBackupLogMessage, ErrorStrategy::Fatal>) -> Self {
    Self { tsfn }
  }
}

impl log::Log for JavascriptLog {
  fn enabled(&self, metadata: &Metadata) -> bool {
    metadata.level() <= Level::Debug
  }

  fn log(&self, record: &Record) {
    if self.enabled(record.metadata()) {
      let message = JsBackupLogMessage {
        progress: Some(JsBackupLog {
          level: record.level().into(),
          context: record.target().to_string(),
          message: record.args().to_string(),
        }),
        error: None,
        complete: false,
      };

      self.tsfn.call(
        message,
        napi::threadsafe_function::ThreadsafeFunctionCallMode::NonBlocking,
      );
    }
  }

  fn flush(&self) {}
}

#[napi]
pub fn init_log(
  context: &JsBackupContext,
  #[napi(ts_arg_type = "(result: JsBackupLogMessage) => void")] callback: JsFunction,
) -> Result<()> {
  let context: Context = context.into();
  let log_level: LevelFilter = context.config.log_level.to_level_filter();
  let tsfn = callback.create_threadsafe_function(0, |ctx| Ok(vec![ctx.value]))?;

  log::set_boxed_logger(Box::new(JavascriptLog::new(tsfn)))
    .map(|()| log::set_max_level(log_level))
    .map_err(|e| Error::from_reason(e.to_string()))?;

  info!("Logging initialized with {log_level}");

  Ok(())
}

In NodeJS i have (using Nest.JS) :

...
interface LogStorage {
  hostname?: string;
  backupNumber?: number;
}

const logAsyncLocalStorage = new AsyncLocalStorage<LogStorage>();

@Injectable()
export class ApplicationLogger implements LoggerService {
  #globalLogger: Logger;

  constructor(
    readonly worker: string,
    private backupsService: BackupsService,
  ) {
    this.#globalLogger = this.#createGlobalLogger(worker);
  }

  #createGlobalLogger(worker: string): Logger {
     ....
  }
  #getLogger(message: Record<string, unknown>): Logger {
    const storage = logAsyncLocalStorage.getStore();

    let hostname = (message.hostname as string | undefined) ?? storage?.hostname;
    let backupNumber = (message.backupNumber as number | undefined) ?? storage?.backupNumber;

    if (hostname !== undefined && backupNumber !== undefined) {
      .... Use host name and backup nummber
    }

    return this.#globalLogger;
  }

  useLogger<R, TArgs extends any[]>(
    hostname: string,
    backupNumber: number,
    callback: (...args: TArgs) => R,
    ...args: TArgs
  ): R {
    return logAsyncLocalStorage.run(
      { hostname, backupNumber },
      (...args) => {
        return callback(...args);
      },
      ...args,
    );
  }

  log(message: string | Record<string, unknown>, context?: string): void {
    const msg = typeof message === 'string' ? { context, message } : { context, ...message };
    this.#getLogger(msg).info(msg);
  }

...
}

Actually this doesn't work : is loose the context.

My question is how to use AsyncLocalStorage ?

When reading https://nodejs.org/api/n-api.html#napi_async_init we see :

The async_resource object needs to be kept alive until napi_async_destroy to keep async_hooks related API acts correctly. In order to retain ABI compatibility with previous versions, napi_async_contexts are not maintaining the strong reference to the async_resource objects to avoid introducing causing memory leaks. However, if the async_resource is garbage collected by JavaScript engine before the napi_async_context was destroyed by napi_async_destroy, calling napi_async_context related APIs like napi_open_callback_scope and napi_make_callback can cause problems like loss of async context when using the AsyncLocalStorage API.

Actually when i use async local storage the context is loose when an api call an napi-rs then call an js callback with ThreadsafeFunction (ie:

Thanks

[phoenix741]

To answer to my question, this code works. If i call call_threadsafe_function and then the callback is called then the stacktrace and local storage is keep.

In my case, i made an error. The init function with the callback is called at program start (with callback registered for future use). The callback is called (in async local storage run) by another thread, but is owned by init not by the process that call it.

To have local storage it's not who call the callback but who create the callback that is important.

So i need to fix my program..

#[napi]
pub fn call_threadsafe_function(callback: JsFunction) -> Result<()> {
  let tsfn: ThreadsafeFunction<u32, ErrorStrategy::CalleeHandled> = callback
    .create_threadsafe_function(0, |ctx| {
      ctx.env.create_uint32(ctx.value + 1).map(|v| vec![v])
    })?;
  for n in 0..100 {
    let tsfn = tsfn.clone();
    tokio::spawn(async move {
      tsfn.call(Ok(n), ThreadsafeFunctionCallMode::Blocking);
    });
  }
  Ok(())
}

import { AsyncLocalStorage } from 'node:async_hooks';
import { callThreadsafeFunction } from './index.js';
 
const asyncLocalStorage = new AsyncLocalStorage();

asyncLocalStorage.run(42, async () => {
  await callThreadsafeFunction(
    () => console.trace('asyncLocalStorage.getStore()', asyncLocalStorage.getStore())
  )
 });