From a1056c51431e4c5d37dcda03bfb6dd00e11c0164 Mon Sep 17 00:00:00 2001
From: James McTavish <mctavish@google.com>
Date: Wed, 10 May 2023 11:41:32 -0400
Subject: [PATCH] Add ability to configure connection timeout (#136)

The connection timeout is currently hardcoded to 2 seconds. It is likely that this is not going to work in all situations, so I'm adding a configuration knob to allow the user to set whatever timeout they'd like.

Depending on feedback, the default timeout may be increased as well. The balance is that if the extension is attempting to discover the database, increases in timeouts will result in slower extension startup.

This PR is related to #135
---
 snapshot_dbg_extension/package.json   |  4 ++++
 snapshot_dbg_extension/src/adapter.ts | 13 ++++++++++---
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/snapshot_dbg_extension/package.json b/snapshot_dbg_extension/package.json
index 04bdc88..f4b9a15 100644
--- a/snapshot_dbg_extension/package.json
+++ b/snapshot_dbg_extension/package.json
@@ -84,6 +84,10 @@
               "debugOutput": {
                 "type": "boolean",
                 "description": "If enabled, will output extension debug messages to the console."
+              },
+              "connectionTimeoutMs": {
+                "type": "number",
+                "description": "Timeout for initial firebase database connection, in milliseconds."
               }
             }
           }
diff --git a/snapshot_dbg_extension/src/adapter.ts b/snapshot_dbg_extension/src/adapter.ts
index afcdaaf..a52719a 100644
--- a/snapshot_dbg_extension/src/adapter.ts
+++ b/snapshot_dbg_extension/src/adapter.ts
@@ -22,6 +22,7 @@ import { debugLog, setDebugLogEnabled } from './debugUtil';
 import { pickDebuggeeId } from './debuggeePicker';
 import { BreakpointManager } from './breakpointManager';
 import { GcloudCredential } from './gcloudCredential';
+import { debuglog } from 'util';
 
 const FIREBASE_APP_NAME = 'snapshotdbg';
 const INITIALIZE_TIME_ALLOWANCE_MS = 2 * 1000; // 2 seconds
@@ -45,6 +46,9 @@ interface IAttachRequestArguments extends DebugProtocol.AttachRequestArguments {
 
     /** Whether to output debug messages to console. */
     debugOutput: boolean;
+
+    /** Timeout for initial firebase database connection, in milliseconds. */
+    connectionTimeoutMs: number | undefined;
 }
 
 
@@ -122,8 +126,9 @@ export class SnapshotDebuggerSession extends DebugSession {
         }
     }
 
-    private async connectToFirebase(credential: GcloudCredential, progress: vscode.Progress<{ message: string, increment: number }>, configuredDatabaseUrl?: string): Promise<Database> {
+    private async connectToFirebase(credential: GcloudCredential, progress: vscode.Progress<{ message: string, increment: number }>, args?: IAttachRequestArguments): Promise<Database> {
         // Build the database URL.
+        const configuredDatabaseUrl = args?.databaseUrl;
         const databaseUrls = [];
         if (configuredDatabaseUrl) {
             databaseUrls.push(configuredDatabaseUrl);
@@ -132,6 +137,8 @@ export class SnapshotDebuggerSession extends DebugSession {
             databaseUrls.push(`https://${this.projectId}-default-rtdb.firebaseio.com`);
         }
 
+        const timeout = args?.connectionTimeoutMs ?? 2000;
+        debugLog(`Timeout set to ${timeout}`);
         for (const databaseUrl of databaseUrls) {
             progress.report({ message: `Connecting to ${databaseUrl}`, increment: 20 });
             this.app = initializeApp({
@@ -144,7 +151,7 @@ export class SnapshotDebuggerSession extends DebugSession {
             // Test the connection by reading the schema version.
             try {
                 const version_snapshot = await withTimeout(
-                    2000,
+                    timeout,
                     db.ref('cdbg/schema_version').get());
                 if (version_snapshot) {
                     const version = version_snapshot.val();
@@ -200,7 +207,7 @@ export class SnapshotDebuggerSession extends DebugSession {
             }
 
             try {
-                this.db = await this.connectToFirebase(credential, progress, args.databaseUrl);
+                this.db = await this.connectToFirebase(credential, progress, args);
             } catch (err) {
                 this.sendErrorResponse(response, 2,
                     'Cannot connect to Firebase.\n\n' +