Make subscriber cleanup race-safe

MODULE.bazel

@@ -1,6 +1,6 @@
 module(
     name = "subspace",
-    version = "3.0.2",
+    version = "3.0.3",
 )
 
 bazel_dep(name = "bazel_skylib", version = "1.9.0")

client/client_test.cc

@@ -18,11 +18,13 @@
 #include <array>
 #include <atomic>
 #include <cerrno>
+#include <condition_variable>
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
 #include <inttypes.h>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <sys/resource.h>
 #if SUBSPACE_SHMEM_MODE == SUBSPACE_SHMEM_MODE_MEMFD
@@ -1153,6 +1155,81 @@ TEST_F(ClientTest, PublishAndReadWithSubscriberQueue) {
   ASSERT_EQ(0, memcmp(msg->buffer, "queued3", 7));
 }
 
+TEST_F(ClientTest, SubscriberJoiningDuringPublishReceivesCommittedMessage) {
+  auto pub_client = EVAL_AND_ASSERT_OK(subspace::Client::Create(Socket()));
+  auto sub_client = EVAL_AND_ASSERT_OK(subspace::Client::Create(Socket()));
+
+  constexpr char kChannel[] = "subscriber_joins_during_publish";
+  auto pub = EVAL_AND_ASSERT_OK(pub_client->CreatePublisher(
+      kChannel,
+      PubOpts(256, 10)
+          .SetChecksum(true)
+          .SetSubscriberQueueArenaSize(subspace::SlotQueueBlockSize(16))));
+
+  std::mutex mutex;
+  std::condition_variable callback_entered_cv;
+  std::condition_variable resume_publish_cv;
+  bool callback_entered = false;
+  bool resume_publish = false;
+  pub.SetChecksumCallback(
+      [&](const std::array<absl::Span<const uint8_t>, 3> &,
+          absl::Span<std::byte> checksum) {
+        std::unique_lock<std::mutex> lock(mutex);
+        callback_entered = true;
+        callback_entered_cv.notify_one();
+        resume_publish_cv.wait(lock, [&] { return resume_publish; });
+        std::fill(checksum.begin(), checksum.end(), std::byte{0});
+      });
+
+  absl::Status publish_status = absl::UnknownError("publish did not run");
+  std::thread publish_thread([&] {
+    absl::StatusOr<void *> buffer = pub.GetMessageBuffer();
+    if (!buffer.ok()) {
+      publish_status = buffer.status();
+      {
+        std::lock_guard<std::mutex> lock(mutex);
+        callback_entered = true;
+      }
+      callback_entered_cv.notify_one();
+      return;
+    }
+    memcpy(*buffer, "joined", 7);
+    publish_status = pub.PublishMessage(7).status();
+  });
+
+  {
+    std::unique_lock<std::mutex> lock(mutex);
+    callback_entered_cv.wait(lock, [&] { return callback_entered; });
+  }
+
+  // Registration seeds the pending publisher-owned generation. The
+  // subscriber must preserve that bit and queue entry until commit.
+  absl::StatusOr<Subscriber> sub_status = sub_client->CreateSubscriber(
+      kChannel, SubOpts().SetSubscriberQueueSize(16).SetChecksum(true));
+  if (sub_status.ok()) {
+    sub_status->SetChecksumCallback(
+        [](const std::array<absl::Span<const uint8_t>, 3> &,
+           absl::Span<std::byte> checksum) {
+          std::fill(checksum.begin(), checksum.end(), std::byte{0});
+        });
+  }
+
+  {
+    std::lock_guard<std::mutex> lock(mutex);
+    resume_publish = true;
+  }
+  resume_publish_cv.notify_one();
+  publish_thread.join();
+  ASSERT_OK(publish_status);
+  ASSERT_OK(sub_status);
+  auto sub = std::move(*sub_status);
+
+  auto message = EVAL_AND_ASSERT_OK(sub.ReadMessage());
+  ASSERT_GT(message.length, 0);
+  EXPECT_EQ(7, message.length);
+  EXPECT_EQ(0, memcmp(message.buffer, "joined", 7));
+}
+
 TEST_F(ClientTest, SubscribersUseDifferentQueueSizes) {
   subspace::Client client;
   ASSERT_OK(client.Init(Socket()));
@@ -4489,6 +4566,83 @@ TEST_F(ClientTest, SubscriberRemovalTriggersServerRetirement) {
   message.Reset();
 }
 
+TEST_F(ClientTest, SubscriberRemovalCanRacePublisherCommit) {
+  auto pub_client = EVAL_AND_ASSERT_OK(subspace::Client::Create(Socket()));
+  auto sub_client = EVAL_AND_ASSERT_OK(subspace::Client::Create(Socket()));
+
+  constexpr char kChannel[] = "server_retirement_during_publish";
+  auto pub = EVAL_AND_ASSERT_OK(pub_client->CreatePublisher(
+      kChannel,
+      PubOpts(256, 10)
+          .SetNotifyRetirement(true)
+          .SetSubscriberQueueArenaSize(subspace::SlotQueueBlockSize(16))));
+  std::optional<Subscriber> sub = EVAL_AND_ASSERT_OK(
+      sub_client->CreateSubscriber(
+          kChannel, SubOpts().SetSubscriberQueueSize(16)));
+
+  subspace::ServerChannel *channel = Server()->FindChannel(kChannel);
+  ASSERT_NE(nullptr, channel);
+
+  int publisher_id = -1;
+  for (const auto &[id, user] : channel->GetUsers()) {
+    if (user != nullptr && user->IsPublisher()) {
+      publisher_id = id;
+      break;
+    }
+  }
+  ASSERT_GE(publisher_id, 0);
+
+  int subscriber_id = -1;
+  channel->GetCcb()->subscribers.Traverse(
+      [&subscriber_id](int id) { subscriber_id = id; });
+  ASSERT_GE(subscriber_id, 0);
+
+  subspace::MessageSlot *slot = nullptr;
+  for (int i = 0; i < channel->NumSlots(); ++i) {
+    subspace::MessageSlot *candidate = &channel->GetCcb()->slots[i];
+    if (candidate->refs.load(std::memory_order_acquire) ==
+        (subspace::kPubOwned | static_cast<uint64_t>(publisher_id))) {
+      slot = candidate;
+      break;
+    }
+  }
+  ASSERT_NE(nullptr, slot);
+
+  const uint64_t cleanup_generation =
+      channel->SubscriberCleanupGenerationFor(-1);
+  slot->ordinal.store(1, std::memory_order_relaxed);
+  slot->vchan_id.store(-1, std::memory_order_relaxed);
+  slot->bridged_slot_id.store(slot->id, std::memory_order_relaxed);
+  channel->GetAvailableSlots(subscriber_id).Set(slot->id);
+  channel->BeginSubscriberQueuePublish(publisher_id);
+
+  // Cleanup must return without waiting for the synthetic in-flight publisher.
+  // Its retirement scan observes kPubOwned and leaves this slot alone.
+  sub.reset();
+  EXPECT_FALSE(channel->RetiredSlots().IsSet(slot->id));
+  EXPECT_NE(cleanup_generation, channel->SubscriberCleanupGenerationFor(-1));
+
+  // Publication commit is the second side of the handshake and therefore
+  // performs the retirement that the server scan could not.
+  slot->refs.store(subspace::BuildRefsBitField(1, -1, 0),
+                   std::memory_order_release);
+  ASSERT_TRUE(channel->TryRetireSlot(slot));
+  channel->EndSubscriberQueuePublish(publisher_id);
+  channel->NotifyPublisherRetirement(slot->id);
+
+  const toolbelt::FileDescriptor &retirement_fd = pub.GetRetirementFd();
+  struct pollfd fd = {
+      .fd = retirement_fd.Fd(),
+      .events = POLLIN,
+  };
+  ASSERT_EQ(1, ::poll(&fd, 1, 1000));
+  int retired_slot = -1;
+  ASSERT_EQ(sizeof(retired_slot),
+            ::read(retirement_fd.Fd(), &retired_slot, sizeof(retired_slot)));
+  EXPECT_EQ(slot->id, retired_slot);
+  EXPECT_EQ(0, ::poll(&fd, 1, 0));
+}
+
 // This tests retirement from the the publisher side using dropped messages.  We
 // have two subscribers, one reads two messages and the other doesn't read any.
 // Since the second subscriber will never see the messages, the publisher will

client/publisher.cc

@@ -293,8 +293,11 @@ void PublisherImpl::RetirePublishedSlotImmediately(MessageSlot *slot) {
   if (slot == nullptr) {
     return;
   }
-  RetiredSlots().Set(slot->id);
-  TriggerRetirement(slot->id);
+  const int32_t retirement_slot_id =
+      slot->bridged_slot_id.load(std::memory_order_relaxed);
+  if (TryRetireSlot(slot)) {
+    TriggerRetirement(retirement_slot_id);
+  }
 }
 
 MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) {
@@ -675,37 +678,48 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother(
     }
   }
 
-  // Set the refs to the ordinal with no refs.
-  slot->refs.store(
-      BuildRefsBitField(slot->ordinal.load(std::memory_order_relaxed),
-                        vchan_id_, 0),
-      std::memory_order_release);
-
-  // Tell all subscribers that the slot is available, BEFORE bumping
-  // total_messages. When subscriber queues are enabled, unreliable C++
-  // subscribers consume the per-subscriber queue first. The available-slot
-  // bitset remains authoritative and provides recovery when queue insertion
-  // fails or entries are evicted.
+  const uint64_t published_ordinal =
+      slot->ordinal.load(std::memory_order_relaxed);
+  const uint64_t published_timestamp =
+      slot->timestamp.load(std::memory_order_relaxed);
+  const int32_t retirement_slot_id =
+      slot->bridged_slot_id.load(std::memory_order_relaxed);
+  const uint64_t cleanup_generation =
+      SubscriberCleanupGenerationFor(vchan_id_);
+
+  // Tell all subscribers that the slot is available while it remains
+  // publisher-owned. The kPubOwned bit is the publication commit barrier:
+  // subscribers preserve the delivery record but cannot claim the slot, and
+  // server cleanup cannot retire it until all delivery records and accounting
+  // below are complete.
   //
-  // SubscriberImpl::NextSlot() uses total_messages as a version stamp
-  // for its cached active_slots_ snapshot: a reliable subscriber that observes
-  // a bumped count must also observe every preceding bits.Set() so its
-  // CollectVisibleSlots() snapshot can't miss the just-published slot.
-  // bits.Set() is relaxed, but the following counter increment is seq_cst, so
-  // the relaxed bit writes are sequenced-before the seq_cst increment and
-  // therefore happens-before any subscriber's seq_cst load of total_messages
-  // that observes the new value.
+  // The available-slot bitset remains authoritative when queue insertion
+  // fails or entries are evicted. A subscriber can disappear after
+  // TraverseSeqCst observes its bit, so recheck membership after setting the
+  // delivery bit. Either this recheck clears a stale write, or the server's
+  // later ClearWasSet observes it.
   SubscriberQueuePublishGuard publish_guard(*this);
   std::vector<InPlaceSlotQueue *> failed_queues;
   ccb_->subscribers.TraverseSeqCst([this, slot, &failed_queues](int sub_id) {
     if (vchan_id_ != -1 && GetSubVchanId(sub_id) != -1 &&
         vchan_id_ != GetSubVchanId(sub_id)) {
       return;
     }
-    // The bitset is the authoritative delivery record. The queue is an
-    // acceleration index and may reject an insertion under contention or
-    // after a peer dies mid-operation.
-    GetAvailableSlots(sub_id).Set(slot->id);
+    InPlaceAtomicBitset &available = GetAvailableSlots(sub_id);
+    available.Set(slot->id);
+    if (!ccb_->subscribers.IsSetSeqCst(sub_id)) {
+      available.Clear(slot->id);
+      // The subscriber ID may have been reused after the first membership
+      // check. Registration publishes membership before seeding this bit. If
+      // the new subscriber is already visible, restore the bit that the stale
+      // cleanup above may have cleared; otherwise its later seed handles the
+      // in-progress generation.
+      if (!ccb_->subscribers.IsSetSeqCst(sub_id)) {
+        return;
+      }
+      available.Set(slot->id);
+    }
+
     InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(sub_id);
     if (queue != nullptr &&
         !queue->Push(slot->id,
@@ -715,7 +729,7 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother(
     }
   });
 
-  // Update counters AFTER notifying subscribers (see above).
+  // Finish all slot and queue bookkeeping before making the slot claimable.
   if (!is_activation) {
     const uint64_t message_size =
         slot->message_size.load(std::memory_order_relaxed);
@@ -724,17 +738,32 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother(
       ccb_->max_message_size = message_size;
     }
   }
-  ccb_->total_messages.fetch_add(1, std::memory_order_seq_cst);
-  // Publish queue failure only after this message's bit and version are
-  // visible. Otherwise a subscriber can consume the failure, take an older
-  // bitset snapshot, leave fallback, and then deliver a newer queue entry
-  // ahead of the failed ordinal.
   for (InPlaceSlotQueue *queue : failed_queues) {
     queue->MarkInsertionFailure();
   }
 
+  // Commit the publication. A subscriber that observes the subsequent
+  // total_messages increment must also observe this release and all preceding
+  // delivery-record writes. PopulateActiveSlots preserves bits for
+  // publisher-owned slots, and queue consumers leave current-generation
+  // entries at the head until this store completes.
+  slot->refs.store(BuildRefsBitField(published_ordinal, vchan_id_, 0),
+                   std::memory_order_release);
+
+  // SubscriberImpl::NextSlot() uses total_messages as a version stamp for its
+  // cached active_slots_ snapshot.
+  ccb_->total_messages.fetch_add(1, std::memory_order_seq_cst);
+
+  // Subscriber removal and publication commit race safely: the operation that
+  // happens second re-evaluates retirement using the current subscriber count.
+  if (!is_activation &&
+      SubscriberCleanupGenerationFor(vchan_id_) != cleanup_generation &&
+      TryRetireSlot(slot)) {
+    TriggerRetirement(retirement_slot_id);
+  }
+
   if (!acquire_next) {
-    return {nullptr, prefix->ordinal, prefix->timestamp};
+    return {nullptr, published_ordinal, published_timestamp};
   }
 
   // A reliable publisher doesn't allocate a slot until it is asked for.
@@ -745,7 +774,7 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother(
   // Find a new slot.x
   MessageSlot *new_slot = FindFreeSlotUnreliable(owner);
 
-  return {new_slot, prefix->ordinal, prefix->timestamp};
+  return {new_slot, published_ordinal, published_timestamp};
 }
 
 } // namespace details