1 /*
2 * Copyright (C) 2018 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "src/tracing/test/mock_producer.h"
18
19 #include "perfetto/ext/tracing/core/client_identity.h"
20 #include "perfetto/ext/tracing/core/trace_writer.h"
21 #include "perfetto/ext/tracing/core/tracing_service.h"
22 #include "perfetto/protozero/scattered_heap_buffer.h"
23 #include "perfetto/tracing/core/data_source_config.h"
24 #include "perfetto/tracing/core/data_source_descriptor.h"
25 #include "protos/perfetto/common/track_event_descriptor.pbzero.h"
26 #include "src/base/test/test_task_runner.h"
27
28 using ::testing::_;
29 using ::testing::Eq;
30 using ::testing::Invoke;
31 using ::testing::InvokeWithoutArgs;
32 using ::testing::Property;
33
34 namespace perfetto {
35
36 namespace {
37
CreateDataSourceDescriptor(const std::initializer_list<std::string> & categories,uint32_t id)38 static DataSourceDescriptor CreateDataSourceDescriptor(
39 const std::initializer_list<std::string>& categories,
40 uint32_t id) {
41 DataSourceDescriptor ds_desc;
42 ds_desc.set_name("track_event");
43 ds_desc.set_id(id);
44
45 protozero::HeapBuffered<protos::pbzero::TrackEventDescriptor> ted;
46 for (auto c : categories) {
47 auto cat = ted->add_available_categories();
48 cat->set_name(c);
49 }
50 ds_desc.set_track_event_descriptor_raw(ted.SerializeAsString());
51 return ds_desc;
52 }
53
54 } // anonymous namespace
55
MockProducer(base::TestTaskRunner * task_runner)56 MockProducer::MockProducer(base::TestTaskRunner* task_runner)
57 : task_runner_(task_runner) {}
58
~MockProducer()59 MockProducer::~MockProducer() {
60 if (!service_endpoint_)
61 return;
62 static int i = 0;
63 auto checkpoint_name = "on_producer_disconnect_" + std::to_string(i++);
64 auto on_disconnect = task_runner_->CreateCheckpoint(checkpoint_name);
65 EXPECT_CALL(*this, OnDisconnect()).WillOnce(Invoke(on_disconnect));
66 service_endpoint_.reset();
67 task_runner_->RunUntilCheckpoint(checkpoint_name);
68 }
69
Connect(TracingService * svc,const std::string & producer_name,uid_t uid,pid_t pid,size_t shared_memory_size_hint_bytes,size_t shared_memory_page_size_hint_bytes,std::unique_ptr<SharedMemory> shm)70 void MockProducer::Connect(TracingService* svc,
71 const std::string& producer_name,
72 uid_t uid,
73 pid_t pid,
74 size_t shared_memory_size_hint_bytes,
75 size_t shared_memory_page_size_hint_bytes,
76 std::unique_ptr<SharedMemory> shm) {
77 producer_name_ = producer_name;
78 service_endpoint_ = svc->ConnectProducer(
79 this, ClientIdentity(uid, pid), producer_name,
80 shared_memory_size_hint_bytes,
81 /*in_process=*/true, TracingService::ProducerSMBScrapingMode::kDefault,
82 shared_memory_page_size_hint_bytes, std::move(shm));
83 auto checkpoint_name = "on_producer_connect_" + producer_name;
84 auto on_connect = task_runner_->CreateCheckpoint(checkpoint_name);
85 EXPECT_CALL(*this, OnConnect()).WillOnce(Invoke(on_connect));
86 task_runner_->RunUntilCheckpoint(checkpoint_name);
87 }
88
RegisterDataSource(const std::string & name,bool ack_stop,bool ack_start,bool handle_incremental_state_clear,bool no_flush)89 void MockProducer::RegisterDataSource(const std::string& name,
90 bool ack_stop,
91 bool ack_start,
92 bool handle_incremental_state_clear,
93 bool no_flush) {
94 DataSourceDescriptor ds_desc;
95 ds_desc.set_name(name);
96 ds_desc.set_will_notify_on_stop(ack_stop);
97 ds_desc.set_will_notify_on_start(ack_start);
98 ds_desc.set_handles_incremental_state_clear(handle_incremental_state_clear);
99 ds_desc.set_no_flush(no_flush);
100 service_endpoint_->RegisterDataSource(ds_desc);
101 }
102
UnregisterDataSource(const std::string & name)103 void MockProducer::UnregisterDataSource(const std::string& name) {
104 service_endpoint_->UnregisterDataSource(name);
105 }
106
RegisterTrackEventDataSource(const std::initializer_list<std::string> & categories,uint32_t id)107 void MockProducer::RegisterTrackEventDataSource(
108 const std::initializer_list<std::string>& categories,
109 uint32_t id) {
110 service_endpoint_->RegisterDataSource(
111 CreateDataSourceDescriptor(categories, id));
112 }
113
UpdateTrackEventDataSource(const std::initializer_list<std::string> & categories,uint32_t id)114 void MockProducer::UpdateTrackEventDataSource(
115 const std::initializer_list<std::string>& categories,
116 uint32_t id) {
117 service_endpoint_->UpdateDataSource(
118 CreateDataSourceDescriptor(categories, id));
119 }
120
RegisterTraceWriter(uint32_t writer_id,uint32_t target_buffer)121 void MockProducer::RegisterTraceWriter(uint32_t writer_id,
122 uint32_t target_buffer) {
123 service_endpoint_->RegisterTraceWriter(writer_id, target_buffer);
124 }
125
UnregisterTraceWriter(uint32_t writer_id)126 void MockProducer::UnregisterTraceWriter(uint32_t writer_id) {
127 service_endpoint_->UnregisterTraceWriter(writer_id);
128 }
129
WaitForTracingSetup()130 void MockProducer::WaitForTracingSetup() {
131 static int i = 0;
132 auto checkpoint_name =
133 "on_shmem_initialized_" + producer_name_ + "_" + std::to_string(i++);
134 auto on_tracing_enabled = task_runner_->CreateCheckpoint(checkpoint_name);
135 EXPECT_CALL(*this, OnTracingSetup()).WillOnce(Invoke(on_tracing_enabled));
136 task_runner_->RunUntilCheckpoint(checkpoint_name);
137 }
138
WaitForDataSourceSetup(const std::string & name)139 void MockProducer::WaitForDataSourceSetup(const std::string& name) {
140 static int i = 0;
141 auto checkpoint_name = "on_ds_setup_" + name + "_" + std::to_string(i++);
142 auto on_ds_start = task_runner_->CreateCheckpoint(checkpoint_name);
143 EXPECT_CALL(*this,
144 SetupDataSource(_, Property(&DataSourceConfig::name, Eq(name))))
145 .WillOnce(Invoke([on_ds_start, this](DataSourceInstanceID ds_id,
146 const DataSourceConfig& cfg) {
147 EXPECT_FALSE(data_source_instances_.count(cfg.name()));
148 auto target_buffer = static_cast<BufferID>(cfg.target_buffer());
149 auto session_id =
150 static_cast<TracingSessionID>(cfg.tracing_session_id());
151 data_source_instances_.emplace(
152 cfg.name(), EnabledDataSource{ds_id, target_buffer, session_id});
153 on_ds_start();
154 }));
155 task_runner_->RunUntilCheckpoint(checkpoint_name);
156 }
157
WaitForDataSourceStart(const std::string & name)158 void MockProducer::WaitForDataSourceStart(const std::string& name) {
159 static int i = 0;
160 auto checkpoint_name = "on_ds_start_" + name + "_" + std::to_string(i++);
161 auto on_ds_start = task_runner_->CreateCheckpoint(checkpoint_name);
162 EXPECT_CALL(*this,
163 StartDataSource(_, Property(&DataSourceConfig::name, Eq(name))))
164 .WillOnce(Invoke([on_ds_start, this](DataSourceInstanceID ds_id,
165 const DataSourceConfig& cfg) {
166 // The data source might have been seen already through
167 // WaitForDataSourceSetup().
168 if (data_source_instances_.count(cfg.name()) == 0) {
169 auto target_buffer = static_cast<BufferID>(cfg.target_buffer());
170 auto session_id =
171 static_cast<TracingSessionID>(cfg.tracing_session_id());
172 data_source_instances_.emplace(
173 cfg.name(), EnabledDataSource{ds_id, target_buffer, session_id});
174 }
175 on_ds_start();
176 }));
177 task_runner_->RunUntilCheckpoint(checkpoint_name);
178 }
179
WaitForDataSourceStop(const std::string & name)180 void MockProducer::WaitForDataSourceStop(const std::string& name) {
181 static int i = 0;
182 auto checkpoint_name = "on_ds_stop_" + name + "_" + std::to_string(i++);
183 auto on_ds_stop = task_runner_->CreateCheckpoint(checkpoint_name);
184 ASSERT_EQ(1u, data_source_instances_.count(name));
185 DataSourceInstanceID ds_id = data_source_instances_[name].id;
186 EXPECT_CALL(*this, StopDataSource(ds_id))
187 .WillOnce(InvokeWithoutArgs(on_ds_stop));
188 task_runner_->RunUntilCheckpoint(checkpoint_name);
189 data_source_instances_.erase(name);
190 }
191
CreateTraceWriter(const std::string & data_source_name)192 std::unique_ptr<TraceWriter> MockProducer::CreateTraceWriter(
193 const std::string& data_source_name) {
194 PERFETTO_DCHECK(data_source_instances_.count(data_source_name));
195 BufferID buf_id = data_source_instances_[data_source_name].target_buffer;
196 return service_endpoint_->CreateTraceWriter(buf_id);
197 }
198
ExpectFlush(TraceWriter * writer_to_flush,bool reply,FlushFlags expected_flags)199 void MockProducer::ExpectFlush(TraceWriter* writer_to_flush,
200 bool reply,
201 FlushFlags expected_flags) {
202 std::vector<TraceWriter*> writers;
203 if (writer_to_flush)
204 writers.push_back(writer_to_flush);
205 ExpectFlush(writers, reply, expected_flags);
206 }
207
ExpectFlush(std::vector<TraceWriter * > writers_to_flush,bool reply,FlushFlags expected_flags)208 void MockProducer::ExpectFlush(std::vector<TraceWriter*> writers_to_flush,
209 bool reply,
210 FlushFlags expected_flags) {
211 auto& expected_call = EXPECT_CALL(*this, Flush(_, _, _, _));
212 expected_call.WillOnce(
213 Invoke([this, writers_to_flush, reply, expected_flags](
214 FlushRequestID flush_req_id, const DataSourceInstanceID*,
215 size_t, FlushFlags actual_flags) {
216 if (expected_flags.flags()) {
217 EXPECT_EQ(actual_flags, expected_flags);
218 }
219 for (auto* writer : writers_to_flush) {
220 writer->Flush();
221 }
222 if (reply) {
223 service_endpoint_->NotifyFlushComplete(flush_req_id);
224 }
225 }));
226 }
227
GetDataSourceInstanceId(const std::string & name)228 DataSourceInstanceID MockProducer::GetDataSourceInstanceId(
229 const std::string& name) {
230 auto it = data_source_instances_.find(name);
231 return it == data_source_instances_.end() ? 0 : it->second.id;
232 }
233
GetDataSourceInstance(const std::string & name)234 const MockProducer::EnabledDataSource* MockProducer::GetDataSourceInstance(
235 const std::string& name) {
236 auto it = data_source_instances_.find(name);
237 return it == data_source_instances_.end() ? nullptr : &it->second;
238 }
239
240 } // namespace perfetto
241