Expose stream handlers in headers to allow dynamic registration for superusers

PiperOrigin-RevId: 556988288

mediapipe/framework/stream_handler/BUILD

@@ -53,8 +53,16 @@ mediapipe_proto_library(
 cc_library(
     name = "barrier_input_stream_handler",
     srcs = ["barrier_input_stream_handler.cc"],
+    hdrs = ["barrier_input_stream_handler.h"],
     deps = [
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
+        "//mediapipe/framework:mediapipe_options_cc_proto",
+        "//mediapipe/framework/tool:tag_map",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
     ],
     alwayslink = 1,
 )
@@ -74,8 +82,15 @@ cc_library(
 cc_library(
     name = "early_close_input_stream_handler",
     srcs = ["early_close_input_stream_handler.cc"],
+    hdrs = ["early_close_input_stream_handler.h"],
     deps = [
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
+        "//mediapipe/framework:mediapipe_options_cc_proto",
+        "//mediapipe/framework/tool:tag_map",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
     ],
     alwayslink = 1,
@@ -84,10 +99,21 @@ cc_library(
 cc_library(
     name = "fixed_size_input_stream_handler",
     srcs = ["fixed_size_input_stream_handler.cc"],
+    hdrs = ["fixed_size_input_stream_handler.h"],
     deps = [
         ":default_input_stream_handler",
         ":fixed_size_input_stream_handler_cc_proto",
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
+        "//mediapipe/framework:mediapipe_options_cc_proto",
+        "//mediapipe/framework:packet",
+        "//mediapipe/framework/tool:tag_map",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/synchronization",
     ],
     alwayslink = 1,
 )
@@ -95,8 +121,18 @@ cc_library(
 cc_library(
     name = "immediate_input_stream_handler",
     srcs = ["immediate_input_stream_handler.cc"],
+    hdrs = ["immediate_input_stream_handler.h"],
     deps = [
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
+        "//mediapipe/framework:mediapipe_options_cc_proto",
+        "//mediapipe/framework/tool:tag_map",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/synchronization",
     ],
     alwayslink = 1,
 )
@@ -122,9 +158,13 @@ cc_library(
 cc_library(
     name = "mux_input_stream_handler",
     srcs = ["mux_input_stream_handler.cc"],
+    hdrs = ["mux_input_stream_handler.h"],
     deps = [
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
-        "//mediapipe/framework/port:logging",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
     ],
@@ -134,16 +174,22 @@ cc_library(
 cc_library(
     name = "sync_set_input_stream_handler",
     srcs = ["sync_set_input_stream_handler.cc"],
+    hdrs = ["sync_set_input_stream_handler.h"],
     deps = [
         ":sync_set_input_stream_handler_cc_proto",
-        "//mediapipe/framework:collection",
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
         "//mediapipe/framework:mediapipe_options_cc_proto",
         "//mediapipe/framework:packet_set",
         "//mediapipe/framework:timestamp",
+        "//mediapipe/framework/port:map_util",
+        "//mediapipe/framework/port:status",
         "//mediapipe/framework/tool:tag_map",
-        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
     ],
     alwayslink = 1,
@@ -152,12 +198,19 @@ cc_library(
 cc_library(
     name = "timestamp_align_input_stream_handler",
     srcs = ["timestamp_align_input_stream_handler.cc"],
+    hdrs = ["timestamp_align_input_stream_handler.h"],
     deps = [
         ":timestamp_align_input_stream_handler_cc_proto",
+        "//mediapipe/framework:calculator_context_manager",
+        "//mediapipe/framework:calculator_framework",
         "//mediapipe/framework:collection_item_id",
         "//mediapipe/framework:input_stream_handler",
+        "//mediapipe/framework:mediapipe_options_cc_proto",
         "//mediapipe/framework:timestamp",
         "//mediapipe/framework/tool:validate_name",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
     ],

mediapipe/framework/stream_handler/barrier_input_stream_handler.cc

@@ -11,46 +11,34 @@
 // 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/stream_handler/barrier_input_stream_handler.h"
 
-#include <algorithm>
+#include <functional>
-#include <memory>
+#include <utility>
-#include <vector>
 
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
 
 namespace mediapipe {
 
-// Implementation of an input stream handler that considers a node as ready for
+void BarrierInputStreamHandler::PrepareForRun(
-// Process() if all input streams have a packet available. This implies it must
-// consider a node as ready for Close() if any input stream is done.
-class BarrierInputStreamHandler : public InputStreamHandler {
- public:
-  BarrierInputStreamHandler() = delete;
-  BarrierInputStreamHandler(
-      std::shared_ptr<tool::TagMap> tag_map,
-      CalculatorContextManager* calculator_context_manager,
-      const MediaPipeOptions& options, bool calculator_run_in_parallel)
-      : InputStreamHandler(std::move(tag_map), calculator_context_manager,
-                           options, calculator_run_in_parallel) {}
-
-  void PrepareForRun(
     std::function<void()> headers_ready_callback,
     std::function<void()> notification_callback,
     std::function<void(CalculatorContext*)> schedule_callback,
-      std::function<void(absl::Status)> error_callback) override {
+    std::function<void(absl::Status)> error_callback) {
   InputStreamHandler::PrepareForRun(
       std::move(headers_ready_callback), std::move(notification_callback),
       std::move(schedule_callback), std::move(error_callback));
   for (auto& stream : input_stream_managers_) {
     stream->DisableTimestamps();
   }
-  }
+}
 
- protected:
+NodeReadiness BarrierInputStreamHandler::GetNodeReadiness(
-  // In BarrierInputStreamHandler, a node is "ready" if:
+    Timestamp* min_stream_timestamp) {
-  // - any stream is done (need to call Close() in this case), or
-  // - all streams have a packet available.
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
   DCHECK(min_stream_timestamp);
   *min_stream_timestamp = Timestamp::Done();
   bool all_available = true;
@@ -72,11 +60,10 @@ class BarrierInputStreamHandler : public InputStreamHandler {
     return NodeReadiness::kReadyForProcess;
   }
   return NodeReadiness::kNotReady;
-  }
+}
 
-  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+void BarrierInputStreamHandler::FillInputSet(Timestamp input_timestamp,
-  void FillInputSet(Timestamp input_timestamp,
+                                             InputStreamShardSet* input_set) {
-                    InputStreamShardSet* input_set) override {
   CHECK(input_timestamp.IsAllowedInStream());
   CHECK(input_set);
   for (CollectionItemId id = input_stream_managers_.BeginId();
@@ -87,8 +74,7 @@ class BarrierInputStreamHandler : public InputStreamHandler {
     AddPacketToShard(&input_set->Get(id), std::move(current_packet),
                      stream_is_done);
   }
-  }
+}
-};
 
 REGISTER_INPUT_STREAM_HANDLER(BarrierInputStreamHandler);
 

mediapipe/framework/stream_handler/barrier_input_stream_handler.h

@@ -0,0 +1,64 @@
+
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_BARRIER_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_BARRIER_INPUT_STREAM_HANDLER_H_
+
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
+#include "mediapipe/framework/tool/tag_map.h"
+
+namespace mediapipe {
+
+// Implementation of an input stream handler that considers a node as ready for
+// Process() if all input streams have a packet available. This implies it must
+// consider a node as ready for Close() if any input stream is done.
+class BarrierInputStreamHandler : public InputStreamHandler {
+ public:
+  BarrierInputStreamHandler() = delete;
+  BarrierInputStreamHandler(
+      std::shared_ptr<tool::TagMap> tag_map,
+      CalculatorContextManager* calculator_context_manager,
+      const mediapipe::MediaPipeOptions& options,
+      bool calculator_run_in_parallel)
+      : InputStreamHandler(std::move(tag_map), calculator_context_manager,
+                           options, calculator_run_in_parallel) {}
+
+  void PrepareForRun(std::function<void()> headers_ready_callback,
+                     std::function<void()> notification_callback,
+                     std::function<void(CalculatorContext*)> schedule_callback,
+                     std::function<void(absl::Status)> error_callback) override;
+
+ protected:
+  // In BarrierInputStreamHandler, a node is "ready" if:
+  // - any stream is done (need to call Close() in this case), or
+  // - all streams have a packet available.
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_BARRIER_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/early_close_input_stream_handler.cc

@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// Copyright 2023 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.
@@ -11,35 +11,25 @@
 // 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/stream_handler/early_close_input_stream_handler.h"
 
 #include <algorithm>
-#include <memory>
-#include <vector>
 
+#include "absl/log/check.h"
 #include "absl/strings/substitute.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
 
 namespace mediapipe {
 
-// Implementation of an input stream handler that considers a node as ready for
+// In EarlyCloseInputStreamHandler, a node is "ready" if:
-// Close() if any input stream is done.
+// - any stream is done (need to call Close() in this case), or
-class EarlyCloseInputStreamHandler : public InputStreamHandler {
+// - the minimum bound (over all empty streams) is greater than the smallest
- public:
+//   timestamp of any stream, which means we have received all the packets
-  EarlyCloseInputStreamHandler() = delete;
+//   that will be available at the next timestamp.
-  EarlyCloseInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
+NodeReadiness EarlyCloseInputStreamHandler::GetNodeReadiness(
-                               CalculatorContextManager* cc_manager,
+    Timestamp* min_stream_timestamp) {
-                               const MediaPipeOptions& options,
-                               bool calculator_run_in_parallel)
-      : InputStreamHandler(std::move(tag_map), cc_manager, options,
-                           calculator_run_in_parallel) {}
-
- protected:
-  // In EarlyCloseInputStreamHandler, a node is "ready" if:
-  // - any stream is done (need to call Close() in this case), or
-  // - the minimum bound (over all empty streams) is greater than the smallest
-  //   timestamp of any stream, which means we have received all the packets
-  //   that will be available at the next timestamp.
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
   DCHECK(min_stream_timestamp);
   *min_stream_timestamp = Timestamp::Done();
   Timestamp min_bound = Timestamp::Done();
@@ -64,11 +54,11 @@ class EarlyCloseInputStreamHandler : public InputStreamHandler {
 
   CHECK_EQ(min_bound, *min_stream_timestamp);
   return NodeReadiness::kNotReady;
-  }
+}
 
-  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
-  void FillInputSet(Timestamp input_timestamp,
+void EarlyCloseInputStreamHandler::FillInputSet(
-                    InputStreamShardSet* input_set) override {
+    Timestamp input_timestamp, InputStreamShardSet* input_set) {
   CHECK(input_timestamp.IsAllowedInStream());
   CHECK(input_set);
   for (CollectionItemId id = input_stream_managers_.BeginId();
@@ -84,8 +74,7 @@ class EarlyCloseInputStreamHandler : public InputStreamHandler {
     AddPacketToShard(&input_set->Get(id), std::move(current_packet),
                      stream_is_done);
   }
-  }
+}
-};
 
 REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler);
 

mediapipe/framework/stream_handler/early_close_input_stream_handler.h

@@ -0,0 +1,56 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_
+
+#include <memory>
+#include <utility>
+
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
+#include "mediapipe/framework/tool/tag_map.h"
+
+namespace mediapipe {
+
+// Implementation of an input stream handler that considers a node as ready for
+// Close() if any input stream is done.
+class EarlyCloseInputStreamHandler : public InputStreamHandler {
+ public:
+  EarlyCloseInputStreamHandler() = delete;
+  EarlyCloseInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
+                               CalculatorContextManager* cc_manager,
+                               const mediapipe::MediaPipeOptions& options,
+                               bool calculator_run_in_parallel)
+      : InputStreamHandler(std::move(tag_map), cc_manager, options,
+                           calculator_run_in_parallel) {}
+
+ protected:
+  // In EarlyCloseInputStreamHandler, a node is "ready" if:
+  // - any stream is done (need to call Close() in this case), or
+  // - the minimum bound (over all empty streams) is greater than the smallest
+  //   timestamp of any stream, which means we have received all the packets
+  //   that will be available at the next timestamp.
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc

@@ -1,4 +1,4 @@
-// Copyright 2019 The MediaPipe Authors.
+// Copyright 2023 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.
@@ -11,50 +11,36 @@
 // 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/stream_handler/fixed_size_input_stream_handler.h"
 
+#include <algorithm>
+#include <list>
 #include <memory>
+#include <utility>
 #include <vector>
 
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
+#include "mediapipe/framework/packet.h"
 #include "mediapipe/framework/stream_handler/default_input_stream_handler.h"
-// TODO: Move protos in another CL after the C++ code migration.
 #include "mediapipe/framework/stream_handler/fixed_size_input_stream_handler.pb.h"
+#include "mediapipe/framework/tool/tag_map.h"
 
 namespace mediapipe {
 
-// Input stream handler that limits each input queue to a maximum of
+FixedSizeInputStreamHandler::FixedSizeInputStreamHandler(
-// target_queue_size packets, discarding older packets as needed.  When a
+    std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
-// timestamp is dropped from a stream, it is dropped from all others as well.
+    const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
-//
-// For example, a calculator node with one input stream and the following input
-// stream handler specs:
-//
-// node {
-//   calculator: "CalculatorRunningAtOneFps"
-//   input_stream: "packets_streaming_in_at_ten_fps"
-//   input_stream_handler {
-//     input_stream_handler: "FixedSizeInputStreamHandler"
-//   }
-// }
-//
-// will always try to keep the newest packet in the input stream.
-//
-// A few details: FixedSizeInputStreamHandler takes action when any stream grows
-// to trigger_queue_size or larger.  It then keeps at most target_queue_size
-// packets in every InputStreamImpl.  Every stream is truncated at the same
-// timestamp, so that each included timestamp delivers the same packets as
-// DefaultInputStreamHandler includes.
-//
-class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
- public:
-  FixedSizeInputStreamHandler() = delete;
-  FixedSizeInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
-                              CalculatorContextManager* cc_manager,
-                              const MediaPipeOptions& options,
-                              bool calculator_run_in_parallel)
     : DefaultInputStreamHandler(std::move(tag_map), cc_manager, options,
                                 calculator_run_in_parallel) {
   const auto& ext =
-        options.GetExtension(FixedSizeInputStreamHandlerOptions::ext);
+      options.GetExtension(mediapipe::FixedSizeInputStreamHandlerOptions::ext);
   trigger_queue_size_ = ext.trigger_queue_size();
   target_queue_size_ = ext.target_queue_size();
   fixed_min_size_ = ext.fixed_min_size();
@@ -63,11 +49,9 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
   // TODO: Either re-enable SetLatePreparation(true) with
   // CalculatorContext::InputTimestamp set correctly, or remove the
   // implementation of SetLatePreparation.
-  }
+}
 
- private:
+void FixedSizeInputStreamHandler::EraseAllSurplus() {
-  // Drops packets if all input streams exceed trigger_queue_size.
-  void EraseAllSurplus() ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
   Timestamp min_timestamp_all_streams = Timestamp::Max();
   for (const auto& stream : input_stream_managers_) {
     // Check whether every InputStreamImpl grew beyond trigger_queue_size.
@@ -85,15 +69,14 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
   for (auto& stream : input_stream_managers_) {
     stream->ErasePacketsEarlierThan(min_timestamp_all_streams);
   }
-  }
+}
 
-  // Returns the latest timestamp allowed before a bound.
+Timestamp FixedSizeInputStreamHandler::PreviousAllowedInStream(
-  Timestamp PreviousAllowedInStream(Timestamp bound) {
+    Timestamp bound) {
   return bound.IsRangeValue() ? bound - 1 : bound;
-  }
+}
 
-  // Returns the lowest timestamp at which a packet may arrive at any stream.
+Timestamp FixedSizeInputStreamHandler::MinStreamBound() {
-  Timestamp MinStreamBound() {
   Timestamp min_bound = Timestamp::Done();
   for (const auto& stream : input_stream_managers_) {
     Timestamp stream_bound = stream->GetMinTimestampAmongNLatest(1);
@@ -105,10 +88,9 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
     min_bound = std::min(min_bound, stream_bound);
   }
   return min_bound;
-  }
+}
 
-  // Returns the lowest timestamp of a packet ready to process.
+Timestamp FixedSizeInputStreamHandler::MinTimestampToProcess() {
-  Timestamp MinTimestampToProcess() {
   Timestamp min_bound = Timestamp::Done();
   for (const auto& stream : input_stream_managers_) {
     bool empty;
@@ -122,13 +104,9 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
     min_bound = std::min(min_bound, stream_timestamp);
   }
   return min_bound;
-  }
+}
 
-  // Keeps only the most recent target_queue_size packets in each stream
+void FixedSizeInputStreamHandler::EraseAnySurplus(bool keep_one) {
-  // exceeding trigger_queue_size.  Also, discards all packets older than the
-  // first kept timestamp on any stream.
-  void EraseAnySurplus(bool keep_one)
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
   // Record the most recent first kept timestamp on any stream.
   for (const auto& stream : input_stream_managers_) {
     int32_t queue_size = (stream->QueueSize() >= trigger_queue_size_)
@@ -149,14 +127,14 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
   for (auto& stream : input_stream_managers_) {
     stream->ErasePacketsEarlierThan(kept_timestamp_);
   }
-  }
+}
 
-  void EraseSurplusPackets(bool keep_one)
+void FixedSizeInputStreamHandler::EraseSurplusPackets(bool keep_one) {
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
   return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
-  }
+}
 
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
+NodeReadiness FixedSizeInputStreamHandler::GetNodeReadiness(
+    Timestamp* min_stream_timestamp) {
   DCHECK(min_stream_timestamp);
   absl::MutexLock lock(&erase_mutex_);
   // kReadyForProcess is returned only once until FillInputSet completes.
@@ -175,32 +153,32 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
   while (*min_stream_timestamp < kept_timestamp_ &&
          result == NodeReadiness::kReadyForProcess) {
     EraseSurplusPackets(false);
-      result =
+    result = DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
-          DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
   }
   pending_ = (result == NodeReadiness::kReadyForProcess);
   return result;
-  }
+}
 
-  void AddPackets(CollectionItemId id,
+void FixedSizeInputStreamHandler::AddPackets(CollectionItemId id,
-                  const std::list<Packet>& packets) override {
+                                             const std::list<Packet>& packets) {
   InputStreamHandler::AddPackets(id, packets);
   absl::MutexLock lock(&erase_mutex_);
   if (!pending_) {
     EraseSurplusPackets(false);
   }
-  }
+}
 
-  void MovePackets(CollectionItemId id, std::list<Packet>* packets) override {
+void FixedSizeInputStreamHandler::MovePackets(CollectionItemId id,
+                                              std::list<Packet>* packets) {
   InputStreamHandler::MovePackets(id, packets);
   absl::MutexLock lock(&erase_mutex_);
   if (!pending_) {
     EraseSurplusPackets(false);
   }
-  }
+}
 
-  void FillInputSet(Timestamp input_timestamp,
+void FixedSizeInputStreamHandler::FillInputSet(Timestamp input_timestamp,
-                    InputStreamShardSet* input_set) override {
+                                               InputStreamShardSet* input_set) {
   CHECK(input_set);
   absl::MutexLock lock(&erase_mutex_);
   if (!pending_) {
@@ -211,19 +189,7 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
   input_timestamp = MinTimestampToProcess();
   DefaultInputStreamHandler::FillInputSet(input_timestamp, input_set);
   pending_ = false;
-  }
+}
-
- private:
-  int32_t trigger_queue_size_;
-  int32_t target_queue_size_;
-  bool fixed_min_size_;
-  // Indicates that GetNodeReadiness has returned kReadyForProcess once, and
-  // the corresponding call to FillInputSet has not yet completed.
-  bool pending_ ABSL_GUARDED_BY(erase_mutex_);
-  // The timestamp used to truncate all input streams.
-  Timestamp kept_timestamp_ ABSL_GUARDED_BY(erase_mutex_);
-  absl::Mutex erase_mutex_;
-};
 
 REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler);
 

mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h

@@ -0,0 +1,108 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_FIXED_SIZE_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_FIXED_SIZE_INPUT_STREAM_HANDLER_H_
+
+#include <cstdint>
+#include <list>
+#include <memory>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/stream_handler/default_input_stream_handler.h"
+
+namespace mediapipe {
+
+// Input stream handler that limits each input queue to a maximum of
+// target_queue_size packets, discarding older packets as needed.  When a
+// timestamp is dropped from a stream, it is dropped from all others as well.
+//
+// For example, a calculator node with one input stream and the following input
+// stream handler specs:
+//
+// node {
+//   calculator: "CalculatorRunningAtOneFps"
+//   input_stream: "packets_streaming_in_at_ten_fps"
+//   input_stream_handler {
+//     input_stream_handler: "FixedSizeInputStreamHandler"
+//   }
+// }
+//
+// will always try to keep the newest packet in the input stream.
+//
+// A few details: FixedSizeInputStreamHandler takes action when any stream grows
+// to trigger_queue_size or larger.  It then keeps at most target_queue_size
+// packets in every InputStreamImpl.  Every stream is truncated at the same
+// timestamp, so that each included timestamp delivers the same packets as
+// DefaultInputStreamHandler includes.
+class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
+ public:
+  FixedSizeInputStreamHandler() = delete;
+  FixedSizeInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
+                              CalculatorContextManager* cc_manager,
+                              const MediaPipeOptions& options,
+                              bool calculator_run_in_parallel);
+
+ private:
+  // Drops packets if all input streams exceed trigger_queue_size.
+  void EraseAllSurplus() ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
+
+  // Returns the latest timestamp allowed before a bound.
+  Timestamp PreviousAllowedInStream(Timestamp bound);
+
+  // Returns the lowest timestamp at which a packet may arrive at any stream.
+  Timestamp MinStreamBound();
+
+  // Returns the lowest timestamp of a packet ready to process.
+  Timestamp MinTimestampToProcess();
+
+  // Keeps only the most recent target_queue_size packets in each stream
+  // exceeding trigger_queue_size.  Also, discards all packets older than the
+  // first kept timestamp on any stream.
+  void EraseAnySurplus(bool keep_one)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
+
+  void EraseSurplusPackets(bool keep_one)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
+
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  void AddPackets(CollectionItemId id,
+                  const std::list<Packet>& packets) override;
+
+  void MovePackets(CollectionItemId id, std::list<Packet>* packets) override;
+
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+
+ private:
+  int32_t trigger_queue_size_;
+  int32_t target_queue_size_;
+  bool fixed_min_size_;
+  // Indicates that GetNodeReadiness has returned kReadyForProcess once, and
+  // the corresponding call to FillInputSet has not yet completed.
+  bool pending_ ABSL_GUARDED_BY(erase_mutex_);
+  // The timestamp used to truncate all input streams.
+  Timestamp kept_timestamp_ ABSL_GUARDED_BY(erase_mutex_);
+  absl::Mutex erase_mutex_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_FIXED_SIZE_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/immediate_input_stream_handler.cc

@@ -11,65 +11,33 @@
 // 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/stream_handler/immediate_input_stream_handler.h"
 
+#include <algorithm>
+#include <functional>
 #include <memory>
 #include <vector>
 
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
+#include "mediapipe/framework/tool/tag_map.h"
 
 namespace mediapipe {
 
 using SyncSet = InputStreamHandler::SyncSet;
 
-// An input stream handler that delivers input packets to the Calculator
-// immediately, with no dependency between input streams.  It also invokes
-// Calculator::Process when any input stream becomes done.
-//
-// NOTE: If packets arrive successively on different input streams with
-// identical or decreasing timestamps, this input stream handler will
-// invoke its Calculator with a sequence of InputTimestamps that is
-// non-increasing.  Its Calculator is responsible for accumulating packets
-// with the required timetamps before processing and delivering output.
-//
-class ImmediateInputStreamHandler : public InputStreamHandler {
- public:
-  ImmediateInputStreamHandler() = delete;
-  ImmediateInputStreamHandler(
-      std::shared_ptr<tool::TagMap> tag_map,
-      CalculatorContextManager* calculator_context_manager,
-      const MediaPipeOptions& options, bool calculator_run_in_parallel);
-
- protected:
-  // Reinitializes this InputStreamHandler before each CalculatorGraph run.
-  void PrepareForRun(std::function<void()> headers_ready_callback,
-                     std::function<void()> notification_callback,
-                     std::function<void(CalculatorContext*)> schedule_callback,
-                     std::function<void(absl::Status)> error_callback) override;
-
-  // Returns kReadyForProcess whenever a Packet is available at any of
-  // the input streams, or any input stream becomes done.
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
-
-  // Selects a packet on each stream with an available packet with the
-  // specified timestamp, leaving other input streams unaffected.
-  void FillInputSet(Timestamp input_timestamp,
-                    InputStreamShardSet* input_set) override;
-
-  // Returns the number of sync-sets maintained by this input-handler.
-  int SyncSetCount() override;
-
-  absl::Mutex mutex_;
-  // The packet-set builder for each input stream.
-  std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
-  // The input timestamp for each kReadyForProcess input stream.
-  std::vector<Timestamp> ready_timestamps_ ABSL_GUARDED_BY(mutex_);
-};
 REGISTER_INPUT_STREAM_HANDLER(ImmediateInputStreamHandler);
 
 ImmediateInputStreamHandler::ImmediateInputStreamHandler(
     std::shared_ptr<tool::TagMap> tag_map,
     CalculatorContextManager* calculator_context_manager,
-    const MediaPipeOptions& options, bool calculator_run_in_parallel)
+    const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
     : InputStreamHandler(tag_map, calculator_context_manager, options,
                          calculator_run_in_parallel) {
   for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) {

mediapipe/framework/stream_handler/immediate_input_stream_handler.h

@@ -0,0 +1,77 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/tool/tag_map.h"
+
+namespace mediapipe {
+
+// An input stream handler that delivers input packets to the Calculator
+// immediately, with no dependency between input streams.  It also invokes
+// Calculator::Process when any input stream becomes done.
+//
+// NOTE: If packets arrive successively on different input streams with
+// identical or decreasing timestamps, this input stream handler will
+// invoke its Calculator with a sequence of InputTimestamps that is
+// non-increasing.  Its Calculator is responsible for accumulating packets
+// with the required timestamps before processing and delivering output.
+class ImmediateInputStreamHandler : public InputStreamHandler {
+ public:
+  ImmediateInputStreamHandler() = delete;
+  ImmediateInputStreamHandler(
+      std::shared_ptr<tool::TagMap> tag_map,
+      CalculatorContextManager* calculator_context_manager,
+      const MediaPipeOptions& options, bool calculator_run_in_parallel);
+
+ protected:
+  // Reinitializes this InputStreamHandler before each CalculatorGraph run.
+  void PrepareForRun(std::function<void()> headers_ready_callback,
+                     std::function<void()> notification_callback,
+                     std::function<void(CalculatorContext*)> schedule_callback,
+                     std::function<void(absl::Status)> error_callback) override;
+
+  // Returns kReadyForProcess whenever a Packet is available at any of
+  // the input streams, or any input stream becomes done.
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Selects a packet on each stream with an available packet with the
+  // specified timestamp, leaving other input streams unaffected.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+
+  // Returns the number of sync-sets maintained by this input-handler.
+  int SyncSetCount() override;
+
+  absl::Mutex mutex_;
+  // The packet-set builder for each input stream.
+  std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
+  // The input timestamp for each kReadyForProcess input stream.
+  std::vector<Timestamp> ready_timestamps_ ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/mux_input_stream_handler.cc

@@ -11,72 +11,46 @@
 // 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/stream_handler/mux_input_stream_handler.h"
 
+#include <utility>
+
+#include "absl/log/check.h"
 #include "absl/strings/substitute.h"
 #include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
-#include "mediapipe/framework/port/logging.h"
 
 namespace mediapipe {
 
-// Implementation of the input stream handler for the MuxCalculator.
+CollectionItemId MuxInputStreamHandler::GetControlStreamId() const {
-//
-// One of the input streams is the control stream; all the other input streams
-// are data streams. To make MuxInputStreamHandler work properly, the tag of the
-// input streams must obey the following rules:
-// Let N be the number of input streams. Data streams must use tag "INPUT" with
-// index 0, ..., N - 2; the control stream must use tag "SELECT".
-//
-// The control stream carries packets of type 'int'. The 'int' value in a
-// control stream packet must be a valid index in the range 0, ..., N - 2 and
-// select the data stream at that index. The selected data stream must have a
-// packet with the same timestamp as the control stream packet.
-//
-// When the control stream is done, GetNodeReadiness() returns
-// NodeReadiness::kReadyForClose.
-//
-// TODO: pass the input stream tags to the MuxInputStreamHandler
-// constructor so that it can refer to input streams by tag. See b/30125118.
-class MuxInputStreamHandler : public InputStreamHandler {
- public:
-  MuxInputStreamHandler() = delete;
-  MuxInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
-                        CalculatorContextManager* cc_manager,
-                        const MediaPipeOptions& options,
-                        bool calculator_run_in_parallel)
-      : InputStreamHandler(std::move(tag_map), cc_manager, options,
-                           calculator_run_in_parallel) {}
-
- private:
-  CollectionItemId GetControlStreamId() const {
   return input_stream_managers_.EndId() - 1;
-  }
+}
-  void RemoveOutdatedDataPackets(Timestamp timestamp) {
+void MuxInputStreamHandler::RemoveOutdatedDataPackets(Timestamp timestamp) {
   const CollectionItemId control_stream_id = GetControlStreamId();
   for (CollectionItemId id = input_stream_managers_.BeginId();
        id < control_stream_id; ++id) {
     input_stream_managers_.Get(id)->ErasePacketsEarlierThan(timestamp);
   }
-  }
+}
 
- protected:
+// In MuxInputStreamHandler, a node is "ready" if:
-  // In MuxInputStreamHandler, a node is "ready" if:
+// - the control stream is done (need to call Close() in this case), or
-  // - the control stream is done (need to call Close() in this case), or
+// - we have received the packets on the control stream and the selected data
-  // - we have received the packets on the control stream and the selected data
+//   stream at the next timestamp.
-  //   stream at the next timestamp.
+NodeReadiness MuxInputStreamHandler::GetNodeReadiness(
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
+    Timestamp* min_stream_timestamp) {
   DCHECK(min_stream_timestamp);
   absl::MutexLock lock(&input_streams_mutex_);
 
-    const auto& control_stream =
+  const auto& control_stream = input_stream_managers_.Get(GetControlStreamId());
-        input_stream_managers_.Get(GetControlStreamId());
   bool empty;
   *min_stream_timestamp = control_stream->MinTimestampOrBound(&empty);
 
   // Data streams may contain some outdated packets which failed to be popped
   // out during "FillInputSet". (This handler doesn't sync input streams,
-    // hence "FillInputSet" can be triggerred before every input stream is
+  // hence "FillInputSet" can be triggered before every input stream is
   // filled with packets corresponding to the same timestamp.)
   RemoveOutdatedDataPackets(*min_stream_timestamp);
   if (empty) {
@@ -115,11 +89,11 @@ class MuxInputStreamHandler : public InputStreamHandler {
   }
   CHECK_EQ(stream_timestamp, *min_stream_timestamp);
   return NodeReadiness::kReadyForProcess;
-  }
+}
 
-  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
-  void FillInputSet(Timestamp input_timestamp,
+void MuxInputStreamHandler::FillInputSet(Timestamp input_timestamp,
-                    InputStreamShardSet* input_set) override {
+                                         InputStreamShardSet* input_set) {
   CHECK(input_timestamp.IsAllowedInStream());
   CHECK(input_set);
   absl::MutexLock lock(&input_streams_mutex_);
@@ -154,13 +128,7 @@ class MuxInputStreamHandler : public InputStreamHandler {
 
   // Discard old packets on data streams.
   RemoveOutdatedDataPackets(input_timestamp.NextAllowedInStream());
-  }
+}
-
- private:
-  // Must be acquired when manipulating the control and data streams to ensure
-  // we have a consistent view of the two streams.
-  absl::Mutex input_streams_mutex_;
-};
 
 REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler);
 

mediapipe/framework/stream_handler/mux_input_stream_handler.h

@@ -0,0 +1,80 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_MUX_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_MUX_INPUT_STREAM_HANDLER_H_
+
+#include <memory>
+#include <utility>
+
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
+#include "mediapipe/framework/input_stream_handler.h"
+
+namespace mediapipe {
+
+// Implementation of the input stream handler for the MuxCalculator.
+//
+// One of the input streams is the control stream; all the other input streams
+// are data streams. To make MuxInputStreamHandler work properly, the tag of the
+// input streams must obey the following rules:
+// Let N be the number of input streams. Data streams must use tag "INPUT" with
+// index 0, ..., N - 2; the control stream must use tag "SELECT".
+//
+// The control stream carries packets of type 'int'. The 'int' value in a
+// control stream packet must be a valid index in the range 0, ..., N - 2 and
+// select the data stream at that index. The selected data stream must have a
+// packet with the same timestamp as the control stream packet.
+//
+// When the control stream is done, GetNodeReadiness() returns
+// NodeReadiness::kReadyForClose.
+//
+// TODO: pass the input stream tags to the MuxInputStreamHandler
+// constructor so that it can refer to input streams by tag. See b/30125118.
+class MuxInputStreamHandler : public InputStreamHandler {
+ public:
+  MuxInputStreamHandler() = delete;
+  MuxInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
+                        CalculatorContextManager* cc_manager,
+                        const MediaPipeOptions& options,
+                        bool calculator_run_in_parallel)
+      : InputStreamHandler(std::move(tag_map), cc_manager, options,
+                           calculator_run_in_parallel) {}
+
+ private:
+  CollectionItemId GetControlStreamId() const;
+  void RemoveOutdatedDataPackets(Timestamp timestamp);
+
+ protected:
+  // In MuxInputStreamHandler, a node is "ready" if:
+  // - the control stream is done (need to call Close() in this case), or
+  // - we have received the packets on the control stream and the selected data
+  //   stream at the next timestamp.
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+
+ private:
+  // Must be acquired when manipulating the control and data streams to ensure
+  // we have a consistent view of the two streams.
+  absl::Mutex input_streams_mutex_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_MUX_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc

@@ -11,91 +11,36 @@
 // 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/stream_handler/sync_set_input_stream_handler.h"
 
-#include <algorithm>
+#include <functional>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
 
-// TODO: Move protos in another CL after the C++ code migration.
+#include "absl/log/check.h"
-#include "absl/strings/substitute.h"
 #include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
-#include "mediapipe/framework/mediapipe_options.pb.h"
 #include "mediapipe/framework/packet_set.h"
+#include "mediapipe/framework/port/map_util.h"
+#include "mediapipe/framework/port/status.h"
 #include "mediapipe/framework/stream_handler/sync_set_input_stream_handler.pb.h"
 #include "mediapipe/framework/timestamp.h"
-#include "mediapipe/framework/tool/tag_map.h"
 
 namespace mediapipe {
 
-// An input stream handler which separates the inputs into sets which
-// are each independently synchronized.  For example, if 5 inputs are
-// present, then the first three can be grouped (and will be synchronized
-// as if they were in a calculator with only those three streams) and the
-// remaining 2 streams can be independently grouped.  The calculator will
-// always be called with all the available packets from a single sync set
-// (never more than one).  The input timestamps seen by the calculator
-// will be ordered sequentially for each sync set but may jump around
-// between sync sets.
-class SyncSetInputStreamHandler : public InputStreamHandler {
- public:
-  SyncSetInputStreamHandler() = delete;
-  SyncSetInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
-                            CalculatorContextManager* cc_manager,
-                            const MediaPipeOptions& extendable_options,
-                            bool calculator_run_in_parallel);
-
-  void PrepareForRun(std::function<void()> headers_ready_callback,
-                     std::function<void()> notification_callback,
-                     std::function<void(CalculatorContext*)> schedule_callback,
-                     std::function<void(absl::Status)> error_callback) override;
-
- protected:
-  // In SyncSetInputStreamHandler, a node is "ready" if any
-  // of its sync sets are ready in the traditional sense (See
-  // DefaultInputStreamHandler).
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
-
-  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
-  // Populates packets for the ready sync-set, and populates timestamp bounds
-  // for all sync-sets.
-  void FillInputSet(Timestamp input_timestamp,
-                    InputStreamShardSet* input_set) override;
-
-  // Populates timestamp bounds for streams outside the ready sync-set.
-  void FillInputBounds(Timestamp input_timestamp,
-                       InputStreamShardSet* input_set)
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
-
-  // Returns the number of sync-sets maintained by this input-handler.
-  int SyncSetCount() override;
-
- private:
-  absl::Mutex mutex_;
-  // The ids of each set of inputs.
-  std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
-  // The index of the ready sync set.  A value of -1 indicates that no
-  // sync sets are ready.
-  int ready_sync_set_index_ ABSL_GUARDED_BY(mutex_) = -1;
-  // The timestamp at which the sync set is ready.  If no sync set is
-  // ready then this variable should be Timestamp::Done() .
-  Timestamp ready_timestamp_ ABSL_GUARDED_BY(mutex_);
-};
-
 REGISTER_INPUT_STREAM_HANDLER(SyncSetInputStreamHandler);
 
-SyncSetInputStreamHandler::SyncSetInputStreamHandler(
-    std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
-    const MediaPipeOptions& extendable_options, bool calculator_run_in_parallel)
-    : InputStreamHandler(std::move(tag_map), cc_manager, extendable_options,
-                         calculator_run_in_parallel) {}
-
 void SyncSetInputStreamHandler::PrepareForRun(
     std::function<void()> headers_ready_callback,
     std::function<void()> notification_callback,
     std::function<void(CalculatorContext*)> schedule_callback,
     std::function<void(absl::Status)> error_callback) {
   const auto& handler_options =
-      options_.GetExtension(SyncSetInputStreamHandlerOptions::ext);
+      options_.GetExtension(mediapipe::SyncSetInputStreamHandlerOptions::ext);
   {
     absl::MutexLock lock(&mutex_);
     sync_sets_.clear();

mediapipe/framework/stream_handler/sync_set_input_stream_handler.h

@@ -0,0 +1,97 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_SYNC_SET_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_SYNC_SET_INPUT_STREAM_HANDLER_H_
+
+#include <functional>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
+#include "mediapipe/framework/packet_set.h"
+#include "mediapipe/framework/stream_handler/sync_set_input_stream_handler.pb.h"
+#include "mediapipe/framework/timestamp.h"
+#include "mediapipe/framework/tool/tag_map.h"
+
+namespace mediapipe {
+
+// An input stream handler which separates the inputs into sets which
+// are each independently synchronized.  For example, if 5 inputs are
+// present, then the first three can be grouped (and will be synchronized
+// as if they were in a calculator with only those three streams) and the
+// remaining 2 streams can be independently grouped.  The calculator will
+// always be called with all the available packets from a single sync set
+// (never more than one).  The input timestamps seen by the calculator
+// will be ordered sequentially for each sync set but may jump around
+// between sync sets.
+class SyncSetInputStreamHandler : public InputStreamHandler {
+ public:
+  SyncSetInputStreamHandler() = delete;
+  SyncSetInputStreamHandler(
+      std::shared_ptr<tool::TagMap> tag_map,
+      CalculatorContextManager* cc_manager,
+      const mediapipe::MediaPipeOptions& extendable_options,
+      bool calculator_run_in_parallel)
+      : InputStreamHandler(std::move(tag_map), cc_manager, extendable_options,
+                           calculator_run_in_parallel) {}
+
+  void PrepareForRun(std::function<void()> headers_ready_callback,
+                     std::function<void()> notification_callback,
+                     std::function<void(CalculatorContext*)> schedule_callback,
+                     std::function<void(absl::Status)> error_callback) override;
+
+ protected:
+  // In SyncSetInputStreamHandler, a node is "ready" if any
+  // of its sync sets are ready in the traditional sense (See
+  // DefaultInputStreamHandler).
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+  // Populates packets for the ready sync-set, and populates timestamp bounds
+  // for all sync-sets.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+
+  // Populates timestamp bounds for streams outside the ready sync-set.
+  void FillInputBounds(Timestamp input_timestamp,
+                       InputStreamShardSet* input_set)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Returns the number of sync-sets maintained by this input-handler.
+  int SyncSetCount() override;
+
+ private:
+  absl::Mutex mutex_;
+  // The ids of each set of inputs.
+  std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
+  // The index of the ready sync set.  A value of -1 indicates that no
+  // sync sets are ready.
+  int ready_sync_set_index_ ABSL_GUARDED_BY(mutex_) = -1;
+  // The timestamp at which the sync set is ready.  If no sync set is
+  // ready then this variable should be Timestamp::Done() .
+  Timestamp ready_timestamp_ ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_SYNC_SET_INPUT_STREAM_HANDLER_H_

mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc

@@ -12,85 +12,39 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h"
+
 #include <algorithm>
+#include <functional>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/log/check.h"
 #include "absl/strings/substitute.h"
 #include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/collection_item_id.h"
 #include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/mediapipe_options.pb.h"
 #include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.pb.h"
 #include "mediapipe/framework/timestamp.h"
 #include "mediapipe/framework/tool/validate_name.h"
 
 namespace mediapipe {
 
-// The input streams must have the same time unit but may have different time
-// origins (also called epochs). The timestamp_base_tag_index option
-// designates an input stream as the timestamp base.
-//
-// TimestampAlignInputStreamHandler operates in two phases:
-//
-// 1. Pre-initialization: In this phase, the input stream handler passes
-// through input packets in the timestamp base input stream, but buffers the
-// input packets in all other input streams. This phase ends when the input
-// stream handler has an input packet in every input stream. It uses the
-// the timestamps of these input packets to calculate the timestamp offset of
-// each input stream with respect to the timestamp base input stream. The
-// timestamp offsets are saved for use in the next phase.
-//
-// 2. Post-initialization: In this phase, the input stream handler behaves
-// like the DefaultInputStreamHandler, except that timestamp offsets are
-// applied to the packet timestamps.
-class TimestampAlignInputStreamHandler : public InputStreamHandler {
- public:
-  TimestampAlignInputStreamHandler() = delete;
-  TimestampAlignInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
-                                   CalculatorContextManager* cc_manager,
-                                   const MediaPipeOptions& options,
-                                   bool calculator_run_in_parallel);
-
-  void PrepareForRun(std::function<void()> headers_ready_callback,
-                     std::function<void()> notification_callback,
-                     std::function<void(CalculatorContext*)> schedule_callback,
-                     std::function<void(absl::Status)> error_callback) override;
-
- protected:
-  // In TimestampAlignInputStreamHandler, a node is "ready" if:
-  // - before the timestamp offsets are initialized: we have received a packet
-  //   in the timestamp base input stream, or
-  // - after the timestamp offsets are initialized: the minimum bound (over
-  //   all empty streams) is greater than the smallest timestamp of any
-  //   stream, which means we have received all the packets that will be
-  //   available at the next timestamp, or
-  // - all streams are done (need to call Close() in this case).
-  // Note that all packet timestamps and timestamp bounds are aligned with the
-  // timestamp base.
-  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
-
-  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
-  void FillInputSet(Timestamp input_timestamp,
-                    InputStreamShardSet* input_set) override;
-
- private:
-  CollectionItemId timestamp_base_stream_id_;
-
-  absl::Mutex mutex_;
-  bool offsets_initialized_ ABSL_GUARDED_BY(mutex_) = false;
-  std::vector<TimestampDiff> timestamp_offsets_;
-};
 REGISTER_INPUT_STREAM_HANDLER(TimestampAlignInputStreamHandler);
 
 TimestampAlignInputStreamHandler::TimestampAlignInputStreamHandler(
     std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
-    const MediaPipeOptions& options, bool calculator_run_in_parallel)
+    const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
     : InputStreamHandler(std::move(tag_map), cc_manager, options,
                          calculator_run_in_parallel),
       timestamp_offsets_(input_stream_managers_.NumEntries()) {
-  const auto& handler_options =
+  const auto& handler_options = options.GetExtension(
-      options.GetExtension(TimestampAlignInputStreamHandlerOptions::ext);
+      mediapipe::TimestampAlignInputStreamHandlerOptions::ext);
   std::string tag;
   int index;
   MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(

mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h

@@ -0,0 +1,91 @@
+// Copyright 2023 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.
+
+#ifndef MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_
+#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/framework/calculator_context_manager.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/collection_item_id.h"
+#include "mediapipe/framework/input_stream_handler.h"
+#include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.pb.h"
+#include "mediapipe/framework/timestamp.h"
+
+namespace mediapipe {
+
+// The input streams must have the same time unit but may have different time
+// origins (also called epochs). The timestamp_base_tag_index option
+// designates an input stream as the timestamp base.
+//
+// TimestampAlignInputStreamHandler operates in two phases:
+//
+// 1. Pre-initialization: In this phase, the input stream handler passes
+// through input packets in the timestamp base input stream, but buffers the
+// input packets in all other input streams. This phase ends when the input
+// stream handler has an input packet in every input stream. It uses the
+// the timestamps of these input packets to calculate the timestamp offset of
+// each input stream with respect to the timestamp base input stream. The
+// timestamp offsets are saved for use in the next phase.
+//
+// 2. Post-initialization: In this phase, the input stream handler behaves
+// like the DefaultInputStreamHandler, except that timestamp offsets are
+// applied to the packet timestamps.
+class TimestampAlignInputStreamHandler : public InputStreamHandler {
+ public:
+  TimestampAlignInputStreamHandler() = delete;
+  TimestampAlignInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
+                                   CalculatorContextManager* cc_manager,
+                                   const mediapipe::MediaPipeOptions& options,
+                                   bool calculator_run_in_parallel);
+
+  void PrepareForRun(std::function<void()> headers_ready_callback,
+                     std::function<void()> notification_callback,
+                     std::function<void(CalculatorContext*)> schedule_callback,
+                     std::function<void(absl::Status)> error_callback) override;
+
+ protected:
+  // In TimestampAlignInputStreamHandler, a node is "ready" if:
+  // - before the timestamp offsets are initialized: we have received a packet
+  //   in the timestamp base input stream, or
+  // - after the timestamp offsets are initialized: the minimum bound (over
+  //   all empty streams) is greater than the smallest timestamp of any
+  //   stream, which means we have received all the packets that will be
+  //   available at the next timestamp, or
+  // - all streams are done (need to call Close() in this case).
+  // Note that all packet timestamps and timestamp bounds are aligned with the
+  // timestamp base.
+  NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
+
+  // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
+  void FillInputSet(Timestamp input_timestamp,
+                    InputStreamShardSet* input_set) override;
+
+ private:
+  CollectionItemId timestamp_base_stream_id_;
+
+  absl::Mutex mutex_;
+  bool offsets_initialized_ ABSL_GUARDED_BY(mutex_) = false;
+  std::vector<TimestampDiff> timestamp_offsets_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_