From 36bf392f6d3192170a97d4340b2917543f255e18 Mon Sep 17 00:00:00 2001
From: Matt Bierner <matb@microsoft.com>
Date: Fri, 20 Sep 2019 14:09:26 -0700
Subject: [PATCH] Make standard tunnel service re-use tunnels when possible

---
 .../services/remote/node/tunnelService.ts     | 44 +++++++++++++++++--
 1 file changed, 41 insertions(+), 3 deletions(-)

diff --git a/src/vs/workbench/services/remote/node/tunnelService.ts b/src/vs/workbench/services/remote/node/tunnelService.ts
index 270f57fcadf..9b669b9cd8f 100644
--- a/src/vs/workbench/services/remote/node/tunnelService.ts
+++ b/src/vs/workbench/services/remote/node/tunnelService.ts
@@ -88,12 +88,20 @@ class NodeRemoteTunnel extends Disposable implements RemoteTunnel {
 export class TunnelService implements ITunnelService {
 	_serviceBrand: undefined;
 
+	private readonly _tunnels = new Map</* port */ number, { refcount: number, readonly value: Promise<RemoteTunnel> }>();
+
 	public constructor(
 		@IWorkbenchEnvironmentService private readonly environmentService: IWorkbenchEnvironmentService,
 		@IRemoteAuthorityResolverService private readonly remoteAuthorityResolverService: IRemoteAuthorityResolverService,
 		@ISignService private readonly signService: ISignService,
-		@ILogService private readonly logService: ILogService
-	) {
+		@ILogService private readonly logService: ILogService,
+	) { }
+
+	dispose(): void {
+		for (const { value } of this._tunnels.values()) {
+			value.then(tunnel => tunnel.dispose());
+		}
+		this._tunnels.clear();
 	}
 
 	openTunnel(remotePort: number): Promise<RemoteTunnel> | undefined {
@@ -102,6 +110,33 @@ export class TunnelService implements ITunnelService {
 			return undefined;
 		}
 
+		const resolvedTunnel = this.retainOrCreateTunnel(remoteAuthority, remotePort);
+		if (!resolvedTunnel) {
+			return resolvedTunnel;
+		}
+
+		return resolvedTunnel.then(tunnel => ({
+			tunnelRemotePort: tunnel.tunnelRemotePort,
+			tunnelLocalPort: tunnel.tunnelLocalPort,
+			dispose: () => {
+				const existing = this._tunnels.get(remotePort);
+				if (existing) {
+					if (--existing.refcount <= 0) {
+						existing.value.then(tunnel => tunnel.dispose());
+						this._tunnels.delete(remotePort);
+					}
+				}
+			}
+		}));
+	}
+
+	private retainOrCreateTunnel(remoteAuthority: string, remotePort: number): Promise<RemoteTunnel> | undefined {
+		const existing = this._tunnels.get(remotePort);
+		if (existing) {
+			++existing.refcount;
+			return existing.value;
+		}
+
 		const options: IConnectionOptions = {
 			commit: product.commit,
 			socketFactory: nodeSocketFactory,
@@ -114,7 +149,10 @@ export class TunnelService implements ITunnelService {
 			signService: this.signService,
 			logService: this.logService
 		};
-		return createRemoteTunnel(options, remotePort);
+
+		const tunnel = createRemoteTunnel(options, remotePort);
+		this._tunnels.set(remotePort, { refcount: 1, value: tunnel });
+		return tunnel;
 	}
 }