From ac0eab10f1960e142f6ae6c66605e3fd8def588b Mon Sep 17 00:00:00 2001
From: Brandon Duffany <brandon@buildbuddy.io>
Date: Wed, 11 Oct 2023 11:05:56 -0400
Subject: [PATCH] Use streaming JSON parse for timing profile to work around
 string size limit

---
 app/invocation/invocation.css             | 20 +++++++
 app/invocation/invocation_timing_card.tsx | 53 ++++++++++++++++--
 app/trace/BUILD                           |  3 +
 app/trace/trace_events.ts                 | 67 +++++++++++++++++------
 4 files changed, 121 insertions(+), 22 deletions(-)

diff --git a/app/invocation/invocation.css b/app/invocation/invocation.css
index 61c54cf4059b..fd50ad0f636d 100644
--- a/app/invocation/invocation.css
+++ b/app/invocation/invocation.css
@@ -1169,3 +1169,23 @@ svg.invocation-query-graph .nodes rect {
   margin-left: auto;
   white-space: nowrap;
 }
+
+.timing.card .progress-bar {
+  height: 8px;
+  border-radius: 4px;
+  max-width: 300px;
+  overflow: clip;
+  background: #e3f2fd;
+}
+
+.timing.card .progress-label {
+  margin-bottom: 4px;
+  font-size: 13px;
+  color: #616161;
+}
+
+.timing.card .progress-bar-inner {
+  height: 100%;
+  background: #2196f3;
+  transform-origin: left;
+}
diff --git a/app/invocation/invocation_timing_card.tsx b/app/invocation/invocation_timing_card.tsx
index a6afdca8a771..386501e1312c 100644
--- a/app/invocation/invocation_timing_card.tsx
+++ b/app/invocation/invocation_timing_card.tsx
@@ -1,7 +1,7 @@
 import React from "react";
 import SetupCodeComponent from "../docs/setup_code";
 import FlameChart from "../flame_chart/flame_chart";
-import { Profile, parseProfile } from "../trace/trace_events";
+import { Profile, readProfile } from "../trace/trace_events";
 import rpcService, { FileEncoding } from "../service/rpc_service";
 import InvocationModel from "./invocation_model";
 import Button from "../components/button/button";
@@ -63,6 +63,8 @@ export default class InvocationTimingCardComponent extends React.Component<Props
     eventPageSize: window.localStorage[eventPageSizeStorageKey] || 100,
   };
 
+  private progressRef = React.createRef<HTMLDivElement>();
+
   componentDidMount() {
     this.fetchProfile();
   }
@@ -81,6 +83,30 @@ export default class InvocationTimingCardComponent extends React.Component<Props
     return Boolean(this.getProfileFile()?.uri?.startsWith("bytestream://"));
   }
 
+  setProgress(bytesLoaded: number, digestSize: number, encoding: FileEncoding) {
+    const container = this.progressRef.current;
+    if (!container) return;
+
+    let approxCompressionRatio = 1;
+    if (encoding === "gzip") {
+      approxCompressionRatio = 11.3;
+    }
+
+    const compressedBytesLoaded = Math.min(bytesLoaded / approxCompressionRatio, digestSize);
+    const progressPercent = 100 * Math.min(1, compressedBytesLoaded / digestSize);
+
+    const spinner = container.querySelector(".loading");
+    spinner?.remove();
+
+    const progressContainer = container.querySelector(".timing-profile-progress")!;
+    progressContainer.removeAttribute("hidden");
+    const progressLabel = progressContainer.querySelector(".progress-label")!;
+    progressLabel.innerHTML = `Loading profile (${format.bytes(compressedBytesLoaded)} / ${format.bytes(digestSize)})`;
+
+    const progressBarInner = progressContainer.querySelector(".progress-bar-inner") as HTMLElement;
+    progressBarInner.style.width = `${progressPercent}%`;
+  }
+
   fetchProfile() {
     if (!this.isTimingEnabled()) {
       this.setState({ loading: false });
@@ -98,15 +124,19 @@ export default class InvocationTimingCardComponent extends React.Component<Props
       storedEncoding = "gzip";
     }
 
+    const digestSize = Number(profileFile.uri.split("/").pop());
+
     this.setState({ loading: true });
-    // Note: we use responseType "text" instead of "json" since the profile is
-    // not always valid JSON (the trailing "]}" may be missing).
     rpcService
-      .fetchBytestreamFile(profileFile.uri, this.props.model.getInvocationId(), "text", {
+      .fetchBytestreamFile(profileFile.uri, this.props.model.getInvocationId(), "stream", {
         // Set the stored encoding header to prevent the server from double-gzipping.
         headers: { "X-Stored-Encoding-Hint": storedEncoding },
       })
-      .then((contents) => this.updateProfile(parseProfile(contents)))
+      .then((body) => {
+        if (body === null) throw new Error("response body is null");
+        return readProfile(body, (n) => this.setProgress(n, digestSize, storedEncoding));
+      })
+      .then((profile) => this.updateProfile(profile))
       .catch((e) => errorService.handleError(e))
       .finally(() => this.setState({ loading: false }));
   }
@@ -207,7 +237,18 @@ export default class InvocationTimingCardComponent extends React.Component<Props
 
   renderEmptyState() {
     if (this.state.loading) {
-      return <div className="loading" />;
+      return (
+        <div ref={this.progressRef}>
+          <div className="loading" />
+          <div className="timing-profile-progress" hidden>
+            <div className="progress-label" />
+            <div className="progress-bar">
+              <div className="progress-bar-inner" />
+              div
+            </div>
+          </div>
+        </div>
+      );
     }
 
     if (!this.props.model.buildToolLogs) {
diff --git a/app/trace/BUILD b/app/trace/BUILD
index e48db018ba96..26d572255a56 100644
--- a/app/trace/BUILD
+++ b/app/trace/BUILD
@@ -55,6 +55,9 @@ ts_library(
 ts_library(
     name = "trace_events",
     srcs = ["trace_events.ts"],
+    deps = [
+        "@npm//tslib",
+    ],
 )
 
 ts_jasmine_node_test(
diff --git a/app/trace/trace_events.ts b/app/trace/trace_events.ts
index c25219e3b6d5..4c5f431e034b 100644
--- a/app/trace/trace_events.ts
+++ b/app/trace/trace_events.ts
@@ -1,3 +1,5 @@
+import { JSONParser, JsonTypes } from "@streamparser/json";
+
 /**
  * This file contains utilities for parsing a raw trace file.
  *
@@ -64,26 +66,59 @@ const TIME_SERIES_EVENT_NAMES_AND_ARG_KEYS: Map<string, string> = new Map([
   ["Network Down usage (total)", "system network down (Mbps)"],
 ]);
 
-export function parseProfile(data: string): Profile {
-  // Note, the trace profile format specifies that the "]" at the end of the
-  // list is optional:
-  // https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.f2f0yd51wi15
-  // Bazel uses the JSON Object Format, which means the trailing "]}" is
-  // optional.
-  if (trailingNonWhitespaceCharacters(data, 2) !== "]}") {
-    data += "]}";
+export async function readProfile(
+  body: ReadableStream<Uint8Array>,
+  progress: (numBytesLoaded: number) => void
+): Promise<Profile> {
+  const reader = body.getReader();
+  const decoder = new TextDecoder("utf-8");
+  let n = 0;
+  let buffer = "";
+  let profile: Profile | null = null;
+
+  while (true) {
+    const { value, done } = await reader.read();
+    if (done) break;
+    // `stream: true` allows us to handle UTF-8 sequences that cross chunk
+    // boundaries (should be relatively rare).
+    const text = decoder.decode(value, { stream: true });
+    buffer += text;
+    n += value.byteLength;
+    progress(n);
+    // Keep accumulating into the buffer until we see the "traceEvents" array.
+    // Each entry in this array is newline-delimited (a special property of
+    // Google's trace event JSON format).
+    if (!profile) {
+      const beginMarker = '"traceEvents":[\n';
+      const index = buffer.indexOf(beginMarker);
+      if (index < 0) continue;
+
+      const before = buffer.substring(0, index + beginMarker.length);
+      const after = buffer.substring(before.length);
+
+      const outerJSON = before + "]}";
+      profile = JSON.parse(outerJSON) as Profile;
+      buffer = after;
+    }
+    if (profile) {
+      buffer = consumeEvents(buffer, profile);
+    }
   }
-  return JSON.parse(data) as Profile;
+  if (!profile) {
+    throw new Error("failed to parse timing profile JSON");
+  }
+  return profile;
 }
 
-function trailingNonWhitespaceCharacters(text: string, numTrailingChars: number) {
-  let out = "";
-  for (let i = text.length - 1; i >= 0; i--) {
-    if (text[i].trim() !== "") out = text[i] + out;
-
-    if (out.length >= numTrailingChars) break;
+function consumeEvents(buffer: string, profile: Profile): string {
+  // Each event entry looks like "   { ... },\n"
+  const parts = buffer.split(",\n");
+  const completeEvents = parts.slice(0, parts.length - 1);
+  for (const rawEvent of completeEvents) {
+    profile!.traceEvents.push(JSON.parse(rawEvent));
   }
-  return out;
+  // If there's a partial event at the end, that becomes the new acc value.
+  return parts[parts.length - 1] || "";
 }
 
 function eventComparator(a: TraceEvent, b: TraceEvent) {