Project import generated by Copybara.

PiperOrigin-RevId: 253489161
2019-06-16 16:06:57 -07:00
commit d68f5e4169
844 changed files with 134997 additions and 0 deletions
@@ -0,0 +1,138 @@
+// Copyright 2019 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "mediapipe/framework/profiler/graph_tracer.h"
+
+#include "absl/time/time.h"
+#include "mediapipe/framework/calculator_context.h"
+#include "mediapipe/framework/calculator_profile.pb.h"
+#include "mediapipe/framework/input_stream_shard.h"
+#include "mediapipe/framework/output_stream_shard.h"
+#include "mediapipe/framework/packet.h"
+#include "mediapipe/framework/profiler/trace_builder.h"
+#include "mediapipe/framework/timestamp.h"
+
+namespace mediapipe {
+
+namespace {
+
+const absl::Duration kDefaultTraceLogInterval = absl::Milliseconds(100);
+
+// Returns a unique identifier for the current thread.
+inline int GetCurrentThreadId() {
+  static int next_thread_id = 0;
+  static thread_local int thread_id = next_thread_id++;
+  return thread_id;
+}
+
+}  // namespace
+
+absl::Duration GraphTracer::GetTraceLogInterval() {
+  return profiler_config_.trace_log_interval_usec()
+             ? absl::Microseconds(profiler_config_.trace_log_interval_usec())
+             : kDefaultTraceLogInterval;
+}
+
+int64 GraphTracer::GetTraceLogCapacity() {
+  return profiler_config_.trace_log_capacity();
+}
+
+GraphTracer::GraphTracer(const ProfilerConfig& profiler_config)
+    : profiler_config_(profiler_config), trace_buffer_(GetTraceLogCapacity()) {
+  event_types_disabled_.resize(static_cast<int>(GraphTrace::EventType_MAX + 1));
+  for (int32 event_type : profiler_config_.trace_event_types_disabled()) {
+    event_types_disabled_[event_type] = true;
+  }
+}
+
+void GraphTracer::LogEvent(TraceEvent event) {
+  if (event_types_disabled_[static_cast<int>(event.event_type)]) {
+    return;
+  }
+  event.set_thread_id(GetCurrentThreadId());
+  trace_buffer_.push_back(event);
+}
+
+void GraphTracer::LogInputEvents(GraphTrace::EventType event_type,
+                                 const CalculatorContext* context,
+                                 absl::Time event_time) {
+  Timestamp input_ts = context->InputTimestamp();
+  for (const InputStreamShard& in_stream : context->Inputs()) {
+    const Packet& packet = in_stream.Value();
+    if (!packet.IsEmpty()) {
+      const std::string* stream_id = &in_stream.Name();
+      LogEvent(TraceEvent(event_type)
+                   .set_event_time(event_time)
+                   .set_is_finish(false)
+                   .set_input_ts(input_ts)
+                   .set_node_id(context->NodeId())
+                   .set_stream_id(stream_id)
+                   .set_packet_ts(packet.Timestamp())
+                   .set_packet_data_id(&packet));
+    }
+  }
+}
+
+void GraphTracer::LogOutputEvents(GraphTrace::EventType event_type,
+                                  const CalculatorContext* context,
+                                  absl::Time event_time) {
+  // For source nodes, the first output timestamp is used as the input_ts.
+  Timestamp input_ts = (context->Inputs().NumEntries() > 0)
+                           ? context->InputTimestamp()
+                           : GetOutputTimestamp(context);
+  for (const OutputStreamShard& out_stream : context->Outputs()) {
+    const std::string* stream_id = &out_stream.Name();
+    for (const Packet& packet : *out_stream.OutputQueue()) {
+      LogEvent(TraceEvent(event_type)
+                   .set_event_time(event_time)
+                   .set_is_finish(true)
+                   .set_input_ts(input_ts)
+                   .set_node_id(context->NodeId())
+                   .set_stream_id(stream_id)
+                   .set_packet_ts(packet.Timestamp())
+                   .set_packet_data_id(&packet));
+    }
+  }
+}
+
+Timestamp GraphTracer::TimestampAfter(absl::Time begin_time) {
+  return TraceBuilder::TimestampAfter(trace_buffer_, begin_time);
+}
+
+void GraphTracer::GetTrace(absl::Time begin_time, absl::Time end_time,
+                           GraphTrace* result) {
+  trace_builder_.CreateTrace(trace_buffer_, begin_time, end_time, result);
+  trace_builder_.Clear();
+}
+
+void GraphTracer::GetLog(absl::Time begin_time, absl::Time end_time,
+                         GraphTrace* result) {
+  trace_builder_.CreateLog(trace_buffer_, begin_time, end_time, result);
+  trace_builder_.Clear();
+}
+
+const TraceBuffer& GraphTracer::GetTraceBuffer() { return trace_buffer_; }
+
+Timestamp GraphTracer::GetOutputTimestamp(const CalculatorContext* context) {
+  for (const OutputStreamShard& out_stream : context->Outputs()) {
+    for (const Packet& packet : *out_stream.OutputQueue()) {
+      if (packet.Timestamp() != Timestamp::Unset()) {
+        return packet.Timestamp();
+      }
+    }
+  }
+  return Timestamp();
+}
+
+}  // namespace mediapipe