src/trace_processor/importers/proto/args_table_utils_unittest.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/importers/proto/args_table_utils.h"

#include "perfetto/ext/base/string_view.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "protos/perfetto/common/descriptor.pbzero.h"
#include "protos/perfetto/trace/interned_data/interned_data.pbzero.h"
#include "protos/perfetto/trace/track_event/source_location.pbzero.h"
#include "src/protozero/test/example_proto/test_messages.descriptor.h"
#include "src/protozero/test/example_proto/test_messages.pbzero.h"
#include "src/trace_processor/args_tracker.h"
#include "src/trace_processor/importers/proto/chrome_compositor_scheduler_state.descriptor.h"
#include "src/trace_processor/trace_storage.h"
#include "test/gtest_and_gmock.h"

namespace perfetto {
namespace trace_processor {
namespace {

constexpr size_t kChunkSize = 42;

using ::testing::_;
using ::testing::Eq;
using ::testing::Invoke;
using ::testing::NiceMock;

class ArgsTableUtilsTest : public ::testing::Test {
 protected:
  ArgsTableUtilsTest() {
    context_.storage.reset(new TraceStorage);
    storage_ = context_.storage.get();
    context_.args_tracker.reset(new ArgsTracker(&context_));
    sequence_state_.reset(new PacketSequenceState(&context_));
  }

  bool HasArg(ArgSetId set_id, const base::StringView& key, Variadic value) {
    const auto& args = storage_->args();
    auto key_id = storage_->string_pool().GetId(key);
    EXPECT_TRUE(key_id);
    auto rows =
        std::equal_range(args.set_ids().begin(), args.set_ids().end(), set_id);
    bool found = false;
    for (; rows.first != rows.second; rows.first++) {
      size_t index = static_cast<size_t>(
          std::distance(args.set_ids().begin(), rows.first));
      if (args.keys()[index] == key_id) {
        EXPECT_EQ(args.flat_keys()[index], key_id);
        if (args.flat_keys()[index] == key_id &&
            args.arg_values()[index] == value) {
          found = true;
          break;
        }
      }
    }
    return found;
  }

  uint32_t arg_set_id_ = 1;
  std::unique_ptr<PacketSequenceState> sequence_state_;
  TraceProcessorContext context_;
  TraceStorage* storage_;
};

TEST_F(ArgsTableUtilsTest, EnsureChromeCompositorStateDescriptorParses) {
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(
      kChromeCompositorSchedulerStateDescriptor.data(),
      kChromeCompositorSchedulerStateDescriptor.size());
  EXPECT_TRUE(status.ok())
      << "Failed to parse kChromeCompositorSchedulerStateDescriptor: "
      << status.message();
}

TEST_F(ArgsTableUtilsTest, EnsureTestMessageProtoParses) {
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  EXPECT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();
}

TEST_F(ArgsTableUtilsTest, BasicSingleLayerProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<EveryField> msg{kChunkSize, kChunkSize};
  msg->set_field_int32(-1);
  msg->set_field_int64(-333123456789ll);
  msg->set_field_uint32(600);
  msg->set_field_uint64(333123456789ll);
  msg->set_field_sint32(-5);
  msg->set_field_sint64(-9000);
  msg->set_field_fixed32(12345);
  msg->set_field_fixed64(444123450000ll);
  msg->set_field_sfixed32(-69999);
  msg->set_field_sfixed64(-200);
  msg->set_field_double(0.5555);
  msg->set_field_bool(true);
  msg->set_small_enum(SmallEnum::TO_BE);
  msg->set_signed_enum(SignedEnum::NEGATIVE);
  msg->set_big_enum(BigEnum::BEGIN);
  msg->set_nested_enum(EveryField::PONG);
  msg->set_field_float(3.14f);
  msg->set_field_string("FizzBuzz");
  msg->add_repeated_int32(1);
  msg->add_repeated_int32(-1);
  msg->add_repeated_int32(100);
  msg->add_repeated_int32(2000000);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.EveryField",
      TraceStorage::CreateRowId(TableId::kTrack, 0));

  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();

  context_.args_tracker->Flush();
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_int32", Variadic::Integer(-1)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_int64",
                     Variadic::Integer(-333123456789ll)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_uint32",
                     Variadic::UnsignedInteger(600)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_uint64",
                     Variadic::UnsignedInteger(333123456789ll)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sint32", Variadic::Integer(-5)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sint64", Variadic::Integer(-9000)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_fixed32", Variadic::Integer(12345)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_fixed64",
                     Variadic::Integer(444123450000ll)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_sfixed32",
                     Variadic::Integer(-69999)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_sfixed64", Variadic::Integer(-200)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_double", Variadic::Real(0.5555)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "field_bool", Variadic::Boolean(true)));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "small_enum",
      Variadic::String(*context_.storage->string_pool().GetId("TO_BE"))));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "signed_enum",
      Variadic::String(*context_.storage->string_pool().GetId("NEGATIVE"))));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "big_enum",
      Variadic::String(*context_.storage->string_pool().GetId("BEGIN"))));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "nested_enum",
             Variadic::String(*context_.storage->string_pool().GetId("PONG"))));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "field_float",
                     Variadic::Real(static_cast<double>(3.14f))));
  ASSERT_TRUE(context_.storage->string_pool().GetId("FizzBuzz"));
  EXPECT_TRUE(HasArg(
      ArgSetId(arg_set_id_), "field_string",
      Variadic::String(*context_.storage->string_pool().GetId("FizzBuzz"))));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "repeated_int32", Variadic::Integer(1)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "repeated_int32", Variadic::Integer(-1)));
}

TEST_F(ArgsTableUtilsTest, NestedProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA",
      TraceStorage::CreateRowId(TableId::kTrack, 0));
  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_c",
                     Variadic::Integer(3)));
}

TEST_F(ArgsTableUtilsTest, CamelCaseFieldsProto) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<CamelCaseFields> msg{kChunkSize, kChunkSize};
  msg->set_barbaz(true);
  msg->set_moomoo(true);
  msg->set___bigbang(true);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.CamelCaseFields",
      TraceStorage::CreateRowId(TableId::kTrack, 0));
  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "barBaz", Variadic::Boolean(true)));
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "MooMoo", Variadic::Boolean(true)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "__bigBang", Variadic::Boolean(true)));
}

TEST_F(ArgsTableUtilsTest, NestedProtoParsingOverrideHandled) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field) {
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        EXPECT_TRUE(state.args_tracker);
        EXPECT_TRUE(state.context);
        EXPECT_TRUE(state.sequence_state);
        auto id = state.context->storage->InternString(
            "super_nested.value_b.replaced");
        int32_t val = field.as_int32();
        state.args_tracker->AddArg(state.row_id, id, id,
                                   Variadic::Integer(val + 1));
        // We've handled this field by adding the desired args.
        return true;
      });
  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA",
      TraceStorage::CreateRowId(TableId::kTrack, 0));
  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_b.replaced",
                     Variadic::Integer(4)));
}

TEST_F(ArgsTableUtilsTest, NestedProtoParsingOverrideSkipped) {
  using namespace protozero::test::protos::pbzero;
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);

  auto binary_proto = msg.SerializeAsArray();

  storage_->mutable_track_table()->Insert({});
  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field) {
        static int val = 0;
        ++val;
        EXPECT_EQ(1, val);
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        EXPECT_TRUE(state.args_tracker);
        EXPECT_TRUE(state.sequence_state);
        EXPECT_TRUE(state.context);
        // By returning false we expect this field to be handled like regular.
        return false;
      });
  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA",
      TraceStorage::CreateRowId(TableId::kTrack, 0));
  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "super_nested.value_c",
                     Variadic::Integer(3)));
}

TEST_F(ArgsTableUtilsTest, LookingUpInternedStateParsingOverride) {
  using namespace protozero::test::protos::pbzero;
  // The test proto, we will use |value_c| as the source_location iid.
  protozero::HeapBuffered<NestedA> msg{kChunkSize, kChunkSize};
  msg->set_super_nested()->set_value_c(3);
  auto binary_proto = msg.SerializeAsArray();

  // The interned source location.
  protozero::HeapBuffered<protos::pbzero::SourceLocation> src_loc{kChunkSize,
                                                                  kChunkSize};
  src_loc->set_iid(3);
  src_loc->set_file_name("test_file_name");
  // We need to update sequence_state to point to it.
  auto binary_data = src_loc.SerializeAsArray();
  std::unique_ptr<uint8_t[]> buffer(new uint8_t[binary_data.size()]);
  for (size_t i = 0; i < binary_data.size(); ++i) {
    buffer.get()[i] = binary_data[i];
  }
  TraceBlobView blob(std::move(buffer), 0, binary_data.size());
  sequence_state_->InternMessage(
      protos::pbzero::InternedData::kSourceLocationsFieldNumber,
      std::move(blob));

  ProtoToArgsTable helper(sequence_state_.get(),
                          sequence_state_->current_generation(), &context_,
                          nullptr, "", 0);
  // Now we override the behaviour of |value_c| so we can expand the iid into
  // multiple args rows.
  helper.AddParsingOverride(
      "super_nested.value_c",
      [](const ProtoToArgsTable::ParsingOverrideState& state,
         const protozero::Field& field) {
        EXPECT_EQ(field.type(), protozero::proto_utils::ProtoWireType::kVarInt);
        auto* decoder = state.sequence_state->LookupInternedMessage<
            protos::pbzero::InternedData::kSourceLocationsFieldNumber,
            protos::pbzero::SourceLocation>(state.sequence_generation,
                                            field.as_uint64());
        if (!decoder) {
          // Lookup failed fall back on default behaviour.
          return false;
        }
        TraceStorage* storage = state.context->storage.get();
        auto file_name_id = storage->InternString("file_name");
        auto line_number_id = storage->InternString("line_number");
        auto file_id = storage->InternString(decoder->file_name());
        state.args_tracker->AddArg(state.row_id, file_name_id, file_name_id,
                                   Variadic::String(file_id));
        state.args_tracker->AddArg(state.row_id, line_number_id, line_number_id,
                                   Variadic::Integer(2));
        return true;
      });

  storage_->mutable_track_table()->Insert({});
  auto status = helper.AddProtoFileDescriptor(kTestMessagesDescriptor.data(),
                                              kTestMessagesDescriptor.size());
  ASSERT_TRUE(status.ok()) << "Failed to parse kTestMessagesDescriptor: "
                           << status.message();

  status = helper.InternProtoIntoArgsTable(
      protozero::ConstBytes{binary_proto.data(), binary_proto.size()},
      ".protozero.test.protos.NestedA",
      TraceStorage::CreateRowId(TableId::kTrack, 0));
  EXPECT_TRUE(status.ok()) << "InternProtoIntoArgsTable failed with error: "
                           << status.message();
  auto file_name_id = storage_->string_pool().GetId("test_file_name");
  ASSERT_TRUE(file_name_id);
  context_.args_tracker->Flush();
  EXPECT_TRUE(HasArg(ArgSetId(arg_set_id_), "file_name",
                     Variadic::String(*file_name_id)));
  EXPECT_TRUE(
      HasArg(ArgSetId(arg_set_id_), "line_number", Variadic::Integer(2)));
}

}  // namespace
}  // namespace trace_processor
}  // namespace perfetto