src/main/java/com/google/devtools/build/lib/profiler/TraceEvent.java

// Copyright 2020 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.google.devtools.build.lib.profiler;

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.gson.stream.JsonReader;
import com.google.gson.stream.JsonToken;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;

/**
 * Represents a single trace event in a JSON profile.
 *
 * <p>Format is documented in
 * https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
 */
@AutoValue
public abstract class TraceEvent {
  public static TraceEvent create(
      @Nullable String category,
      String name,
      @Nullable String type,
      @Nullable Duration timestamp,
      @Nullable Duration duration,
      long processId,
      long threadId,
      @Nullable ImmutableMap<String, Object> args,
      @Nullable String primaryOutputPath,
      @Nullable String targetLabel,
      @Nullable String mnemonic) {
    return new AutoValue_TraceEvent(
        category,
        name,
        type,
        timestamp,
        duration,
        processId,
        threadId,
        args,
        primaryOutputPath,
        targetLabel,
        mnemonic);
  }

  @Nullable
  public abstract String category();

  public abstract String name();

  @Nullable
  public abstract String type();

  @Nullable
  public abstract Duration timestamp();

  @Nullable
  public abstract Duration duration();

  public abstract long processId();

  public abstract long threadId();

  @Nullable
  public abstract ImmutableMap<String, Object> args();

  // Only applicable to action-related TraceEvents.
  @Nullable
  public abstract String primaryOutputPath();

  @Nullable
  public abstract String targetLabel();

  @Nullable
  public abstract String mnemonic();

  private static TraceEvent createFromJsonReader(JsonReader reader) throws IOException {
    String category = null;
    String name = null;
    Duration timestamp = null;
    Duration duration = null;
    long processId = -1;
    long threadId = -1;
    String primaryOutputPath = null;
    String targetLabel = null;
    String mnemonic = null;
    String type = null;
    ImmutableMap<String, Object> args = null;

    reader.beginObject();
    while (reader.hasNext()) {
      switch (reader.nextName()) {
        case "cat":
          category = reader.nextString();
          break;
        case "name":
          name = reader.nextString();
          break;
        case "ph":
          type = reader.nextString();
          break;
        case "ts":
          // Duration has no microseconds :-/.
          timestamp = Duration.ofNanos(reader.nextLong() * 1000);
          break;
        case "dur":
          duration = Duration.ofNanos(reader.nextLong() * 1000);
          break;
        case "pid":
          processId = reader.nextLong();
          break;
        case "tid":
          threadId = reader.nextLong();
          break;
        case "out":
          primaryOutputPath = reader.nextString();
          break;
        case "args":
          args = parseMap(reader);
          Object target = args.get("target");
          targetLabel = target instanceof String ? (String) target : null;
          Object mnemonicValue = args.get("mnemonic");
          mnemonic = mnemonicValue instanceof String ? (String) mnemonicValue : null;
          break;
        default:
          reader.skipValue();
      }
    }
    reader.endObject();
    return TraceEvent.create(
        category,
        name,
        type,
        timestamp,
        duration,
        processId,
        threadId,
        args,
        primaryOutputPath,
        targetLabel,
        mnemonic);
  }

  private static ImmutableMap<String, Object> parseMap(JsonReader reader) throws IOException {
    ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();

    reader.beginObject();
    while (reader.peek() != JsonToken.END_OBJECT) {
      String name = reader.nextName();
      Object val = parseSingleValueRecursively(reader);
      builder.put(name, val);
    }
    reader.endObject();

    return builder.buildOrThrow();
  }

  private static ImmutableList<Object> parseArray(JsonReader reader) throws IOException {
    ImmutableList.Builder<Object> builder = ImmutableList.builder();

    reader.beginArray();
    while (reader.peek() != JsonToken.END_ARRAY) {
      Object val = parseSingleValueRecursively(reader);
      builder.add(val);
    }
    reader.endArray();

    return builder.build();
  }

  @Nullable
  private static Object parseSingleValueRecursively(JsonReader reader) throws IOException {
    JsonToken nextToken = reader.peek();
    switch (nextToken) {
      case BOOLEAN:
        return reader.nextBoolean();
      case NULL:
        reader.nextNull();
        return null;
      case NUMBER:
        // Json's only numeric type is number, using Double to accommodate all types
        return reader.nextDouble();
      case STRING:
        return reader.nextString();
      case BEGIN_OBJECT:
        return parseMap(reader);
      case BEGIN_ARRAY:
        return parseArray(reader);
      default:
        throw new IOException("Unexpected token " + nextToken.name());
    }
  }

  public static List<TraceEvent> parseTraceEvents(JsonReader reader) throws IOException {
    List<TraceEvent> traceEvents = new ArrayList<>();
    reader.beginArray();
    while (reader.hasNext()) {
      traceEvents.add(createFromJsonReader(reader));
    }
    reader.endArray();
    return traceEvents;
  }
}