From ab3a5f0fbf1883c4d1dfe1df2db80a7045a390c4 Mon Sep 17 00:00:00 2001 From: MediaPipe Team 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( R"pb( @@ -480,15 +480,11 @@ TEST(MuxCalculatorTest, CrashesOnCloseWithDeafultInputStreamHandler) { MP_ASSERT_OK(graph.AddPacketToInputStream( "value_0", MakePacket(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