Fixes for container auto termination

packages/client/src/client.ts

@@ -157,6 +157,10 @@ export class NetworkClientImpl implements NetworkClient {
       case opNames.sendContainer:
         await send(await agent.request(agentParams[0], agentParams[1], agentParams[2]))
         break
+      case opNames.terminate:
+        await agent.terminate(agentParams[0] as ContainerUuid)
+        await send('')
+        break
       default:
         throw new Error('Unknown method')
     }
@@ -193,8 +197,9 @@ export class NetworkClientImpl implements NetworkClient {
   }
 
   async onRegister (): Promise<void> {
+    // TODO: Add retry in container requests on re-connect to new network.
     for (const [uuid, ref] of this.references.entries()) {
-      await this.handleRefUpdate(uuid, await this.getContainerRef(uuid, ref.request))
+      await this.handleRefUpdate(uuid, await this.retryGetContainerRef(uuid, ref.request))
     }
     this.registered = true
     // We need to re-register all our managed agents
@@ -241,7 +246,7 @@ export class NetworkClientImpl implements NetworkClient {
         lastVisit: container.lastVisit
       } satisfies ContainerRecord)
     }
-    const toClean = await this.client.request<ContainerEndpointRef[]>(opNames.register, {
+    const toClean = await this.client.request<ContainerUuid[]>(opNames.register, {
       uuid: agent.uuid,
       containers,
       kinds: agent.kinds,
@@ -335,6 +340,23 @@ export class NetworkClientImpl implements NetworkClient {
     return await this.client.request<ContainerEndpointRef>(opNames.getContainer, { uuid, request })
   }
 
+  async retryGetContainerRef (uuid: ContainerUuid, request: ContainerRequest): Promise<ContainerEndpointRef> {
+    let waitTimeout: number = 1
+    while (true) {
+      try {
+        const ref = await this.getContainerRef(uuid, request)
+        if (waitTimeout > 1) {
+          console.log(`Successfully got container ref for ${uuid} after ${waitTimeout - 1} retries.`)
+        }
+        return ref
+      } catch (err) {
+        console.warn(`Error getting container ref for ${uuid}. Will retry...`)
+        await this.tickMgr.waitTick(waitTimeout)
+        waitTimeout++
+      }
+    }
+  }
+
   async release (uuid: ContainerUuid): Promise<void> {
     await this.client.request<any>(opNames.releaseContainer, { uuid })
   }

packages/core/src/agent.ts

@@ -24,7 +24,7 @@ export class AgentImpl implements NetworkAgent {
   // Own, managed containers
   private readonly _byId = new Map<ContainerUuid, ContainerRecordImpl | Promise<ContainerRecordImpl>>()
 
-  private readonly _containers = new Map<ContainerEndpointRef, ContainerRecordImpl>()
+  private readonly _containers = new Map<ContainerUuid, ContainerRecordImpl>()
 
   endpoint?: AgentEndpointRef | undefined
 
@@ -94,17 +94,22 @@ export class AgentImpl implements NetworkAgent {
     )
     this._byId.set(uuid, container)
     container = await container
-    this._containers.set(container.endpoint, container)
+    this._containers.set(uuid, container)
     this._byId.set(uuid, container)
 
     return container.endpoint
   }
 
-  async terminate (endpoint: ContainerEndpointRef): Promise<void> {
-    const current = this._containers.get(endpoint)
+  async terminate (uuid: ContainerUuid): Promise<void> {
+    const current = this._byId.get(uuid)
     if (current !== undefined) {
-      this._containers.delete(endpoint)
-      await current.container.terminate()
+      this._containers.delete(uuid)
+      this._byId.delete(uuid)
+      if (current instanceof Promise) {
+        await (await current).container.terminate() // Await promise before terminating
+      } else {
+        await current.container.terminate()
+      }
     }
   }
 

packages/core/src/api/agent.ts

@@ -53,5 +53,5 @@ export interface NetworkAgent {
   request: (target: ContainerUuid, operation: string, data?: any) => Promise<any>
 
   // ask for immediate termination for container
-  terminate: (container: ContainerEndpointRef) => Promise<void>
+  terminate: (container: ContainerUuid) => Promise<void>
 }

packages/core/src/api/network.ts

@@ -18,7 +18,7 @@ export interface Network {
    * Register or reregister agent in network.
    * On every network restart agent should reconnect to network.
    */
-  register: (record: AgentRecord, agent: NetworkAgent) => Promise<ContainerEndpointRef[]>
+  register: (record: AgentRecord, agent: NetworkAgent) => Promise<ContainerUuid[]>
 
   // Unregister an agent from the network.
   // Will call terminate for every connection/references.

packages/core/src/network.ts

@@ -109,7 +109,7 @@ export class NetworkImpl implements Network, NetworkWithClients {
     return await agent.api.request(target, operation, data)
   }
 
-  async register (record: AgentRecord, agent: NetworkAgent): Promise<ContainerEndpointRef[]> {
+  async register (record: AgentRecord, agent: NetworkAgent): Promise<ContainerUuid[]> {
     const containers: ContainerRecord[] = record.containers
     const newContainers = new Map<ContainerUuid, ContainerRecordImpl>(
       containers.map((record) => [
@@ -152,17 +152,21 @@ export class NetworkImpl implements Network, NetworkWithClients {
       }
     }
 
-    const containersToShutdown: ContainerEndpointRef[] = []
+    const containersToShutdown: ContainerUuid[] = []
 
     // Update active container registry.
     for (const rec of containers) {
       const oldAgentId = this._containers.get(rec.uuid)
       if (oldAgentId === undefined) {
         containerEvent.added.push(rec)
+        const containerImpl = newContainers.get(rec.uuid)
+        if (containerImpl !== undefined) {
+          this._orphanedContainers.set(rec.endpoint, containerImpl)
+        }
         this._containers.set(rec.uuid, record.agentId)
       }
-      if (oldAgentId !== record.agentId) {
-        containersToShutdown.push(rec.endpoint)
+      if (oldAgentId !== undefined && oldAgentId !== record.agentId) {
+        containersToShutdown.push(rec.uuid)
       }
     }
 
@@ -357,6 +361,7 @@ export class NetworkImpl implements Network, NetworkWithClients {
 
   async terminate (container: ContainerRecordImpl): Promise<void> {
     this._containers.delete(container.record.uuid) // Remove from active container registry
+    this._orphanedContainers.delete(container.record.endpoint)
     this.eventQueue.push({
       added: [],
       deleted: [container.record],
@@ -365,11 +370,7 @@ export class NetworkImpl implements Network, NetworkWithClients {
     const agent = this._agents.get(container.record.agentId)
     agent?.containers.delete(container.record.uuid)
 
-    let endpoint = container.endpoint
-    if (endpoint instanceof Promise) {
-      endpoint = await endpoint
-    }
-    await agent?.api.terminate(endpoint)
+    await agent?.api.terminate(container.record.uuid)
   }
 
   /**

packages/server/src/server.ts

@@ -37,8 +37,8 @@ class AgentCallbackHandler implements NetworkAgent {
     return await this.rpcServer.request(this.client, opNames.sendContainer, [this.uuid, [target, operation, data]])
   }
 
-  async terminate (): Promise<void> {
-    // Ignore
+  async terminate (containerUuid: ContainerUuid): Promise<void> {
+    return await this.rpcServer.request(this.client, opNames.terminate, [this.uuid, [containerUuid]])
   }
 
   async getContainer (uuid: ContainerUuid): Promise<Container | undefined> {

pods/network-tool/src/benchmark.ts

@@ -48,6 +48,7 @@ export function registerBenchmark (): void {
       []
     )
     .action(async (cmd: { network: string, label: string[], endpoint: string }) => {
+      console.log('Starting benchmark agent')
       const network = process.env.NETWORK_HOST ?? cmd.network
       const client = createNetworkClient(network)
 

pods/network-tool/src/utils.ts

@@ -29,6 +29,11 @@ export function registerShutdown (): void {
     void shutdown()
   })
 
+  // Handle Ctrl+C in console
+  process.on('SIGBREAK', (): void => {
+    void shutdown()
+  })
+
   process.on('exit', (): void => {
     void shutdown()
   })