From b9d020cb7d32e936943b963c401cc3aeb9f88407 Mon Sep 17 00:00:00 2001
From: Hadon Nash <hadon@google.com>
Date: Tue, 13 Dec 2022 16:58:12 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 495163109
---
 mediapipe/framework/scheduler.cc | 11 ++++++++---
 mediapipe/framework/scheduler.h  |  2 +-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/mediapipe/framework/scheduler.cc b/mediapipe/framework/scheduler.cc
index afef4f383..854c10fd5 100644
--- a/mediapipe/framework/scheduler.cc
+++ b/mediapipe/framework/scheduler.cc
@@ -117,7 +117,7 @@ void Scheduler::SubmitWaitingTasksOnQueues() {
 // Note: state_mutex_ is held when this function is entered or
 // exited.
 void Scheduler::HandleIdle() {
-  if (handling_idle_) {
+  if (++handling_idle_ > 1) {
     // Someone is already inside this method.
     // Note: This can happen in the sections below where we unlock the mutex
     // and make more nodes runnable: the nodes can run and become idle again
@@ -127,7 +127,6 @@ void Scheduler::HandleIdle() {
     VLOG(2) << "HandleIdle: already in progress";
     return;
   }
-  handling_idle_ = true;
 
   while (IsIdle() && (state_ == STATE_RUNNING || state_ == STATE_CANCELLING)) {
     // Remove active sources that are closed.
@@ -165,11 +164,17 @@ void Scheduler::HandleIdle() {
       }
     }
 
+    // If HandleIdle has been called again, then continue scheduling.
+    if (handling_idle_ > 1) {
+      handling_idle_ = 1;
+      continue;
+    }
+
     // Nothing left to do.
     break;
   }
 
-  handling_idle_ = false;
+  handling_idle_ = 0;
 }
 
 // Note: state_mutex_ is held when this function is entered or exited.
diff --git a/mediapipe/framework/scheduler.h b/mediapipe/framework/scheduler.h
index dd1572d99..b59467b9f 100644
--- a/mediapipe/framework/scheduler.h
+++ b/mediapipe/framework/scheduler.h
@@ -302,7 +302,7 @@ class Scheduler {
   // - We need it to be reentrant, which Mutex does not support.
   // - We want simultaneous calls to return immediately instead of waiting,
   //   and Mutex's TryLock is not guaranteed to work.
-  bool handling_idle_ ABSL_GUARDED_BY(state_mutex_) = false;
+  int handling_idle_ ABSL_GUARDED_BY(state_mutex_) = 0;
 
   // Mutex for the scheduler state and related things.
   // Note: state_ is declared as atomic so that its getter methods don't need