Fixed closure issue with state decorators when running multiple instances

Author: rpanic

packages/module/src/runtime/Runtime.ts

@@ -311,9 +311,7 @@ export class Runtime<Modules extends RuntimeModulesRecord>
   }
 
   public get stateServiceProvider(): StateServiceProvider {
-    return this.dependencyContainer.resolve<StateServiceProvider>(
-      "StateServiceProvider"
-    );
+    return this.container.resolve<StateServiceProvider>("StateServiceProvider");
   }
 
   public get stateService(): SimpleAsyncStateService {

packages/protocol/src/state/protocol/ProtocolState.ts

@@ -27,42 +27,6 @@ export interface StatefulModule {
   };
 }
 
-export function createStateGetter<TargetModule extends StatefulModule>(
-  target: TargetModule,
-  propertyKey: string,
-  valueReference: Reference<State<unknown> | undefined>,
-  prefix: string,
-  debugInfo: { parentName: string; baseModuleNames: string }
-) {
-  return function getter(this: TargetModule) {
-    // const self = this;
-    const { value } = valueReference;
-    // Short-circuit this to return the state in case its already initialized
-    if (value !== undefined && value.path !== undefined) {
-      return value;
-    }
-
-    if (this.name === undefined) {
-      throw errors.missingName(this.constructor.name);
-    }
-
-    if (!this.parent) {
-      throw errors.missingParent(
-        this.constructor.name,
-        debugInfo.parentName,
-        debugInfo.baseModuleNames
-      );
-    }
-
-    const path = Path.fromProperty(this.name, propertyKey, prefix);
-    if (value) {
-      value.path = path;
-      value.stateServiceProvider = this.parent.stateServiceProvider;
-    }
-    return value;
-  };
-}
-
 /**
  * Decorates a runtime module property as state, passing down some
  * underlying values to improve developer experience.
@@ -72,31 +36,62 @@ export function state() {
     target: TargetTransitioningModule,
     propertyKey: string
   ) => {
-    const stateReference = createReference<State<unknown> | undefined>(
-      undefined
-    );
-
-    const isProtocol = target instanceof TransitioningProtocolModule;
-    const statePrefix = isProtocol
-      ? PROTOKIT_PREFIXES.STATE_PROTOCOL
-      : PROTOKIT_PREFIXES.STATE_RUNTIME;
-    const debugInfo = isProtocol
-      ? { parentName: "protocol", baseModuleNames: "...Hook" }
-      : { parentName: "runtime", baseModuleNames: "RuntimeModule" };
-
     Object.defineProperty(target, propertyKey, {
       enumerable: true,
 
-      get: createStateGetter(
-        target,
-        propertyKey,
-        stateReference,
-        statePrefix,
-        debugInfo
-      ),
+      get: function get(this: TargetTransitioningModule) {
+        // The reason for why we store the state value in this weird way is that
+        // in the decorator on the prototype of the class. This means that if there
+        // are multiple instances of this class, any closure that this getter shares
+        // will be the same for all instances.
+        // Therefore, we need to somehow save the set instance on the instance itself
+
+        // eslint-disable-next-line max-len
+        // eslint-disable-next-line @typescript-eslint/consistent-type-assertions,@typescript-eslint/no-unsafe-assignment
+        const reference: Reference<State<unknown>> | undefined = (this as any)[
+          `protokit_state_cache_${propertyKey}`
+        ];
+
+        // Short-circuit this to return the state in case its already initialized
+        if (reference !== undefined && reference.value.path !== undefined) {
+          return reference.value;
+        }
+
+        if (this.name === undefined) {
+          throw errors.missingName(this.constructor.name);
+        }
+
+        const isProtocol = target instanceof TransitioningProtocolModule;
+
+        if (!this.parent) {
+          const debugInfo = isProtocol
+            ? { parentName: "protocol", baseModuleNames: "...Hook" }
+            : { parentName: "runtime", baseModuleNames: "RuntimeModule" };
+
+          throw errors.missingParent(
+            this.constructor.name,
+            debugInfo.parentName,
+            debugInfo.baseModuleNames
+          );
+        }
+
+        const statePrefix = isProtocol
+          ? PROTOKIT_PREFIXES.STATE_PROTOCOL
+          : PROTOKIT_PREFIXES.STATE_RUNTIME;
+        const path = Path.fromProperty(this.name, propertyKey, statePrefix);
+        if (reference) {
+          const { value } = reference;
+          value.path = path;
+          value.stateServiceProvider = this.parent.stateServiceProvider;
+        }
+        return reference?.value;
+      },
 
-      set: (newValue: State<unknown>) => {
-        stateReference.value = newValue;
+      set: function set(
+        this: TargetTransitioningModule & any,
+        newValue: State<unknown>
+      ) {
+        this[`protokit_state_cache_${propertyKey}`] = createReference(newValue);
       },
     });
   };

packages/sequencer/src/sequencer/executor/Sequencer.ts

@@ -120,6 +120,8 @@ export class Sequencer<Modules extends SequencerModulesRecord>
       // eslint-disable-next-line no-await-in-loop
       await sequencerModule.start();
     }
+
+    // TODO This currently also warns for client appchains
     if (!moduleClassNames.includes("SequencerStartupModule")) {
       log.warn("SequencerStartupModule is not defined.");
     }

packages/stack/test/graphql/Post.ts

@@ -1,10 +1,5 @@
-import {
-  RuntimeModule,
-  runtimeMethod,
-  runtimeModule,
-  state,
-} from "@proto-kit/module";
-import { StateMap } from "@proto-kit/protocol";
+import { RuntimeModule, runtimeMethod, runtimeModule } from "@proto-kit/module";
+import { StateMap, state } from "@proto-kit/protocol";
 import {
   CircuitString,
   Field,

packages/stack/test/graphql/graphql.test.ts

@@ -20,6 +20,7 @@ import {
   GraphqlNetworkStateTransportModule,
 } from "@proto-kit/sdk";
 import { beforeAll } from "@jest/globals";
+import { container } from "tsyringe";
 
 import { startServer, TestBalances } from "../../src/scripts/graphql/server";
 
@@ -103,7 +104,7 @@ describe("graphql client test", () => {
 
     appChain = prepareClient();
 
-    await appChain.start();
+    await appChain.start(false, container.createChildContainer());
 
     trigger = server.sequencer.resolveOrFail(
       "BlockTrigger",