xref: /aosp_15_r20/external/federated-compute/fcp/tracing/test_tracing_recorder_impl.h (revision 14675a029014e728ec732f129a32e299b2da0601)
1 // Copyright 2019 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #ifndef FCP_TRACING_TEST_TRACING_RECORDER_IMPL_H_
16 #define FCP_TRACING_TEST_TRACING_RECORDER_IMPL_H_
17 
18 #include <memory>
19 #include <vector>
20 
21 #include "absl/container/flat_hash_map.h"
22 #include "fcp/tracing/tracing_recorder_impl.h"
23 
24 namespace fcp {
25 namespace tracing_internal {
26 
27 class TestTracingSpanImpl;
28 
29 class TestTracingRecorderImpl : public TracingRecorderImpl {
30  public:
31   // Allows listening for/handling traces as they appear.
32   class TraceListener {
33    public:
34     virtual ~TraceListener() = default;
35     // Called exactly once when the root span is created.
36     virtual void OnRoot(TracingSpanId id, flatbuffers::DetachedBuffer data) = 0;
37     // Called when a new span or event is created.
38     // @param id Will only be called once per
39     // unique ID (although it may be called multiple times with the same
40     // parent span ID.) id
41     virtual void OnTrace(TracingSpanId parent_id, TracingSpanId id,
42                          flatbuffers::DetachedBuffer data) = 0;
43   };
44   explicit TestTracingRecorderImpl(TraceListener* trace_listener);
45   ~TestTracingRecorderImpl() override;
46   TracingSpanImpl* GetRootSpan() override;
47   void TraceImpl(TracingSpanId id, flatbuffers::DetachedBuffer&& buf,
48                  const TracingTraitsBase& traits) override;
49   std::unique_ptr<TracingSpanImpl> CreateChildSpan(
50       TracingSpanId parent_span_id, flatbuffers::DetachedBuffer&& buf,
51       const TracingTraitsBase& traits) override;
52 
53  private:
54   TraceListener* trace_listener_;
55   std::unique_ptr<TestTracingSpanImpl> root_span_;
56 };
57 
58 }  // namespace tracing_internal
59 }  // namespace fcp
60 
61 #endif  // FCP_TRACING_TEST_TRACING_RECORDER_IMPL_H_
62