From 116480137693fa4182072dfad0fe5c7a7e6d0779 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Mon, 9 Sep 2019 16:41:04 -0500
Subject: [PATCH] Keep a maximum number of connections instead of a timeout

There's no way to actually know if those clients have gone away, so it
seems it might be better to base it on whether the user has connected
again with new clients to determine if the old clients are now invalid.
---
 src/connection.ts | 74 ++++++++++++++++++++++++-----------------------
 src/server.ts     | 12 ++++++++
 2 files changed, 50 insertions(+), 36 deletions(-)

diff --git a/src/connection.ts b/src/connection.ts
index 15bf1c01..f872ff6c 100644
--- a/src/connection.ts
+++ b/src/connection.ts
@@ -12,48 +12,54 @@ import { uriTransformerPath } from "vs/server/src/util";
 import { IExtHostReadyMessage } from "vs/workbench/services/extensions/common/extensionHostProtocol";
 
 export abstract class Connection {
-	protected readonly _onClose = new Emitter<void>();
+	private readonly _onClose = new Emitter<void>();
 	public readonly onClose = this._onClose.event;
-	protected disposed: boolean = false;
-	public constructor(protected protocol: Protocol) {}
+	private disposed = false;
+	private _offline: number | undefined;
+
+	public constructor(protected protocol: Protocol) {
+		protocol.onClose(() => this.dispose()); // Explicit close.
+		protocol.onSocketClose(() => this._offline = Date.now()); // Might reconnect.
+	}
+
+	public get offline(): number | undefined {
+		return this._offline;
+	}
+
+	public reconnect(socket: ISocket, buffer: VSBuffer): void {
+		this._offline = undefined;
+		this.doReconnect(socket, buffer);
+	}
+
+	public dispose(): void {
+		if (!this.disposed) {
+			this.disposed = true;
+			this.doDispose();
+			this._onClose.fire();
+		}
+	}
+
 	/**
 	 * Set up the connection on a new socket.
 	 */
-	public abstract reconnect(socket: ISocket, buffer: VSBuffer): void;
-	protected abstract dispose(): void;
+	protected abstract doReconnect(socket: ISocket, buffer: VSBuffer): void;
+	protected abstract doDispose(): void;
 }
 
 /**
  * Used for all the IPC channels.
  */
 export class ManagementConnection extends Connection {
-	private timeout: NodeJS.Timeout | undefined;
-	private readonly wait = 1000 * 60;
-
-	public constructor(protocol: Protocol) {
-		super(protocol);
-		protocol.onClose(() => this.dispose());
-		protocol.onSocketClose(() => {
-			this.timeout = setTimeout(() => this.dispose(), this.wait);
-		});
+	protected doDispose(): void {
+		this.protocol.sendDisconnect();
+		this.protocol.dispose();
+		this.protocol.getSocket().end();
 	}
 
-	public reconnect(socket: ISocket, buffer: VSBuffer): void {
-		clearTimeout(this.timeout as any); // Not sure why the type doesn't work.
+	protected doReconnect(socket: ISocket, buffer: VSBuffer): void {
 		this.protocol.beginAcceptReconnection(socket, buffer);
 		this.protocol.endAcceptReconnection();
 	}
-
-	protected dispose(): void {
-		if (!this.disposed) {
-			clearTimeout(this.timeout as any); // Not sure why the type doesn't work.
-			this.disposed = true;
-			this.protocol.sendDisconnect();
-			this.protocol.dispose();
-			this.protocol.getSocket().end();
-			this._onClose.fire();
-		}
-	}
 }
 
 export class ExtensionHostConnection extends Connection {
@@ -70,18 +76,14 @@ export class ExtensionHostConnection extends Connection {
 		this.protocol.getUnderlyingSocket().pause();
 	}
 
-	protected dispose(): void {
-		if (!this.disposed) {
-			this.disposed = true;
-			if (this.process) {
-				this.process.kill();
-			}
-			this.protocol.getSocket().end();
-			this._onClose.fire();
+	protected doDispose(): void {
+		if (this.process) {
+			this.process.kill();
 		}
+		this.protocol.getSocket().end();
 	}
 
-	public reconnect(socket: ISocket, buffer: VSBuffer): void {
+	protected doReconnect(socket: ISocket, buffer: VSBuffer): void {
 		// This is just to set the new socket.
 		this.protocol.beginAcceptReconnection(socket, null);
 		this.protocol.dispose();
diff --git a/src/server.ts b/src/server.ts
index b778141c..f07e41f8 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -426,6 +426,7 @@ export class MainServer extends Server {
 	public readonly onDidClientConnect = this._onDidClientConnect.event;
 	private readonly ipc = new IPCServer(this.onDidClientConnect);
 
+	private readonly maxOfflineConnections = 5;
 	private readonly connections = new Map<ConnectionType, Map<string, Connection>>();
 
 	private readonly services = new ServiceCollection();
@@ -586,6 +587,7 @@ export class MainServer extends Server {
 					);
 				}
 				connections.set(token, connection);
+				this.disposeOldOfflineConnections();
 				connection.onClose(() => connections.delete(token));
 				break;
 			case ConnectionType.Tunnel: return protocol.tunnel();
@@ -593,6 +595,16 @@ export class MainServer extends Server {
 		}
 	}
 
+	private disposeOldOfflineConnections(): void {
+		this.connections.forEach((connections) => {
+			const offline = Array.from(connections.values())
+				.filter((connection) => typeof connection.offline !== "undefined");
+			for (let i = 0, max = offline.length - this.maxOfflineConnections; i < max; ++i) {
+				offline[i].dispose();
+			}
+		});
+	}
+
 	private async initializeServices(args: ParsedArgs): Promise<void> {
 		const environmentService = new EnvironmentService(args, process.execPath);
 		const logService = new SpdLogService(RemoteExtensionLogFileName, environmentService.logsPath, getLogLevel(environmentService));