| 86 | constexpr char kUserMetadataMarker = '#'; |
| 87 | |
| 88 | Status HostTracer::CollectData(RunMetadata* run_metadata) { |
| 89 | if (recording_) { |
| 90 | return Status(error::INTERNAL, "TraceMeRecorder not stopped"); |
| 91 | } |
| 92 | // Pair up start and end events, and add complete events to trace_entries. |
| 93 | absl::flat_hash_map<uint64, uint64> end_times; |
| 94 | for (const auto& thread : events_) { |
| 95 | for (const auto& event : thread.events) { |
| 96 | if (event.end_time && !event.start_time) { |
| 97 | end_times.emplace(event.activity_id, event.end_time); |
| 98 | } |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | StepStatsCollector step_stats_collector(run_metadata->mutable_step_stats()); |
| 103 | |
| 104 | const string cpu_name = "/host:CPU"; |
| 105 | for (auto& thread : events_) { |
| 106 | step_stats_collector.SaveThreadName(cpu_name, thread.thread.tid, |
| 107 | thread.thread.name); |
| 108 | for (auto& event : thread.events) { |
| 109 | if (!event.end_time) { |
| 110 | auto it = end_times.find(event.activity_id); |
| 111 | if (it != end_times.end()) event.end_time = it->second; |
| 112 | } |
| 113 | if (event.start_time && event.end_time) { |
| 114 | NodeExecStats* ns = new NodeExecStats; |
| 115 | if (event.name.back() != kUserMetadataMarker) { |
| 116 | ns->set_node_name(std::move(event.name)); |
| 117 | } else { |
| 118 | // Expect the format will be "<name>#<metadata>#" |
| 119 | std::vector<absl::string_view> parts = |
| 120 | absl::StrSplit(event.name, kUserMetadataMarker); |
| 121 | if (parts.size() >= 2) { |
| 122 | ns->set_node_name(string(parts[0])); |
| 123 | ns->set_timeline_label(string(parts[1])); |
| 124 | } else { |
| 125 | ns->set_node_name(std::move(event.name)); |
| 126 | } |
| 127 | } |
| 128 | ns->set_all_start_micros(event.start_time / EnvTime::kMicrosToNanos); |
| 129 | ns->set_all_end_rel_micros((event.end_time - event.start_time) / |
| 130 | EnvTime::kMicrosToNanos); |
| 131 | ns->set_thread_id(thread.thread.tid); |
| 132 | ns->set_correlation_id(event.activity_id); |
| 133 | ns->set_parent_id(event.parent_id); |
| 134 | // TODO(fishx): Add thread name to RunMetadata |
| 135 | step_stats_collector.Save(cpu_name, ns); |
| 136 | } |
| 137 | } |
| 138 | } |
| 139 | events_.clear(); |
| 140 | step_stats_collector.Finalize(); |
| 141 | return Status::OK(); |
| 142 | } |
| 143 | } // namespace |
| 144 | |
| 145 | // Not in anonymous namespace for testing purposes. |