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
22namespace KINETO_NAMESPACE {
23
24using namespace libkineto;
25
26class 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