[core] Threaded actors get stuck forever if they receive two exit signals (#51582)

If a threaded actor receives exit signals twice as shown in the above
screenshot, it will execute
[task_receiver_->Stop()](https://github.com/ray-project/ray/blob/6bb9cef9257046ae31f78f6c52015a8ebf009f81/src/ray/core_worker/core_worker.cc#L1224)
twice. However, the second call to `task_receiver_->Stop()` will get
stuck forever when executing the
[releaser](https://github.com/ray-project/ray/blob/6bb9cef9257046ae31f78f6c52015a8ebf009f81/src/ray/core_worker/transport/concurrency_group_manager.cc#L135).

### Reproduction

```sh
ray start --head --include-dashboard=True --num-cpus=1
# https://gist.github.com/kevin85421/7a42ac3693537c2148fa554065bb5223
python3 test.py

# Some actors are still ALIVE. If all actors are DEAD, increase the number of actors.
ray list actors
```

---------

Signed-off-by: kaihsun <kaihsun@anyscale.com>

Author: kevin85421

python/ray/_private/state_api_test_utils.py

@@ -29,6 +29,7 @@
     PredicateType,
     SupportedFilterType,
 )
+import ray._private.test_utils as test_utils
 
 
 @dataclass
@@ -113,6 +114,16 @@ def invoke_state_api(
     return res
 
 
+def invoke_state_api_n(*args, **kwargs):
+    def verify():
+        NUM_API_CALL_SAMPLES = 10
+        for _ in range(NUM_API_CALL_SAMPLES):
+            invoke_state_api(*args, **kwargs)
+        return True
+
+    test_utils.wait_for_condition(verify, retry_interval_ms=2000, timeout=30)
+
+
 def aggregate_perf_results(state_stats: StateAPIStats = GLOBAL_STATE_STATS):
     """Aggregate stats of state API calls
 

python/ray/tests/test_actor.py

@@ -20,6 +20,10 @@
 from ray._private.test_utils import SignalActor
 from ray.core.generated import gcs_pb2
 from ray._private.utils import hex_to_binary
+from ray._private.state_api_test_utils import invoke_state_api, invoke_state_api_n
+
+from ray.util.state import list_actors
+
 
 # NOTE: We have to import setproctitle after ray because we bundle setproctitle
 # with ray.
@@ -283,6 +287,44 @@ def get_val(self):
     assert ray.get(g.remote()) == num_remote_functions - 1
 
 
+@pytest.mark.parametrize("enable_concurrency_group", [False, True])
+def test_exit_actor(ray_start_regular, enable_concurrency_group):
+    concurrency_groups = {"io": 1} if enable_concurrency_group else None
+
+    @ray.remote(concurrency_groups=concurrency_groups)
+    class TestActor:
+        def exit(self):
+            ray.actor.exit_actor()
+
+    num_actors = 30
+    actor_class_name = TestActor.__ray_metadata__.class_name
+
+    actors = [TestActor.remote() for _ in range(num_actors)]
+    ray.get([actor.__ray_ready__.remote() for actor in actors])
+    invoke_state_api(
+        lambda res: len(res) == num_actors,
+        list_actors,
+        filters=[("state", "=", "ALIVE"), ("class_name", "=", actor_class_name)],
+        limit=1000,
+    )
+
+    ray.wait([actor.exit.remote() for actor in actors], timeout=10.0)
+
+    invoke_state_api_n(
+        lambda res: len(res) == 0,
+        list_actors,
+        filters=[("state", "=", "ALIVE"), ("class_name", "=", actor_class_name)],
+        limit=1000,
+    )
+
+    invoke_state_api(
+        lambda res: len(res) == num_actors,
+        list_actors,
+        filters=[("state", "=", "DEAD"), ("class_name", "=", actor_class_name)],
+        limit=1000,
+    )
+
+
 @pytest.mark.skipif(client_test_enabled(), reason="internal api")
 def test_actor_method_metadata_cache(ray_start_regular):
     class Actor(object):

release/nightly_tests/stress_tests/test_state_api_scale.py

@@ -7,6 +7,7 @@
     StateAPIMetric,
     aggregate_perf_results,
     invoke_state_api,
+    invoke_state_api_n,
     GLOBAL_STATE_STATS,
 )
 
@@ -38,16 +39,6 @@
 MiB = 1024 * 1024
 
 
-def invoke_state_api_n(*args, **kwargs):
-    def verify():
-        NUM_API_CALL_SAMPLES = 10
-        for _ in range(NUM_API_CALL_SAMPLES):
-            invoke_state_api(*args, **kwargs)
-        return True
-
-    test_utils.wait_for_condition(verify, retry_interval_ms=2000, timeout=30)
-
-
 def test_many_tasks(num_tasks: int):
     TASK_NAME_TEMPLATE = "pi4_sample_{num_tasks}"
     if num_tasks == 0:

src/ray/core_worker/core_worker.cc

@@ -1174,6 +1174,13 @@ void CoreWorker::Exit(
     const rpc::WorkerExitType exit_type,
     const std::string &detail,
     const std::shared_ptr<LocalMemoryBuffer> &creation_task_exception_pb_bytes) {
+  if (is_exit_) {
+    RAY_LOG(INFO) << "Exit signal received, but the core worker has already received a "
+                     "signal and is shutting down.";
+    return;
+  }
+  is_exit_ = true;
+
   RAY_LOG(INFO) << "Exit signal received, this process will exit after all outstanding "
                    "tasks have finished"
                 << ", exit_type=" << rpc::WorkerExitType_Name(exit_type)

src/ray/core_worker/core_worker.h

@@ -1893,6 +1893,15 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
 
   std::atomic<bool> is_shutdown_ = false;
 
+  /// Whether the `Exit` function has been called, to avoid executing the exit
+  /// process multiple times.
+  ///
+  /// TODO(kevin85421): Currently, there are two public functions, `Exit` and `Shutdown`,
+  /// to terminate the core worker gracefully. We should unify them into `Exit()` so we
+  /// don't need `is_shutdown_` in the future. See
+  /// https://github.com/ray-project/ray/issues/51642 for more details.
+  std::atomic<bool> is_exit_ = false;
+
   int64_t max_direct_call_object_size_;
 
   friend class CoreWorkerTest;