Retry up to 5 times the initial authority resolving and also retry up to 5 times the initial connection (#156444)

Fixes #127565: Retry up to 5 times the initial authority resolving and retry up to 5 times the initial connection

2 files changed, 66 insertions, 27 deletions

diff --git a/src/vs/platform/remote/common/remoteAgentConnection.ts b/src/vs/platform/remote/common/remoteAgentConnection.ts
index 2335efbde29..fde30c2d8a6 100644
--- a/src/vs/platform/remote/common/remoteAgentConnection.ts
+++ b/src/vs/platform/remote/common/remoteAgentConnection.ts
@@ -417,30 +417,48 @@ export interface IAddressProvider {
 }
 
 export async function connectRemoteAgentManagement(options: IConnectionOptions, remoteAuthority: string, clientId: string): Promise<ManagementPersistentConnection> {
-	try {
-		const reconnectionToken = generateUuid();
-		const simpleOptions = await resolveConnectionOptions(options, reconnectionToken, null);
-		const { protocol } = await doConnectRemoteAgentManagement(simpleOptions, CancellationToken.None);
-		return new ManagementPersistentConnection(options, remoteAuthority, clientId, reconnectionToken, protocol);
-	} catch (err) {
-		options.logService.error(`[remote-connection] An error occurred in the very first connect attempt, it will be treated as a permanent error! Error:`);
-		options.logService.error(err);
-		PersistentConnection.triggerPermanentFailure(0, 0, RemoteAuthorityResolverError.isHandled(err));
-		throw err;
-	}
+	return createInitialConnection(
+		options,
+		async (simpleOptions) => {
+			const { protocol } = await doConnectRemoteAgentManagement(simpleOptions, CancellationToken.None);
+			return new ManagementPersistentConnection(options, remoteAuthority, clientId, simpleOptions.reconnectionToken, protocol);
+		}
+	);
 }
 
 export async function connectRemoteAgentExtensionHost(options: IConnectionOptions, startArguments: IRemoteExtensionHostStartParams): Promise<ExtensionHostPersistentConnection> {
-	try {
-		const reconnectionToken = generateUuid();
-		const simpleOptions = await resolveConnectionOptions(options, reconnectionToken, null);
-		const { protocol, debugPort } = await doConnectRemoteAgentExtensionHost(simpleOptions, startArguments, CancellationToken.None);
-		return new ExtensionHostPersistentConnection(options, startArguments, reconnectionToken, protocol, debugPort);
-	} catch (err) {
-		options.logService.error(`[remote-connection] An error occurred in the very first connect attempt, it will be treated as a permanent error! Error:`);
-		options.logService.error(err);
-		PersistentConnection.triggerPermanentFailure(0, 0, RemoteAuthorityResolverError.isHandled(err));
-		throw err;
+	return createInitialConnection(
+		options,
+		async (simpleOptions) => {
+			const { protocol, debugPort } = await doConnectRemoteAgentExtensionHost(simpleOptions, startArguments, CancellationToken.None);
+			return new ExtensionHostPersistentConnection(options, startArguments, simpleOptions.reconnectionToken, protocol, debugPort);
+		}
+	);
+}
+
+/**
+ * Will attempt to connect 5 times. If it fails 5 consecutive times, it will give up.
+ */
+async function createInitialConnection<T extends PersistentConnection>(options: IConnectionOptions, connectionFactory: (simpleOptions: ISimpleConnectionOptions) => Promise<T>): Promise<T> {
+	const MAX_ATTEMPTS = 5;
+
+	for (let attempt = 1; ; attempt++) {
+		try {
+			const reconnectionToken = generateUuid();
+			const simpleOptions = await resolveConnectionOptions(options, reconnectionToken, null);
+			const result = await connectionFactory(simpleOptions);
+			return result;
+		} catch (err) {
+			if (attempt < MAX_ATTEMPTS) {
+				options.logService.error(`[remote-connection][attempt ${attempt}] An error occurred in initial connection! Will retry... Error:`);
+				options.logService.error(err);
+			} else {
+				options.logService.error(`[remote-connection][attempt ${attempt}]  An error occurred in initial connection! It will be treated as a permanent error. Error:`);
+				options.logService.error(err);
+				PersistentConnection.triggerPermanentFailure(0, 0, RemoteAuthorityResolverError.isHandled(err));
+				throw err;
+			}
+		}
 	}
 }
 
diff --git a/src/vs/workbench/services/extensions/electron-sandbox/electronExtensionService.ts b/src/vs/workbench/services/extensions/electron-sandbox/electronExtensionService.ts
index 34cb13522aa..ee051e3d1f1 100644
--- a/src/vs/workbench/services/extensions/electron-sandbox/electronExtensionService.ts
+++ b/src/vs/workbench/services/extensions/electron-sandbox/electronExtensionService.ts
@@ -413,6 +413,32 @@ export abstract class ElectronExtensionService extends AbstractExtensionService
 		throw new Error(`Cannot get canonical URI because no extension is installed to resolve ${getRemoteAuthorityPrefix(remoteAuthority)}`);
 	}
 
+	private async _resolveAuthorityInitial(remoteAuthority: string): Promise<ResolverResult> {
+		const MAX_ATTEMPTS = 5;
+
+		for (let attempt = 1; ; attempt++) {
+			const sw = StopWatch.create(false);
+			this._logService.info(`[attempt ${attempt}] Invoking resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)})`);
+			try {
+				const resolverResult = await this._resolveAuthority(remoteAuthority);
+				this._logService.info(`[attempt ${attempt}] resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)}) returned '${resolverResult.authority.host}:${resolverResult.authority.port}' after ${sw.elapsed()} ms`);
+				return resolverResult;
+			} catch (err) {
+				this._logService.error(`[attempt ${attempt}] resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)}) returned an error after ${sw.elapsed()} ms`, err);
+
+				if (RemoteAuthorityResolverError.isNoResolverFound(err)) {
+					// There is no point in retrying if there is no resolver found
+					throw err;
+				}
+
+				if (attempt >= MAX_ATTEMPTS) {
+					// Too many failed attempts, give up
+					throw err;
+				}
+			}
+		}
+	}
+
 	private async _resolveAuthorityAgain(): Promise<void> {
 		const remoteAuthority = this._environmentService.remoteAuthority;
 		if (!remoteAuthority) {
@@ -473,14 +499,9 @@ export abstract class ElectronExtensionService extends AbstractExtensionService
 			}
 
 			let resolverResult: ResolverResult;
-
-			const sw = StopWatch.create(false);
-			this._logService.info(`Invoking resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)})`);
 			try {
-				resolverResult = await this._resolveAuthority(remoteAuthority);
-				this._logService.info(`resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)}) returned '${resolverResult.authority.host}:${resolverResult.authority.port}' after ${sw.elapsed()} ms`);
+				resolverResult = await this._resolveAuthorityInitial(remoteAuthority);
 			} catch (err) {
-				this._logService.error(`resolveAuthority(${getRemoteAuthorityPrefix(remoteAuthority)}) returned an error after ${sw.elapsed()} ms`, err);
 				if (RemoteAuthorityResolverError.isNoResolverFound(err)) {
 					err.isHandled = await this._handleNoResolverFound(remoteAuthority);
 				} else {