reduce task leaking (#2868)

* reduce task leaking

* fix onLeaveContext

container-runtime/src/Adapters/EffectCreator.ts

@@ -104,7 +104,18 @@ const rpcRoundFor =
 export function makeEffects(context: EffectContext): Effects {
   const rpcRound = rpcRoundFor(context.procedureId)
   const self: Effects = {
+    child: (name) =>
+      makeEffects({ ...context, callbacks: context.callbacks?.child(name) }),
     constRetry: context.constRetry,
+    isInContext: !!context.callbacks,
+    onLeaveContext:
+      context.callbacks?.onLeaveContext?.bind(context.callbacks) ||
+      (() => {
+        console.warn(
+          "no context for this effects object",
+          new Error().stack?.replace(/^Error/, ""),
+        )
+      }),
     clearCallbacks(...[options]: Parameters<T.Effects["clearCallbacks"]>) {
       return rpcRound("clear-callbacks", {
         ...options,
@@ -313,5 +324,14 @@ export function makeEffects(context: EffectContext): Effects {
       >
     },
   }
+  self.onLeaveContext(() => {
+    self.isInContext = false
+    self.onLeaveContext = () => {
+      console.warn(
+        "this effects object is already out of context",
+        new Error().stack?.replace(/^Error/, ""),
+      )
+    }
+  })
   return self
 }

container-runtime/src/Adapters/RpcListener.ts

@@ -238,21 +238,6 @@ export class RpcListener {
     return this._system
   }
 
-  private callbackHolders: Map<string, CallbackHolder> = new Map()
-  private removeCallbackHolderFor(procedure: string) {
-    const prev = this.callbackHolders.get(procedure)
-    if (prev) {
-      this.callbackHolders.delete(procedure)
-      this.callbacks?.removeChild(prev)
-    }
-  }
-  private callbackHolderFor(procedure: string): CallbackHolder {
-    this.removeCallbackHolderFor(procedure)
-    const callbackHolder = this.callbacks!.child()
-    this.callbackHolders.set(procedure, callbackHolder)
-    return callbackHolder
-  }
-
   callCallback(callback: number, args: any[]): void {
     if (this.callbacks) {
       this.callbacks
@@ -302,7 +287,7 @@ export class RpcListener {
         return null
       })
       .when(startType, async ({ id }) => {
-        const callbacks = this.callbackHolderFor("main")
+        const callbacks = this.callbacks?.child("main")
         const effects = makeEffects({
           procedureId: null,
           callbacks,
@@ -313,7 +298,7 @@ export class RpcListener {
         )
       })
       .when(stopType, async ({ id }) => {
-        this.removeCallbackHolderFor("main")
+        this.callbacks?.removeChild("main")
         return handleRpc(
           id,
           this.system.stop().then((result) => ({ result })),
@@ -338,7 +323,7 @@ export class RpcListener {
                   procedureId: null,
                 }),
               )
-              const callbacks = this.callbackHolderFor("containerInit")
+              const callbacks = this.callbacks.child("containerInit")
               await system.containerInit(
                 makeEffects({
                   procedureId: null,
@@ -420,7 +405,7 @@ export class RpcListener {
     ): { result: any } => {
       return { result }
     }
-    const callbacks = this.callbackHolderFor(procedure)
+    const callbacks = this.callbacks?.child(procedure)
     const effects = makeEffects({
       procedureId,
       callbacks,

container-runtime/src/Models/CallbackHolder.ts

@@ -14,7 +14,8 @@ export class CallbackHolder {
   constructor(private effects?: T.Effects) {}
 
   private callbacks = new Map<number, Function>()
-  private children: WeakRef<CallbackHolder>[] = []
+  private onLeaveContextCallbacks: Function[] = []
+  private children: Map<string, CallbackHolder> = new Map()
   private newId() {
     return CallbackIdCell.inc++
   }
@@ -32,23 +33,25 @@ export class CallbackHolder {
       })
     return id
   }
-  child(): CallbackHolder {
+  child(name: string): CallbackHolder {
+    this.removeChild(name)
     const child = new CallbackHolder()
-    this.children.push(new WeakRef(child))
+    this.children.set(name, child)
     return child
   }
-  removeChild(child: CallbackHolder) {
-    this.children = this.children.filter((c) => {
-      const ref = c.deref()
-      return ref && ref !== child
-    })
+  removeChild(name: string) {
+    const child = this.children.get(name)
+    if (child) {
+      child.leaveContext()
+      this.children.delete(name)
+    }
   }
   private getCallback(index: number): Function | undefined {
     let callback = this.callbacks.get(index)
     if (callback) this.callbacks.delete(index)
     else {
-      for (let i = 0; i < this.children.length; i++) {
-        callback = this.children[i].deref()?.getCallback(index)
+      for (let [_, child] of this.children) {
+        callback = child.getCallback(index)
         if (callback) return callback
       }
     }
@@ -59,4 +62,21 @@ export class CallbackHolder {
     if (!callback) return Promise.resolve()
     return Promise.resolve().then(() => callback(...args))
   }
+  onLeaveContext(fn: Function) {
+    this.onLeaveContextCallbacks.push(fn)
+  }
+  leaveContext() {
+    for (let [_, child] of this.children) {
+      child.leaveContext()
+    }
+    this.children = new Map()
+    for (let fn of this.onLeaveContextCallbacks) {
+      try {
+        fn()
+      } catch (e) {
+        console.warn(e)
+      }
+    }
+    this.onLeaveContextCallbacks = []
+  }
 }