1 | /* |
---|---|
2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
3 | * All rights reserved. |
4 | * |
5 | * This source code is licensed under the BSD-style license found in the |
6 | * LICENSE file in the root directory of this source tree. |
7 | */ |
8 | |
9 | #pragma once |
10 | |
11 | #if !USE_GOOGLE_LOG |
12 | |
13 | #include <atomic> |
14 | #include <map> |
15 | #include <set> |
16 | #include <vector> |
17 | |
18 | // TODO(T90238193) |
19 | // @lint-ignore-every CLANGTIDY facebook-hte-RelativeInclude |
20 | #include "ILoggerObserver.h" |
21 | |
22 | namespace KINETO_NAMESPACE { |
23 | |
24 | using namespace libkineto; |
25 | |
26 | class LoggerCollector : public ILoggerObserver { |
27 | public: |
28 | LoggerCollector() : buckets_() {} |
29 | |
30 | void write(const std::string& message, LoggerOutputType ot = ERROR) override { |
31 | // Skip STAGE output type which is only used by USTLoggerCollector. |
32 | // Skip VERBOSE output type which may bloat metadata section of traces. |
33 | if (ot == STAGE || ot == VERBOSE) { |
34 | return; |
35 | } |
36 | buckets_[ot].push_back(message); |
37 | } |
38 | |
39 | const std::map<LoggerOutputType, std::vector<std::string>> extractCollectorMetadata() override { |
40 | return buckets_; |
41 | } |
42 | |
43 | void reset() override { |
44 | trace_duration_ms = 0; |
45 | event_count = 0; |
46 | destinations.clear(); |
47 | } |
48 | |
49 | void addDevice(const int64_t device) override { |
50 | devices.insert(device); |
51 | } |
52 | |
53 | void setTraceDurationMS(const int64_t duration) override { |
54 | trace_duration_ms = duration; |
55 | } |
56 | |
57 | void addEventCount(const int64_t count) override { |
58 | event_count += count; |
59 | } |
60 | |
61 | void addDestination(const std::string& dest) override { |
62 | if (!dest.empty()) { |
63 | destinations.insert(dest); |
64 | } |
65 | } |
66 | |
67 | void addMetadata(const std::string& key, const std::string& value) override {}; |
68 | |
69 | protected: |
70 | std::map<LoggerOutputType, std::vector<std::string>> buckets_; |
71 | |
72 | // These are useful metadata to collect from CUPTIActivityProfiler for internal tracking. |
73 | std::set<int64_t> devices; |
74 | int64_t trace_duration_ms{0}; |
75 | std::atomic<uint64_t> event_count{0}; |
76 | std::set<std::string> destinations; |
77 | |
78 | }; |
79 | |
80 | } // namespace KINETO_NAMESPACE |
81 | |
82 | #endif // !USE_GOOGLE_LOG |
83 |