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#include <fstream>
12#include <map>
13#include <ostream>
14#include <thread>
15#include <unordered_map>
16
17#include "GenericTraceActivity.h"
18#include "output_base.h"
19
20namespace KINETO_NAMESPACE {
21 // Previous declaration of TraceSpan is struct. Must match the same here.
22 struct TraceSpan;
23}
24
25namespace KINETO_NAMESPACE {
26
27class Config;
28
29class ChromeTraceLogger : public libkineto::ActivityLogger {
30 public:
31 explicit ChromeTraceLogger(const std::string& traceFileName);
32
33 // Note: the caller of these functions should handle concurrency
34 // i.e., we these functions are not thread-safe
35 void handleDeviceInfo(
36 const DeviceInfo& info,
37 uint64_t time) override;
38
39 void handleOverheadInfo(const OverheadInfo& info, int64_t time) override;
40
41 void handleResourceInfo(const ResourceInfo& info, int64_t time) override;
42
43 void handleTraceSpan(const TraceSpan& span) override;
44
45 void handleActivity(const ITraceActivity& activity) override;
46 void handleGenericActivity(const GenericTraceActivity& activity) override;
47
48 void handleTraceStart(
49 const std::unordered_map<std::string, std::string>& metadata) override;
50
51 void finalizeTrace(
52 const Config& config,
53 std::unique_ptr<ActivityBuffers> buffers,
54 int64_t endTime,
55 std::unordered_map<std::string, std::vector<std::string>>& metadata) override;
56
57 std::string traceFileName() const {
58 return fileName_;
59 }
60
61 private:
62
63 // Create a flow event (arrow)
64 void handleLink(
65 char type,
66 const ITraceActivity& e,
67 int64_t id,
68 const std::string& name);
69
70 void addIterationMarker(const TraceSpan& span);
71
72 void openTraceFile();
73
74 void handleGenericInstantEvent(const ITraceActivity& op);
75
76 void handleGenericLink(const ITraceActivity& activity);
77
78 void metadataToJSON(
79 const std::unordered_map<std::string, std::string>& metadata);
80
81 std::string& sanitizeStrForJSON(std::string& value);
82
83 std::string fileName_;
84 std::string tempFileName_;
85 std::ofstream traceOf_;
86};
87
88} // namespace KINETO_NAMESPACE
89