How to use node.js async thread local context for HTTP and GraphQL

async-local-storage-context-example.js

const { AsyncLocalStorage } = require('node:async_hooks')
const get = require('lodash/get')

const httpStorage = new AsyncLocalStorage()
const gqlStorage = new AsyncLocalStorage()
const KEY = Symbol('_context')

// RESOLVERS //

async function resolveAllUsers (context, path) {
    await sleep()
    logGql(context, path, 'resolveAllUsers')
}

async function resolveAllTickets (context, path) {
    await sleep()
    logGql(context, path, 'resolveAllTickets')
}

async function resolveAllProperties (context, path) {
    await sleep()
    logGql(context, path, 'resolveAllProperties')
    await executeGql(context, { users: resolveAllUsers })
}

async function resolveAllTicketRelations (context, path) {
    await sleep()
    logGql(context, path, 'resolveAllTicketRelations')
}

// END RESOLVERS //

function logGql () {
    const contextName = httpStorage.getStore()?.req
    const contextQuery = gqlStorage.getStore()?.query
    const argumentName = arguments[0]
    const argumentQuery = arguments[1]
    const resolverName = arguments[2]
    if (argumentName !== contextName) throw new Error('wrong context!')
    console.log(`resolver: http=${argumentName} path=${argumentQuery.join('->')}->${resolverName}(...); gql-context-value!=${contextQuery};`)
}

function sleep () {
    return new Promise((res) => {
        setTimeout(res, 300 + Math.floor(Math.random() * 100))
    })
}

function addGqlContextWrapper (fn) {
    return async (context, path) => {
        const threadContext = httpStorage.getStore()?.q
        if (!threadContext) throw new Error('no http.q context')
        const x = gqlStorage.getStore()
        if (x) {
            // inside sub query
            return gqlStorage.run(x, () => fn(context, path))
        } else {
            // not inside sub query
            const queryContext = get(threadContext, path)[KEY]
            if (!queryContext) throw new Error('no http.q context for ' + path.join('->'))

            return gqlStorage.run({ query: queryContext }, () => fn(context, path))
        }
    }
}

function prepareHttpContext (gql, parent = undefined) {
    const result = (parent) ? { [KEY]: parent } : {}
    if (typeof gql === 'object') {
        Object.keys(gql).forEach((key) => {
            result[key] = prepareHttpContext(gql[key], (parent) ? parent : key)
        })
    }
    return result
}

function addHttpContextWrapper (fn) {
    return (name, gql) => {
        httpStorage.enterWith({ 'req': name, 'q': prepareHttpContext(gql) })
        try {
            return fn(name, gql)
        } finally {
            // httpStorage.disable()
        }
    }
}

async function executeGql (context, gql, path = undefined) {
    // gql: { name: resolver }
    if (!gql) return
    const result = {}
    if (typeof gql === 'object') {
        await Promise.all(Object.keys(gql).map((key) => {
            return executeGql(context, gql[key], (path) ? [...path, key] : [key])
                .then((value) => { result[key] = value })
        }))
        return result
    } else {
        return await addGqlContextWrapper(gql)(context, path)
    }
}

function expressRequestProcessor (requestId, body) {
    function graphQLRequestProcessor (requestId, query) { return executeGql(requestId, query) }

    return addHttpContextWrapper(graphQLRequestProcessor)(requestId, body)
}

async function main () {
    // Run parallel HTTP requests
    Promise.all([
        expressRequestProcessor('r1', { allTickets: { resolveAllTickets, bar: { foo: { buz: resolveAllTicketRelations } } }, allProperty: { resolveAllProperties, resolveAllTickets } }),
        expressRequestProcessor('r2', { allTickets: { resolveAllTickets, bar: { foo: { buz: resolveAllTicketRelations } } }, allProperty: { resolveAllProperties, resolveAllTickets } }),
        expressRequestProcessor('r3', { allTickets: { resolveAllTickets, bar: { foo: { buz: resolveAllTicketRelations } } }, allProperty: { resolveAllProperties, resolveAllTickets } }),
    ])
}

main()