diff --git a/packages/client/src/client.ts b/packages/client/src/client.ts index cfa0c95afe..5914c988bf 100644 --- a/packages/client/src/client.ts +++ b/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 { + // 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(opNames.register, { + const toClean = await this.client.request(opNames.register, { uuid: agent.uuid, containers, kinds: agent.kinds, @@ -335,6 +340,23 @@ export class NetworkClientImpl implements NetworkClient { return await this.client.request(opNames.getContainer, { uuid, request }) } + async retryGetContainerRef (uuid: ContainerUuid, request: ContainerRequest): Promise { + 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 { await this.client.request(opNames.releaseContainer, { uuid }) } diff --git a/packages/core/src/agent.ts b/packages/core/src/agent.ts index b1a0b3d2d5..2ad41739db 100644 --- a/packages/core/src/agent.ts +++ b/packages/core/src/agent.ts @@ -24,7 +24,7 @@ export class AgentImpl implements NetworkAgent { // Own, managed containers private readonly _byId = new Map>() - private readonly _containers = new Map() + private readonly _containers = new Map() 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 { - const current = this._containers.get(endpoint) + async terminate (uuid: ContainerUuid): Promise { + 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() + } } } diff --git a/packages/core/src/api/agent.ts b/packages/core/src/api/agent.ts index 1f57d1658a..bb1b9ee903 100644 --- a/packages/core/src/api/agent.ts +++ b/packages/core/src/api/agent.ts @@ -53,5 +53,5 @@ export interface NetworkAgent { request: (target: ContainerUuid, operation: string, data?: any) => Promise // ask for immediate termination for container - terminate: (container: ContainerEndpointRef) => Promise + terminate: (container: ContainerUuid) => Promise } diff --git a/packages/core/src/api/network.ts b/packages/core/src/api/network.ts index c19adab485..456122a3d2 100644 --- a/packages/core/src/api/network.ts +++ b/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 + register: (record: AgentRecord, agent: NetworkAgent) => Promise // Unregister an agent from the network. // Will call terminate for every connection/references. diff --git a/packages/core/src/network.ts b/packages/core/src/network.ts index e4beb1165b..0c6763d9fe 100644 --- a/packages/core/src/network.ts +++ b/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 { + async register (record: AgentRecord, agent: NetworkAgent): Promise { const containers: ContainerRecord[] = record.containers const newContainers = new Map( 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 { 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) } /** diff --git a/packages/server/src/server.ts b/packages/server/src/server.ts index 9d7e40780e..7c837f4edc 100644 --- a/packages/server/src/server.ts +++ b/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 { - // Ignore + async terminate (containerUuid: ContainerUuid): Promise { + return await this.rpcServer.request(this.client, opNames.terminate, [this.uuid, [containerUuid]]) } async getContainer (uuid: ContainerUuid): Promise { diff --git a/pods/network-tool/src/benchmark.ts b/pods/network-tool/src/benchmark.ts index 1e03e0b2f0..1374d0edf9 100644 --- a/pods/network-tool/src/benchmark.ts +++ b/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) diff --git a/pods/network-tool/src/utils.ts b/pods/network-tool/src/utils.ts index 3d120c04c9..ec400b86e2 100644 --- a/pods/network-tool/src/utils.ts +++ b/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() })