From 510a269dde893d3fd7eb4dafdccecf5532071d97 Mon Sep 17 00:00:00 2001
From: Koen Vlaswinkel <koesie10@github.com>
Date: Tue, 22 Oct 2024 15:37:35 +0200
Subject: [PATCH] Add id to streaming comparison

This avoids a bug where comparisons could potentially overlap if the
user opens a new comparison while the previous one has not yet finished
loading.
---
 .../ql-vscode/src/common/interface-types.ts       |  4 ++++
 extensions/ql-vscode/src/compare/compare-view.ts  |  6 ++++++
 extensions/ql-vscode/src/view/compare/Compare.tsx | 15 +++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/extensions/ql-vscode/src/common/interface-types.ts b/extensions/ql-vscode/src/common/interface-types.ts
index 11b4ab9f7..81ef4e612 100644
--- a/extensions/ql-vscode/src/common/interface-types.ts
+++ b/extensions/ql-vscode/src/common/interface-types.ts
@@ -424,6 +424,8 @@ export type InterpretedQueryCompareResult = {
 
 export interface StreamingComparisonSetupMessage {
   readonly t: "streamingComparisonSetup";
+  // The id of this streaming comparison
+  readonly id: string;
   readonly currentResultSetName: string;
   readonly message: string | undefined;
   // The from and to fields will only contain a chunk of the results
@@ -432,12 +434,14 @@ export interface StreamingComparisonSetupMessage {
 
 interface StreamingComparisonAddResultsMessage {
   readonly t: "streamingComparisonAddResults";
+  readonly id: string;
   // The from and to fields will only contain a chunk of the results
   readonly result: QueryCompareResult;
 }
 
 interface StreamingComparisonCompleteMessage {
   readonly t: "streamingComparisonComplete";
+  readonly id: string;
 }
 
 /**
diff --git a/extensions/ql-vscode/src/compare/compare-view.ts b/extensions/ql-vscode/src/compare/compare-view.ts
index 2692edcbe..c4c25df54 100644
--- a/extensions/ql-vscode/src/compare/compare-view.ts
+++ b/extensions/ql-vscode/src/compare/compare-view.ts
@@ -34,6 +34,7 @@ import {
 } from "./result-set-names";
 import { compareInterpretedResults } from "./interpreted-results";
 import { isCanary } from "../config";
+import { nanoid } from "nanoid";
 
 interface ComparePair {
   from: CompletedLocalQueryInfo;
@@ -206,6 +207,8 @@ export class CompareView extends AbstractWebview<
       return;
     }
 
+    const id = nanoid();
+
     // Streaming itself is implemented like this:
     // - 1 setup message which contains the first 1,000 results
     // - n "add results" messages which contain 1,000 results each
@@ -213,6 +216,7 @@ export class CompareView extends AbstractWebview<
 
     await this.postMessage({
       t: "streamingComparisonSetup",
+      id,
       result: this.chunkResults(result, 0, 1000),
       currentResultSetName,
       message,
@@ -226,12 +230,14 @@ export class CompareView extends AbstractWebview<
 
       await this.postMessage({
         t: "streamingComparisonAddResults",
+        id,
         result: chunk,
       });
     }
 
     await this.postMessage({
       t: "streamingComparisonComplete",
+      id,
     });
   }
 
diff --git a/extensions/ql-vscode/src/view/compare/Compare.tsx b/extensions/ql-vscode/src/view/compare/Compare.tsx
index ac226d6dd..b64b39b33 100644
--- a/extensions/ql-vscode/src/view/compare/Compare.tsx
+++ b/extensions/ql-vscode/src/view/compare/Compare.tsx
@@ -74,6 +74,13 @@ export function Compare(_: Record<string, never>): React.JSX.Element {
               break;
             }
 
+            if (prev.id !== msg.id) {
+              console.warn(
+                'Received "streamingComparisonAddResults" with different id, ignoring',
+              );
+              break;
+            }
+
             let result: QueryCompareResult;
             switch (prev.result.kind) {
               case "raw":
@@ -121,6 +128,14 @@ export function Compare(_: Record<string, never>): React.JSX.Element {
               setComparison(null);
               break;
             }
+
+            if (streamingComparisonRef.current.id !== msg.id) {
+              console.warn(
+                'Received "streamingComparisonComplete" with different id, ignoring',
+              );
+              break;
+            }
+
             setComparison({
               ...streamingComparisonRef.current,
               t: "setComparisons",