preserve hook context reference across hooks

Signed-off-by: Adam Wootton <[email protected]>

Author: ajwootto

libs/providers/multi-provider/src/lib/hook-executor.ts

@@ -7,18 +7,17 @@ export class HookExecutor {
   constructor(private logger: Logger) {}
 
   async beforeHooks(hooks: Hook[] | undefined, hookContext: HookContext, hints: HookHints) {
-    const newContext = { ...hookContext.context };
     for (const hook of hooks ?? []) {
       // freeze the hookContext
       Object.freeze(hookContext);
 
-      Object.assign(newContext, {
+      Object.assign(hookContext.context, {
         ...(await hook?.before?.(hookContext, Object.freeze(hints))),
       });
     }
 
     // after before hooks, freeze the EvaluationContext.
-    return Object.freeze(newContext);
+    return Object.freeze(hookContext.context);
   }
 
   async afterHooks(

libs/providers/multi-provider/src/lib/multi-provider.spec.ts

@@ -267,11 +267,14 @@ describe('MultiProvider', () => {
           logger: logger,
         };
 
+        let weakMap = new WeakMap();
+
         provider1.hooks = [
           {
             before: async (context) => {
               hook1Called = true;
-              expect(context).toBe(hookContext);
+              expect(context).toEqual(hookContext);
+              weakMap.set(context, 'test');
               return { ...context.context, hook1: true };
             },
             after: async (context) => {
@@ -280,13 +283,18 @@ describe('MultiProvider', () => {
                 hook1: true,
                 hook2: true,
               });
+              expect(weakMap.get(context)).toEqual('test');
               after1Called = true;
             },
           },
           {
             before: async (context) => {
               hook2Called = true;
-              expect(context).toBe(hookContext);
+              expect(weakMap.get(context)).toEqual('test');
+              expect(context.context).toEqual({
+                test: true,
+                hook1: true,
+              });
               return { ...context.context, hook2: true };
             },
           },
@@ -295,6 +303,7 @@ describe('MultiProvider', () => {
         provider2.hooks = [
           {
             after: async (context) => {
+              expect(weakMap.get(context)).toBeFalsy();
               expect(context.context).toEqual({
                 test: true,
               });
@@ -339,7 +348,7 @@ describe('MultiProvider', () => {
         expect(after2Called).toBe(true);
       });
 
-      it('runs error hook and finally hook with modified context', async () => {
+      it('runs error hook and finally hook with modified context using same object reference', async () => {
         const provider1 = new TestProvider();
         let hook1Called = false;
         let error1Called = false;
@@ -359,18 +368,22 @@ describe('MultiProvider', () => {
           logger: logger,
         };
 
+        let weakMap = new WeakMap();
+
         provider1.hooks = [
           {
             before: async (context) => {
               hook1Called = true;
-              expect(context).toBe(hookContext);
+              weakMap.set(context, 'exists');
+              expect(context).toEqual(hookContext);
               return { ...context.context, hook1: true };
             },
             error: async (context) => {
               expect(context.context).toEqual({
                 test: true,
                 hook1: true,
               });
+              expect(weakMap.get(context)).toEqual('exists');
               error1Called = true;
               throw new Error('error hook error');
             },
@@ -379,6 +392,7 @@ describe('MultiProvider', () => {
                 test: true,
                 hook1: true,
               });
+              expect(weakMap.get(context)).toEqual('exists');
               finally1Called = true;
             },
           },
@@ -392,6 +406,7 @@ describe('MultiProvider', () => {
 
         provider1.resolveBooleanEvaluation.mockRejectedValue(new Error('test error'));
 
+        // call the multiprovider before hook to set up the hookcontext
         await multiProvider.hooks[0].before!(hookContext);
         await expect(() => multiProvider.resolveBooleanEvaluation('flag', false, context)).rejects.toThrow();
         expect(hook1Called).toBe(true);

libs/providers/multi-provider/src/lib/multi-provider.ts

@@ -241,8 +241,12 @@ export class MultiProvider implements Provider {
     let providerContext: EvaluationContext | undefined = undefined;
     let evaluationResult: ResolutionDetails<T>;
 
+    // create a copy of the shared hook context because we're going to mutate the evaluation context
+    const hookContextCopy = { ...hookContext, context: { ...hookContext.context } };
+
     try {
-      providerContext = await this.hookExecutor.beforeHooks(provider.hooks, hookContext, hookHints);
+      // return the modified provider context and mutate the hook context to contain it
+      providerContext = await this.hookExecutor.beforeHooks(provider.hooks, hookContextCopy, hookHints);
 
       evaluationResult = (await this.callProviderResolve(
         provider,
@@ -257,27 +261,13 @@ export class MultiProvider implements Provider {
         flagKey,
       };
 
-      await this.hookExecutor.afterHooks(
-        provider.hooks,
-        { ...hookContext, context: providerContext },
-        afterHookEvalDetails,
-        hookHints,
-      );
+      await this.hookExecutor.afterHooks(provider.hooks, hookContextCopy, afterHookEvalDetails, hookHints);
       return evaluationResult;
     } catch (error: unknown) {
-      await this.hookExecutor.errorHooks(
-        provider.hooks,
-        { ...hookContext, context: providerContext ?? hookContext.context },
-        error,
-        hookHints,
-      );
+      await this.hookExecutor.errorHooks(provider.hooks, hookContextCopy, error, hookHints);
       throw error;
     } finally {
-      await this.hookExecutor.finallyHooks(
-        provider.hooks,
-        { ...hookContext, context: providerContext ?? hookContext.context },
-        hookHints,
-      );
+      await this.hookExecutor.finallyHooks(provider.hooks, hookContextCopy, hookHints);
     }
   }