remove error msg based check and add test

Signed-off-by: Sagar Sumit <[email protected]>

Author: codope

src/ray/core_worker/core_worker.cc

@@ -538,9 +538,7 @@ void CoreWorker::Shutdown() {
   RAY_LOG(INFO) << "Core worker ready to be deallocated.";
 }
 
-bool CoreWorker::IsShuttingDown() const {
-  return is_shutdown_.load();
-}
+bool CoreWorker::IsShuttingDown() const { return is_shutdown_.load(); }
 
 void CoreWorker::ConnectToRayletInternal() {
   // Tell the raylet the port that we are listening on.

src/ray/core_worker/core_worker_process.cc

@@ -423,33 +423,30 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
       /*put_in_local_plasma_callback=*/
       [this](const RayObject &object, const ObjectID &object_id) {
         auto core_worker = GetCoreWorker();
-        
+
         // Check if the core worker is shutting down before attempting plasma operations.
         // During shutdown, the plasma store connection may already be broken, so we
         // should avoid plasma operations entirely.
         if (core_worker->IsShuttingDown()) {
-          RAY_LOG(INFO) << "Skipping plasma store operation for error object " << object_id 
-                       << " because core worker is shutting down.";
+          RAY_LOG(INFO) << "Skipping plasma store operation for error object "
+                        << object_id << " because core worker is shutting down.";
           return;
         }
-        
-        auto status = core_worker->PutInLocalPlasmaStore(object, object_id, /*pin_object=*/true);
+
+        auto status =
+            core_worker->PutInLocalPlasmaStore(object, object_id, /*pin_object=*/true);
         if (!status.ok()) {
-          if (status.IsIOError() && 
-              (status.message().find("Broken pipe") != std::string::npos ||
-               status.message().find("Connection reset") != std::string::npos ||
-               status.message().find("Bad file descriptor") != std::string::npos)) {
-            // This is likely a shutdown race where the plasma store
-            // connection was closed before we could complete the operation.
-            // Log as warning since this is expected during shutdown scenarios.
-            RAY_LOG(WARNING) << "Failed to put error object " << object_id 
-                            << " in plasma store due to connection error (likely shutdown): " 
-                            << status.ToString();
+          // During shutdown, plasma store connections can be closed causing IOError.
+          // We only tolerate IOError during shutdown to avoid masking real errors.
+          if (status.IsIOError() && core_worker->IsShuttingDown()) {
+            // Double-check shutdown state - this handles the race where shutdown
+            // began after our first check but before the plasma operation.
+            RAY_LOG(WARNING) << "Failed to put error object " << object_id
+                             << " in plasma store during shutdown: " << status.ToString();
           } else {
-            // For other types of errors, maintain the original
-            // behavior with RAY_CHECK_OK to catch real issues.
-            RAY_CHECK_OK(status) << "Failed to put error object " << object_id 
-                                << " in plasma store: " << status.ToString();
+            // For any other error, or IOError when not shutting down, this indicates
+            // a real problem that should crash the process.
+            RAY_CHECK_OK(status);
           }
         }
       },

src/ray/core_worker/test/task_manager_test.cc

@@ -2692,6 +2692,74 @@ TEST_F(TaskManagerTest, TestTaskRetriedOnNodePreemption) {
   // Cleanup
   manager_.FailPendingTask(spec.TaskId(), rpc::ErrorType::WORKER_DIED);
 }
+
+class PlasmaShutdownRaceTest : public ::testing::Test {
+ public:
+  PlasmaShutdownRaceTest() : is_shutting_down_(false) {}
+
+  Status SimulatePlasmaCallback(const ObjectID &object_id, bool simulate_failure) {
+    if (is_shutting_down_) {
+      skipped_operations_.insert(object_id);
+      return Status::OK();
+    }
+
+    if (simulate_failure) {
+      auto status = Status::IOError("Broken pipe");
+      if (status.IsIOError() && is_shutting_down_) {
+        tolerated_operations_.insert(object_id);
+        return Status::OK();
+      } else {
+        failed_operations_.insert(object_id);
+        return status;
+      }
+    }
+
+    successful_operations_.insert(object_id);
+    return Status::OK();
+  }
+
+  void SetShuttingDown(bool shutting_down) { is_shutting_down_ = shutting_down; }
+
+ protected:
+  bool is_shutting_down_;
+  std::unordered_set<ObjectID> skipped_operations_;
+  std::unordered_set<ObjectID> tolerated_operations_;
+  std::unordered_set<ObjectID> successful_operations_;
+  std::unordered_set<ObjectID> failed_operations_;
+};
+
+// Test plasma callback behavior during shutdown to prevent RAY_CHECK crashes
+TEST_F(PlasmaShutdownRaceTest, PlasmaCallbackHandlesShutdownRaceCondition) {
+  auto object_id = ObjectID::FromRandom();
+
+  SetShuttingDown(false);
+  ASSERT_TRUE(SimulatePlasmaCallback(object_id, false).ok());
+  ASSERT_EQ(successful_operations_.count(object_id), 1);
+
+  auto object_id2 = ObjectID::FromRandom();
+  auto status = SimulatePlasmaCallback(object_id2, true);
+  ASSERT_FALSE(status.ok());
+  ASSERT_TRUE(status.IsIOError());
+  ASSERT_EQ(failed_operations_.count(object_id2), 1);
+
+  auto object_id3 = ObjectID::FromRandom();
+  SetShuttingDown(true);
+  ASSERT_TRUE(SimulatePlasmaCallback(object_id3, false).ok());
+  ASSERT_EQ(skipped_operations_.count(object_id3), 1);
+
+  auto object_id4 = ObjectID::FromRandom();
+  SetShuttingDown(false);
+  auto status4 = Status::IOError("Broken pipe");
+  SetShuttingDown(true);
+
+  if (status4.IsIOError() && is_shutting_down_) {
+    tolerated_operations_.insert(object_id4);
+  } else {
+    failed_operations_.insert(object_id4);
+  }
+  ASSERT_EQ(tolerated_operations_.count(object_id4), 1);
+}
+
 }  // namespace core
 }  // namespace ray