[socket-transfer]: IsLastSemaphore requires poisoning as well or else it will write to the atm after an error has been set.

PiperOrigin-RevId: 763624413

Author: pschuh

xla/python/transfer/BUILD

@@ -153,6 +153,7 @@ xla_cc_test(
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
         "@com_google_googletest//:gtest_main",

xla/python/transfer/streaming_ifrt.cc

@@ -221,6 +221,7 @@ class DmaDestination : public ChunkDestination {
   }
 
   void Poison(absl::Status s) override {
+    semaphore_.Poison();
     atm_->SetBufferError(buffer_index_, std::move(s));
   }
 

xla/python/transfer/streaming_ifrt.h

@@ -183,27 +183,40 @@ class IsLastSemaphore {
       : guard_counter_(value), counter_(value) {}
 
   template <typename T>
-  auto DoWork(size_t value, T&& cb) -> decltype(cb(false)) {
-    bool is_last = guard_counter_.fetch_sub(value) - value == 0;
-    if (is_last && counter_.fetch_sub(value) - value != 0) {
-      // Wait if we happen to slip in between guard_counter and counter.
+  auto DoWork(size_t value, T&& cb) -> absl::Status {
+    bool is_last;
+    {
       absl::MutexLock l(&mu_);
-      auto cond = [this]() { return counter_.load() == 0; };
-      mu_.Await(absl::Condition(&cond));
+      if (is_poisoned_) {
+        return absl::OkStatus();
+      }
+      guard_counter_ -= value;
+      is_last = guard_counter_ == 0;
+      if (is_last) {
+        // Wait if we happen to slip in between guard_counter and counter.
+        auto cond = [this, value]() { return counter_ == value; };
+        mu_.Await(absl::Condition(&cond));
+      }
     }
     auto cleanup = absl::MakeCleanup([&]() {
-      if (!is_last && (counter_.fetch_sub(value) - value) == 0) {
-        // Wake any waiters.
-        absl::MutexLock l(&mu_);
-      }
+      absl::MutexLock l(&mu_);
+      counter_ -= value;
     });
     return cb(is_last);
   }
 
+  void Poison() {
+    absl::MutexLock l(&mu_);
+    is_poisoned_ = true;
+    auto cond = [this]() { return counter_ == guard_counter_; };
+    mu_.Await(absl::Condition(&cond));
+  }
+
  private:
   absl::Mutex mu_;
-  std::atomic<ssize_t> guard_counter_;
-  std::atomic<ssize_t> counter_;
+  bool is_poisoned_ = false;
+  ssize_t guard_counter_;
+  ssize_t counter_;
 };
 
 }  // namespace internal

xla/python/transfer/streaming_ifrt_test.cc

@@ -25,6 +25,7 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 #include "absl/log/check.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
@@ -151,9 +152,15 @@ TEST(PremappedCopierState, RoundTrip) {
 TEST(Semaphore, Basic) {
   internal::IsLastSemaphore semaphore(15);
   for (size_t i = 0; i < 10; ++i) {
-    semaphore.DoWork(1, [&](bool is_last) { EXPECT_FALSE(is_last); });
+    CHECK_OK(semaphore.DoWork(1, [&](bool is_last) -> absl::Status {
+      EXPECT_FALSE(is_last);
+      return absl::OkStatus();
+    }));
   }
-  semaphore.DoWork(5, [&](bool is_last) { EXPECT_TRUE(is_last); });
+  CHECK_OK(semaphore.DoWork(5, [&](bool is_last) -> absl::Status {
+    EXPECT_TRUE(is_last);
+    return absl::OkStatus();
+  }));
 }
 
 TEST(Semaphore, Async) {
@@ -177,24 +184,26 @@ TEST(Semaphore, Async) {
       tsl::Env::Default()->StartThread({}, "t1", [&]() {
         for (size_t i = 0; i < 8; ++i) {
           thread_wait_flip(0);
-          o_semaphore.DoWork(1, [&](bool is_last) {
+          CHECK_OK(o_semaphore.DoWork(1, [&](bool is_last) -> absl::Status {
             thread_flip(0);
             EXPECT_FALSE(is_last);
-          });
+            return absl::OkStatus();
+          }));
         }
       }));
   std::unique_ptr<tsl::Thread> t2(
       tsl::Env::Default()->StartThread({}, "t2", [&]() {
         for (size_t i = 0; i < 8; ++i) {
           thread_wait_flip(1);
-          o_semaphore.DoWork(1, [&](bool is_last) {
+          CHECK_OK(o_semaphore.DoWork(1, [&](bool is_last) -> absl::Status {
             thread_flip(1);
             if (i == 7) {
               EXPECT_TRUE(is_last);
             } else {
               EXPECT_FALSE(is_last);
             }
-          });
+            return absl::OkStatus();
+          }));
         }
       }));
 }