From ab3a5f0fbf1883c4d1dfe1df2db80a7045a390c4 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Thu, 17 Nov 2022 16:28:08 -0800
Subject: [PATCH] Make MuxCalculator with DefaultInputStreamHandler to handle
 graph closure gracefully

PiperOrigin-RevId: 489336722
---
 mediapipe/calculators/core/mux_calculator.cc     |  4 ++++
 .../calculators/core/mux_calculator_test.cc      | 16 ++++++----------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/mediapipe/calculators/core/mux_calculator.cc b/mediapipe/calculators/core/mux_calculator.cc
index a0ce2ae34..88b04a32b 100644
--- a/mediapipe/calculators/core/mux_calculator.cc
+++ b/mediapipe/calculators/core/mux_calculator.cc
@@ -41,6 +41,10 @@ class MuxCalculator : public Node {
                           StreamHandler("MuxInputStreamHandler"));
 
   absl::Status Process(CalculatorContext* cc) final {
+    if (kSelect(cc).IsStream() && kSelect(cc).IsEmpty()) {
+      return absl::OkStatus();
+    }
+
     int select = *kSelect(cc);
     RET_CHECK(0 <= select && select < kIn(cc).Count());
     if (!kIn(cc)[select].IsEmpty()) {
diff --git a/mediapipe/calculators/core/mux_calculator_test.cc b/mediapipe/calculators/core/mux_calculator_test.cc
index a3ac8a27a..6b9434be9 100644
--- a/mediapipe/calculators/core/mux_calculator_test.cc
+++ b/mediapipe/calculators/core/mux_calculator_test.cc
@@ -439,7 +439,7 @@ TEST(MuxCalculatorTest, HandlesCloseGracefully) {
   EXPECT_TRUE(output_packets.empty());
 }
 
-TEST(MuxCalculatorTest, CrashesOnCloseWithDeafultInputStreamHandler) {
+TEST(MuxCalculatorTest, HandlesCloseGracefullyWithDeafultInputStreamHandler) {
   CalculatorGraphConfig config =
       mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
           R"pb(
@@ -480,15 +480,11 @@ TEST(MuxCalculatorTest, CrashesOnCloseWithDeafultInputStreamHandler) {
   MP_ASSERT_OK(graph.AddPacketToInputStream(
       "value_0", MakePacket<int>(0).At(Timestamp(1000))));
   MP_ASSERT_OK(graph.WaitUntilIdle());
-  // Currently MuxCalculator crashes with a correct packet set from
-  // DefaultInputStreamHandler.  The SELECT packet is missing at Timestamp 1000,
-  // and an empty packet is the correct representation of that.
-  EXPECT_DEATH(
-      {
-        (void)graph.CloseAllInputStreams();
-        (void)graph.WaitUntilDone();
-      },
-      "Check failed: payload_");
+  MP_ASSERT_OK(graph.CloseAllInputStreams());
+  MP_ASSERT_OK(graph.WaitUntilDone());
+
+  ASSERT_EQ(output_packets.size(), 1);
+  EXPECT_TRUE(output_packets[0].IsEmpty());
 }
 
 }  // namespace