Use streaming JSON parse for timing profile to work around string siz…

…e limit

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";
@@ -81,6 +81,10 @@ export default class InvocationTimingCardComponent extends React.Component<Props
     return Boolean(this.getProfileFile()?.uri?.startsWith("bytestream://"));
   }
 
+  setProgress(n: number, digestSize: number, encoding: FileEncoding) {
+    // TODO: show a progress indicator
+  }
+
   fetchProfile() {
     if (!this.isTimingEnabled()) {
       this.setState({ loading: false });
@@ -98,15 +102,21 @@ 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 }));
   }

app/trace/BUILD

@@ -55,6 +55,10 @@ ts_library(
 ts_library(
     name = "trace_events",
     srcs = ["trace_events.ts"],
+    deps = [
+        "@npm//@streamparser/json",
+        "@npm//tslib",
+    ],
 )
 
 ts_jasmine_node_test(

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,43 @@ 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 parser = new JSONParser();
+  // Keep track of the last 2 tokens. If the trailing tokens aren't "]}" then we
+  // have an incomplete stream and need to manually close the sequence.
+  const last2Tokens: string[] = [];
+  parser.onToken = (parsedToken) => {
+    last2Tokens.push(parsedToken.value?.toString() || "");
+    if (last2Tokens.length > 2) last2Tokens.shift();
+  };
+  // Values are visited in child => parent order, so the timing profile should
+  // be the last visited value after parsing all elements.
+  let lastValue: JsonTypes.JsonPrimitive | JsonTypes.JsonStruct | null = null;
+  parser.onValue = (parsedElement) => {
+    lastValue = parsedElement.value;
+  };
+  const reader = body.getReader();
+  let n = 0;
+  while (true) {
+    const { value, done } = await reader.read();
+    if (done) break;
+    parser.write(new Uint8Array(value.buffer));
+    n += value.byteLength;
+    progress(n);
   }
-  return JSON.parse(data) as 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;
+  if (last2Tokens.join("") !== "]}") {
+    parser.write("]}");
+  }
+  if (!lastValue) {
+    throw new Error("failed to parse timing profile: final parsed value is null");
+  }
+  if (typeof lastValue !== "object") {
+    throw new Error(`unsupported profile type "${typeof lastValue}"`);
   }
-  return out;
+  return lastValue as Profile;
 }
 
 function eventComparator(a: TraceEvent, b: TraceEvent) {

package.json

@@ -10,6 +10,7 @@
     "@bazel/esbuild": "^5.4.0",
     "@bazel/jasmine": "^5.4.0",
     "@bazel/typescript": "^5.4.0",
+    "@streamparser/json": "^0.0.17",
     "@types/d3-scale": "^3.0.0",
     "@types/d3-time": "^3.0.0",
     "@types/dagre-d3": "^0.4.39",

yarn.lock

@@ -110,6 +110,11 @@
   resolved "https://registry.yarnpkg.com/@protobufjs/utf8/-/utf8-1.1.0.tgz#a777360b5b39a1a2e5106f8e858f2fd2d060c570"
   integrity sha1-p3c2C1s5oaLlEG+OhY8v0tBgxXA=
 
+"@streamparser/json@^0.0.17":
+  version "0.0.17"
+  resolved "https://registry.yarnpkg.com/@streamparser/json/-/json-0.0.17.tgz#b68742ebb49eec9c1fcc76cfa730dd74a5131382"
+  integrity sha512-mW54K6CTVJVLwXRB6kSS1xGWPmtTuXAStWnlvtesmcySgtop+eFPWOywBFPpJO4UD173epYsPSP6HSW8kuqN8w==
+
 "@types/d3-array@^3.0.3":
   version "3.0.4"
   resolved "https://registry.yarnpkg.com/@types/d3-array/-/d3-array-3.0.4.tgz#44eebe40be57476cad6a0cd6a85b0f57d54185a2"