diff --git a/MODULE.bazel b/MODULE.bazel index 020fa03..295651a 100644 --- a/MODULE.bazel +++ b/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") diff --git a/client/client_test.cc b/client/client_test.cc index eb06537..e349ecf 100644 --- a/client/client_test.cc +++ b/client/client_test.cc @@ -18,11 +18,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #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, 3> &, + absl::Span checksum) { + std::unique_lock 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 buffer = pub.GetMessageBuffer(); + if (!buffer.ok()) { + publish_status = buffer.status(); + { + std::lock_guard lock(mutex); + callback_entered = true; + } + callback_entered_cv.notify_one(); + return; + } + memcpy(*buffer, "joined", 7); + publish_status = pub.PublishMessage(7).status(); + }); + + { + std::unique_lock 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 sub_status = sub_client->CreateSubscriber( + kChannel, SubOpts().SetSubscriberQueueSize(16).SetChecksum(true)); + if (sub_status.ok()) { + sub_status->SetChecksumCallback( + [](const std::array, 3> &, + absl::Span checksum) { + std::fill(checksum.begin(), checksum.end(), std::byte{0}); + }); + } + + { + std::lock_guard 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 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(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 diff --git a/client/publisher.cc b/client/publisher.cc index ee72140..528abfd 100644 --- a/client/publisher.cc +++ b/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,26 +678,26 @@ 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 failed_queues; ccb_->subscribers.TraverseSeqCst([this, slot, &failed_queues](int sub_id) { @@ -702,10 +705,21 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother( 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 diff --git a/client/subscriber.cc b/client/subscriber.cc index 8a8e84d..501d8a3 100644 --- a/client/subscriber.cc +++ b/client/subscriber.cc @@ -74,11 +74,11 @@ void SubscriberImpl::RemoveActiveMessage(MessageSlot *slot) { // std::cerr << this << " remove active message " << slot->id << " " // << slot->ordinal << " refs " << std::hex << slot->refs.load() << // std::dec << "\n"; - slot->sub_owners.Clear(subscriber_id_); - AtomicIncRefCount(slot, IsReliable(), -1, - slot->ordinal.load(std::memory_order_relaxed), - slot->vchan_id.load(std::memory_order_relaxed), true, - [this, slot]() { + if (slot->sub_owners.ClearWasSet(subscriber_id_)) { + AtomicIncRefCount(slot, IsReliable(), -1, + slot->ordinal.load(std::memory_order_relaxed), + slot->vchan_id.load(std::memory_order_relaxed), true, + [this, slot]() { // When a slot retires we want to use the slot id that was // originally used for the message. If the message came // in from a bridge we want to notify the original sender @@ -95,9 +95,10 @@ void SubscriberImpl::RemoveActiveMessage(MessageSlot *slot) { // slot->bridged_slot_id, slot->ordinal, // slot->vchan_id); // std::cerr << details; - TriggerRetirement( - slot->bridged_slot_id.load(std::memory_order_relaxed)); - }); + TriggerRetirement(slot->bridged_slot_id.load( + std::memory_order_relaxed)); + }); + } if (--num_active_messages_ < options_.MaxActiveMessages()) { Trigger(); if (IsReliable()) { @@ -110,15 +111,21 @@ void SubscriberImpl::PopulateActiveSlots(InPlaceAtomicBitset &bits) { uint64_t num_messages = 0; do { num_messages = ccb_->total_messages.load(std::memory_order_seq_cst); - bits.ClearAll(); for (int i = 0; i < NumSlots(); i++) { MessageSlot *s = &ccb_->slots[i]; uint64_t refs = s->refs.load(std::memory_order_acquire); + if ((refs & kPubOwned) != 0) { + // A publisher may already have installed this generation's delivery + // bit while still preparing queue entries and accounting. Preserve the + // bit until the publisher commits and advances total_messages. + continue; + } if (VirtualChannelIdMatch(s, vchan_id_) && - s->ordinal.load(std::memory_order_relaxed) != 0 && - (refs & kPubOwned) == 0) { + s->ordinal.load(std::memory_order_relaxed) != 0) { bits.Set(i); + } else { + bits.Clear(i); } } } while (num_messages != @@ -214,7 +221,9 @@ void SubscriberImpl::ClaimSlot(MessageSlot *slot, int vchan_id, } void SubscriberImpl::UnreadSlot(MessageSlot *slot) { - DecrementSlotRef(slot, false); + if (slot->sub_owners.ClearWasSet(subscriber_id_)) { + DecrementSlotRef(slot, false); + } // A queued hint has already been consumed by NextSlot(). If delivery is // rejected (for example at max_active_messages), the slot remains unread in // the authoritative bitset but is no longer present in the queue. Stay on @@ -283,9 +292,18 @@ SubscriberImpl::FindNextQueuedSlot(uint64_t max_queue_position) { if (!queue->TryPeek(queued)) { return std::nullopt; } - if (queued.slot_id < 0 || queued.slot_id >= NumSlots()) { - queue->DropFront(); - continue; + if (queued.slot_id >= 0 && queued.slot_id < NumSlots()) { + MessageSlot *peeked_slot = &ccb_->slots[queued.slot_id]; + const uint64_t peeked_ordinal = + peeked_slot->ordinal.load(std::memory_order_relaxed); + if ((peeked_slot->refs.load(std::memory_order_acquire) & kPubOwned) != + 0 && + peeked_ordinal == queued.ordinal) { + // The queue entry belongs to the generation currently being + // published. Leave it at the head until the publisher's release store + // makes the slot claimable. + return std::nullopt; + } } QueuedSlot popped; if (!queue->TryPop(popped)) { @@ -512,6 +530,10 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, if (!stable_poll_drain) { next_slot_cache_valid_ = false; } + // Record ownership before returning the pinned slot. If the process + // dies before ClientImpl can allocate and claim the ActiveMessage, the + // server must still be able to identify and release this reference. + new_slot->sub_owners.Set(subscriber_id_); return new_slot; } // Push() may fail after a peer dies or loses a bounded CAS race. The @@ -653,6 +675,9 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, // Successful claim. Advance the cursor so the next NextSlot() call // picks up the next ordinal in the cached, sorted list. ++next_slot_cursor_; + // Record ownership before returning the pinned slot. This closes the + // crash window between AtomicIncRefCount and ClaimSlot. + new_slot->slot->sub_owners.Set(subscriber_id_); return new_slot->slot; } // CAS failed: another subscriber raced us, or the slot was retired and @@ -725,6 +750,9 @@ MessageSlot *SubscriberImpl::LastSlot(MessageSlot *slot, bool reliable, new_slot->vchan_id, false); continue; } + // ReadNewest also returns a pinned slot through ClientImpl before + // ClaimSlot runs, so make that transient reference server-visible. + new_slot->slot->sub_owners.Set(subscriber_id_); return new_slot->slot; } newest_snapshot_.clear(); diff --git a/common/atomic_bitset.h b/common/atomic_bitset.h index e26f2c0..1335a5e 100644 --- a/common/atomic_bitset.h +++ b/common/atomic_bitset.h @@ -43,6 +43,21 @@ template class AtomicBitSet { bits_[word].fetch_or(1ULL << offset, std::memory_order_relaxed); } + void SetSeqCst(size_t bit) { + size_t word = bit / 64; + size_t offset = bit % 64; + bits_[word].fetch_or(1ULL << offset, std::memory_order_seq_cst); + } + + // Atomically set a bit and return whether it was previously clear. + bool SetWasClear(size_t bit) { + size_t word = bit / 64; + size_t offset = bit % 64; + uint64_t mask = 1ULL << offset; + uint64_t old = bits_[word].fetch_or(mask, std::memory_order_release); + return (old & mask) == 0; + } + void Clear(size_t bit) { size_t word = bit / 64; size_t offset = bit % 64; @@ -77,6 +92,12 @@ template class AtomicBitSet { return bits_[word].load(std::memory_order_relaxed) & (1ULL << offset); } + bool IsSetSeqCst(size_t bit) const { + size_t word = bit / 64; + size_t offset = bit % 64; + return bits_[word].load(std::memory_order_seq_cst) & (1ULL << offset); + } + void ClearAll() { for (size_t i = 0; i < BitsToWords(num_bits_); i++) { bits_[i].store(0, std::memory_order_relaxed); diff --git a/common/channel.cc b/common/channel.cc index 5e223d0..8bec312 100644 --- a/common/channel.cc +++ b/common/channel.cc @@ -261,9 +261,9 @@ bool Channel::AtomicIncRefCount(MessageSlot *slot, bool reliable, int inc, // "%d: AtomicIncRefCount: %s slot %d ordinal %d retired_refs: %d NumSubscribers: %d retire: %d\n", getpid(), Name(), slot->id, ordinal, retired_refs, NumSubscribers(ref_vchan_id), retire); // std::cerr << details; if (retire && new_refs == 0 && new_reliable_refs == 0 && - retired_refs >= NumSubscribers(ref_vchan_id)) { + retired_refs >= NumSubscribers(ref_vchan_id) && + RetiredSlots().SetWasClear(slot->id)) { // All subscribers have seen the slot, retire it. - RetiredSlots().Set(slot->id); if (retire_callback) { // std::cerr << "Calling retire callback for slot " << slot->id // << std::endl; @@ -361,9 +361,8 @@ uint64_t Channel::GetVirtualMemoryUsage() const { return size; } -void Channel::CleanupSlots( - int owner, bool reliable, bool is_pub, int vchan_id, - std::function retire_callback) { +void Channel::CleanupSlots(int owner, bool reliable, bool is_pub, + int vchan_id) { if (is_pub) { // Clear every slot owned by this publisher. Explicit multi-slot leases can // leave more than one slot publisher-owned when a process exits. @@ -388,28 +387,52 @@ void Channel::CleanupSlots( ccb_->subscribers.Clear(owner); ccb_->num_subs.RemoveSubscriber(vchan_id); - // Go through all the slots and remove the owner from the owners bitset. + InPlaceAtomicBitset &available = GetAvailableSlots(owner); for (int i = 0; i < NumSlots(); i++) { MessageSlot *slot = &ccb_->slots[i]; - if (slot->sub_owners.IsSet(owner)) { - slot->sub_owners.Clear(owner); - std::function notify_retirement; - if (retire_callback && - (slot->flags.load(std::memory_order_relaxed) & - kMessageIsActivation) == 0) { - const int32_t slot_id = - slot->bridged_slot_id.load(std::memory_order_relaxed); - notify_retirement = [&retire_callback, slot_id]() { - retire_callback(slot_id); - }; - } - AtomicIncRefCount(slot, reliable, -1, 0, 0, true, - std::move(notify_retirement)); + + // The available-slot bitset is the authoritative delivery record. A + // process may die before reading a published slot, in which case there + // is no sub_owners entry to clean up. Removing the subscriber changes + // the retirement threshold, so every unread slot must be re-evaluated. + available.ClearWasSet(i); + + if (slot->sub_owners.ClearWasSet(owner)) { + // The subscriber has already been removed from NumSubscribers above, + // so lowering the retirement threshold accounts for this owner. + AtomicIncRefCount(slot, reliable, -1, 0, 0, false); } } } } +bool Channel::TryRetireSlot(MessageSlot *slot) { + if (slot->ordinal.load(std::memory_order_relaxed) == 0) { + return false; + } + + const uint64_t refs = slot->refs.load(std::memory_order_acquire); + if ((refs & kPubOwned) != 0) { + return false; + } + + const uint64_t ref_count = refs & kRefCountMask; + const uint64_t reliable_ref_count = + (refs >> kReliableRefCountShift) & kRefCountMask; + const uint64_t retired_refs = + (refs >> kRetiredRefsShift) & kRetiredRefsMask; + int ref_vchan_id = (refs >> kVchanIdShift) & kVchanIdMask; + if (ref_vchan_id == kVchanIdMask) { + ref_vchan_id = -1; + } + + if (ref_count != 0 || reliable_ref_count != 0 || + retired_refs < static_cast(NumSubscribers(ref_vchan_id))) { + return false; + } + return RetiredSlots().SetWasClear(slot->id); +} + #if SUBSPACE_SHMEM_MODE == SUBSPACE_SHMEM_MODE_POSIX absl::StatusOr Channel::PosixSharedMemoryName(const std::string &shadow_file) { diff --git a/common/channel.h b/common/channel.h index d21d6e8..41715a7 100644 --- a/common/channel.h +++ b/common/channel.h @@ -134,7 +134,7 @@ constexpr int kDefaultSubscriberQueueSize = 16; constexpr uint64_t kDefaultSubscriberQueueArenaSize = 64'000; constexpr size_t kDefaultMaxAvailableSlotQueueCapacity = 1024; constexpr size_t kMaxSlotQueueCasAttempts = 64; -constexpr uint32_t kChannelControlBlockVersion = 4; +constexpr uint32_t kChannelControlBlockVersion = 5; constexpr size_t kMaxChannelControlBlockSize = 1ULL << 30; // This limits the number of virtual channels. Each virtual channel @@ -575,23 +575,119 @@ GetMetadataSpan(const MessagePrefix *prefix, int32_t checksum_size, // This counts the number of subscribers given a virtual channel id. class SubscriberCounter { public: - void AddSubscriber(int vchan_id) { num_subs_[vchan_id + 1]++; } + SubscriberCounter() { + sequence_.store(0, std::memory_order_relaxed); + ResetCounts(); + } - void RemoveSubscriber(int vchan_id) { num_subs_[vchan_id + 1]--; } + SubscriberCounter(const SubscriberCounter &other) { + sequence_.store(0, std::memory_order_relaxed); + for (size_t i = 0; i < num_subs_.size(); ++i) { + num_subs_[i].store( + other.num_subs_[i].load(std::memory_order_acquire), + std::memory_order_relaxed); + } + } + + SubscriberCounter &operator=(const SubscriberCounter &other) { + if (this == &other) { + return *this; + } + const uint64_t sequence = + sequence_.load(std::memory_order_relaxed) & ~uint64_t{1}; + sequence_.store(sequence + 1, std::memory_order_release); + for (size_t i = 0; i < num_subs_.size(); ++i) { + num_subs_[i].store( + other.num_subs_[i].load(std::memory_order_acquire), + std::memory_order_relaxed); + } + sequence_.store(sequence + 2, std::memory_order_release); + return *this; + } + + void Reset() { + BeginWrite(); + ResetCounts(); + EndWrite(); + } + + void AddSubscriber(int vchan_id) { + BeginWrite(); + num_subs_[vchan_id + 1].fetch_add(1, std::memory_order_relaxed); + EndWrite(); + } + + void RemoveSubscriber(int vchan_id) { + BeginWrite(); + num_subs_[vchan_id + 1].fetch_sub(1, std::memory_order_relaxed); + EndWrite(); + } // If vchan_id is valid we also count the number of subscribers to the // multiplexer itself. int NumSubscribers(int vchan_id) { - int n = num_subs_[0]; - if (vchan_id == -1) { - return n; + for (;;) { + const uint64_t before = sequence_.load(std::memory_order_acquire); + if ((before & 1) != 0) { + // A server process can die between BeginWrite and EndWrite. Treat an + // interrupted update conservatively instead of making surviving + // clients spin forever; the recovered server rebuilds the counter. + return kMaxSlotOwners; + } + const int mux_count = num_subs_[0].load(std::memory_order_relaxed); + const int count = + vchan_id == -1 + ? mux_count + : mux_count + + num_subs_[vchan_id + 1].load(std::memory_order_relaxed); + if (sequence_.load(std::memory_order_acquire) == before) { + return count; + } } - return n + num_subs_[vchan_id + 1]; } private: + void BeginWrite() { + sequence_.fetch_add(1, std::memory_order_acq_rel); + } + + void EndWrite() { sequence_.fetch_add(1, std::memory_order_release); } + + void ResetCounts() { + for (auto &count : num_subs_) { + count.store(0, std::memory_order_relaxed); + } + } + // Vchan ID -1 means invalid vchan ID so we just use element 0 for that. - std::array num_subs_ = {}; + std::atomic sequence_; + std::array, kMaxVchanId + 1> num_subs_; +}; + +class SubscriberCleanupGeneration { +public: + SubscriberCleanupGeneration() { + for (auto &generation : generations_) { + generation.store(0, std::memory_order_relaxed); + } + } + + void Increment(int vchan_id) { + generations_[vchan_id + 1].fetch_add(1, std::memory_order_release); + } + + uint64_t Get(int vchan_id) const { + const uint64_t mux_generation = + generations_[0].load(std::memory_order_acquire); + if (vchan_id == -1) { + return mux_generation; + } + return mux_generation + + generations_[vchan_id + 1].load(std::memory_order_acquire); + } + +private: + std::array, kMaxVchanId + 1> generations_; }; class OrdinalAccumulator { @@ -641,6 +737,10 @@ struct ChannelControlBlock { // a.k.a CCB std::array sub_vchan_ids; SubscriberCounter num_subs; + // Incremented by the server after removing a subscriber and before its + // all-slot retirement scan. Publishers use this to determine whether + // subscriber cleanup raced their publication commit. + SubscriberCleanupGeneration subscriber_cleanup_generation; // Statistics counters. std::atomic total_bytes; @@ -837,18 +937,25 @@ class Channel : public std::enable_shared_from_this { void RegisterSubscriber(int sub_id, int vchan_id, bool is_new) { ccb_->sub_vchan_ids[sub_id] = vchan_id; - if (is_new && !IsPlaceholder()) { - GetAvailableSlots(sub_id).ClearAll(); - } - ccb_->subscribers.Set(sub_id); - if (is_new && !IsPlaceholder()) { - SeedAvailableSlotQueue(sub_id, vchan_id); - } + const bool was_registered = ccb_->subscribers.IsSet(sub_id); + const bool register_membership = is_new || !was_registered; SubscriberCounter num_subs; ccb_->subscribers.Traverse([this, &num_subs](size_t id) { num_subs.AddSubscriber(ccb_->sub_vchan_ids[id]); }); + if (register_membership && !was_registered) { + num_subs.AddSubscriber(vchan_id); + if (!IsPlaceholder()) { + GetAvailableSlots(sub_id).ClearAll(); + } + // Publish the increased retirement threshold before making the + // subscriber visible to publishers. + } ccb_->num_subs = num_subs; + ccb_->subscribers.SetSeqCst(sub_id); + if (register_membership && !was_registered && !IsPlaceholder()) { + SeedAvailableSlotQueue(sub_id, vchan_id); + } } int GetSubVchanId(int32_t i) const { return ccb_->sub_vchan_ids[i]; } @@ -856,11 +963,12 @@ class Channel : public std::enable_shared_from_this { void SeedAvailableSlotQueue(int sub_id, int vchan_id) { InPlaceAtomicBitset &bits = GetAvailableSlots(sub_id); InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(sub_id); + // Include a non-zero publisher-owned generation. This closes the inverse + // registration race: a publisher may have snapshotted subscribers before + // this subscriber is registered, while the registration scan happens + // before that publisher commits. Consumers preserve the bit/queue entry + // until kPubOwned is cleared. auto visible = [vchan_id](MessageSlot &slot) { - const uint64_t refs = slot.refs.load(std::memory_order_acquire); - if ((refs & kPubOwned) != 0) { - return false; - } const uint64_t ordinal = slot.ordinal.load(std::memory_order_relaxed); const int buffer_index = slot.buffer_index.load(std::memory_order_relaxed); @@ -961,14 +1069,18 @@ class Channel : public std::enable_shared_from_this { } std::string SlotType() const { return type_; } - void CleanupSlots( - int owner, bool reliable, bool is_pub, int vchan_id, - std::function retire_callback = {}); + void CleanupSlots(int owner, bool reliable, bool is_pub, int vchan_id); + + bool TryRetireSlot(MessageSlot *slot); int NumSubscribers(int vchan_id) const { return ccb_->num_subs.NumSubscribers(vchan_id); } + uint64_t SubscriberCleanupGenerationFor(int vchan_id) const { + return ccb_->subscriber_cleanup_generation.Get(vchan_id); + } + int GetChannelId() const { return channel_id_; } int NumUpdates() const { return num_updates_; } diff --git a/docs/server-architecture.md b/docs/server-architecture.md index 0737fd0..bb9b0d8 100644 --- a/docs/server-architecture.md +++ b/docs/server-architecture.md @@ -74,7 +74,8 @@ Each channel requires three shared memory regions, created via `shm_open()` (POS - One per channel. - Contains: channel name, num_slots, ordinals, activation tracker. -- CCB version 4 uses atomic slot metadata. `total_messages` advances for every +- CCB version 5 uses atomic slot metadata and subscriber counters. + `total_messages` advances for every completed publication, including activation messages, and also versions subscriber delivery snapshots. - Variable-length: `MessageSlot` array, retired/free/available bitsets, a @@ -93,6 +94,16 @@ Each channel requires three shared memory regions, created via `shm_open()` (POS active. Shadow recovery reconciles subscriber offsets with allocated blocks, conservatively retires orphan blocks, and only reclaims them after their recorded publisher hazards have quiesced. +- A slot remains publisher-owned while its available-slot bits and queue hints + are prepared. Subscribers preserve those records but cannot claim the slot + until the publisher commits it with a release store and advances + `total_messages`. Subscriber registration also seeds non-zero + publisher-owned generations so a join racing publication cannot miss the + message. +- Subscriber removal does not wait for in-progress publishers. It releases the + dead subscriber's references and re-evaluates every slot. Publication commit + performs the same retirement check when a cleanup generation changed, so the + operation that finishes second safely completes retirement. ### Buffer Control Block (BCB) diff --git a/rust_client/src/bitset.rs b/rust_client/src/bitset.rs index 098d6ea..086113b 100644 --- a/rust_client/src/bitset.rs +++ b/rust_client/src/bitset.rs @@ -33,18 +33,38 @@ impl AtomicBitSet { self.bits[word].fetch_or(1u64 << offset, Ordering::Relaxed); } + pub fn set_seq_cst(&self, bit: usize) { + let word = bit / 64; + let offset = bit % 64; + self.bits[word].fetch_or(1u64 << offset, Ordering::SeqCst); + } + pub fn clear(&self, bit: usize) { let word = bit / 64; let offset = bit % 64; self.bits[word].fetch_and(!(1u64 << offset), Ordering::Relaxed); } + pub fn clear_was_set(&self, bit: usize) -> bool { + let word = bit / 64; + let offset = bit % 64; + self.bits[word].fetch_and(!(1u64 << offset), Ordering::Acquire) + & (1u64 << offset) + != 0 + } + pub fn is_set(&self, bit: usize) -> bool { let word = bit / 64; let offset = bit % 64; self.bits[word].load(Ordering::Relaxed) & (1u64 << offset) != 0 } + pub fn is_set_seq_cst(&self, bit: usize) -> bool { + let word = bit / 64; + let offset = bit % 64; + self.bits[word].load(Ordering::SeqCst) & (1u64 << offset) != 0 + } + pub fn clear_all(&self) { for w in &self.bits { w.store(0, Ordering::Relaxed); @@ -157,6 +177,15 @@ impl InPlaceAtomicBitSet { .fetch_or(1u64 << offset, Ordering::Relaxed); } + pub fn set_was_clear(&self, bit: usize) -> bool { + let word_idx = bit / 64; + let offset = bit % 64; + self.word(word_idx) + .fetch_or(1u64 << offset, Ordering::Release) + & (1u64 << offset) + == 0 + } + pub fn clear(&self, bit: usize) { let word_idx = bit / 64; let offset = bit % 64; @@ -164,6 +193,15 @@ impl InPlaceAtomicBitSet { .fetch_and(!(1u64 << offset), Ordering::Relaxed); } + pub fn clear_was_set(&self, bit: usize) -> bool { + let word_idx = bit / 64; + let offset = bit % 64; + self.word(word_idx) + .fetch_and(!(1u64 << offset), Ordering::Acquire) + & (1u64 << offset) + != 0 + } + pub fn is_set(&self, bit: usize) -> bool { let word_idx = bit / 64; let offset = bit % 64; diff --git a/rust_client/src/channel.rs b/rust_client/src/channel.rs index 7663c67..8f0a320 100644 --- a/rust_client/src/channel.rs +++ b/rust_client/src/channel.rs @@ -32,7 +32,7 @@ pub const MAX_CHANNELS: usize = 1024; pub const MAX_SLOT_OWNERS: usize = 1024; pub const MAX_AVAILABLE_SLOT_QUEUE_CAPACITY: usize = 1024; const MAX_SLOT_QUEUE_CAS_ATTEMPTS: usize = 64; -pub const CHANNEL_CONTROL_BLOCK_VERSION: u32 = 4; +pub const CHANNEL_CONTROL_BLOCK_VERSION: u32 = 5; pub const MAX_VCHAN_ID: usize = 1023; pub const MAX_CHANNEL_NAME: usize = 64; pub const MAX_BUFFERS: usize = 1024; @@ -376,6 +376,23 @@ impl SlotQueueHeader { self.insertion_failed.store(true, Ordering::Release); } + pub fn try_peek(&self) -> Option<(i32, u64)> { + if self.capacity == 0 { + return None; + } + + let head = self.head.load(Ordering::Acquire); + let entry = + unsafe { &*self.entries().add((head % self.capacity as u64) as usize) }; + if entry.sequence.load(Ordering::Acquire) != head + 1 { + return None; + } + Some(( + entry.slot_id.load(Ordering::Relaxed), + entry.ordinal.load(Ordering::Relaxed), + )) + } + pub fn try_pop(&self) -> Option<(i32, u64)> { if self.capacity == 0 { return None; @@ -498,24 +515,71 @@ impl ActivationTracker { #[repr(C)] pub struct SubscriberCounter { - num_subs: [i32; MAX_VCHAN_ID + 1], + sequence: AtomicU64, + num_subs: [AtomicI32; MAX_VCHAN_ID + 1], } impl SubscriberCounter { - pub fn add_subscriber(&mut self, vchan_id: i32) { - self.num_subs[(vchan_id + 1) as usize] += 1; + fn replace(&self, counts: &[i32; MAX_VCHAN_ID + 1]) { + let sequence = self.sequence.load(Ordering::Relaxed) & !1; + self.sequence.store(sequence + 1, Ordering::Release); + for (count, value) in self.num_subs.iter().zip(counts.iter()) { + count.store(*value, Ordering::Relaxed); + } + self.sequence.store(sequence + 2, Ordering::Release); } - pub fn remove_subscriber(&mut self, vchan_id: i32) { - self.num_subs[(vchan_id + 1) as usize] -= 1; + pub fn add_subscriber(&self, vchan_id: i32) { + self.sequence.fetch_add(1, Ordering::AcqRel); + self.num_subs[(vchan_id + 1) as usize].fetch_add(1, Ordering::Relaxed); + self.sequence.fetch_add(1, Ordering::Release); + } + + pub fn remove_subscriber(&self, vchan_id: i32) { + self.sequence.fetch_add(1, Ordering::AcqRel); + self.num_subs[(vchan_id + 1) as usize].fetch_sub(1, Ordering::Relaxed); + self.sequence.fetch_add(1, Ordering::Release); } pub fn num_subscribers(&self, vchan_id: i32) -> i32 { - let n = self.num_subs[0]; + loop { + let before = self.sequence.load(Ordering::Acquire); + if (before & 1) != 0 { + // A server may die during a write. Conservatively prevent + // retirement until recovery rebuilds the counter. + return MAX_SLOT_OWNERS as i32; + } + let mux_count = self.num_subs[0].load(Ordering::Relaxed); + let count = if vchan_id == -1 { + mux_count + } else { + mux_count + + self.num_subs[(vchan_id + 1) as usize].load(Ordering::Relaxed) + }; + if self.sequence.load(Ordering::Acquire) == before { + return count; + } + } + } +} + +#[repr(C)] +pub struct SubscriberCleanupGeneration { + generations: [AtomicU64; MAX_VCHAN_ID + 1], +} + +impl SubscriberCleanupGeneration { + pub fn increment(&self, vchan_id: i32) { + self.generations[(vchan_id + 1) as usize].fetch_add(1, Ordering::Release); + } + + pub fn get(&self, vchan_id: i32) -> u64 { + let mux_generation = self.generations[0].load(Ordering::Acquire); if vchan_id == -1 { - n + mux_generation } else { - n + self.num_subs[(vchan_id + 1) as usize] + mux_generation + + self.generations[(vchan_id + 1) as usize].load(Ordering::Acquire) } } } @@ -538,6 +602,7 @@ pub struct ChannelControlBlock { pub sub_vchan_ids: [i16; MAX_SLOT_OWNERS], pub num_subs: SubscriberCounter, + pub subscriber_cleanup_generation: SubscriberCleanupGeneration, pub total_bytes: AtomicU64, pub total_messages: AtomicU64, @@ -1011,14 +1076,21 @@ impl Channel { pub fn register_subscriber(&self, sub_id: usize, vchan_id: i32, is_new: bool) { let ccb = self.ccb(); - ccb.subscribers.set(sub_id); + let was_registered = ccb.subscribers.is_set(sub_id); + let register_membership = is_new || !was_registered; unsafe { let ccb_mut = &mut *self.ccb; ccb_mut.sub_vchan_ids[sub_id] = vchan_id as i16; - if is_new { - ccb_mut.num_subs.add_subscriber(vchan_id); - } } + let mut counts = [0i32; MAX_VCHAN_ID + 1]; + ccb.subscribers.traverse(|id| { + counts[(ccb.sub_vchan_ids[id] + 1) as usize] += 1; + }); + if register_membership && !was_registered { + counts[(vchan_id + 1) as usize] += 1; + } + ccb.num_subs.replace(&counts); + ccb.subscribers.set_seq_cst(sub_id); } /// Atomically increment/decrement the ref count on a slot. @@ -1088,8 +1160,8 @@ impl Channel { && new_refs == 0 && new_reliable_refs == 0 && retired_refs >= self.num_subscribers(ref_vchan_id) + && self.retired_slots().set_was_clear(slot.id as usize) { - self.retired_slots().set(slot.id as usize); if let Some(cb) = retire_callback { cb(); } @@ -1099,6 +1171,36 @@ impl Channel { } } + pub fn try_retire_slot(&self, slot_idx: usize) -> bool { + let slot = self.slot_ref(slot_idx); + if slot.ordinal() == 0 { + return false; + } + + let refs = slot.refs.load(Ordering::Acquire); + if (refs & PUB_OWNED) != 0 { + return false; + } + + let ref_count = refs & REF_COUNT_MASK; + let reliable_ref_count = (refs >> RELIABLE_REF_COUNT_SHIFT) & REF_COUNT_MASK; + let retired_refs = (refs >> RETIRED_REFS_SHIFT) & RETIRED_REFS_MASK; + let encoded_vchan_id = (refs >> VCHAN_ID_SHIFT) & VCHAN_ID_MASK; + let ref_vchan_id = if encoded_vchan_id == VCHAN_ID_MASK { + -1 + } else { + encoded_vchan_id as i32 + }; + + if ref_count != 0 + || reliable_ref_count != 0 + || retired_refs < self.num_subscribers(ref_vchan_id) as u64 + { + return false; + } + self.retired_slots().set_was_clear(slot_idx) + } + /// Get the buffer address for a slot, accounting for prefix. pub fn get_buffer_address(&self, slot_idx: usize) -> *mut u8 { let slot = self.slot_ref(slot_idx); @@ -1241,16 +1343,25 @@ impl Channel { } else { let ccb = self.ccb(); ccb.subscribers.clear(owner as usize); - unsafe { - (*self.ccb).num_subs.remove_subscriber(vchan_id); - } + ccb.num_subs.remove_subscriber(vchan_id); for i in 0..self.num_slots as usize { let slot = self.slot_ref(i); - if slot.sub_owners.is_set(owner as usize) { - slot.sub_owners.clear(owner as usize); - self.atomic_inc_ref_count::(i, reliable, -1, 0, 0, true, None); + self.get_available_slots(owner as usize).clear_was_set(i); + if slot.sub_owners.clear_was_set(owner as usize) { + self.atomic_inc_ref_count::(i, reliable, -1, 0, 0, false, None); + } + } + ccb.subscriber_cleanup_generation.increment(vchan_id); + for i in 0..self.num_slots as usize { + let slot = self.slot_ref(i); + if (slot.flags() & MESSAGE_IS_ACTIVATION) != 0 { + continue; + } + if vchan_id != -1 && i32::from(slot.vchan_id()) != vchan_id { + continue; } + self.try_retire_slot(i); } } } diff --git a/rust_client/src/publisher.rs b/rust_client/src/publisher.rs index 8f337e3..f515fcf 100644 --- a/rust_client/src/publisher.rs +++ b/rust_client/src/publisher.rs @@ -281,8 +281,10 @@ impl PublisherImpl { } pub fn retire_published_slot_immediately(&self, slot_idx: usize) { - self.channel.retired_slots().set(slot_idx); - self.trigger_retirement(slot_idx); + let retirement_slot_id = self.channel.slot_ref(slot_idx).bridged_slot_id(); + if self.channel.try_retire_slot(slot_idx) { + self.trigger_retirement(retirement_slot_id as usize); + } } pub fn find_free_slot_unreliable( @@ -670,14 +672,22 @@ impl PublisherImpl { } } - // Release the slot: store refs with ordinal, no PUB_OWNED. let ordinal = slot.ordinal(); - slot.refs.store( - build_refs_bit_field(ordinal, vchan_id, 0), - Ordering::Release, - ); + let retirement_slot_id = slot.bridged_slot_id(); + let (return_ordinal, return_timestamp) = if !prefix.is_null() { + unsafe { ((*prefix).ordinal, (*prefix).timestamp) } + } else { + (0, 0) + }; + let cleanup_generation = self + .channel + .ccb() + .subscriber_cleanup_generation + .get(vchan_id); - // Tell all subscribers the slot is available. + // Prepare every delivery record and statistic while the slot remains + // PUB_OWNED. Subscribers preserve the bit/queue entry but cannot claim + // the slot, and server cleanup cannot retire it before commit. let ccb = self.channel.ccb(); { let _publish_guard = SubscriberQueuePublishGuard::new( @@ -693,7 +703,19 @@ impl PublisherImpl { { return; } - self.channel.get_available_slots(sub_id).set(slot_idx); + let available = self.channel.get_available_slots(sub_id); + available.set(slot_idx); + if !ccb.subscribers.is_set_seq_cst(sub_id) { + available.clear(slot_idx); + // Registration publishes membership before seeding the + // in-progress generation. Restore a bit cleared for the + // previous occupant if this subscriber ID was reused. + if !ccb.subscribers.is_set_seq_cst(sub_id) { + return; + } + available.set(slot_idx); + } + let queue = self.channel.get_available_slot_queue(sub_id); if let Some(queue) = queue { if !queue.push( @@ -705,52 +727,62 @@ impl PublisherImpl { } } }); - ccb.total_messages.fetch_add(1, Ordering::SeqCst); + + if !is_activation { + let message_size = slot.message_size(); + ccb.total_bytes.fetch_add(message_size, Ordering::Relaxed); + let msg_size = message_size as u32; + let mut old_max = ccb.max_message_size.load(Ordering::Relaxed); + while msg_size > old_max { + match ccb.max_message_size.compare_exchange_weak( + old_max, + msg_size, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(v) => old_max = v, + } + } + } + for queue in failed_queues { unsafe { (&*queue).mark_insertion_failure() }; } - } - if !is_activation { - let message_size = slot.message_size(); - self.channel - .ccb() - .total_bytes - .fetch_add(message_size, Ordering::Relaxed); - let msg_size = message_size as u32; - let mut old_max = self.channel.ccb().max_message_size.load(Ordering::Relaxed); - while msg_size > old_max { - match self.channel.ccb().max_message_size.compare_exchange_weak( - old_max, - msg_size, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(v) => old_max = v, - } + // Commit only after all delivery state is ready. + slot.refs.store( + build_refs_bit_field(ordinal, vchan_id, 0), + Ordering::Release, + ); + ccb.total_messages.fetch_add(1, Ordering::SeqCst); + + // Whichever happens second, subscriber cleanup or this commit, + // completes retirement using the current subscriber count. + if !is_activation + && ccb + .subscriber_cleanup_generation + .get(vchan_id) + != cleanup_generation + && self.channel.try_retire_slot(slot_idx) + { + self.trigger_retirement(retirement_slot_id as usize); } } - let (ordinal, timestamp) = if !prefix.is_null() { - unsafe { ((*prefix).ordinal, (*prefix).timestamp) } - } else { - (0, 0) - }; - if !acquire_next { return PublishedMessage { new_slot: None, - ordinal, - timestamp, + ordinal: return_ordinal, + timestamp: return_timestamp, }; } if reliable { return PublishedMessage { new_slot: None, - ordinal, - timestamp, + ordinal: return_ordinal, + timestamp: return_timestamp, }; } @@ -758,8 +790,8 @@ impl PublisherImpl { PublishedMessage { new_slot, - ordinal, - timestamp, + ordinal: return_ordinal, + timestamp: return_timestamp, } } diff --git a/rust_client/src/subscriber.rs b/rust_client/src/subscriber.rs index 71f9121..9745df0 100644 --- a/rust_client/src/subscriber.rs +++ b/rust_client/src/subscriber.rs @@ -280,23 +280,27 @@ impl SubscriberImpl { pub fn remove_active_message(&self, slot_idx: usize) { let slot = self.channel.slot_ref(slot_idx); - slot.sub_owners.clear(self.subscriber_id as usize); - let ordinal = slot.ordinal(); - let vchan_id = slot.vchan_id() as i32; - let bridged_slot_id = slot.bridged_slot_id(); let reliable = self.options.reliable; - self.channel.atomic_inc_ref_count( - slot_idx, - reliable, - -1, - ordinal, - vchan_id, - true, - Some(|| { - self.trigger_retirement(bridged_slot_id as usize); - }), - ); + if slot + .sub_owners + .clear_was_set(self.subscriber_id as usize) + { + let ordinal = slot.ordinal(); + let vchan_id = slot.vchan_id() as i32; + let bridged_slot_id = slot.bridged_slot_id(); + self.channel.atomic_inc_ref_count( + slot_idx, + reliable, + -1, + ordinal, + vchan_id, + true, + Some(|| { + self.trigger_retirement(bridged_slot_id as usize); + }), + ); + } let new_count = self.num_active_messages.fetch_sub(1, Ordering::Relaxed) - 1; if new_count < self.options.max_active_messages { self.trigger(); @@ -313,16 +317,22 @@ impl SubscriberImpl { pub fn populate_active_slots(&self, bits: &crate::bitset::InPlaceAtomicBitSet) { loop { let total = self.total_messages(); - bits.clear_all(); for i in 0..self.channel.num_slots as usize { let s = self.channel.slot_ref(i); let refs = s.refs.load(Ordering::Acquire); + if (refs & PUB_OWNED) != 0 { + // Preserve a delivery bit installed by an in-progress + // publisher. The total_messages commit will force a fresh + // snapshot after PUB_OWNED is cleared. + continue; + } if virtual_channel_id_match(s.vchan_id(), self.channel.vchan_id) && s.ordinal() != 0 - && (refs & PUB_OWNED) == 0 { bits.set(i); + } else { + bits.clear(i); } } @@ -428,16 +438,28 @@ impl SubscriberImpl { return None; } loop { - let queue_at_boundary = match self + let queue = match self .channel .get_available_slot_queue(self.subscriber_id as usize) { - Some(queue) => queue.head() >= max_queue_position, - None => true, + Some(queue) => queue, + None => break, }; - if queue_at_boundary { + if queue.head() >= max_queue_position { break; } + if let Some((slot_id, ordinal)) = queue.try_peek() { + if slot_id >= 0 && (slot_id as usize) < self.channel.num_slots as usize { + let slot = self.channel.slot_ref(slot_id as usize); + if (slot.refs.load(Ordering::Acquire) & PUB_OWNED) != 0 + && slot.ordinal() == ordinal + { + // Leave the current generation at the queue head until + // the publisher commits it. + break; + } + } + } let Some((slot_id, ordinal)) = self .channel .get_available_slot_queue(self.subscriber_id as usize) @@ -594,6 +616,10 @@ impl SubscriberImpl { .saturating_add(concurrent_drops as i32); } } + self.channel + .slot_ref(slot_idx) + .sub_owners + .set(self.subscriber_id as usize); return Some(slot_idx); } } @@ -753,6 +779,10 @@ impl SubscriberImpl { ); continue; } + self.channel + .slot_ref(active.slot_index) + .sub_owners + .set(self.subscriber_id as usize); return Some(active.slot_index); } } @@ -856,6 +886,10 @@ impl SubscriberImpl { ); continue; } + self.channel + .slot_ref(active.slot_index) + .sub_owners + .set(self.subscriber_id as usize); return Some(active.slot_index); } } @@ -918,7 +952,14 @@ impl SubscriberImpl { } pub fn unread_slot(&mut self, slot_idx: usize, ordinal: u64, vchan_id: i32) { - self.decrement_slot_ref(slot_idx, ordinal, vchan_id, false); + if self + .channel + .slot_ref(slot_idx) + .sub_owners + .clear_was_set(self.subscriber_id as usize) + { + self.decrement_slot_ref(slot_idx, ordinal, vchan_id, false); + } if self .channel .get_available_slot_queue(self.subscriber_id as usize) diff --git a/rust_client/tests/client_test.rs b/rust_client/tests/client_test.rs index 231dfca..5a29f19 100644 --- a/rust_client/tests/client_test.rs +++ b/rust_client/tests/client_test.rs @@ -1727,6 +1727,75 @@ fn integration_custom_checksum_callback() { assert!(msg2.checksum_error); } +#[test] +fn integration_subscriber_joining_during_publish_receives_message() { + use std::sync::{Arc, Condvar, Mutex}; + + let pub_client = new_client("test_join_publish_p"); + let sub_client = new_client("test_join_publish_s"); + let pub_opts = PublisherOptions::new() + .set_slot_size(256) + .set_num_slots(10) + .set_checksum(true) + .set_subscriber_queue_arena_size(DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE); + let publisher = pub_client + .create_publisher("rust_join_during_publish", &pub_opts) + .unwrap(); + + let state = Arc::new((Mutex::new((false, false)), Condvar::new())); + let callback_state = Arc::clone(&state); + publisher.set_checksum_callback(move |spans: &[&[u8]], checksum: &mut [u8]| { + let (lock, cv) = &*callback_state; + let mut state = lock.lock().unwrap(); + state.0 = true; + cv.notify_all(); + while !state.1 { + state = cv.wait(state).unwrap(); + } + drop(state); + calculate_crc32_checksum(spans, checksum); + }); + + let publisher_thread = publisher.clone(); + let publish_thread = std::thread::spawn(move || { + let payload = b"joined"; + let (buffer, _) = publisher_thread + .get_message_buffer(payload.len() as i32) + .unwrap() + .unwrap(); + unsafe { + std::ptr::copy_nonoverlapping(payload.as_ptr(), buffer, payload.len()); + } + publisher_thread + .publish_message(payload.len() as i64) + .unwrap(); + }); + + let (lock, cv) = &*state; + let mut state_guard = lock.lock().unwrap(); + while !state_guard.0 { + state_guard = cv.wait(state_guard).unwrap(); + } + drop(state_guard); + + let sub_opts = SubscriberOptions::new() + .set_subscriber_queue_size(16) + .set_checksum(true); + let subscriber = sub_client + .create_subscriber("rust_join_during_publish", &sub_opts) + .unwrap(); + + let mut state_guard = lock.lock().unwrap(); + state_guard.1 = true; + cv.notify_all(); + drop(state_guard); + publish_thread.join().unwrap(); + + let message = subscriber.read_message(ReadMode::ReadNext).unwrap(); + assert_eq!(message.length, 6); + assert_eq!(unsafe { message.as_slice() }, b"joined"); +} + // ── Checksum + metadata tests ──────────────────────────────────────────────── #[test] diff --git a/server/server_channel.cc b/server/server_channel.cc index 7eb8548..2d4903e 100644 --- a/server/server_channel.cc +++ b/server/server_channel.cc @@ -338,7 +338,8 @@ ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, return p.status(); } ccb_ = reinterpret_cast(*p); - ccb_->num_subs = SubscriberCounter(); + new (&ccb_->num_subs) SubscriberCounter(); + new (&ccb_->subscriber_cleanup_generation) SubscriberCleanupGeneration(); // Create buffer control block. p = CreateSharedMemory(channel_id_, "bcb", sizeof(BufferControlBlock), @@ -552,6 +553,31 @@ std::vector ServerChannel::GetRetirementFds() const { } return r; } + +void ServerChannel::NotifyPublisherRetirement(int32_t slot_id) { + for (auto &[id, user] : users_) { + if (user == nullptr || !user->IsPublisher()) { + continue; + } + auto &fd = + static_cast(user.get())->GetRetirementFdWriter(); + if (!fd.Valid()) { + continue; + } + absl::StatusOr written = fd.Write(&slot_id, sizeof(slot_id)); + if (!written.ok()) { + logger_.Log(toolbelt::LogLevel::kError, + "Failed to trigger retirement for slot %d: %s", slot_id, + written.status().ToString().c_str()); + } else if (*written != sizeof(slot_id)) { + logger_.Log(toolbelt::LogLevel::kError, + "Failed to trigger retirement for slot %d: wrote %zd " + "bytes, expected %zu bytes", + slot_id, *written, sizeof(slot_id)); + } + } +} + // User ids are allocated from the multiplexer as all virtual channels // on the mux share the same CCB. absl::StatusOr ServerChannel::AllocateUserId(const char *type) { @@ -931,32 +957,32 @@ void ServerChannel::CleanupSlots(int owner, bool reliable, bool is_pub, } ccb_->subscribers.ClearSeqCst(owner); - Channel::CleanupSlots( - owner, reliable, is_pub, vchan_id, [this](int32_t slot_id) { - for (auto &[id, user] : users_) { - if (user == nullptr || !user->IsPublisher()) { - continue; - } - auto &fd = - static_cast(user.get())->GetRetirementFdWriter(); - if (!fd.Valid()) { - continue; - } - absl::StatusOr written = - fd.Write(&slot_id, sizeof(slot_id)); - if (!written.ok()) { - logger_.Log(toolbelt::LogLevel::kError, - "Failed to trigger retirement for slot %d: %s", slot_id, - written.status().ToString().c_str()); - } else if (*written != sizeof(slot_id)) { - logger_.Log(toolbelt::LogLevel::kError, - "Failed to trigger retirement for slot %d: wrote %zd " - "bytes, expected %zu bytes", - slot_id, *written, sizeof(slot_id)); - } - } - }); RetireSubscriberQueue(owner); + + Channel::CleanupSlots(owner, reliable, is_pub, vchan_id); + ccb_->subscriber_cleanup_generation.Increment(vchan_id); + + // Re-evaluate every published slot after lowering the subscriber count. + // Publication commit performs the same check. If this scan encounters a + // publisher-owned slot it leaves it alone; the publisher's later commit + // observes the new count and completes retirement. If commit happened first, + // this scan completes retirement. + for (int i = 0; i < NumSlots(); ++i) { + MessageSlot *slot = &ccb_->slots[i]; + if ((slot->flags.load(std::memory_order_relaxed) & + kMessageIsActivation) != 0) { + continue; + } + if (vchan_id != -1 && + slot->vchan_id.load(std::memory_order_relaxed) != vchan_id) { + continue; + } + + if (TryRetireSlot(slot)) { + NotifyPublisherRetirement( + slot->bridged_slot_id.load(std::memory_order_relaxed)); + } + } } std::vector ServerChannel::RegisterExistingSubscribers() { @@ -1004,6 +1030,13 @@ std::vector ChannelMultiplexer::RegisterExistingSubscribers() { return warnings; } +void ChannelMultiplexer::NotifyPublisherRetirement(int32_t slot_id) { + ServerChannel::NotifyPublisherRetirement(slot_id); + for (VirtualChannel *vchan : virtual_channels_) { + vchan->NotifyPublisherRetirement(slot_id); + } +} + void ServerChannel::TriggerAllSubscribers() { for (auto &[id, user] : users_) { if (user == nullptr) { diff --git a/server/server_channel.h b/server/server_channel.h index 60b5311..fd0f318 100644 --- a/server/server_channel.h +++ b/server/server_channel.h @@ -283,6 +283,7 @@ class ServerChannel : public Channel { std::vector GetReliablePublisherTriggerFds() const; std::vector GetRetirementFds() const; + virtual void NotifyPublisherRetirement(int32_t slot_id); // Translate a user id into a User pointer. The pointer ownership // is kept by the ServerChannel. @@ -541,6 +542,7 @@ class ChannelMultiplexer : public ServerChannel { bool IsMux() const override { return true; } bool HasPublisherOwnedBy(const ClientHandler *handler) const override; std::vector RegisterExistingSubscribers() override; + void NotifyPublisherRetirement(int32_t slot_id) override; bool IsEmpty() const override { return virtual_channels_.empty() && ServerChannel::IsEmpty(); } diff --git a/shadow/shadow_test.cc b/shadow/shadow_test.cc index 67d8507..a867a7f 100644 --- a/shadow/shadow_test.cc +++ b/shadow/shadow_test.cc @@ -959,8 +959,18 @@ TEST_F(ShadowRecoveryTest, RecoversMuxSubscriberQueueTopology) { ASSERT_TRUE(WaitForShadowState([this]() { return shadow_->WithChannels([](auto &channels) { - return channels.contains("/queue_recovery/*") && - channels.contains("/queue_recovery/0"); + auto mux = channels.find("/queue_recovery/*"); + auto vchan = channels.find("/queue_recovery/0"); + if (mux == channels.end() || vchan == channels.end() || + !mux->second.has_max_subscribers || + mux->second.max_subscribers != 2 || + mux->second.subscribers.size() != 1 || + vchan->second.publishers.size() != 1) { + return false; + } + return mux->second.subscribers.begin()->second.subscriber_queue_size == 4 && + vchan->second.publishers.begin() + ->second.max_outstanding_slot_leases == 3; }); }));