src/trace_processor/rpc/rpc.cc

/*
 * Copyright (C) 2019 The Android Open Source Project
 *
 * 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.
 */

#include "src/trace_processor/rpc/rpc.h"

#include <cinttypes>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "perfetto/base/logging.h"
#include "perfetto/base/status.h"
#include "perfetto/base/time.h"
#include "perfetto/ext/base/file_utils.h"
#include "perfetto/ext/base/scoped_file.h"
#include "perfetto/ext/base/status_macros.h"
#include "perfetto/ext/base/temp_file.h"
#include "perfetto/ext/base/utils.h"
#include "perfetto/ext/base/version.h"
#include "perfetto/ext/protozero/proto_ring_buffer.h"
#include "perfetto/ext/trace_processor/rpc/query_result_serializer.h"
#include "perfetto/protozero/field.h"
#include "perfetto/protozero/proto_utils.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "perfetto/public/compiler.h"
#include "perfetto/trace_processor/basic_types.h"
#include "perfetto/trace_processor/trace_blob.h"
#include "perfetto/trace_processor/trace_blob_view.h"
#include "perfetto/trace_processor/trace_processor.h"
#include "src/trace_processor/tp_metatrace.h"

#include "protos/perfetto/trace_processor/metatrace_categories.pbzero.h"
#include "protos/perfetto/trace_processor/trace_processor.pbzero.h"

namespace perfetto::trace_processor {

// Takes ownership of a Stream's responses and hands them to the transport's
// RpcResponseFunction one slice at a time.
class Rpc::Outbox {
 public:
  explicit Outbox(RpcResponseFunction response_fn)
      : response_fn_(std::move(response_fn)) {}

  void Send(std::vector<protozero::ScatteredHeapBuffer::Slice> slices) {
    for (const auto& slice : slices) {
      auto range = slice.GetUsedRange();
      response_fn_(range.begin, static_cast<uint32_t>(range.size()));
    }
  }

  void Close() { response_fn_(nullptr, 0); }

 private:
  RpcResponseFunction response_fn_;
};

namespace {
// Writes a "Loading trace ..." update every N bytes.
constexpr size_t kProgressUpdateBytes = 50ul * 1000 * 1000;
using TraceProcessorRpcStream = protos::pbzero::TraceProcessorRpcStream;
using RpcProto = protos::pbzero::TraceProcessorRpc;

// Most RPC messages are either very small or a query results.
// QueryResultSerializer splits rows into batches of approximately 128KB. Try
// avoid extra heap allocations for the nominal case.
constexpr auto kSliceSize =
    QueryResultSerializer::kDefaultBatchSplitThreshold + 4096;

// Holds a trace_processor::TraceProcessorRpc pbzero message. Avoids extra
// copies by handing the fragmented heap buffer itself to the Outbox. It also
// takes care of prefixing each message with the proto preamble and varint size.
class Response {
 public:
  Response(int64_t seq, int method);
  Response(const Response&) = delete;
  Response& operator=(const Response&) = delete;
  RpcProto* operator->() { return msg_; }
  void Send(Rpc::Outbox&);

 private:
  RpcProto* msg_ = nullptr;

  // The reason why we use TraceProcessorRpcStream as root message is because
  // the RPC wire protocol expects each message to be prefixed with a proto
  // preamble and varint size. This happens to be the same serialization of a
  // repeated field (this is really the same trick we use between
  // Trace and TracePacket in trace.proto)
  protozero::HeapBuffered<TraceProcessorRpcStream> buf_;
};

Response::Response(int64_t seq, int method) : buf_(kSliceSize, kSliceSize) {
  msg_ = buf_->add_msg();
  msg_->set_seq(seq);
  msg_->set_response(static_cast<RpcProto::TraceProcessorMethod>(method));
}

void Response::Send(Rpc::Outbox& outbox) {
  outbox.Send(buf_.TakeSlices());
}

// Statement-cursor metadata stamped on the first message of a
// TPM_STATEMENT_STREAMING response; null for TPM_QUERY_STREAMING.
struct StatementStreamHeader {
  uint32_t tail_offset = 0;
  bool statement_executed = false;
};

// Emits the StatementResult envelope of a TPM_STATEMENT_STREAMING reply,
// stamping the cursor metadata when |header| is non-null. The single place
// where the cursor fields are written.
protos::pbzero::QueryResult* SetStatementResult(
    Response& resp,
    const StatementStreamHeader* header) {
  auto* stmt_res = resp->set_statement_result();
  if (header) {
    stmt_res->set_tail_offset(header->tail_offset);
    stmt_res->set_statement_executed(header->statement_executed);
  }
  return stmt_res->set_result();
}

// Returns the QueryResult submessage of |resp| for the endpoint being served:
// the bare query_result for TPM_QUERY_STREAMING (|header| == nullptr) or the
// one nested in a StatementResult for TPM_STATEMENT_STREAMING, stamping
// |header| on the first message of the stream.
protos::pbzero::QueryResult* SetResult(Response& resp,
                                       const StatementStreamHeader* header,
                                       bool first_message) {
  if (!header)
    return resp->set_query_result();
  return SetStatementResult(resp, first_message ? header : nullptr);
}

// Streams every message of |serializer|'s result through |outbox|, one
// Response per message. Shared tail of TPM_QUERY_STREAMING and
// TPM_STATEMENT_STREAMING.
PERFETTO_NO_INLINE void StreamSerializerResponses(
    QueryResultSerializer* serializer,
    int req_type,
    int64_t* tx_seq_id,
    Rpc::Outbox& outbox,
    const StatementStreamHeader* header) {
  for (bool has_more = true, first = true; has_more; first = false) {
    const int64_t seq_id = (*tx_seq_id)++;
    Response resp(seq_id, req_type);
    has_more = serializer->Serialize(SetResult(resp, header, first));
    const uint32_t resp_size = resp->Finalize();
    if (resp_size < protozero::proto_utils::kMaxMessageLength) {
      // This is the nominal case.
      resp.Send(outbox);
      continue;
    }
    // In rare cases a query can end up with a batch which is too big.
    // Normally batches are automatically split before hitting the limit,
    // but one can come up with a query where a single cell is > 256MB.
    // If this happens, just bail out gracefully rather than creating an
    // unparsable proto which will cause a RPC framing error.
    // If we hit this, we have to discard `resp` because it's
    // unavoidably broken (due to have overflown the 4-bytes size) and
    // can't be parsed. Instead create a new response with the error.
    Response err_resp(seq_id, req_type);
    auto* qres = SetResult(err_resp, header, first);
    qres->add_batch()->set_is_last_batch(true);
    qres->set_error("The query ended up with a response that is too big (" +
                    std::to_string(resp_size) +
                    " bytes). This usually happens when a single row is >= 256 "
                    "MiB. Consider writing large values to a file instead.");
    err_resp.Send(outbox);
    break;
  }
}

// Sends the single-message TPM_STATEMENT_STREAMING replies which don't stream
// a result set: errors detected before execution (|error| != nullptr) and the
// no-statement-left case (|header| with statement_executed == false).
PERFETTO_NO_INLINE void SendSingleStatementResponse(
    int req_type,
    int64_t* tx_seq_id,
    Rpc::Outbox& outbox,
    const StatementStreamHeader* header,
    const char* error) {
  Response resp((*tx_seq_id)++, req_type);
  auto* qres = SetStatementResult(resp, header);
  if (error)
    qres->set_error(error);
  qres->add_batch()->set_is_last_batch(true);
  resp.Send(outbox);
}

class RpcExportOutput : public TraceProcessor::ExportOutput {
 public:
  // |path|, when provided, makes the export write to that file instead of
  // streaming through Write().
  explicit RpcExportOutput(const Rpc::ExportCallback& callback,
                           std::optional<std::string> path = std::nullopt)
      : callback_(callback), path_(std::move(path)) {}

  base::Status Write(const void* data, size_t size) override {
    return callback_(static_cast<const uint8_t*>(data), size,
                     /*has_more=*/true);
  }

  std::optional<std::string> GetFilePath() const override { return path_; }

 private:
  const Rpc::ExportCallback& callback_;
  std::optional<std::string> path_;
};

}  // namespace

Rpc::Rpc(std::unique_ptr<TraceProcessor> preloaded_instance,
         bool has_preloaded_eof,
         Config default_config,
         TraceProcessor::PlatformInterface* platform,
         std::function<void(TraceProcessor*)> on_trace_processor_created)
    : default_config_(default_config),
      platform_(platform),
      on_trace_processor_created_(std::move(on_trace_processor_created)),
      current_config_(std::move(default_config)),
      trace_processor_(std::move(preloaded_instance)),
      eof_(trace_processor_ ? has_preloaded_eof : false) {
  if (!trace_processor_) {
    ResetTraceProcessorInternal(default_config_);
  }
}

Rpc::Rpc() : Rpc(nullptr, false, Config(), nullptr, {}) {}
Rpc::~Rpc() = default;

void Rpc::ResetTraceProcessorInternal(const Config& config) {
  current_config_ = config;
  bytes_parsed_ = bytes_last_progress_ = 0;
  t_parse_started_ = base::GetWallTimeNs().count();

  // Clear all summarizers before destroying the old trace processor, as they
  // hold pointers to it.
  summarizers_.Clear();

  trace_processor_ = TraceProcessor::CreateInstance(config, platform_);
  if (on_trace_processor_created_) {
    on_trace_processor_created_(trace_processor_.get());
  }

  // Deliberately not touching the RPC channel state: the Streams' tokenizers
  // and Rpc's own {tx,rx} sequence ids.
  // This is invoked from the same client to clear the current trace state
  // before loading a new one. The IPC channel is orthogonal to that and the
  // message numbering continues regardless of the reset.
}

Rpc::Stream::Stream(Rpc& rpc, RpcResponseFunction response_fn)
    : rpc_(rpc), outbox_(std::make_unique<Outbox>(std::move(response_fn))) {}

Rpc::Stream::~Stream() = default;

Rpc::RequestHandle::~RequestHandle() {
  PERFETTO_CHECK(stream_ == nullptr);
}

Rpc::RequestHandle::RequestHandle(RequestHandle&& other) noexcept
    : stream_(other.stream_), write_(std::move(other.write_)) {
  other.stream_ = nullptr;
}

Rpc::RequestHandle& Rpc::RequestHandle::operator=(
    RequestHandle&& other) noexcept {
  this->~RequestHandle();  // CHECKs that any reservation held was consumed.
  new (this) RequestHandle(std::move(other));
  return *this;
}

void Rpc::RequestHandle::EndRequest(size_t size_written) {
  PERFETTO_CHECK(stream_ != nullptr);
  Stream* stream = std::exchange(stream_, nullptr);
  write_.EndWrite(size_written);
  stream->FinishRequest();
  stream->rpc_.DrainStream(*stream);
}

void Rpc::RequestHandle::AbortRequest() {
  PERFETTO_CHECK(stream_ != nullptr);
  write_.AbortWrite();
  std::exchange(stream_, nullptr)->FinishRequest();
}

Rpc::RequestHandle Rpc::Stream::BeginRequest(size_t size) {
  // The handle keeps the caller from losing a reservation, but not from
  // holding two at once. The ring buffer has the same guard for its own
  // handles; this one is here so the CHECK fires at the Stream API boundary.
  PERFETTO_CHECK(!request_in_flight_);
  request_in_flight_ = true;
  return RequestHandle(this, rxbuf_.BeginWrite(size));
}

void Rpc::Stream::FinishRequest() {
  PERFETTO_CHECK(request_in_flight_);
  request_in_flight_ = false;
}

void Rpc::DrainStream(Stream& stream) {
  for (;;) {
    auto msg = stream.rxbuf_.ReadMessage();
    if (!msg.valid()) {
      if (msg.fatal_framing_error()) {
        protozero::HeapBuffered<TraceProcessorRpcStream> err_msg;
        err_msg->add_msg()->set_fatal_error("RPC framing error");
        stream.outbox_->Send(err_msg.TakeSlices());
        stream.outbox_->Close();
      }
      break;
    }
    ParseRpcRequest(*stream.outbox_, std::move(msg));
  }
}

namespace {
using ProtoEnum = protos::pbzero::MetatraceCategories;
TraceProcessor::MetatraceCategories MetatraceCategoriesToPublicEnum(
    ProtoEnum categories) {
  TraceProcessor::MetatraceCategories result =
      TraceProcessor::MetatraceCategories::NONE;
  if (categories & ProtoEnum::QUERY_TIMELINE) {
    result = static_cast<TraceProcessor::MetatraceCategories>(
        result | TraceProcessor::MetatraceCategories::QUERY_TIMELINE);
  }
  if (categories & ProtoEnum::QUERY_DETAILED) {
    result = static_cast<TraceProcessor::MetatraceCategories>(
        result | TraceProcessor::MetatraceCategories::QUERY_DETAILED);
  }
  if (categories & ProtoEnum::FUNCTION_CALL) {
    result = static_cast<TraceProcessor::MetatraceCategories>(
        result | TraceProcessor::MetatraceCategories::FUNCTION_CALL);
  }
  if (categories & ProtoEnum::DB) {
    result = static_cast<TraceProcessor::MetatraceCategories>(
        result | TraceProcessor::MetatraceCategories::DB);
  }
  if (categories & ProtoEnum::API_TIMELINE) {
    result = static_cast<TraceProcessor::MetatraceCategories>(
        result | TraceProcessor::MetatraceCategories::API_TIMELINE);
  }
  return result;
}

}  // namespace

// [data, len] here is a tokenized TraceProcessorRpc proto message, without the
// size header.
void Rpc::ParseRpcRequest(Outbox& outbox,
                          protozero::ProtoRingBuffer::Message message) {
  RpcProto::Decoder req(message.data(), message.size());
  // Captured up front: TPM_APPEND_TRACE_DATA hands |message| to the parser.
  const size_t len = message.size();

  // We allow restarting the sequence from 0. This happens when refreshing the
  // browser while using the external trace_processor_shell --httpd.
  if (req.seq() != 0 && rx_seq_id_ != 0 && req.seq() != rx_seq_id_ + 1) {
    char err_str[255];
    // "(ERR:rpc_seq)" is intercepted by error_dialog.ts in the UI.
    snprintf(err_str, sizeof(err_str),
             "RPC request out of order. Expected %" PRId64 ", got %" PRId64
             " (ERR:rpc_seq)",
             rx_seq_id_ + 1, req.seq());
    PERFETTO_ELOG("%s", err_str);
    protozero::HeapBuffered<TraceProcessorRpcStream> err_msg;
    err_msg->add_msg()->set_fatal_error(err_str);
    outbox.Send(err_msg.TakeSlices());
    outbox.Close();
    return;
  }
  rx_seq_id_ = req.seq();

  // The static cast is to prevent that the compiler breaks future proofness.
  const int req_type = static_cast<int>(req.request());
  static const char kErrFieldNotSet[] = "RPC error: request field not set";
  switch (req_type) {
    case RpcProto::TPM_APPEND_TRACE_DATA: {
      Response resp(tx_seq_id_++, req_type);
      auto* result = resp->set_append_result();
      if (!req.has_append_trace_data()) {
        result->set_error(kErrFieldNotSet);
      } else {
        protozero::ConstBytes byte_range = req.append_trace_data();
        base::Status res = Parse(TraceBlobView(TraceBlob::Adopt(
            byte_range.data, byte_range.size, std::move(message))));
        if (!res.ok()) {
          result->set_error(res.message());
        }
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_FINALIZE_TRACE_DATA: {
      Response resp(tx_seq_id_++, req_type);
      auto* result = resp->set_finalize_data_result();
      base::Status res = NotifyEndOfFile();
      if (!res.ok()) {
        result->set_error(res.message());
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_QUERY_STREAMING: {
      if (!req.has_query_args()) {
        Response resp(tx_seq_id_++, req_type);
        auto* result = resp->set_query_result();
        result->set_error(kErrFieldNotSet);
        resp.Send(outbox);
      } else {
        protozero::ConstBytes args = req.query_args();
        protos::pbzero::QueryArgs::Decoder query(args.data, args.size);
        std::string sql = query.sql_query().ToStdString();

        PERFETTO_TP_TRACE(metatrace::Category::API_TIMELINE, "RPC_QUERY",
                          [&](metatrace::Record* r) {
                            r->AddArg("SQL", sql);
                            if (query.has_tag()) {
                              r->AddArg("tag", query.tag());
                            }
                          });

        const auto t_start = base::GetWallTimeNs();
        auto it = trace_processor_->ExecuteQuery(sql);

        QueryResultSerializer serializer(std::move(it), t_start);
        StreamSerializerResponses(&serializer, req_type, &tx_seq_id_, outbox,
                                  /*header=*/nullptr);
      }
      break;
    }
    case RpcProto::TPM_STATEMENT_STREAMING: {
      if (!req.has_statement_args()) {
        SendSingleStatementResponse(req_type, &tx_seq_id_, outbox,
                                    /*header=*/nullptr, kErrFieldNotSet);
        break;
      }
      protozero::ConstBytes args = req.statement_args();
      protos::pbzero::StatementArgs::Decoder stmt_args(args.data, args.size);
      std::string sql = stmt_args.sql().ToStdString();
      uint32_t offset = stmt_args.start_offset();

      PERFETTO_TP_TRACE(metatrace::Category::API_TIMELINE, "RPC_STATEMENT",
                        [&](metatrace::Record* r) {
                          r->AddArg("SQL", sql);
                          r->AddArg("start_offset", std::to_string(offset));
                          if (stmt_args.has_tag()) {
                            r->AddArg("tag", stmt_args.tag());
                          }
                        });

      const auto t_start = base::GetWallTimeNs();
      // Offsets are defined against the server's normalized SQL, which is
      // never longer than |sql| as sent, so this only rejects offsets that
      // are invalid in both forms. Offsets in the gap between the two sizes
      // are rejected by ExecuteNextStatement's own range check below.
      if (offset > sql.size()) {
        SendSingleStatementResponse(req_type, &tx_seq_id_, outbox,
                                    /*header=*/nullptr,
                                    "StatementArgs.start_offset out of range");
        break;
      }
      std::optional<Iterator> it =
          trace_processor_->ExecuteNextStatement(sql, &offset);
      if (!it.has_value()) {
        StatementStreamHeader header{offset, /*statement_executed=*/false};
        SendSingleStatementResponse(req_type, &tx_seq_id_, outbox, &header,
                                    /*error=*/nullptr);
        break;
      }

      // A pre-execution failure (bad offset, statement failed to prepare)
      // surfaces as an iterator carrying an error: nothing executed and
      // |offset| is unchanged, so the header must not claim otherwise.
      // Runtime errors (surfacing during serialization) did execute.
      const bool statement_executed = it->Status().ok();
      QueryResultSerializer serializer(std::move(*it), t_start);
      StatementStreamHeader header{offset, statement_executed};
      StreamSerializerResponses(&serializer, req_type, &tx_seq_id_, outbox,
                                &header);
      break;
    }
    case RpcProto::TPM_COMPUTE_METRIC: {
      Response resp(tx_seq_id_++, req_type);
      auto* result = resp->set_metric_result();
      if (!req.has_compute_metric_args()) {
        result->set_error(kErrFieldNotSet);
      } else {
        protozero::ConstBytes args = req.compute_metric_args();
        ComputeMetricInternal(args.data, args.size, result);
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_SUMMARIZE_TRACE: {
      Response resp(tx_seq_id_++, req_type);
      auto* result = resp->set_trace_summary_result();
      if (!req.has_trace_summary_args()) {
        result->set_error(kErrFieldNotSet);
      } else {
        protozero::ConstBytes args = req.trace_summary_args();
        ComputeTraceSummaryInternal(args.data, args.size, result);
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_GET_METRIC_DESCRIPTORS: {
      Response resp(tx_seq_id_++, req_type);
      auto descriptor_set = trace_processor_->GetMetricDescriptors();
      auto* result = resp->set_metric_descriptors();
      result->AppendRawProtoBytes(descriptor_set.data(), descriptor_set.size());
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_RESTORE_INITIAL_TABLES: {
      trace_processor_->RestoreInitialTables();
      Response resp(tx_seq_id_++, req_type);
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_ENABLE_METATRACE: {
      using protos::pbzero::MetatraceCategories;
      protozero::ConstBytes args = req.enable_metatrace_args();
      EnableMetatrace(args.data, args.size);

      Response resp(tx_seq_id_++, req_type);
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_DISABLE_AND_READ_METATRACE: {
      Response resp(tx_seq_id_++, req_type);
      DisableAndReadMetatraceInternal(resp->set_metatrace());
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_GET_STATUS: {
      Response resp(tx_seq_id_++, req_type);
      std::vector<uint8_t> status = GetStatus();
      resp->set_status()->AppendRawProtoBytes(status.data(), status.size());
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_RESET_TRACE_PROCESSOR: {
      Response resp(tx_seq_id_++, req_type);
      protozero::ConstBytes args = req.reset_trace_processor_args();
      ResetTraceProcessor(args.data, args.size);
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_REGISTER_SQL_PACKAGE: {
      Response resp(tx_seq_id_++, req_type);
      base::Status status = RegisterSqlPackage(req.register_sql_package_args());
      auto* res = resp->set_register_sql_package_result();
      if (!status.ok()) {
        res->set_error(status.message());
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_CREATE_SUMMARIZER: {
      Response resp(tx_seq_id_++, req_type);
      protozero::ConstBytes args = req.create_summarizer_args();
      protos::pbzero::CreateSummarizerArgs::Decoder decoder(args.data,
                                                            args.size);
      std::string summarizer_id = decoder.summarizer_id().ToStdString();

      auto* result = resp->set_create_summarizer_result();
      if (summarizers_.Find(summarizer_id)) {
        result->set_error("Summarizer already exists: " + summarizer_id);
      } else {
        std::unique_ptr<Summarizer> summarizer;
        base::Status status = trace_processor_->CreateSummarizer(&summarizer);
        if (!status.ok()) {
          result->set_error(status.message());
        } else {
          summarizers_.Insert(summarizer_id, std::move(summarizer));
          result->set_summarizer_id(summarizer_id);
        }
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_UPDATE_SUMMARIZER_SPEC: {
      Response resp(tx_seq_id_++, req_type);
      protozero::ConstBytes args = req.update_summarizer_spec_args();
      protos::pbzero::UpdateSummarizerSpecArgs::Decoder decoder(args.data,
                                                                args.size);
      std::string summarizer_id = decoder.summarizer_id().ToStdString();

      auto* result = resp->set_update_summarizer_spec_result();
      auto* summarizer_ptr = summarizers_.Find(summarizer_id);
      if (!summarizer_ptr) {
        result->set_error("Summarizer not found: " + summarizer_id);
      } else {
        Summarizer* summarizer = summarizer_ptr->get();
        protozero::ConstBytes spec_bytes = decoder.spec();
        SummarizerUpdateSpecResult update_result;
        base::Status status = summarizer->UpdateSpec(
            spec_bytes.data, spec_bytes.size, &update_result);
        if (!status.ok()) {
          result->set_error(status.message());
        }
        for (const auto& sync_info : update_result.queries) {
          auto* query = result->add_queries();
          query->set_query_id(sync_info.query_id);
          if (sync_info.error) {
            query->set_error(*sync_info.error);
          }
          query->set_was_updated(sync_info.was_updated);
          query->set_was_dropped(sync_info.was_dropped);
        }
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_QUERY_SUMMARIZER: {
      Response resp(tx_seq_id_++, req_type);
      protozero::ConstBytes args = req.query_summarizer_args();
      protos::pbzero::QuerySummarizerArgs::Decoder decoder(args.data,
                                                           args.size);
      std::string summarizer_id = decoder.summarizer_id().ToStdString();
      std::string query_id = decoder.query_id().ToStdString();

      auto* result = resp->set_query_summarizer_result();
      auto* summarizer_ptr = summarizers_.Find(summarizer_id);
      if (!summarizer_ptr) {
        result->set_error("Summarizer not found: " + summarizer_id);
      } else {
        Summarizer* summarizer = summarizer_ptr->get();
        SummarizerQueryResult query_result;
        base::Status status = summarizer->Query(query_id, &query_result);
        result->set_exists(query_result.exists);
        if (!status.ok()) {
          result->set_error(status.message());
        } else if (query_result.exists) {
          result->set_table_name(query_result.table_name);
          result->set_row_count(query_result.row_count);
          for (const auto& col : query_result.columns) {
            result->add_columns(col);
          }
          result->set_duration_ms(query_result.duration_ms);
          result->set_sql(query_result.sql);
          result->set_textproto(query_result.textproto);
          result->set_standalone_sql(query_result.standalone_sql);
        }
      }
      resp.Send(outbox);
      break;
    }
    case RpcProto::TPM_EXPORT: {
      protos::pbzero::ExportArgs::Decoder args(req.export_args());
      std::optional<TraceProcessor::ExportFormat> format =
          ParseExportFormat(args.format());
      base::Status status =
          format ? Export(*format,
                          [&](const uint8_t* chunk, size_t chunk_len,
                              bool has_more) {
                            Response resp(tx_seq_id_++, req_type);
                            auto* result = resp->set_export_result();
                            if (chunk && chunk_len > 0) {
                              result->set_data(chunk, chunk_len);
                            }
                            result->set_has_more(has_more);
                            resp.Send(outbox);
                            return base::OkStatus();
                          })
                 : base::ErrStatus("Export format is required");
      if (!status.ok()) {
        Response resp(tx_seq_id_++, req_type);
        auto* result = resp->set_export_result();
        result->set_error(status.message());
        result->set_has_more(false);
        resp.Send(outbox);
      }
      break;
    }
    case RpcProto::TPM_DESTROY_SUMMARIZER: {
      Response resp(tx_seq_id_++, req_type);
      protozero::ConstBytes args = req.destroy_summarizer_args();
      protos::pbzero::DestroySummarizerArgs::Decoder decoder(args.data,
                                                             args.size);
      std::string summarizer_id = decoder.summarizer_id().ToStdString();

      auto* result = resp->set_destroy_summarizer_result();
      if (!summarizers_.Find(summarizer_id)) {
        result->set_error("Summarizer not found: " + summarizer_id);
      } else {
        // Erasing will call the Summarizer destructor which drops all tables.
        summarizers_.Erase(summarizer_id);
      }
      resp.Send(outbox);
      break;
    }
    default: {
      // This can legitimately happen if the client is newer. We reply with a
      // generic "unknown request" response, so the client can do feature
      // detection
      PERFETTO_DLOG("[RPC] Unknown request type (%d), size=%zu", req_type, len);
      Response resp(tx_seq_id_++, req_type);
      resp->set_invalid_request(
          static_cast<RpcProto::TraceProcessorMethod>(req_type));
      resp.Send(outbox);
      break;
    }
  }  // switch(req_type)
}

base::Status Rpc::Parse(TraceBlobView blob) {
  const size_t len = blob.size();
  PERFETTO_TP_TRACE(
      metatrace::Category::API_TIMELINE, "RPC_PARSE",
      [&](metatrace::Record* r) { r->AddArg("length", std::to_string(len)); });
  if (eof_) {
    // Reset the trace processor state if another trace has been previously
    // loaded. Use the same TraceProcessor Config.
    ResetTraceProcessorInternal(current_config_);
  }

  eof_ = false;
  bytes_parsed_ += len;
  MaybePrintProgress();

  if (len == 0)
    return base::OkStatus();

  return trace_processor_->Parse(std::move(blob));
}

base::Status Rpc::NotifyEndOfFile() {
  PERFETTO_TP_TRACE(metatrace::Category::API_TIMELINE,
                    "RPC_NOTIFY_END_OF_FILE");

  eof_ = true;
  RETURN_IF_ERROR(trace_processor_->NotifyEndOfFile());
  MaybePrintProgress();
  return base::OkStatus();
}

void Rpc::ResetTraceProcessor(const uint8_t* args, size_t len) {
  protos::pbzero::ResetTraceProcessorArgs::Decoder reset_trace_processor_args(
      args, len);
  Config config = default_config_;
  if (reset_trace_processor_args.has_drop_track_event_data_before()) {
    config.drop_track_event_data_before =
        reset_trace_processor_args.drop_track_event_data_before() ==
                protos::pbzero::ResetTraceProcessorArgs::
                    TRACK_EVENT_RANGE_OF_INTEREST
            ? DropTrackEventDataBefore::kTrackEventRangeOfInterest
            : DropTrackEventDataBefore::kNoDrop;
  }
  if (reset_trace_processor_args.has_ingest_ftrace_in_raw_table()) {
    config.ingest_ftrace_in_raw_table =
        reset_trace_processor_args.ingest_ftrace_in_raw_table();
  }
  if (reset_trace_processor_args.has_analyze_trace_proto_content()) {
    config.analyze_trace_proto_content =
        reset_trace_processor_args.analyze_trace_proto_content();
  }
  if (reset_trace_processor_args.has_ftrace_drop_until_all_cpus_valid()) {
    config.soft_drop_ftrace_data_before =
        reset_trace_processor_args.ftrace_drop_until_all_cpus_valid()
            ? SoftDropFtraceDataBefore::kAllPerCpuBuffersValid
            : SoftDropFtraceDataBefore::kNoDrop;
  }
  using Args = protos::pbzero::ResetTraceProcessorArgs;
  switch (reset_trace_processor_args.parsing_mode()) {
    case Args::ParsingMode::DEFAULT:
      config.parsing_mode = ParsingMode::kDefault;
      break;
    case Args::ParsingMode::TOKENIZE_ONLY:
      config.parsing_mode = ParsingMode::kTokenizeOnly;
      break;
    case Args::ParsingMode::TOKENIZE_AND_SORT:
      config.parsing_mode = ParsingMode::kTokenizeAndSort;
      break;
  }
  switch (reset_trace_processor_args.sorting_mode()) {
    case Args::SortingMode::DEFAULT_HEURISTICS:
      config.sorting_mode = SortingMode::kDefaultHeuristics;
      break;
    case Args::SortingMode::FORCE_FULL_SORT:
      config.sorting_mode = SortingMode::kForceFullSort;
      break;
  }
  for (auto it = reset_trace_processor_args.extra_parsing_descriptors(); it;
       ++it) {
    protozero::ConstBytes bytes = it->as_bytes();
    config.extra_parsing_descriptors.emplace_back(
        reinterpret_cast<const char*>(bytes.data), bytes.size);
  }
  ResetTraceProcessorInternal(config);
}

base::Status Rpc::RegisterSqlPackage(protozero::ConstBytes bytes) {
  protos::pbzero::RegisterSqlPackageArgs::Decoder args(bytes);
  SqlPackage package;
  package.name = args.package_name().ToStdString();
  package.allow_override = args.allow_override();
  for (auto it = args.modules(); it; ++it) {
    protos::pbzero::RegisterSqlPackageArgs::Module::Decoder m(*it);
    package.modules.emplace_back(m.name().ToStdString(), m.sql().ToStdString());
  }
  return trace_processor_->RegisterSqlPackage(package);
}

void Rpc::MaybePrintProgress() {
  if (eof_ || bytes_parsed_ - bytes_last_progress_ > kProgressUpdateBytes) {
    bytes_last_progress_ = bytes_parsed_;
    auto t_load_s =
        static_cast<double>(base::GetWallTimeNs().count() - t_parse_started_) /
        1e9;
    fprintf(stderr, "\rLoading trace %.2f MB (%.1f MB/s)%s",
            static_cast<double>(bytes_parsed_) / 1e6,
            static_cast<double>(bytes_parsed_) / 1e6 / t_load_s,
            (eof_ ? "\n" : ""));
    fflush(stderr);
  }
}

void Rpc::Query(const uint8_t* args,
                size_t len,
                const QueryResultBatchCallback& result_callback) {
  protos::pbzero::QueryArgs::Decoder query(args, len);
  std::string sql = query.sql_query().ToStdString();
  PERFETTO_TP_TRACE(metatrace::Category::API_TIMELINE, "RPC_QUERY",
                    [&](metatrace::Record* r) {
                      r->AddArg("SQL", sql);
                      if (query.has_tag()) {
                        r->AddArg("tag", query.tag());
                      }
                    });

  const auto t_start = base::GetWallTimeNs();
  auto it = trace_processor_->ExecuteQuery(sql);

  QueryResultSerializer serializer(std::move(it), t_start);

  protozero::HeapBuffered<protos::pbzero::QueryResult> buffered(kSliceSize,
                                                                kSliceSize);
  for (bool has_more = true; has_more;) {
    has_more = serializer.Serialize(buffered.get());
    const auto& res = buffered.GetSlices();
    for (uint32_t i = 0; i < res.size(); ++i) {
      auto used = res[i].GetUsedRange();
      result_callback(used.begin, used.size(), has_more || i < res.size() - 1);
    }
    if (res.size() == 0 && !has_more) {
      result_callback(nullptr, 0, false);
    }
    buffered.Reset();
  }
}

std::optional<TraceProcessor::ExportFormat> Rpc::ParseExportFormat(
    int32_t format) {
  if (format == protos::pbzero::ExportArgs::PERFETTO) {
    return TraceProcessor::ExportFormat::kPerfetto;
  }
  if (format == protos::pbzero::ExportArgs::ARROW_TAR) {
    return TraceProcessor::ExportFormat::kArrowTar;
  }
  if (format == protos::pbzero::ExportArgs::SQLITE) {
    return TraceProcessor::ExportFormat::kSqlite;
  }
  return std::nullopt;
}

base::Status Rpc::Export(TraceProcessor::ExportFormat format,
                         const ExportCallback& callback) {
  if (format == TraceProcessor::ExportFormat::kSqlite) {
    return ExportSqlite(callback);
  }
  RpcExportOutput output(callback);
  RETURN_IF_ERROR(trace_processor_->Export(format, &output));
  return callback(nullptr, 0, /*has_more=*/false);
}

base::Status Rpc::ExportSqlite(const ExportCallback& callback) {
  // SQLite export needs random-access output, so the server materializes the
  // database in a server-controlled temporary file and streams the bytes back
  // to the client. The client never names a path on the server, so this does
  // not expose the server's filesystem.
  auto file = base::TempFile::Create();
  if (!file.ok()) {
    return base::ErrStatus(
        "Failed to create temporary file for SQLite export: %s",
        file.status().c_message());
  }

  RpcExportOutput output(callback, file->path());
  RETURN_IF_ERROR(
      trace_processor_->Export(TraceProcessor::ExportFormat::kSqlite, &output));

  base::ScopedFile fd = base::OpenFile(file->path(), O_RDONLY);
  if (!fd) {
    return base::ErrStatus("Failed to open exported SQLite database");
  }
  std::vector<uint8_t> buffer(64 * 1024);
  for (;;) {
    ssize_t read = base::Read(fd.get(), buffer.data(), buffer.size());
    if (read < 0) {
      return base::ErrStatus("Failed to read exported SQLite database");
    }
    if (read == 0) {
      break;
    }
    RETURN_IF_ERROR(callback(buffer.data(), static_cast<size_t>(read),
                             /*has_more=*/true));
  }
  return callback(nullptr, 0, /*has_more=*/false);
}

void Rpc::RestoreInitialTables() {
  trace_processor_->RestoreInitialTables();
}

std::vector<uint8_t> Rpc::ComputeMetric(const uint8_t* args, size_t len) {
  protozero::HeapBuffered<protos::pbzero::ComputeMetricResult> result;
  ComputeMetricInternal(args, len, result.get());
  return result.SerializeAsArray();
}

std::vector<uint8_t> Rpc::ComputeTraceSummary(const uint8_t* args, size_t len) {
  protozero::HeapBuffered<protos::pbzero::TraceSummaryResult> result;
  ComputeTraceSummaryInternal(args, len, result.get());
  return result.SerializeAsArray();
}

void Rpc::ComputeMetricInternal(const uint8_t* data,
                                size_t len,
                                protos::pbzero::ComputeMetricResult* result) {
  protos::pbzero::ComputeMetricArgs::Decoder args(data, len);
  std::vector<std::string> metric_names;
  for (auto it = args.metric_names(); it; ++it) {
    metric_names.emplace_back(it->as_std_string());
  }

  PERFETTO_TP_TRACE(metatrace::Category::API_TIMELINE, "RPC_COMPUTE_METRIC",
                    [&](metatrace::Record* r) {
                      for (const auto& metric : metric_names) {
                        r->AddArg("Metric", metric);
                        r->AddArg("Format", std::to_string(args.format()));
                      }
                    });

  PERFETTO_DLOG("[RPC] ComputeMetrics(%zu, %s), format=%d", metric_names.size(),
                metric_names.empty() ? "" : metric_names.front().c_str(),
                args.format());
  switch (args.format()) {
    case protos::pbzero::ComputeMetricArgs::BINARY_PROTOBUF: {
      std::vector<uint8_t> metrics_proto;
      base::Status status =
          trace_processor_->ComputeMetric(metric_names, &metrics_proto);
      if (status.ok()) {
        result->set_metrics(metrics_proto.data(), metrics_proto.size());
      } else {
        result->set_error(status.message());
      }
      break;
    }
    case protos::pbzero::ComputeMetricArgs::TEXTPROTO: {
      std::string metrics_string;
      base::Status status = trace_processor_->ComputeMetricText(
          metric_names, TraceProcessor::MetricResultFormat::kProtoText,
          &metrics_string);
      if (status.ok()) {
        result->set_metrics_as_prototext(metrics_string);
      } else {
        result->set_error(status.message());
      }
      break;
    }
    case protos::pbzero::ComputeMetricArgs::JSON: {
      std::string metrics_string;
      base::Status status = trace_processor_->ComputeMetricText(
          metric_names, TraceProcessor::MetricResultFormat::kJson,
          &metrics_string);
      if (status.ok()) {
        result->set_metrics_as_json(metrics_string);
      } else {
        result->set_error(status.message());
      }
      break;
    }
  }
}

void Rpc::ComputeTraceSummaryInternal(
    const uint8_t* data,
    size_t len,
    protos::pbzero::TraceSummaryResult* result) {
  protos::pbzero::TraceSummaryArgs::Decoder args(data, len);
  if (!args.has_proto_specs() && !args.has_textproto_specs()) {
    result->set_error("TraceSummary missing trace_summary_spec");
    return;
  }
  if (!args.has_output_format()) {
    result->set_error("TraceSummary missing format");
    return;
  }

  auto comp_spec = args.computation_spec();
  protos::pbzero::TraceSummaryArgs::ComputationSpec::Decoder comp_spec_decoder(
      comp_spec.data, comp_spec.size);

  TraceSummaryComputationSpec computation_spec;

  if (comp_spec_decoder.has_run_all_metrics() &&
      comp_spec_decoder.run_all_metrics() == true) {
    computation_spec.v2_metric_ids = std::nullopt;
  } else {
    computation_spec.v2_metric_ids = std::vector<std::string>();
    for (auto it = comp_spec_decoder.metric_ids(); it; ++it) {
      computation_spec.v2_metric_ids->push_back(it->as_std_string());
    }
  }

  if (comp_spec_decoder.has_metadata_query_id()) {
    computation_spec.metadata_query_id =
        comp_spec_decoder.metadata_query_id().ToStdString();
  }

  std::vector<TraceSummarySpecBytes> summary_specs;
  for (auto it = args.proto_specs(); it; ++it) {
    TraceSummarySpecBytes spec;
    spec.ptr = it->data();
    spec.size = it->size();
    spec.format = TraceSummarySpecBytes::Format::kBinaryProto;
    summary_specs.push_back(spec);
  }
  for (auto it = args.textproto_specs(); it; ++it) {
    TraceSummarySpecBytes spec;
    spec.ptr = it->data();
    spec.size = it->size();
    spec.format = TraceSummarySpecBytes::Format::kTextProto;
    summary_specs.push_back(spec);
  }
  TraceSummaryOutputSpec output_spec;
  switch (args.output_format()) {
    case protos::pbzero::TraceSummaryArgs::BINARY_PROTOBUF:
      output_spec.format = TraceSummaryOutputSpec::Format::kBinaryProto;
      break;
    case protos::pbzero::TraceSummaryArgs::TEXTPROTO:
      output_spec.format = TraceSummaryOutputSpec::Format::kTextProto;
      break;
    default:
      PERFETTO_FATAL("Unknown format");
  }
  std::vector<uint8_t> output;
  base::Status status = trace_processor_->Summarize(
      computation_spec, summary_specs, &output, output_spec);
  if (!status.ok()) {
    result->set_error(status.message());
    return;
  }
  switch (output_spec.format) {
    case TraceSummaryOutputSpec::Format::kBinaryProto: {
      result->set_proto_summary(output.data(), output.size());
      break;
    }
    case TraceSummaryOutputSpec::Format::kTextProto: {
      std::string textproto_output(output.begin(), output.end());
      result->set_textproto_summary(textproto_output);
      break;
    }
  }
}

void Rpc::EnableMetatrace(const uint8_t* data, size_t len) {
  using protos::pbzero::MetatraceCategories;
  TraceProcessor::MetatraceConfig config;
  protos::pbzero::EnableMetatraceArgs::Decoder args(data, len);
  config.categories = MetatraceCategoriesToPublicEnum(
      static_cast<MetatraceCategories>(args.categories()));
  trace_processor_->EnableMetatrace(config);
}

std::vector<uint8_t> Rpc::DisableAndReadMetatrace() {
  protozero::HeapBuffered<protos::pbzero::DisableAndReadMetatraceResult> result;
  DisableAndReadMetatraceInternal(result.get());
  return result.SerializeAsArray();
}

void Rpc::DisableAndReadMetatraceInternal(
    protos::pbzero::DisableAndReadMetatraceResult* result) {
  std::vector<uint8_t> trace_proto;
  base::Status status = trace_processor_->DisableAndReadMetatrace(&trace_proto);
  if (status.ok()) {
    result->set_metatrace(trace_proto.data(), trace_proto.size());
  } else {
    result->set_error(status.message());
  }
}

std::vector<uint8_t> Rpc::GetStatus() {
  protozero::HeapBuffered<protos::pbzero::StatusResult> status;
  status->set_loaded_trace_name(trace_processor_->GetCurrentTraceName());
  status->set_human_readable_version(base::GetVersionString());
  if (const char* version_code = base::GetVersionCode(); version_code) {
    status->set_version_code(version_code);
  }
  status->set_api_version(protos::pbzero::TRACE_PROCESSOR_CURRENT_API_VERSION);
  return status.SerializeAsArray();
}

}  // namespace perfetto::trace_processor