Request Context: AsyncLocalStorage and context.Context

import { AsyncLocalStorage } from "node:async_hooks";
 
const requestContext = new AsyncLocalStorage();
 
// The deepest-level logger: it does not take requestId as a parameter.
function log(message) {
  const store = requestContext.getStore(); // undefined outside of run()
  console.log(`[${store?.requestId ?? "-"}] ${message}`);
}
 
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
 
async function loadUser(id) {
  await sleep(10);
  log(`loaded user ${id}`);
  return { id };
}
 
function handleRequest(requestId, userId) {
  // Everything inside the callback, even after await, sees the same store.
  return requestContext.run({ requestId }, async () => {
    log("start");
    await loadUser(userId);
    log("done");
  });
}
 
// Two requests interleave but never mix up their requestId.
await Promise.all([handleRequest("req-1", 1), handleRequest("req-2", 2)]);
log("outside");
// → [req-1] start
// → [req-2] start
// → [req-1] loaded user 1
// → [req-1] done
// → [req-2] loaded user 2
// → [req-2] done
// → [-] outside

package main
 
import (
	"context"
	"fmt"
	"sync"
	"time"
)
 
// Unexported key type: no other package can create a colliding key.
type requestIDKey struct{}
 
func withRequestID(ctx context.Context, id string) context.Context {
	return context.WithValue(ctx, requestIDKey{}, id) // returns a new ctx; the old one is unchanged
}
 
// Even the deepest-level logger has to take ctx as a parameter.
func log(ctx context.Context, message string) {
	id, ok := ctx.Value(requestIDKey{}).(string) // Value returns any: needs a type assertion
	if !ok {
		id = "-"
	}
	fmt.Printf("[%s] %s\n", id, message)
}
 
func loadUser(ctx context.Context, id int) {
	time.Sleep(10 * time.Millisecond)
	log(ctx, fmt.Sprintf("loaded user %d", id))
}
 
func handleRequest(ctx context.Context, requestID string, userID int) {
	ctx = withRequestID(ctx, requestID)
	log(ctx, "start")
	loadUser(ctx, userID) // by convention, ctx is always the first parameter
	log(ctx, "done")
}
 
func main() {
	ctx := context.Background()
 
	var wg sync.WaitGroup
	wg.Go(func() { handleRequest(ctx, "req-1", 1) })
	wg.Go(func() { handleRequest(ctx, "req-2", 2) })
	wg.Wait()
 
	log(ctx, "outside")
}
 
// → [req-1] start
// → [req-2] start
// → [req-1] loaded user 1
// → [req-1] done
// → [req-2] loaded user 2
// → [req-2] done
// → [-] outside
// (the order between req-1 and req-2 can change from run to run)

import { AsyncLocalStorage } from "node:async_hooks";
import { EventEmitter } from "node:events";
 
const requestContext = new AsyncLocalStorage();
const currentId = () => requestContext.getStore()?.requestId ?? "-";
 
requestContext.run({ requestId: "req-1" }, () => {
  // Callbacks scheduled inside run() carry the store, even though they run after run() has returned.
  setTimeout(() => console.log("timeout:", currentId()), 10);
  Promise.resolve().then(() => console.log("then:", currentId()));
});
 
// Scheduled outside run(): no store.
setTimeout(() => console.log("outside:", currentId()), 20);
 
// EventEmitter listeners run in the context of emit(), not of on().
const emitter = new EventEmitter();
requestContext.run({ requestId: "req-1" }, () => {
  emitter.on("job", () => console.log("listener:", currentId()));
});
requestContext.run({ requestId: "req-2" }, () => emitter.emit("job"));
// → listener: req-2
// → then: req-1
// → timeout: req-1
// → outside: -

package main
 
import (
	"context"
	"fmt"
	"sync"
	"time"
)
 
type requestIDKey struct{}
 
func requestID(ctx context.Context) string {
	if id, ok := ctx.Value(requestIDKey{}).(string); ok {
		return id
	}
	return "-"
}
 
// A goroutine inherits nothing from the goroutine that started it:
// to see the request ID it must receive ctx as a parameter.
func worker(ctx context.Context, name string) {
	fmt.Printf("%s: %s\n", name, requestID(ctx))
}
 
func main() {
	ctx := context.WithValue(context.Background(), requestIDKey{}, "req-1")
 
	var wg sync.WaitGroup
	wg.Go(func() { worker(ctx, "goroutine") })                   // pass ctx explicitly
	wg.Go(func() { worker(context.Background(), "forgot ctx") }) // forgot to pass it: request ID lost
	wg.Wait()
 
	// Same for time.AfterFunc: the callback only sees ctx if the closure captures it.
	done := make(chan struct{})
	time.AfterFunc(10*time.Millisecond, func() {
		worker(ctx, "timer")
		close(done)
	})
	<-done
}
 
// → goroutine: req-1
// → forgot ctx: -
// → timer: req-1
// (the first two lines may swap places)

import { AsyncLocalStorage } from "node:async_hooks";
 
const requestContext = new AsyncLocalStorage();
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
 
async function query(sql) {
  requestContext.getStore().queries.push(sql); // the store is a plain object: mutate it directly
  await sleep(1);
}
 
async function handleRequest() {
  await query("SELECT * FROM users");
  await Promise.all([query("SELECT * FROM orders"), query("SELECT * FROM items")]);
}
 
const store = { requestId: "req-1", queries: [] };
await requestContext.run(store, handleRequest);
 
// Changes made inside are visible outside, since it is the same object.
console.log(`${store.requestId} ran ${store.queries.length} queries`);
// → req-1 ran 3 queries

package main
 
import (
	"context"
	"fmt"
	"sync"
	"time"
)
 
// Context values cannot be modified (WithValue always creates a new ctx),
// so to collect per-request data, store a pointer.
type queryLog struct {
	mu      sync.Mutex // goroutines of the same request may write concurrently
	queries []string
}
 
type queryLogKey struct{}
 
func query(ctx context.Context, sql string) {
	ql := ctx.Value(queryLogKey{}).(*queryLog) // panics if missing: use the `v, ok :=` form if unsure
	ql.mu.Lock()
	ql.queries = append(ql.queries, sql)
	ql.mu.Unlock()
	time.Sleep(time.Millisecond)
}
 
func handleRequest(ctx context.Context) {
	query(ctx, "SELECT * FROM users")
 
	var wg sync.WaitGroup
	wg.Go(func() { query(ctx, "SELECT * FROM orders") })
	wg.Go(func() { query(ctx, "SELECT * FROM items") })
	wg.Wait()
}
 
func main() {
	ql := &queryLog{}
	ctx := context.WithValue(context.Background(), queryLogKey{}, ql)
	handleRequest(ctx)
 
	fmt.Printf("req-1 ran %d queries\n", len(ql.queries)) // → req-1 ran 3 queries
}