Skip to content

Commit

Permalink
Use streaming JSON parse for timing profile to work around string siz…
Browse files Browse the repository at this point in the history
…e limit
  • Loading branch information
bduffany committed Oct 11, 2023
1 parent 9aaf060 commit ac0eab1
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 22 deletions.
20 changes: 20 additions & 0 deletions app/invocation/invocation.css
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
53 changes: 47 additions & 6 deletions app/invocation/invocation_timing_card.tsx
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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();
}
Expand All @@ -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 });
Expand All @@ -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 }));
}
Expand Down Expand Up @@ -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) {
Expand Down
3 changes: 3 additions & 0 deletions app/trace/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ ts_library(
ts_library(
name = "trace_events",
srcs = ["trace_events.ts"],
deps = [
"@npm//tslib",
],
)

ts_jasmine_node_test(
Expand Down
67 changes: 51 additions & 16 deletions app/trace/trace_events.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import { JSONParser, JsonTypes } from "@streamparser/json";

/**
* This file contains utilities for parsing a raw trace file.
*
Expand Down Expand Up @@ -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) {
Expand Down

0 comments on commit ac0eab1

Please sign in to comment.