svelte/packages/svelte/src/internal/server/context.js

/** @import { ALSContext, SSRContext } from '#server' */
/** @import { AsyncLocalStorage } from 'node:async_hooks' */
/** @import { Hydratable } from '#shared' */
import { DEV } from 'esm-env';
import * as e from './errors.js';

/** @type {SSRContext | null} */
export var ssr_context = null;

/** @param {SSRContext | null} v */
export function set_ssr_context(v) {
	ssr_context = v;
}

/**
 * @template T
 * @param {any} key
 * @returns {T}
 */
export function getContext(key) {
	const context_map = get_or_init_context_map('getContext');
	const result = /** @type {T} */ (context_map.get(key));

	return result;
}

/**
 * @template T
 * @param {any} key
 * @param {T} context
 * @returns {T}
 */
export function setContext(key, context) {
	get_or_init_context_map('setContext').set(key, context);
	return context;
}

/**
 * @param {any} key
 * @returns {boolean}
 */
export function hasContext(key) {
	return get_or_init_context_map('hasContext').has(key);
}

/** @returns {Map<any, any>} */
export function getAllContexts() {
	return get_or_init_context_map('getAllContexts');
}

/**
 * @param {string} name
 * @returns {Map<unknown, unknown>}
 */
function get_or_init_context_map(name) {
	if (ssr_context === null) {
		e.lifecycle_outside_component(name);
	}

	return (ssr_context.c ??= new Map(get_parent_context(ssr_context) || undefined));
}

/**
 * @param {Function} [fn]
 */
export function push(fn) {
	ssr_context = { p: ssr_context, c: null, r: null };

	if (DEV) {
		ssr_context.function = fn;
		ssr_context.element = ssr_context.p?.element;
	}
}

export function pop() {
	ssr_context = /** @type {SSRContext} */ (ssr_context).p;
}

/**
 * @param {SSRContext} ssr_context
 * @returns {Map<unknown, unknown> | null}
 */
function get_parent_context(ssr_context) {
	let parent = ssr_context.p;

	while (parent !== null) {
		const context_map = parent.c;
		if (context_map !== null) {
			return context_map;
		}
		parent = parent.p;
	}

	return null;
}

/**
 * Wraps an `await` expression in such a way that the component context that was
 * active before the expression evaluated can be reapplied afterwards —
 * `await a + b()` becomes `(await $.save(a))() + b()`, meaning `b()` will have access
 * to the context of its component.
 * @template T
 * @param {Promise<T>} promise
 * @returns {Promise<() => T>}
 */
export async function save(promise) {
	var previous_context = ssr_context;
	var previous_sync_store = sync_store;
	var value = await promise;

	return () => {
		ssr_context = previous_context;
		sync_store = previous_sync_store;
		return value;
	};
}

/**
 * @template T
 * @type {Hydratable<T>}
 */
export async function hydratable(key, fn, { transport } = {}) {
	const store = await get_render_store();

	if (store.hydratables.has(key)) {
		// TODO error
		throw new Error("can't have two hydratables with the same key");
	}

	const result = fn();
	store.hydratables.set(key, { value: result, transport });
	return result;
}

/** @type {ALSContext | null} */
export let sync_store = null;

/** @param {ALSContext | null} store */
export function set_sync_store(store) {
	sync_store = store;
}

/** @type {Promise<AsyncLocalStorage<ALSContext | null> | null>} */
const als = import('node:async_hooks')
	.then((hooks) => new hooks.AsyncLocalStorage())
	.catch(() => {
		// can't use ALS but can still use manual context preservation
		return null;
	});

/** @returns {Promise<ALSContext | null>} */
async function try_get_render_store() {
	return sync_store ?? (await als)?.getStore() ?? null;
}

/** @returns {Promise<ALSContext>} */
export async function get_render_store() {
	const store = await try_get_render_store();

	if (!store) {
		// TODO make this a proper e.error
		let message = 'Could not get rendering context.';

		if (await als) {
			message += ' This is an internal error.';
		} else {
			message +=
				' In environments without `AsyncLocalStorage`, `hydratable` must be accessed synchronously, not after an `await`.' +
				' If it was accessed synchronously then this is an internal error.';
		}

		throw new Error(message);
	}

	return store;
}

/**
 * @template T
 * @param {ALSContext} store
 * @param {() => Promise<T>} fn
 * @returns {Promise<T>}
 */
export async function with_render_store(store, fn) {
	try {
		sync_store = store;
		const storage = await als;
		return storage ? storage.run(store, fn) : fn();
	} finally {
		sync_store = null;
	}
}