1 // Copyright 2020 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_PRODUCER_TEST_UTILS_H_
6 #define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_PRODUCER_TEST_UTILS_H_
7 
8 #include <memory>
9 #include <vector>
10 
11 #include "base/callback_forward.h"
12 #include "base/optional.h"
13 #include "services/tracing/public/cpp/perfetto/producer_client.h"
14 #include "services/tracing/public/cpp/perfetto/task_runner.h"
15 #include "third_party/perfetto/include/perfetto/ext/tracing/core/trace_writer.h"
16 #include "third_party/perfetto/include/perfetto/protozero/root_message.h"
17 #include "third_party/perfetto/include/perfetto/protozero/scattered_stream_null_delegate.h"
18 #include "third_party/perfetto/protos/perfetto/trace/chrome/chrome_trace_event.pb.h"
19 #include "third_party/perfetto/protos/perfetto/trace/trace_packet.pb.h"
20 #include "third_party/perfetto/protos/perfetto/trace/trace_packet.pbzero.h"
21 
22 namespace tracing {
23 
24 // Test producer client for data source tests.
25 class TestProducerClient : public ProducerClient {
26  public:
27   explicit TestProducerClient(
28       std::unique_ptr<PerfettoTaskRunner> main_thread_task_runner,
29       bool log_only_main_thread = true);
30   ~TestProducerClient() override;
31 
32   // ProducerClient implementation:
33   std::unique_ptr<perfetto::TraceWriter> CreateTraceWriter(
34       perfetto::BufferID target_buffer,
35       perfetto::BufferExhaustedPolicy =
36           perfetto::BufferExhaustedPolicy::kDefault) override;
37 
38   void FlushPacketIfPossible();
39 
40   perfetto::protos::pbzero::TracePacket* NewTracePacket();
41 
42   size_t GetFinalizedPacketCount();
43 
44   const perfetto::protos::TracePacket* GetFinalizedPacket(
45       size_t packet_index = 0);
46 
47   const google::protobuf::RepeatedPtrField<perfetto::protos::ChromeMetadata>*
48   GetChromeMetadata(size_t packet_index = 0);
49 
50   const perfetto::protos::ChromeMetadataPacket* GetProtoChromeMetadata(
51       size_t packet_index = 0);
52 
53   const std::vector<std::unique_ptr<perfetto::protos::TracePacket>>&
finalized_packets()54   finalized_packets() const {
55     return finalized_packets_;
56   }
57 
58   TestProducerClient(TestProducerClient&&) = delete;
59   TestProducerClient& operator=(TestProducerClient&&) = delete;
60 
61  private:
62   std::vector<std::unique_ptr<perfetto::protos::TracePacket>>
63       finalized_packets_;
64   std::vector<std::unique_ptr<perfetto::protos::TracePacket>>
65       legacy_metadata_packets_;
66   std::vector<std::unique_ptr<perfetto::protos::TracePacket>>
67       proto_metadata_packets_;
68   protozero::RootMessage<perfetto::protos::pbzero::TracePacket> trace_packet_;
69   protozero::ScatteredStreamWriterNullDelegate delegate_;
70   protozero::ScatteredStreamWriter stream_;
71   std::unique_ptr<PerfettoTaskRunner> main_thread_task_runner_;
72   bool log_only_main_thread_;
73 };
74 
75 class TestTraceWriter : public perfetto::TraceWriter {
76  public:
77   using TracePacketCallback = base::RepeatingCallback<void(
78       std::unique_ptr<perfetto::protos::TracePacket>)>;
79   explicit TestTraceWriter(TestProducerClient* producer_client);
80 
81   perfetto::TraceWriter::TracePacketHandle NewTracePacket() override;
82   void Flush(std::function<void()> callback = {}) override {}
83   perfetto::WriterID writer_id() const override;
84   uint64_t written() const override;
85 
86   TestTraceWriter(TestTraceWriter&&) = delete;
87   TestTraceWriter& operator=(TestTraceWriter&&) = delete;
88 
89  private:
90   TestProducerClient* producer_client_;
91 };
92 
93 }  // namespace tracing
94 
95 #endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_PRODUCER_TEST_UTILS_H_
96