From af66bce16c6411e787ad9a40e98c6b6499fcba63 Mon Sep 17 00:00:00 2001
From: ido <ido.pluto@w-on.uk>
Date: Sun, 30 Jun 2024 20:20:32 +0300
Subject: [PATCH] fix(context): per request lock

---
 packages/context/src/index.ts | 63 ++++++++++++++++++++++-------------
 1 file changed, 40 insertions(+), 23 deletions(-)

diff --git a/packages/context/src/index.ts b/packages/context/src/index.ts
index 8fe40a7..1fed1cf 100644
--- a/packages/context/src/index.ts
+++ b/packages/context/src/index.ts
@@ -8,49 +8,66 @@ declare global {
     }
 }
 
-type ContextAstro = AstroGlobal | { request: Request, locals: any, props: any; };
-function getContextHistory(astro: ContextAstro, name: string) {
-    const contexts: Map<string, any[]> = astro.locals.amContext ??= new Map();
-    contexts.set(name, contexts.get(name) ?? []);
-    return contexts.get(name);
+type ContextAstro = AstroGlobal | {
+    request: Request,
+    locals: any,
+    props: any;
+};
+
+type AMContext = {
+    lock: Map<string, any>;
+    history: Map<string, any[]>;
+};
+
+function getAMContextFromAstro(astro: ContextAstro, name: string) {
+    const amContext = astro.locals.amContext ??= {
+        lock: new Map(),
+        history: new Map()
+    };
+
+    amContext.history.set(name, amContext.history.get(name) ?? []);
+    return amContext;
+}
+
+async function getContextHistoryAfterLock(astro: ContextAstro, name: string, lock?: string) {
+    const contexts: AMContext = getAMContextFromAstro(astro, name);
+
+    while (contexts.lock.get(lock)) {
+        await contexts.lock.get(lock);
+    }
+
+    return {
+        value: contexts.history.get(name),
+        lock: contexts.lock
+    };
 }
 
 export default function getContext(astro: ContextAstro, name = "default") {
-    return getContextHistory(astro, name).at(-1) ?? {};
+    const contexts: AMContext = getAMContextFromAstro(astro, name);
+    return contexts.history.get(name).at(-1) ?? {};
 }
 
 type AsyncContextOptions = { name?: string, context?: any, lock?: string; };
-const activeLock: Map<string, Map<string, Promise<void>>> = new Map();
 
 export async function asyncContext<T>(promise: () => Promise<T>, astro: ContextAstro, { name = "default", context = null, lock }: AsyncContextOptions = {}): Promise<T> {
-    activeLock.set(name, activeLock.get(name) ?? new Map());
-    const lockContext = activeLock.get(name);
-
-    while (lockContext.get(lock)) {
-        await lockContext.get(lock);
-    }
-
-    const contextHistory = getContextHistory(astro, name);
+    const contextHistory = await getContextHistoryAfterLock(astro, name);
 
-    contextHistory.push({
-        ...contextHistory.at(-1),
+    contextHistory.value.push({
+        ...contextHistory.value.at(-1),
         ...(context ?? astro.props)
     });
 
     let resolver: () => void | null;
     if (lock) {
-        lockContext.set(lock, new Promise<void>(resolve => resolver = resolve));
+        contextHistory.lock.set(lock, new Promise<void>(resolve => resolver = resolve));
     }
 
     try {
         const response = await promise();
-        contextHistory.pop();
+        contextHistory.value.pop();
         return response;
     } finally {
-        lockContext.delete(lock);
-        if(lockContext.size === 0) {
-            activeLock.delete(name);
-        }
+        contextHistory.lock.delete(lock);
         resolver?.();
     }
 }