Merge pull request #3771 from github/koesie10/streaming-comparison

Implement streaming for compare view

Author: koesie10

extensions/ql-vscode/src/common/interface-types.ts

@@ -360,6 +360,9 @@ interface ChangeCompareMessage {
 export type ToCompareViewMessage =
   | SetComparisonQueryInfoMessage
   | SetComparisonsMessage
+  | StreamingComparisonSetupMessage
+  | StreamingComparisonAddResultsMessage
+  | StreamingComparisonCompleteMessage
   | SetUserSettingsMsg;
 
 /**
@@ -419,6 +422,28 @@ export type InterpretedQueryCompareResult = {
   to: Result[];
 };
 
+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
+  readonly result: QueryCompareResult;
+}
+
+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;
+}
+
 /**
  * Extract the name of the default result. Prefer returning
  * 'alerts', or '#select'. Otherwise return the first in the list.

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;
@@ -183,13 +184,97 @@ export class CompareView extends AbstractWebview<
         message = getErrorMessage(e);
       }
 
+      await this.streamResults(result, currentResultSetDisplayName, message);
+    }
+  }
+
+  private async streamResults(
+    result: QueryCompareResult | undefined,
+    currentResultSetName: string,
+    message: string | undefined,
+  ) {
+    // Since there is a string limit of 1GB in Node.js, the comparison is send as a JSON.stringified string to the webview
+    // and some comparisons may be larger than that, we sometimes need to stream results. This uses a heuristic of 2,000 results
+    // to determine if we should stream results.
+
+    if (!this.shouldStreamResults(result)) {
       await this.postMessage({
         t: "setComparisons",
         result,
-        currentResultSetName: currentResultSetDisplayName,
+        currentResultSetName,
         message,
       });
+      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
+    // - 1 complete message which just tells the webview that we're done
+
+    await this.postMessage({
+      t: "streamingComparisonSetup",
+      id,
+      result: this.chunkResults(result, 0, 1000),
+      currentResultSetName,
+      message,
+    });
+
+    const { from, to } = result;
+
+    const maxResults = Math.max(from.length, to.length);
+    for (let i = 1000; i < maxResults; i += 1000) {
+      const chunk = this.chunkResults(result, i, i + 1000);
+
+      await this.postMessage({
+        t: "streamingComparisonAddResults",
+        id,
+        result: chunk,
+      });
+    }
+
+    await this.postMessage({
+      t: "streamingComparisonComplete",
+      id,
+    });
+  }
+
+  private shouldStreamResults(
+    result: QueryCompareResult | undefined,
+  ): result is QueryCompareResult {
+    if (result === undefined) {
+      return false;
+    }
+
+    // We probably won't run into limits if we have less than 2,000 total results
+    const totalResults = result.from.length + result.to.length;
+    return totalResults > 2000;
+  }
+
+  private chunkResults(
+    result: QueryCompareResult,
+    start: number,
+    end: number,
+  ): QueryCompareResult {
+    if (result.kind === "raw") {
+      return {
+        ...result,
+        from: result.from.slice(start, end),
+        to: result.to.slice(start, end),
+      };
+    }
+
+    if (result.kind === "interpreted") {
+      return {
+        ...result,
+        from: result.from.slice(start, end),
+        to: result.to.slice(start, end),
+      };
+    }
+
+    assertNever(result);
   }
 
   protected getPanelConfig(): WebviewPanelConfig {

extensions/ql-vscode/src/view/compare/Compare.tsx

@@ -1,11 +1,13 @@
-import { useState, useEffect } from "react";
+import { useState, useEffect, useRef } from "react";
 import { styled } from "styled-components";
 
 import type {
   ToCompareViewMessage,
   SetComparisonsMessage,
   SetComparisonQueryInfoMessage,
   UserSettings,
+  StreamingComparisonSetupMessage,
+  QueryCompareResult,
 } from "../../common/interface-types";
 import { DEFAULT_USER_SETTINGS } from "../../common/interface-types";
 import CompareSelector from "./CompareSelector";
@@ -37,6 +39,12 @@ export function Compare(_: Record<string, never>): React.JSX.Element {
     DEFAULT_USER_SETTINGS,
   );
 
+  // This is a ref because we don't need to re-render when we get a new streaming comparison message
+  // and we don't want to change the listener every time we get a new message
+  const streamingComparisonRef = useRef<StreamingComparisonSetupMessage | null>(
+    null,
+  );
+
   const message = comparison?.message || "Empty comparison";
   const hasRows =
     comparison?.result &&
@@ -53,6 +61,87 @@ export function Compare(_: Record<string, never>): React.JSX.Element {
           case "setComparisons":
             setComparison(msg);
             break;
+          case "streamingComparisonSetup":
+            setComparison(null);
+            streamingComparisonRef.current = msg;
+            break;
+          case "streamingComparisonAddResults": {
+            const prev = streamingComparisonRef.current;
+            if (prev === null) {
+              console.warn(
+                'Received "streamingComparisonAddResults" before "streamingComparisonSetup"',
+              );
+              break;
+            }
+
+            if (prev.id !== msg.id) {
+              console.warn(
+                'Received "streamingComparisonAddResults" with different id, ignoring',
+              );
+              break;
+            }
+
+            let result: QueryCompareResult;
+            switch (prev.result.kind) {
+              case "raw":
+                if (msg.result.kind !== "raw") {
+                  throw new Error(
+                    "Streaming comparison: expected raw results, got interpreted results",
+                  );
+                }
+
+                result = {
+                  ...prev.result,
+                  from: [...prev.result.from, ...msg.result.from],
+                  to: [...prev.result.to, ...msg.result.to],
+                };
+                break;
+              case "interpreted":
+                if (msg.result.kind !== "interpreted") {
+                  throw new Error(
+                    "Streaming comparison: expected interpreted results, got raw results",
+                  );
+                }
+
+                result = {
+                  ...prev.result,
+                  from: [...prev.result.from, ...msg.result.from],
+                  to: [...prev.result.to, ...msg.result.to],
+                };
+                break;
+              default:
+                throw new Error("Unexpected comparison result kind");
+            }
+
+            streamingComparisonRef.current = {
+              ...prev,
+              result,
+            };
+
+            break;
+          }
+          case "streamingComparisonComplete":
+            if (streamingComparisonRef.current === null) {
+              console.warn(
+                'Received "streamingComparisonComplete" before "streamingComparisonSetup"',
+              );
+              setComparison(null);
+              break;
+            }
+
+            if (streamingComparisonRef.current.id !== msg.id) {
+              console.warn(
+                'Received "streamingComparisonComplete" with different id, ignoring',
+              );
+              break;
+            }
+
+            setComparison({
+              ...streamingComparisonRef.current,
+              t: "setComparisons",
+            });
+            streamingComparisonRef.current = null;
+            break;
           case "setUserSettings":
             setUserSettings(msg.userSettings);
             break;