Skip to content

Commit

Permalink
[Performance] Read Bazel profile as a stream (#165)
Browse files Browse the repository at this point in the history
Instead of reading the Bazel profile all at once, use a stream to reduce
memory consumption.

Benchmark with a Bazel profile of unzipped size 357MB
Using YourKit as a profiler and checking the shallow size at the end of
the `BazelProfile` constructor:

Before: 3.6 GB
After: 349 MB

Contributes to #163

---------

Signed-off-by: Sara Adams <sara.e.adams@gmail.com>
saraadams authored Dec 21, 2023
1 parent 8019030 commit ef5c4d5
Showing 1 changed file with 64 additions and 54 deletions.
Original file line number Diff line number Diff line change
@@ -27,14 +27,15 @@
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.google.gson.stream.JsonReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -77,69 +78,78 @@ public static BazelProfile createFromPath(String path) throws IllegalArgumentExc

public static BazelProfile createFromInputStream(InputStream inputStream)
throws IllegalArgumentException {
JsonObject bazelProfile;
try {
bazelProfile = JsonParser.parseReader(new InputStreamReader(inputStream)).getAsJsonObject();
} catch (IllegalStateException e) {
throw new IllegalArgumentException("Could not parse Bazel profile.", e);
}
return new BazelProfile(bazelProfile);
return new BazelProfile(
new JsonReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8)));
}

private final BazelVersion bazelVersion;
private final Map<String, String> otherData = new HashMap<>();
private final Map<ThreadId, ProfileThread> threads = new HashMap<>();

private BazelProfile(JsonObject profile) {
if (!profile.has(TraceEventFormatConstants.SECTION_OTHER_DATA)
|| !profile.has(TraceEventFormatConstants.SECTION_TRACE_EVENTS)) {
throw new IllegalArgumentException(
String.format(
"Invalid profile, JSON file missing \"%s\" and/or \"%s\"",
TraceEventFormatConstants.SECTION_OTHER_DATA,
TraceEventFormatConstants.SECTION_TRACE_EVENTS));
}
private BazelProfile(JsonReader profileReader) {
try {
profile
.get(TraceEventFormatConstants.SECTION_OTHER_DATA)
.getAsJsonObject()
.entrySet()
.forEach(entry -> otherData.put(entry.getKey(), entry.getValue().getAsString()));
this.bazelVersion =
BazelVersion.parse(otherData.get(BazelProfileConstants.OTHER_DATA_BAZEL_VERSION));

profile
.get(TraceEventFormatConstants.SECTION_TRACE_EVENTS)
.getAsJsonArray()
.forEach(
element -> {
JsonObject object = element.getAsJsonObject();
int pid;
int tid;
try {
pid = object.get(TraceEventFormatConstants.EVENT_PROCESS_ID).getAsInt();
tid = object.get(TraceEventFormatConstants.EVENT_THREAD_ID).getAsInt();
} catch (Exception e) {
// Skip events that do not have a valid pid or tid.
return;
}
ThreadId threadId = new ThreadId(pid, tid);
ProfileThread profileThread =
threads.compute(
threadId,
(key, t) -> {
if (t == null) {
t = new ProfileThread(threadId);
}
return t;
});
// TODO: Use success response to take action on errant events.
profileThread.addEvent(object);
});
} catch (IllegalStateException e) {
boolean hasOtherData = false;
boolean hasTraceEvents = false;
profileReader.beginObject();
while (profileReader.hasNext()) {
switch (profileReader.nextName()) {
case TraceEventFormatConstants.SECTION_OTHER_DATA:
hasOtherData = true;
profileReader.beginObject();
while (profileReader.hasNext()) {
otherData.put(profileReader.nextName(), profileReader.nextString());
}
profileReader.endObject();
break;
case TraceEventFormatConstants.SECTION_TRACE_EVENTS:
hasTraceEvents = true;
profileReader.beginArray();
while (profileReader.hasNext()) {
var traceEvent = JsonParser.parseReader(profileReader).getAsJsonObject();
int pid;
int tid;
try {
pid = traceEvent.get(TraceEventFormatConstants.EVENT_PROCESS_ID).getAsInt();
tid = traceEvent.get(TraceEventFormatConstants.EVENT_THREAD_ID).getAsInt();
} catch (Exception e) {
// Skip events that do not have a valid pid or tid.
continue;
}
ThreadId threadId = new ThreadId(pid, tid);
ProfileThread profileThread =
threads.compute(
threadId,
(key, t) -> {
if (t == null) {
t = new ProfileThread(threadId);
}
return t;
});
// TODO: Use success response to take action on errant events.
profileThread.addEvent(traceEvent);
}
profileReader.endArray();
break;
default:
// We only care about otherData and traceEvents.
profileReader.skipValue();
}
}
profileReader.endObject();
if (!hasOtherData || !hasTraceEvents) {
throw new IllegalArgumentException(
String.format(
"Invalid profile, JSON file missing \"%s\" and/or \"%s\"",
TraceEventFormatConstants.SECTION_OTHER_DATA,
TraceEventFormatConstants.SECTION_TRACE_EVENTS));
}
} catch (IllegalStateException | IOException e) {
throw new IllegalArgumentException("Could not parse Bazel profile.", e);
}

this.bazelVersion =
BazelVersion.parse(otherData.get(BazelProfileConstants.OTHER_DATA_BAZEL_VERSION));

if (!containsMainThread()) {
throw new IllegalArgumentException(
String.format(

0 comments on commit ef5c4d5

Please sign in to comment.