diff --git a/c_client/client_test.cc b/c_client/client_test.cc index 3ac5e2bc..13e505f5 100644 --- a/c_client/client_test.cc +++ b/c_client/client_test.cc @@ -282,6 +282,7 @@ TEST_F(ClientTest, CreatePublisherThenSubscriber) { ASSERT_NE(nullptr, client.client); SubspacePublisherOptions pub_opts = CPublisherOptionsDefault(256, 10); + ASSERT_EQ(0, pub_opts.subscriber_queue_arena_size); pub_opts.type.type = "foo"; pub_opts.type.type_length = strlen(pub_opts.type.type); SubspacePublisher pub = subspace_create_publisher(client, "dave1", pub_opts); @@ -300,6 +301,9 @@ TEST_F(ClientTest, CreatePublisherThenSubscriber) { subspace_create_subscriber(client, "dave1", CSubscriberOptionsDefault()); ASSERT_NE(nullptr, sub.subscriber); ASSERT_FALSE(subspace_has_error()); + ASSERT_EQ(0, subspace_get_publisher_queue_size(pub)); + ASSERT_EQ(0, subspace_get_publisher_queue_arena_size(pub)); + ASSERT_EQ(0, subspace_get_subscriber_queue_size(sub)); ASSERT_TRUE(subspace_remove_subscriber(&sub)); ASSERT_TRUE(subspace_remove_publisher(&pub)); @@ -781,11 +785,13 @@ TEST_F(ClientTest, ClientPublisherSubscriberIntrospection) { pub_opts.mux = mux; pub_opts.mux_length = strlen(mux); pub_opts.metadata_size = 8; + pub_opts.subscriber_queue_arena_size = 12'000; SubspacePublisher pub = subspace_create_publisher(client, "c_introspection", pub_opts); ASSERT_NE(nullptr, pub.publisher); SubspaceSubscriberOptions sub_opts = CSubscriberOptionsDefault(); + sub_opts.subscriber_queue_size = 4; sub_opts.type.type = type; sub_opts.type.type_length = strlen(type); sub_opts.mux = mux; @@ -836,6 +842,8 @@ TEST_F(ClientTest, ClientPublisherSubscriberIntrospection) { ASSERT_FALSE(subspace_is_publisher_for_tunnel(pub)); ASSERT_EQ(192, subspace_get_publisher_slot_size(pub)); ASSERT_EQ(6, subspace_get_publisher_num_slots(pub)); + ASSERT_EQ(16, subspace_get_publisher_queue_size(pub)); + ASSERT_EQ(12'000, subspace_get_publisher_queue_arena_size(pub)); ASSERT_TRUE(SubspaceStringEquals(subspace_get_publisher_name(pub), "c_introspection")); ASSERT_TRUE(SubspaceStringEquals(subspace_get_publisher_type(pub), type)); @@ -860,6 +868,7 @@ TEST_F(ClientTest, ClientPublisherSubscriberIntrospection) { ASSERT_EQ(0, subspace_get_subscriber_num_active_messages(sub)); ASSERT_EQ(8, subspace_get_subscriber_metadata_size(sub)); ASSERT_EQ(4, subspace_get_subscriber_checksum_size(sub)); + ASSERT_EQ(4, subspace_get_subscriber_queue_size(sub)); ASSERT_GE(subspace_get_subscriber_prefix_size(sub), 64); ASSERT_GE(subspace_get_subscriber_virtual_memory_usage(sub), 0U); @@ -1398,8 +1407,11 @@ TEST_F(ClientTest, InvalidArgumentsReportErrors) { ASSERT_EQ(-1, subspace_get_publisher_retirement_fd(invalid_publisher)); ASSERT_EQ(0, subspace_get_subscriber_slot_size(invalid_subscriber)); ASSERT_EQ(0, subspace_get_subscriber_num_slots(invalid_subscriber)); + ASSERT_EQ(0, subspace_get_subscriber_queue_size(invalid_subscriber)); ASSERT_EQ(0, subspace_get_publisher_slot_size(invalid_publisher)); ASSERT_EQ(0, subspace_get_publisher_num_slots(invalid_publisher)); + ASSERT_EQ(0, subspace_get_publisher_queue_size(invalid_publisher)); + ASSERT_EQ(0, subspace_get_publisher_queue_arena_size(invalid_publisher)); ASSERT_EQ(0, subspace_get_publisher_metadata_size(invalid_publisher)); ASSERT_EQ(0, subspace_get_subscriber_metadata_size(invalid_subscriber)); ASSERT_EQ(0, subspace_get_publisher_prefix_size(invalid_publisher)); diff --git a/c_client/subspace.cc b/c_client/subspace.cc index df475041..a75de3ab 100644 --- a/c_client/subspace.cc +++ b/c_client/subspace.cc @@ -127,6 +127,9 @@ SubspaceChannelInfo ToCChannelInfo(const subspace::ChannelInfo &info, .type = ToCString(type), .slot_size = info.slot_size, .num_slots = info.num_slots, + .subscriber_queue_size = info.subscriber_queue_size, + .subscriber_queue_arena_size = + info.subscriber_queue_arena_size, .reliable = info.reliable}; } @@ -493,6 +496,7 @@ bool subspace_get_all_channel_stats(SubspaceClient client, SubspaceSubscriberOptions subspace_subscriber_options_default(void) { SubspaceSubscriberOptions options = {}; options.max_active_messages = 1; + options.detect_dropped_messages = true; options.vchan_id = -1; return options; } @@ -502,6 +506,7 @@ SubspacePublisherOptions subspace_publisher_options_default(int32_t slot_size, SubspacePublisherOptions options = { slot_size, num_slots, + 0, false, false, false, @@ -530,6 +535,7 @@ subspace_create_subscriber(SubspaceClient client, const char *channel_name, SubspaceSubscriberOptions options) { subspace::SubscriberOptions subspace_options; subspace_options.SetReliable(options.reliable) + .SetSubscriberQueueSize(options.subscriber_queue_size) .SetBridge(options.bridge) .SetForTunnel(options.for_tunnel) .SetType(StringFromPointer(options.type.type, options.type.type_length)) @@ -541,6 +547,7 @@ subspace_create_subscriber(SubspaceClient client, const char *channel_name, .SetChecksum(options.checksum) .SetPassChecksumErrors(options.pass_checksum_errors) .SetKeepActiveMessage(options.keep_active_message) + .SetDetectDroppedMessages(options.detect_dropped_messages) .SetSplitBufferCallbacks(ToCppSplitCallbacks(options.split_callbacks)); subspace_options.SetLogDroppedMessages(options.log_dropped_messages); subspace_clear_error(); @@ -581,6 +588,7 @@ SubspacePublisher subspace_create_publisher(SubspaceClient client, .SetChecksum(options.checksum) .SetChecksumSize(options.checksum_size) .SetMetadataSize(options.metadata_size) + .SetSubscriberQueueArenaSize(options.subscriber_queue_arena_size) .SetPreferRetiredSlots(options.prefer_retired_slots) .SetMaxPublishers(options.max_publishers) .SetUseSplitBuffers(options.use_split_buffers) @@ -1448,6 +1456,21 @@ int32_t subspace_get_publisher_num_slots(SubspacePublisher publisher) { return (*PublisherPtr(publisher))->NumSlots(); } +int32_t subspace_get_publisher_queue_size(SubspacePublisher publisher) { + if (publisher.publisher == nullptr) { + return 0; + } + return (*PublisherPtr(publisher))->SubscriberQueueSize(); +} + +uint64_t +subspace_get_publisher_queue_arena_size(SubspacePublisher publisher) { + if (publisher.publisher == nullptr) { + return 0; + } + return (*PublisherPtr(publisher))->SubscriberQueueArenaSize(); +} + SubspaceString subspace_get_publisher_name(SubspacePublisher publisher) { if (publisher.publisher == nullptr) { return {}; @@ -1602,6 +1625,14 @@ int subspace_get_subscriber_num_slots(SubspaceSubscriber subscriber) { return (*sub_ptr)->NumSlots(); } +int32_t +subspace_get_subscriber_queue_size(SubspaceSubscriber subscriber) { + if (subscriber.subscriber == nullptr) { + return 0; + } + return (*SubscriberPtr(subscriber))->SubscriberQueueSize(); +} + int64_t subspace_get_subscriber_current_ordinal(SubspaceSubscriber subscriber) { if (subscriber.subscriber == nullptr) { return -1; @@ -1850,13 +1881,19 @@ bool subspace_snapshot_message_slot(SubspaceMessageSlot slot, } auto *message_slot = reinterpret_cast(slot.slot); *snapshot = {.id = message_slot->id, - .ordinal = message_slot->ordinal, - .message_size = message_slot->message_size, - .buffer_index = message_slot->buffer_index, - .vchan_id = message_slot->vchan_id, - .timestamp = message_slot->timestamp, - .flags = message_slot->flags, - .bridged_slot_id = message_slot->bridged_slot_id}; + .ordinal = + message_slot->ordinal.load(std::memory_order_relaxed), + .message_size = + message_slot->message_size.load(std::memory_order_relaxed), + .buffer_index = + message_slot->buffer_index.load(std::memory_order_relaxed), + .vchan_id = + message_slot->vchan_id.load(std::memory_order_relaxed), + .timestamp = + message_slot->timestamp.load(std::memory_order_relaxed), + .flags = message_slot->flags.load(std::memory_order_relaxed), + .bridged_slot_id = message_slot->bridged_slot_id.load( + std::memory_order_relaxed)}; return true; } diff --git a/c_client/subspace.h b/c_client/subspace.h index b5c92ac5..27c8f818 100644 --- a/c_client/subspace.h +++ b/c_client/subspace.h @@ -92,6 +92,8 @@ typedef struct { SubspaceString type; uint64_t slot_size; int num_slots; + int subscriber_queue_size; + uint64_t subscriber_queue_arena_size; bool reliable; } SubspaceChannelInfo; @@ -215,6 +217,9 @@ typedef struct { typedef struct { const int32_t slot_size; // Initial size of slots (might be resized). const int num_slots; // Number of slots (never changes) + // Total bytes reserved for packed per-subscriber queues in the CCB. The + // options factory selects zero, disabling queues in favor of the bitset path. + uint64_t subscriber_queue_arena_size; bool local; // If true, messages stay local to this machine. bool reliable; // Reliable publisher. bool bridge; // This publisher is for the bridge. @@ -257,12 +262,16 @@ typedef struct { typedef struct { bool reliable; // Reliable subscriber. + // Capacity of this subscriber's CCB slot queue. 0 uses the publisher + // default. + int32_t subscriber_queue_size; bool bridge; // This subscriber is for the bridge. bool for_tunnel; // Mark subscriptions for external tunnels. SubspaceTypeInfo type; // Type of the message. This is an opaque string. int max_active_messages; // Max number of message that can be active at once. bool pass_activation; // Pass activation message in read. bool log_dropped_messages; // Log dropped messages to stderr. + bool detect_dropped_messages; // Detect and count ordinal gaps internally. bool read_write; // Map buffers writable for this subscriber. const char *mux; // Optional mux channel name for virtual channels. size_t mux_length; @@ -402,6 +411,8 @@ int subspace_get_subscriber_fd(SubspaceSubscriber subscriber); // is received. int32_t subspace_get_subscriber_slot_size(SubspaceSubscriber subscriber); int subspace_get_subscriber_num_slots(SubspaceSubscriber subscriber); +int32_t +subspace_get_subscriber_queue_size(SubspaceSubscriber subscriber); // This is a shortcut to wait for a message to be available. It will block // until a message is available. @@ -536,6 +547,9 @@ bool subspace_is_publisher_for_tunnel(SubspacePublisher publisher); bool subspace_publisher_uses_split_buffers(SubspacePublisher publisher); int32_t subspace_get_publisher_slot_size(SubspacePublisher publisher); int32_t subspace_get_publisher_num_slots(SubspacePublisher publisher); +int32_t subspace_get_publisher_queue_size(SubspacePublisher publisher); +uint64_t +subspace_get_publisher_queue_arena_size(SubspacePublisher publisher); SubspaceString subspace_get_publisher_name(SubspacePublisher publisher); SubspaceString subspace_get_publisher_type(SubspacePublisher publisher); SubspaceString subspace_get_publisher_mux(SubspacePublisher publisher); diff --git a/client/client.cc b/client/client.cc index 2df7bcb7..bc734185 100644 --- a/client/client.cc +++ b/client/client.cc @@ -433,7 +433,8 @@ ClientImpl::CreatePublisher(const std::string &channel_name, }; std::shared_ptr channel = std::make_shared( - channel_name, opts.num_slots, pub_resp.channel_id(), + channel_name, opts.num_slots, pub_resp.subscriber_queue_size(), + pub_resp.subscriber_queue_arena_size(), pub_resp.channel_id(), pub_resp.publisher_id(), pub_resp.vchan_id(), session_id_, pub_resp.type(), opts, [this](Channel *c) { @@ -568,9 +569,12 @@ ClientImpl::CreateSubscriber(const std::string &channel_name, subscriber_options.use_split_buffers = sub_resp.use_split_buffers(); std::shared_ptr channel = std::make_shared( - channel_name, sub_resp.num_slots(), sub_resp.channel_id(), - sub_resp.subscriber_id(), sub_resp.vchan_id(), session_id_, - sub_resp.type(), subscriber_options, + channel_name, sub_resp.num_slots(), + sub_resp.default_subscriber_queue_size(), + sub_resp.subscriber_queue_arena_size(), + sub_resp.subscriber_queue_size(), sub_resp.channel_id(), + sub_resp.subscriber_id(), sub_resp.vchan_id(), session_id_, sub_resp.type(), + subscriber_options, [this](Channel *c) { return CheckReload(static_cast(c)); }, @@ -582,6 +586,7 @@ ClientImpl::CreateSubscriber(const std::string &channel_name, }); channel->SetNumSlots(sub_resp.num_slots()); + channel->SetEffectiveSubscriberQueueSize(sub_resp.subscriber_queue_size()); { int32_t cs = sub_resp.checksum_size() > 0 ? sub_resp.checksum_size() : 4; int32_t ms = sub_resp.metadata_size() > 0 ? sub_resp.metadata_size() : 0; @@ -767,7 +772,7 @@ ClientImpl::PublishMessageInternal(PublisherImpl *publisher, if (debug_) { if (old_slot != nullptr) { printf("publish old slot: %d: %" PRId64 "\n", old_slot->id, - old_slot->ordinal); + old_slot->ordinal.load(std::memory_order_relaxed)); } } @@ -797,7 +802,7 @@ ClientImpl::PublishMessageInternal(PublisherImpl *publisher, if (debug_) { printf("publish new slot: %d: %" PRId64 "\n", msg.new_slot->id, - msg.new_slot->ordinal); + msg.new_slot->ordinal.load(std::memory_order_relaxed)); } return Message(message_size, nullptr, msg.ordinal, msg.timestamp, @@ -1125,7 +1130,7 @@ ClientImpl::ReadMessageInternal(SubscriberImpl *subscriber, ReadMode mode, MessageSlot *old_slot = subscriber->CurrentSlot(); int64_t last_ordinal = -1; if (old_slot != nullptr) { - last_ordinal = old_slot->ordinal; + last_ordinal = old_slot->ordinal.load(std::memory_order_relaxed); if (debug_) { printf("read old slot: %d: %" PRId64 "\n", old_slot->id, last_ordinal); } @@ -1154,27 +1159,13 @@ ClientImpl::ReadMessageInternal(SubscriberImpl *subscriber, ReadMode mode, return Message(); } subscriber->SetSlot(new_slot); + int64_t delivered_message_size = + static_cast( + new_slot->message_size.load(std::memory_order_relaxed)); if (debug_) { - printf("read new_slot: %d: %" PRId64 "\n", new_slot->id, new_slot->ordinal); - } - - if (mode == ReadMode::kReadNext && last_ordinal != -1) { - int drops = subscriber->DetectDrops(new_slot->vchan_id); - if (drops > 0) { - // We dropped a message. If we have a callback registered for this - // channel, call it with the number of dropped messages. - auto it = dropped_message_callbacks_.find(subscriber); - if (it != dropped_message_callbacks_.end()) { - it->second(subscriber, drops); - } - subscriber->RecordDroppedMessages(drops); - if (subscriber->options_.log_dropped_messages) { - logger_.Log(toolbelt::LogLevel::kWarning, - "Dropped %d message%s on channel %s", drops, - drops == 1 ? "" : "s", subscriber->Name().c_str()); - } - } + printf("read new_slot: %d: %" PRId64 "\n", new_slot->id, + new_slot->ordinal.load(std::memory_order_relaxed)); } MessagePrefix *prefix = subscriber->Prefix(new_slot); @@ -1185,7 +1176,7 @@ ClientImpl::ReadMessageInternal(SubscriberImpl *subscriber, ReadMode mode, if (prefix->HasChecksum()) { auto data = GetMessageChecksumData(prefix, subscriber->GetCurrentBufferAddress(), - new_slot->message_size, + delivered_message_size, subscriber->ChecksumSize(), subscriber->MetadataSize()); absl::Span cksum = @@ -1208,13 +1199,17 @@ ClientImpl::ReadMessageInternal(SubscriberImpl *subscriber, ReadMode mode, // Call the on receive callback. if (subscriber->on_receive_callback_ != nullptr) { absl::StatusOr status_or_size = subscriber->on_receive_callback_( - subscriber->GetCurrentBufferAddress(), new_slot->message_size); + subscriber->GetCurrentBufferAddress(), delivered_message_size); if (!status_or_size.ok()) { + subscriber->UnreadSlot(new_slot); + subscriber->SetSlot(nullptr); return status_or_size.status(); } - new_slot->message_size = status_or_size.value(); + delivered_message_size = status_or_size.value(); } - if (new_slot->message_size <= 0) { + if (delivered_message_size <= 0) { + subscriber->UnreadSlot(new_slot); + subscriber->SetSlot(nullptr); return Message(); } // We have a new slot, clear the subscriber's slot. @@ -1222,19 +1217,47 @@ ClientImpl::ReadMessageInternal(SubscriberImpl *subscriber, ReadMode mode, // Allocate a new active message for the slot. auto msg = subscriber->SetActiveMessage( - new_slot->message_size, new_slot, subscriber->GetCurrentBufferAddress(), + delivered_message_size, new_slot, subscriber->GetCurrentBufferAddress(), subscriber->CurrentOrdinal(), subscriber->Timestamp(new_slot), - new_slot->vchan_id, is_activation, checksum_error); + new_slot->vchan_id.load(std::memory_order_relaxed), is_activation, + checksum_error); // If we are unable to allocate a new message (due to message limits) // restore the slot so that we pick it up next time. if (msg->length == 0) { subscriber->UnreadSlot(new_slot); // Subscriber does not have a slot now but the slot it had is still active. + // Do not wrap the reusable ActiveMessage in an empty Message. A caller may + // retain that empty handle while this slot is retried, which would keep an + // extra reference after the ActiveMessage becomes valid and prevent its + // active-message count from being released. + return Message(); } else { + if (mode == ReadMode::kReadNext && + subscriber->options_.DetectDroppedMessages()) { + int drops = subscriber->ConsumeQueueDrops(); + if (last_ordinal != -1) { + drops = std::max( + drops, subscriber->DetectDrops( + new_slot->vchan_id.load(std::memory_order_relaxed))); + } + if (drops > 0) { + auto it = dropped_message_callbacks_.find(subscriber); + if (it != dropped_message_callbacks_.end()) { + it->second(subscriber, drops); + } + subscriber->RecordDroppedMessages(drops); + if (subscriber->options_.log_dropped_messages) { + logger_.Log(toolbelt::LogLevel::kWarning, + "Dropped %d message%s on channel %s", drops, + drops == 1 ? "" : "s", subscriber->Name().c_str()); + } + } + } // We have a slot, claim it. - subscriber->ClaimSlot(new_slot, subscriber->VirtualChannelId(), - mode == ReadMode::kReadNewest); + subscriber->ClaimSlot( + new_slot, new_slot->vchan_id.load(std::memory_order_relaxed), + mode == ReadMode::kReadNewest); } auto ret_msg = Message(msg); if (subscriber->IsBridge()) { @@ -1290,7 +1313,8 @@ ClientImpl::FindMessageInternal(SubscriberImpl *subscriber, // Not found. return Message(); } - return Message(new_slot->message_size, subscriber->GetCurrentBufferAddress(), + return Message(new_slot->message_size.load(std::memory_order_relaxed), + subscriber->GetCurrentBufferAddress(), subscriber->CurrentOrdinal(), subscriber->Timestamp(), subscriber->VirtualChannelId(), false, new_slot->id, false); } @@ -1355,7 +1379,7 @@ int64_t ClientImpl::GetCurrentOrdinal(SubscriberImpl *sub) { if (slot == nullptr) { return -1; } - return slot->ordinal; + return slot->ordinal.load(std::memory_order_relaxed); } bool ClientImpl::CheckReload(ClientChannel *channel) { @@ -1389,8 +1413,6 @@ absl::Status ClientImpl::ReloadSubscriber(SubscriberImpl *subscriber) { if (subscriber->NumUpdates() == updates) { return absl::OkStatus(); } - subscriber->SetNumUpdates(updates); - if (absl::Status status = CheckConnected(); !status.ok()) { return status; } @@ -1399,6 +1421,8 @@ absl::Status ClientImpl::ReloadSubscriber(SubscriberImpl *subscriber) { cmd->set_channel_name(subscriber->Name()); cmd->set_subscriber_id(subscriber->GetSubscriberId()); cmd->set_mux(subscriber->options_.mux); + cmd->set_subscriber_queue_size( + subscriber->options_.SubscriberQueueSize()); // Send request to server and wait for response. Response resp; @@ -1414,14 +1438,26 @@ absl::Status ClientImpl::ReloadSubscriber(SubscriberImpl *subscriber) { return absl::InternalError(sub_resp.error()); } - // Unmap the channel memory. - subscriber->Unmap(); + // A subscriber-created placeholder is the only case where the server + // replaces the CCB. Once num_slots is non-zero, publisher updates retain the + // existing CCB and only require refreshed descriptors and buffers. + const bool remap_ccb = subscriber->NumSlots() == 0; + if (remap_ccb) { + subscriber->ResetDeliveryState(); + subscriber->Unmap(); + } if (!sub_resp.type().empty()) { subscriber->SetType(sub_resp.type()); } subscriber->options_.use_split_buffers = sub_resp.use_split_buffers(); subscriber->SetNumSlots(sub_resp.num_slots()); + subscriber->SetSubscriberQueueSize( + sub_resp.default_subscriber_queue_size()); + subscriber->SetSubscriberQueueArenaSize( + sub_resp.subscriber_queue_arena_size()); + subscriber->SetEffectiveSubscriberQueueSize( + sub_resp.subscriber_queue_size()); { int32_t cs = sub_resp.checksum_size() > 0 ? sub_resp.checksum_size() : 4; int32_t ms = sub_resp.metadata_size() > 0 ? sub_resp.metadata_size() : 0; @@ -1431,15 +1467,15 @@ absl::Status ClientImpl::ReloadSubscriber(SubscriberImpl *subscriber) { subscriber->AllocateChecksumBuffer(); } - SharedMemoryFds channel_fds(std::move(fds[sub_resp.ccb_fd_index()]), - std::move(fds[sub_resp.bcb_fd_index()])); - // subscriber->SetSlots(sub_resp.slot_size(), sub_resp.num_slots()); - - if (absl::Status status = subscriber->Map(std::move(channel_fds), scb_fd_); - !status.ok()) { - return status; + if (remap_ccb) { + SharedMemoryFds channel_fds(std::move(fds[sub_resp.ccb_fd_index()]), + std::move(fds[sub_resp.bcb_fd_index()])); + if (absl::Status status = subscriber->Map(std::move(channel_fds), scb_fd_); + !status.ok()) { + return status; + } + subscriber->InitActiveMessages(); } - subscriber->InitActiveMessages(); if (absl::Status status = subscriber->AttachBuffers(); !status.ok()) { return status; @@ -1459,6 +1495,7 @@ absl::Status ClientImpl::ReloadSubscriber(SubscriberImpl *subscriber) { subscriber->AddRetirementTrigger(fds[size_t(index)]); } + subscriber->SetNumUpdates(updates); // subscriber->Dump(); return absl::OkStatus(); } @@ -1565,7 +1602,7 @@ absl::Status ClientImpl::ActivateReliableChannel(PublisherImpl *publisher) { return absl::InternalError( absl::StrFormat("Channel %s has no buffer", publisher->Name())); } - slot->message_size = 1; + slot->message_size.store(1, std::memory_order_relaxed); publisher->ActivateSlotAndGetAnother( /*reliable=*/true, @@ -1588,7 +1625,7 @@ absl::Status ClientImpl::ActivateChannel(PublisherImpl *publisher) { absl::StrFormat("3 Channel %s has no buffer", publisher->Name())); } MessageSlot *slot = publisher->CurrentSlot(); - slot->message_size = 1; + slot->message_size.store(1, std::memory_order_relaxed); Channel::PublishedMessage msg = publisher->ActivateSlotAndGetAnother( /*reliable=*/false, @@ -1715,6 +1752,9 @@ ClientImpl::GetChannelInfo(const std::string &channel) { result.type = info.type(); result.slot_size = info.slot_size(); result.num_slots = info.num_slots(); + result.subscriber_queue_size = info.subscriber_queue_size(); + result.subscriber_queue_arena_size = + info.subscriber_queue_arena_size(); return result; } @@ -1751,6 +1791,9 @@ absl::StatusOr> ClientImpl::GetChannelInfo() { result.type = info.type(); result.slot_size = info.slot_size(); result.num_slots = info.num_slots(); + result.subscriber_queue_size = info.subscriber_queue_size(); + result.subscriber_queue_arena_size = + info.subscriber_queue_arena_size(); r.push_back(result); } return r; @@ -1875,6 +1918,7 @@ void ClientImpl::FillCreatePublisherRequest(CreatePublisherRequest *cmd, cmd->set_max_publishers(opts.MaxPublishers()); cmd->set_use_split_buffers(opts.UseSplitBuffers()); cmd->set_split_buffers_over_bridge(opts.SplitBuffersOverBridge()); + cmd->set_subscriber_queue_arena_size(opts.SubscriberQueueArenaSize()); cmd->set_process_id(static_cast(getpid())); } @@ -1915,6 +1959,7 @@ void ClientImpl::FillCreateSubscriberRequest(CreateSubscriberRequest *cmd, cmd->set_max_active_messages(opts.MaxActiveMessages()); cmd->set_mux(opts.Mux()); cmd->set_vchan_id(opts.VchanId()); + cmd->set_subscriber_queue_size(opts.SubscriberQueueSize()); cmd->set_process_id(static_cast(getpid())); } @@ -1993,6 +2038,8 @@ absl::Status ClientImpl::ReregisterPublisher(PublisherImpl *publisher) { FillCreatePublisherRequest(req.mutable_create_publisher(), publisher->Name(), publisher->options_, publisher->GetPublisherId()); + req.mutable_create_publisher()->set_active_queue_publish_depth( + publisher->ActiveQueuePublishDepth()); Response resp; std::vector fds; diff --git a/client/client.h b/client/client.h index a83492e1..3e6b8257 100644 --- a/client/client.h +++ b/client/client.h @@ -69,6 +69,8 @@ struct ChannelInfo { std::string type; uint64_t slot_size; int num_slots; + int subscriber_queue_size; + uint64_t subscriber_queue_arena_size; bool reliable; }; @@ -967,6 +969,10 @@ class Publisher { int32_t SlotSize() const { return impl_->SlotSize(); } int32_t NumSlots() const { return impl_->NumSlots(); } + int32_t SubscriberQueueSize() const { return impl_->SubscriberQueueSize(); } + uint64_t SubscriberQueueArenaSize() const { + return impl_->SubscriberQueueArenaSize(); + } const std::vector> &GetBuffers() const { return client_->GetBuffers(impl_.get()); @@ -1403,6 +1409,7 @@ class Subscriber { int32_t SlotSize() const { return impl_->SlotSize(); } int32_t NumSlots() const { return impl_->NumSlots(); } + int32_t SubscriberQueueSize() const { return impl_->SubscriberQueueSize(); } const std::vector> &GetBuffers() const { return client_->GetBuffers(impl_.get()); @@ -1505,8 +1512,10 @@ class Subscriber { bool AtomicIncRefCount(int slot_id, int inc) { MessageSlot *slot = impl_->GetSlot(slot_id); if (slot != nullptr) { - return impl_->AtomicIncRefCount(slot, IsReliable(), inc, slot->ordinal, - slot->vchan_id, false); + return impl_->AtomicIncRefCount( + slot, IsReliable(), inc, + slot->ordinal.load(std::memory_order_relaxed), + slot->vchan_id.load(std::memory_order_relaxed), false); } return false; } diff --git a/client/client_channel.cc b/client/client_channel.cc index bbcc4bbf..9566b023 100644 --- a/client/client_channel.cc +++ b/client/client_channel.cc @@ -121,27 +121,45 @@ ClientChannel::CreatePosixSharedMemoryFile(const std::string &filename, absl::Status ClientChannel::Map(SharedMemoryFds fds, const toolbelt::FileDescriptor &scb_fd) { + absl::StatusOr checked_ccb_size = + CheckedCcbSize(num_slots_, subscriber_queue_arena_size_); + if (!checked_ccb_size.ok()) { + return checked_ccb_size.status(); + } scb_ = reinterpret_cast(MapMemory( scb_fd.Fd(), sizeof(SystemControlBlock), PROT_READ | PROT_WRITE, "SCB")); if (scb_ == MAP_FAILED) { + scb_ = nullptr; return absl::InternalError(absl::StrFormat( "Failed to map SystemControlBlock: %s (scb_fd=%d, ccb_fd=%d, " "bcb_fd=%d, scb_size=%zu, ccb_size=%zu, bcb_size=%zu)", strerror(errno), scb_fd.Fd(), fds.ccb.Fd(), fds.bcb.Fd(), - sizeof(SystemControlBlock), CcbSize(num_slots_), + sizeof(SystemControlBlock), *checked_ccb_size, sizeof(BufferControlBlock))); } - ccb_ = reinterpret_cast(MapMemory( - fds.ccb.Fd(), CcbSize(num_slots_), PROT_READ | PROT_WRITE, "CCB")); + ccb_ = reinterpret_cast( + MapMemory(fds.ccb.Fd(), *checked_ccb_size, PROT_READ | PROT_WRITE, "CCB")); if (ccb_ == MAP_FAILED) { int mmap_errno = errno; UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + scb_ = nullptr; + ccb_ = nullptr; return absl::InternalError(absl::StrFormat( "Failed to map ChannelControlBlock: %s (scb_fd=%d, ccb_fd=%d, " "bcb_fd=%d, ccb_size=%zu)", strerror(mmap_errno), scb_fd.Fd(), fds.ccb.Fd(), fds.bcb.Fd(), - CcbSize(num_slots_))); + *checked_ccb_size)); + } + if (ccb_->version != kChannelControlBlockVersion) { + const uint32_t version = ccb_->version; + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + scb_ = nullptr; + ccb_ = nullptr; + return absl::FailedPreconditionError(absl::StrFormat( + "unsupported channel control block version %u (expected %u)", + version, kChannelControlBlockVersion)); } bcb_ = reinterpret_cast(MapMemory( @@ -149,7 +167,10 @@ absl::Status ClientChannel::Map(SharedMemoryFds fds, if (bcb_ == MAP_FAILED) { int mmap_errno = errno; UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); - UnmapMemory(ccb_, CcbSize(num_slots_), "CCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + scb_ = nullptr; + ccb_ = nullptr; + bcb_ = nullptr; return absl::InternalError(absl::StrFormat( "Failed to map BufferControlBlock: %s (scb_fd=%d, ccb_fd=%d, " "bcb_fd=%d, bcb_size=%zu)", @@ -248,16 +269,18 @@ void ClientChannel::UnmapSplitBufferSet(size_t buffer_index, } bool ClientChannel::ValidateSlotBuffer(MessageSlot *slot) { - if (slot->buffer_index < 0) { + const int buffer_index = + slot->buffer_index.load(std::memory_order_relaxed); + if (buffer_index < 0) { return true; } - if (static_cast(slot->buffer_index) < buffers_.size() && - buffers_[slot->buffer_index]->IsSplitBuffers()) { + if (static_cast(buffer_index) < buffers_.size() && + buffers_[buffer_index]->IsSplitBuffers()) { return slot->id >= 0 && static_cast(slot->id) < - buffers_[slot->buffer_index]->split_slot_buffers.size() && - buffers_[slot->buffer_index]->split_slot_buffers[slot->id] != + buffers_[buffer_index]->split_slot_buffers.size() && + buffers_[buffer_index]->split_slot_buffers[slot->id] != nullptr; } @@ -379,8 +402,10 @@ uint64_t ClientChannel::GetVirtualMemoryUsage() const { return Channel::GetVirtualMemoryUsage(); } - uint64_t size = sizeof(SystemControlBlock) + CcbSize(num_slots_) + - sizeof(BufferControlBlock); + uint64_t size = + sizeof(SystemControlBlock) + + CcbSize(num_slots_, subscriber_queue_arena_size_) + + sizeof(BufferControlBlock); for (int i = 0; i < ccb_->num_buffers; i++) { if (bcb_->refs[i].load(std::memory_order_relaxed) <= 0) { continue; @@ -1058,7 +1083,8 @@ void ClientChannel::TriggerRetirement(int slot_id) { return; } MessageSlot *slot = GetSlot(slot_id); - if ((slot->flags & kMessageIsActivation) != 0) { + if ((slot->flags.load(std::memory_order_relaxed) & + kMessageIsActivation) != 0) { // Don't retire activation messages. return; } diff --git a/client/client_channel.h b/client/client_channel.h index 9ef27a83..164ab5bf 100644 --- a/client/client_channel.h +++ b/client/client_channel.h @@ -81,10 +81,13 @@ struct BufferSet { // a publisher or a subscriber, as defined as the subclasses. class ClientChannel : public Channel { public: - ClientChannel(const std::string &name, int num_slots, int channel_id, + ClientChannel(const std::string &name, int num_slots, + int subscriber_queue_size, + uint64_t subscriber_queue_arena_size, int channel_id, int vchan_id, uint64_t session_id, std::string type, std::function reload, int user_id, int group_id) - : Channel(name, num_slots, channel_id, std::move(type), + : Channel(name, num_slots, channel_id, subscriber_queue_size, + subscriber_queue_arena_size, std::move(type), std::move(reload)), vchan_id_(vchan_id), session_id_(std::move(session_id)), user_id_(user_id), group_id_(group_id) { active_slots_.reserve(num_slots); @@ -117,7 +120,8 @@ class ClientChannel : public Channel { // What is the address of the message buffer (after the prefix area) // for the slot given a slot id. void *GetBufferAddress(int slot_id) { - int buffer_index = ccb_->slots[slot_id].buffer_index; + const int buffer_index = + ccb_->slots[slot_id].buffer_index.load(std::memory_order_relaxed); if (buffer_index >= 0 && static_cast(buffer_index) < buffers_.size() && buffers_[buffer_index]->IsSplitBuffers()) { @@ -133,7 +137,8 @@ class ClientChannel : public Channel { if (slot == nullptr) { return nullptr; } - int buffer_index = ccb_->slots[slot->id].buffer_index; + const int buffer_index = + ccb_->slots[slot->id].buffer_index.load(std::memory_order_relaxed); if (buffer_index >= 0 && static_cast(buffer_index) < buffers_.size() && buffers_[buffer_index]->IsSplitBuffers()) { @@ -151,7 +156,8 @@ class ClientChannel : public Channel { if (slot == nullptr) { return nullptr; } - int buffer_index = ccb_->slots[slot->id].buffer_index; + const int buffer_index = + ccb_->slots[slot->id].buffer_index.load(std::memory_order_relaxed); if (buffer_index >= 0 && static_cast(buffer_index) < buffers_.size() && buffers_[buffer_index]->IsSplitBuffers()) { @@ -174,13 +180,13 @@ class ClientChannel : public Channel { // Get the size associated with the given slot id. int SlotSize(int slot_id) const { - if (ccb_->slots[slot_id].buffer_index < 0 || - static_cast(ccb_->slots[slot_id].buffer_index) >= buffers_.size()) { + const int buffer_index = + ccb_->slots[slot_id].buffer_index.load(std::memory_order_relaxed); + if (buffer_index < 0 || + static_cast(buffer_index) >= buffers_.size()) { return 0; } - return buffers_.empty() - ? 0 - : buffers_[ccb_->slots[slot_id].buffer_index]->slot_size; + return buffers_.empty() ? 0 : buffers_[buffer_index]->slot_size; } int SlotSize(MessageSlot *slot) const { @@ -190,11 +196,13 @@ class ClientChannel : public Channel { if (buffers_.empty()) { return 0; } - if (ccb_->slots[slot->id].buffer_index < 0 || - static_cast(ccb_->slots[slot->id].buffer_index) >= buffers_.size()) { + const int buffer_index = + ccb_->slots[slot->id].buffer_index.load(std::memory_order_relaxed); + if (buffer_index < 0 || + static_cast(buffer_index) >= buffers_.size()) { return 0; } - return buffers_[ccb_->slots[slot->id].buffer_index]->slot_size; + return buffers_[buffer_index]->slot_size; } // Get the biggest slot size for the channel. int SlotSize() const { @@ -215,8 +223,9 @@ class ClientChannel : public Channel { constexpr int kMaxRetries = 1000; int retries = 0; while (retries < kMaxRetries) { - size_t index = ccb_->slots[slot_id].buffer_index; - if (index != -1ULL && index < buffers_.size()) { + const int index = + ccb_->slots[slot_id].buffer_index.load(std::memory_order_relaxed); + if (index >= 0 && static_cast(index) < buffers_.size()) { return buffers_.empty() ? nullptr : (buffers_[index]->buffer); } CheckReload(); @@ -226,7 +235,8 @@ class ClientChannel : public Channel { if (abort_on_range) { // If the index is out of range, we have a problem. // This should never happen. - int index = ccb_->slots[slot_id].buffer_index; + const int index = + ccb_->slots[slot_id].buffer_index.load(std::memory_order_relaxed); std::cerr << this << " Invalid buffer index for slot " << slot_id << ": " << index << " there are " << buffers_.size() << " buffers" << std::endl; @@ -344,7 +354,7 @@ class ClientChannel : public Channel { bool ValidateSlotBuffer(MessageSlot *slot); void SetMessageSize(int64_t message_size) { - slot_->message_size = message_size; + slot_->message_size.store(message_size, std::memory_order_relaxed); } bool IsVirtual() const { return vchan_id_ != -1; } diff --git a/client/client_test.cc b/client/client_test.cc index fd927896..44ef22e8 100644 --- a/client/client_test.cc +++ b/client/client_test.cc @@ -83,7 +83,8 @@ uint64_t AlignPage(uint64_t size) { uint64_t ExpectedSplitBufferVirtualMemoryUsage(int num_slots, uint64_t slot_size, uint64_t prefix_size) { - return sizeof(subspace::SystemControlBlock) + subspace::CcbSize(num_slots) + + return sizeof(subspace::SystemControlBlock) + + subspace::CcbSize(num_slots, /*subscriber_queue_arena_size=*/0) + sizeof(subspace::BufferControlBlock) + AlignPage(prefix_size * static_cast(num_slots)) + AlignPage(slot_size) * static_cast(num_slots); @@ -281,6 +282,12 @@ TEST(AndroidBufferRegistrationTest, FailedRegistrationRollsBackNumBuffers) { absl::StatusOr ccb_fd = CreateTestMemfd("subspace_test_ccb", subspace::CcbSize(kNumSlots)); ASSERT_OK(ccb_fd); + auto *ccb = reinterpret_cast( + subspace::MapMemory(ccb_fd->Fd(), subspace::CcbSize(kNumSlots), + PROT_READ | PROT_WRITE, "test CCB")); + ASSERT_NE(MAP_FAILED, ccb); + ccb->version = subspace::kChannelControlBlockVersion; + subspace::UnmapMemory(ccb, subspace::CcbSize(kNumSlots), "test CCB"); absl::StatusOr bcb_fd = CreateTestMemfd( "subspace_test_bcb", sizeof(subspace::BufferControlBlock)); ASSERT_OK(bcb_fd); @@ -293,7 +300,9 @@ TEST(AndroidBufferRegistrationTest, FailedRegistrationRollsBackNumBuffers) { subspace::PublisherOptions options; options.SetUseSplitBuffers(false); subspace::details::PublisherImpl publisher( - "android_registration_rollback", kNumSlots, /*channel_id=*/0, + "android_registration_rollback", kNumSlots, + /*subscriber_queue_size=*/0, /*subscriber_queue_arena_size=*/0, + /*channel_id=*/0, /*publisher_id=*/0, /*vchan_id=*/-1, /*session_id=*/123, "", options, [](subspace::Channel *) { return false; }, /*user_id=*/0, /*group_id=*/0); @@ -881,6 +890,619 @@ TEST_F(ClientTest, PublishSingleMessageAndRead) { ASSERT_EQ(0, msg->length); } +TEST_F(ClientTest, PublishAndReadWithSubscriberQueue) { + subspace::Client pub_client; + subspace::Client sub_client; + ASSERT_OK(pub_client.Init(Socket())); + ASSERT_OK(sub_client.Init(Socket())); + + absl::StatusOr pub = pub_client.CreatePublisher( + "subscriber_queue_read", + subspace::PublisherOptions() + .SetSlotSize(256) + .SetNumSlots(40) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize)); + ASSERT_OK(pub); + + absl::StatusOr sub = + sub_client.CreateSubscriber("subscriber_queue_read"); + ASSERT_OK(sub); + + absl::StatusOr buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + memcpy(*buffer, "queued1", 7); + absl::StatusOr pub_status = pub->PublishMessage(7); + ASSERT_OK(pub_status); + + absl::StatusOr msg = sub->ReadMessage(); + ASSERT_OK(msg); + ASSERT_EQ(7, msg->length); + ASSERT_EQ(0, memcmp(msg->buffer, "queued1", 7)); + msg->Reset(); + + buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + memcpy(*buffer, "queued2", 7); + absl::StatusOr pub_status2 = pub->PublishMessage(7); + ASSERT_OK(pub_status2); + + buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + memcpy(*buffer, "queued3", 7); + absl::StatusOr pub_status3 = pub->PublishMessage(7); + ASSERT_OK(pub_status3); + + msg = sub->ReadMessage(subspace::ReadMode::kReadNewest); + ASSERT_OK(msg); + ASSERT_EQ(7, msg->length); + ASSERT_EQ(0, memcmp(msg->buffer, "queued3", 7)); +} + +TEST_F(ClientTest, SubscribersUseDifferentQueueSizes) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "different_subscriber_queue_sizes", + subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(40) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + auto small = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + "different_subscriber_queue_sizes", + subspace::SubscriberOptions().SetSubscriberQueueSize(2))); + auto defaults = EVAL_AND_ASSERT_OK( + client.CreateSubscriber("different_subscriber_queue_sizes")); + + EXPECT_EQ(2, small.SubscriberQueueSize()); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + defaults.SubscriberQueueSize()); + + for (uint8_t value = 1; value <= 4; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + Message small_message = EVAL_AND_ASSERT_OK(small.ReadMessage()); + ASSERT_EQ(1, small_message.length); + EXPECT_EQ(3, *static_cast(small_message.buffer)); + small_message.Reset(); + + Message default_message = EVAL_AND_ASSERT_OK(defaults.ReadMessage()); + ASSERT_EQ(1, default_message.length); + EXPECT_EQ(1, *static_cast(default_message.buffer)); +} + +TEST_F(ClientTest, PublisherQueueArenaRemainsFixedWithoutPublishers) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "publisher_queue_default_without_publishers"; + std::unique_ptr subscriber; + { + auto publisher = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize(4096))); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + publisher.SubscriberQueueSize()); + subscriber = std::make_unique( + EVAL_AND_ASSERT_OK(client.CreateSubscriber(kChannel))); + } + + auto mismatched = client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize(8192)); + ASSERT_FALSE(mismatched.ok()); + EXPECT_THAT(mismatched.status().message(), + ::testing::HasSubstr( + "subscriber queue arena size is 4096, not 8192")); +} + +TEST_F(ClientTest, PublisherQueueArenaMatchesAcrossVirtualChannels) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kMux[] = "publisher_queue_default_mux"; + auto first = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "publisher_queue_default_vchan_a", + subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(16) + .SetSubscriberQueueArenaSize(4096) + .SetMux(kMux))); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + first.SubscriberQueueSize()); + auto second_vchan_subscriber = + EVAL_AND_ASSERT_OK(client.CreateSubscriber( + "publisher_queue_default_vchan_b", + subspace::SubscriberOptions().SetMux(kMux))); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + second_vchan_subscriber.SubscriberQueueSize()); + + auto mismatched = client.CreatePublisher( + "publisher_queue_default_vchan_b", + subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(16) + .SetSubscriberQueueArenaSize(8192) + .SetMux(kMux)); + ASSERT_FALSE(mismatched.ok()); + EXPECT_THAT(mismatched.status().message(), + ::testing::HasSubstr( + "subscriber queue arena size is 4096, not 8192")); +} + +TEST_F(ClientTest, FailedSubscriberQueuePushFallsBackToBitset) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_push_fallback"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + auto sub = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(2))); + + subspace::ServerChannel *server_channel = Server()->FindChannel(kChannel); + ASSERT_NE(nullptr, server_channel); + int sub_id = -1; + server_channel->GetCcb()->subscribers.Traverse( + [&sub_id](int id) { sub_id = id; }); + ASSERT_GE(sub_id, 0); + subspace::InPlaceSlotQueue *queue = + server_channel->GetAvailableSlotQueueAddress(sub_id); + ASSERT_NE(nullptr, queue); + auto *entries = reinterpret_cast( + reinterpret_cast(queue) + + sizeof(subspace::InPlaceSlotQueue)); + // Model a consumer that advanced head but died before releasing the entry. + entries[0].sequence.store(1, std::memory_order_release); + + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + memcpy(buffer, "fallback", 8); + ASSERT_OK(pub.PublishMessage(8)); + + Message message = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(8, message.length); + EXPECT_EQ(0, memcmp(message.buffer, "fallback", 8)); +} + +TEST_F(ClientTest, ConcurrentQueueReservationOrderDoesNotDropOlderOrdinal) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_out_of_order"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + auto sub = EVAL_AND_ASSERT_OK(client.CreateSubscriber(kChannel)); + + for (uint8_t value = 1; value <= 2; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + subspace::ServerChannel *server_channel = Server()->FindChannel(kChannel); + ASSERT_NE(nullptr, server_channel); + int sub_id = -1; + server_channel->GetCcb()->subscribers.Traverse( + [&sub_id](int id) { sub_id = id; }); + ASSERT_GE(sub_id, 0); + subspace::InPlaceSlotQueue *queue = + server_channel->GetAvailableSlotQueueAddress(sub_id); + ASSERT_NE(nullptr, queue); + + std::vector data_slots; + for (int i = 0; i < server_channel->NumSlots(); ++i) { + subspace::MessageSlot *slot = &server_channel->GetCcb()->slots[i]; + if (slot->message_size.load(std::memory_order_relaxed) == 1) { + data_slots.push_back(slot); + } + } + ASSERT_EQ(2u, data_slots.size()); + std::sort( + data_slots.begin(), data_slots.end(), [](const auto *a, const auto *b) { + return a->ordinal.load(std::memory_order_relaxed) < + b->ordinal.load(std::memory_order_relaxed); + }); + + // Concurrent publishers reserve queue positions independently of ordinal + // assignment. Recreate the resulting newer-before-older hint order while + // retaining the authoritative bits written by PublishMessage(). + queue->DiscardAll(); + ASSERT_TRUE(queue->Push( + data_slots[1]->id, + data_slots[1]->ordinal.load(std::memory_order_relaxed))); + ASSERT_TRUE(queue->Push( + data_slots[0]->id, + data_slots[0]->ordinal.load(std::memory_order_relaxed))); + + for (uint8_t expected = 1; expected <= 2; ++expected) { + Message message = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, message.length); + EXPECT_EQ(expected, *static_cast(message.buffer)); + } +} + +TEST_F(ClientTest, SubscriberQueueOverflowReportsDroppedMessages) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_overflow_reporting"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + auto sub = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(2))); + int64_t reported_drops = 0; + ASSERT_OK(sub.RegisterDroppedMessageCallback( + [&reported_drops](Subscriber *, int64_t drops) { + reported_drops += drops; + })); + + for (uint8_t value = 1; value <= 4; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + Message message = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, message.length); + EXPECT_EQ(3, *static_cast(message.buffer)); + EXPECT_EQ(2, reported_drops); +} + +TEST_F(ClientTest, QueueMessageSurvivesMaxActiveMessageRejection) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_max_active"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + subspace::SubscriberOptions options; + options.SetSubscriberQueueSize(4).SetMaxActiveMessages(1); + auto sub = + EVAL_AND_ASSERT_OK(client.CreateSubscriber(kChannel, options)); + + for (uint8_t value = 1; value <= 3; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + Message first = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, first.length); + EXPECT_EQ(1, *static_cast(first.buffer)); + + Message blocked = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + EXPECT_EQ(0, blocked.length); + first.Reset(); + + Message recovered = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, recovered.length); + EXPECT_EQ(2, *static_cast(recovered.buffer)); + recovered.Reset(); + + Message next = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, next.length); + EXPECT_EQ(3, *static_cast(next.buffer)); +} + +TEST_F(ClientTest, SubscriberQueuePollDrainHandlesActivationOrdinals) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_poll_activation"; + auto sub = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(4))); + ASSERT_GE(sub.GetPollFd().fd, 0); + + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize) + .SetActivate(true))); + subspace::ServerChannel *server_channel = Server()->FindChannel(kChannel); + ASSERT_NE(nullptr, server_channel); + EXPECT_EQ(1, server_channel->GetCcb()->total_messages.load()); + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + memcpy(buffer, "visible", 7); + ASSERT_OK(pub.PublishMessage(7)); + EXPECT_EQ(2, server_channel->GetCcb()->total_messages.load()); + + Message message = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(7, message.length); + EXPECT_EQ(0, memcmp(message.buffer, "visible", 7)); +} + +TEST_F(ClientTest, SubscriberQueueOverrideExhaustingArenaIsRejected) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "subscriber_queue_arena_exhaustion", + subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(40) + .SetSubscriberQueueArenaSize( + 2 * subspace::SlotQueueBlockSize(1024) + + subspace::SlotQueueBlockSize( + subspace::kDefaultSubscriberQueueSize)))); + + std::vector large_subscribers; + bool exhausted = false; + for (int i = 0; i < 32; ++i) { + auto subscriber = client.CreateSubscriber( + "subscriber_queue_arena_exhaustion", + subspace::SubscriberOptions().SetSubscriberQueueSize(1024)); + if (!subscriber.ok()) { + EXPECT_THAT(subscriber.status().message(), + ::testing::HasSubstr("does not fit")); + exhausted = true; + break; + } + large_subscribers.push_back(std::move(*subscriber)); + } + ASSERT_TRUE(exhausted); + + // Retiring one queue makes its arena block available to the next subscriber. + large_subscribers.pop_back(); + auto replacement = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + "subscriber_queue_arena_exhaustion", + subspace::SubscriberOptions().SetSubscriberQueueSize(1024))); + EXPECT_EQ(1024, replacement.SubscriberQueueSize()); + + auto defaults = EVAL_AND_ASSERT_OK( + client.CreateSubscriber("subscriber_queue_arena_exhaustion")); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + defaults.SubscriberQueueSize()); +} + +TEST_F(ClientTest, SubscriberQueueReuseWaitsForPublisherTraversal) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_hazard_reuse"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(16) + .SetSubscriberQueueArenaSize( + 3 * subspace::SlotQueueBlockSize(1024)))); + subspace::ServerChannel *channel = Server()->FindChannel(kChannel); + ASSERT_NE(nullptr, channel); + + int publisher_id = -1; + for (const auto &entry : channel->GetUsers()) { + if (entry.second->IsPublisher()) { + publisher_id = entry.first; + } + } + ASSERT_GE(publisher_id, 0); + + uint64_t first_offset = 0; + { + auto first = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, + subspace::SubscriberOptions().SetSubscriberQueueSize(1024))); + int subscriber_id = -1; + channel->GetCcb()->subscribers.Traverse( + [&subscriber_id](int id) { subscriber_id = id; }); + ASSERT_GE(subscriber_id, 0); + first_offset = channel->GetAvailableSlotQueueIndexAddress() + ->offsets[subscriber_id] + .load(std::memory_order_acquire); + channel->BeginSubscriberQueuePublish(publisher_id); + } + + uint64_t second_offset = 0; + { + auto second = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, + subspace::SubscriberOptions().SetSubscriberQueueSize(1024))); + int subscriber_id = -1; + channel->GetCcb()->subscribers.Traverse( + [&subscriber_id](int id) { subscriber_id = id; }); + ASSERT_GE(subscriber_id, 0); + second_offset = channel->GetAvailableSlotQueueIndexAddress() + ->offsets[subscriber_id] + .load(std::memory_order_acquire); + EXPECT_NE(first_offset, second_offset); + } + + channel->EndSubscriberQueuePublish(publisher_id); + auto reclaimed = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(1024))); + int reclaimed_id = -1; + channel->GetCcb()->subscribers.Traverse( + [&reclaimed_id](int id) { reclaimed_id = id; }); + ASSERT_GE(reclaimed_id, 0); + const uint64_t reclaimed_offset = + channel->GetAvailableSlotQueueIndexAddress() + ->offsets[reclaimed_id] + .load(std::memory_order_acquire); + EXPECT_EQ(first_offset, reclaimed_offset); +} + +TEST_F(ClientTest, SubscriberQueueArenaCoalescesAdjacentBlocks) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_coalesce"; + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(16) + .SetSubscriberQueueArenaSize( + 2 * subspace::SlotQueueBlockSize(512)))); + subspace::ServerChannel *channel = Server()->FindChannel(kChannel); + ASSERT_NE(nullptr, channel); + + auto first = std::make_unique( + EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, + subspace::SubscriberOptions().SetSubscriberQueueSize(512)))); + auto second = std::make_unique( + EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, + subspace::SubscriberOptions().SetSubscriberQueueSize(512)))); + std::vector queue_offsets; + channel->GetCcb()->subscribers.Traverse([channel, &queue_offsets](int id) { + queue_offsets.push_back(channel->GetAvailableSlotQueueIndexAddress() + ->offsets[id] + .load(std::memory_order_acquire)); + }); + ASSERT_EQ(2, queue_offsets.size()); + const uint64_t first_offset = queue_offsets[0]; + const uint64_t second_offset = queue_offsets[1]; + const uint64_t lower_offset = std::min(first_offset, second_offset); + first.reset(); + second.reset(); + + auto coalesced = EVAL_AND_ASSERT_OK(client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(1024))); + int coalesced_id = -1; + channel->GetCcb()->subscribers.Traverse( + [&coalesced_id](int id) { coalesced_id = id; }); + ASSERT_GE(coalesced_id, 0); + const uint64_t coalesced_offset = + channel->GetAvailableSlotQueueIndexAddress() + ->offsets[coalesced_id] + .load(std::memory_order_acquire); + EXPECT_EQ(lower_offset, coalesced_offset); +} + +TEST_F(ClientTest, SubscriberFirstQueueOverrideSurvivesPlaceholderRemap) { + subspace::Client pub_client; + subspace::Client sub_client; + ASSERT_OK(pub_client.Init(Socket())); + ASSERT_OK(sub_client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_first_queue_override"; + auto sub = EVAL_AND_ASSERT_OK(sub_client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(2))); + EXPECT_TRUE(sub.IsPlaceholder()); + EXPECT_EQ(0, sub.SubscriberQueueSize()); + + auto pub = EVAL_AND_ASSERT_OK(pub_client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(32) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + for (uint8_t value = 1; value <= 4; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + Message message = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(1, message.length); + EXPECT_EQ(3, *static_cast(message.buffer)); + EXPECT_FALSE(sub.IsPlaceholder()); + EXPECT_EQ(2, sub.SubscriberQueueSize()); +} + +TEST_F(ClientTest, SubscriberFirstOversizedQueueFallsBackToBitset) { + subspace::Client pub_client; + subspace::Client sub_client; + ASSERT_OK(pub_client.Init(Socket())); + ASSERT_OK(sub_client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_first_oversized_queue"; + std::vector subscribers; + for (int i = 0; i < 16; ++i) { + subscribers.push_back(EVAL_AND_ASSERT_OK(sub_client.CreateSubscriber( + kChannel, subspace::SubscriberOptions().SetSubscriberQueueSize(1024)))); + ASSERT_TRUE(subscribers.back().IsPlaceholder()); + } + + auto pub = EVAL_AND_ASSERT_OK(pub_client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(32) + .SetSubscriberQueueArenaSize( + 8 * subspace::SlotQueueBlockSize(1024)))); + for (uint8_t value = 1; value <= 4; ++value) { + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + *static_cast(buffer) = value; + ASSERT_OK(pub.PublishMessage(1)); + } + + int queued = 0; + int bitset = 0; + for (Subscriber &sub : subscribers) { + Message message = + EVAL_AND_ASSERT_OK(sub.ReadMessage(subspace::ReadMode::kReadNewest)); + ASSERT_EQ(1, message.length); + EXPECT_EQ(4, *static_cast(message.buffer)); + EXPECT_FALSE(sub.IsPlaceholder()); + if (sub.SubscriberQueueSize() == 0) { + ++bitset; + } else { + EXPECT_EQ(1024, sub.SubscriberQueueSize()); + ++queued; + } + } + EXPECT_GT(queued, 0); + EXPECT_GT(bitset, 0); +} + +TEST_F(ClientTest, SubscriberQueueChurnKeepsQueuesIndependent) { + subspace::Client pub_client; + subspace::Client sub_client; + ASSERT_OK(pub_client.Init(Socket())); + ASSERT_OK(sub_client.Init(Socket())); + + constexpr char kChannel[] = "subscriber_queue_churn"; + auto pub = EVAL_AND_ASSERT_OK(pub_client.CreatePublisher( + kChannel, subspace::PublisherOptions() + .SetSlotSize(64) + .SetNumSlots(64) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + + for (int iteration = 1; iteration <= 1100; ++iteration) { + const int queue_size = 1 + iteration % 4; + auto sub = EVAL_AND_ASSERT_OK(sub_client.CreateSubscriber( + kChannel, + subspace::SubscriberOptions().SetSubscriberQueueSize(queue_size))); + EXPECT_EQ(queue_size, sub.SubscriberQueueSize()); + + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + memcpy(buffer, &iteration, sizeof(iteration)); + ASSERT_OK(pub.PublishMessage(sizeof(iteration))); + + Message message = + EVAL_AND_ASSERT_OK(sub.ReadMessage(subspace::ReadMode::kReadNewest)); + ASSERT_EQ(sizeof(iteration), message.length); + EXPECT_EQ(iteration, *static_cast(message.buffer)); + } +} + TEST_F(ClientTest, SplitBuffersPublishWithHandlesAndSeparatePrefix) { subspace::Client pub_client; subspace::Client sub_client; @@ -1803,10 +2425,73 @@ TEST_F(ClientTest, PublishConcurrentlyFromOneClientToOneSubscriber) { ASSERT_OK(pub_client.Init(Socket())); for (int i = 0; i < kNumPublishers; ++i) { absl::StatusOr pub = pub_client.CreatePublisher( - channel_name, PubOpts(256, 2 * kNumPublishers + 16)); + channel_name, + PubOpts(256, 2 * kNumPublishers + 16) + .SetSubscriberQueueArenaSize(0)); + ASSERT_OK(pub) << pub.status(); + pubs.emplace_back(std::move(*pub)); + } + ASSERT_EQ(0, sub.SubscriberQueueSize()); + + std::vector pub_threads; + pub_threads.reserve(kNumPublishers); + for (int i = 0; i < kNumPublishers; ++i) { + pub_threads.emplace_back(std::thread([&pubs, i]() { + std::array msg = {}; + auto size = std::snprintf(msg.data(), msg.size(), "M%d", i); + auto buffer = pubs[i].GetMessageBuffer(size); + ASSERT_OK(buffer) << buffer.status(); + ASSERT_NE(nullptr, *buffer); + std::memcpy(*buffer, msg.data(), size); + ASSERT_OK(pubs[i].PublishMessage(size)); + })); + } + + for (auto &t : pub_threads) { + t.join(); + } + + std::vector all_recv_msgs; + all_recv_msgs.reserve(kNumPublishers); + while (true) { + auto message = *sub.ReadMessage(); + size_t size = message.length; + if (size == 0) { + break; + } + all_recv_msgs.emplace_back(std::string( + reinterpret_cast(message.buffer), message.length)); + } + EXPECT_EQ(all_recv_msgs.size(), kNumPublishers); + std::sort(all_recv_msgs.begin(), all_recv_msgs.end()); + auto last_uniq = std::unique(all_recv_msgs.begin(), all_recv_msgs.end()); + EXPECT_EQ(last_uniq - all_recv_msgs.begin(), kNumPublishers); +} + +TEST_F(ClientTest, PublishConcurrentlyFromOneClientToOneQueuedSubscriber) { + std::string channel_name = "checkin_channel_queued"; + subspace::Client sub_client; + ASSERT_OK(sub_client.Init(Socket())); + + const int kNumPublishers = + absl::GetFlag(FLAGS_use_split_buffers) ? 16 : 100; + std::vector pubs; + pubs.reserve(kNumPublishers); + subspace::Client pub_client; + InitClient(pub_client); + for (int i = 0; i < kNumPublishers; ++i) { + absl::StatusOr pub = pub_client.CreatePublisher( + channel_name, + PubOpts(256, 2 * kNumPublishers + 16) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize)); ASSERT_OK(pub) << pub.status(); pubs.emplace_back(std::move(*pub)); } + auto sub = EVAL_AND_ASSERT_OK(sub_client.CreateSubscriber( + channel_name, + SubOpts().SetSubscriberQueueSize(kNumPublishers))); + ASSERT_EQ(kNumPublishers, sub.SubscriberQueueSize()); std::vector pub_threads; pub_threads.reserve(kNumPublishers); @@ -1875,7 +2560,9 @@ TEST_F(ClientTest, PublishConcurrentlyToOneSubscriber) { } ASSERT_TRUE(connected); absl::StatusOr pub = pub_client.CreatePublisher( - channel_name, PubOpts(256, 2 * kNumPublishers + 16)); + channel_name, + PubOpts(256, 2 * kNumPublishers + 16) + .SetSubscriberQueueArenaSize(0)); ASSERT_OK(pub) << pub.status(); std::array msg = {}; auto size = std::snprintf(msg.data(), msg.size(), "M%d", i); @@ -1894,6 +2581,106 @@ TEST_F(ClientTest, PublishConcurrentlyToOneSubscriber) { for (auto &t : pub_threads) { t.join(); } + ASSERT_EQ(0, sub.SubscriberQueueSize()); + + std::vector all_recv_msgs; + all_recv_msgs.reserve(kNumPublishers); + while (true) { + auto message = *sub.ReadMessage(); + size_t size = message.length; + if (size == 0) { + break; + } + all_recv_msgs.emplace_back(std::string( + reinterpret_cast(message.buffer), message.length)); + } + EXPECT_EQ(all_recv_msgs.size(), kNumPublishers); + std::sort(all_recv_msgs.begin(), all_recv_msgs.end()); + auto last_uniq = std::unique(all_recv_msgs.begin(), all_recv_msgs.end()); + EXPECT_EQ(last_uniq - all_recv_msgs.begin(), kNumPublishers); +} + +TEST_F(ClientTest, PublishConcurrentlyToOneQueuedSubscriber) { + std::string channel_name = "checkin_channel_multi_client_queued"; + subspace::Client sub_client; + ASSERT_OK(sub_client.Init(Socket())); + + std::vector pub_threads; +#ifdef __APPLE__ + constexpr int kNumPublishers = 16; +#else + const int kNumPublishers = + absl::GetFlag(FLAGS_use_split_buffers) ? 16 : 100; +#endif + auto channel_publisher = EVAL_AND_ASSERT_OK(sub_client.CreatePublisher( + channel_name, + PubOpts(256, 2 * kNumPublishers + 16) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize))); + ASSERT_EQ(subspace::kDefaultSubscriberQueueArenaSize, + channel_publisher.SubscriberQueueArenaSize()); + auto sub = EVAL_AND_ASSERT_OK(sub_client.CreateSubscriber( + channel_name, + SubOpts().SetSubscriberQueueSize(kNumPublishers))); + ASSERT_EQ(kNumPublishers, sub.SubscriberQueueSize()); + + pub_threads.reserve(kNumPublishers); + std::atomic publishers_finished{0}; + for (int i = 0; i < kNumPublishers; ++i) { + pub_threads.emplace_back(std::thread( + [&channel_name, &publishers_finished, kNumPublishers, i]() { + // Keep every publisher alive until all messages have been published. + subspace::Client pub_client; + absl::StatusOr pub = + absl::UnknownError("publisher not created"); + [&]() { + bool connected = false; + for (int attempt = 0; attempt < 100; ++attempt) { + if (pub_client.Init(Socket()).ok()) { + connected = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + if (!connected) { + ADD_FAILURE() << "Failed to connect publisher " << i; + return; + } + pub = pub_client.CreatePublisher( + channel_name, + PubOpts(256, 2 * kNumPublishers + 16) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize)); + if (!pub.ok()) { + ADD_FAILURE() << pub.status(); + return; + } + std::array msg = {}; + auto size = std::snprintf(msg.data(), msg.size(), "M%d", i); + auto buffer = pub->GetMessageBuffer(size); + if (!buffer.ok() || *buffer == nullptr) { + ADD_FAILURE() << buffer.status(); + return; + } + std::memcpy(*buffer, msg.data(), size); + auto publish_status = pub->PublishMessage(size); + if (!publish_status.ok()) { + ADD_FAILURE() << publish_status.status(); + return; + } + }(); + publishers_finished.fetch_add(1, std::memory_order_release); + while (publishers_finished.load(std::memory_order_acquire) < + kNumPublishers) { + std::this_thread::yield(); + } + })); + } + + for (auto &t : pub_threads) { + t.join(); + } + ASSERT_EQ(kNumPublishers, sub.SubscriberQueueSize()); std::vector all_recv_msgs; all_recv_msgs.reserve(kNumPublishers); @@ -2289,6 +3076,35 @@ TEST_F(ClientTest, ReliablePublisher1) { machine.Run(); } +TEST_F(ClientTest, ReliablePublisherDoesNotBlockOnUnreliableSubscriber) { + subspace::Client client; + InitClient(client); + + constexpr int kNumSlots = 5; + absl::StatusOr pub = client.CreatePublisher( + "rel_pub_unrel_sub", 32, kNumSlots, + subspace::PublisherOptions().SetReliable(true)); + ASSERT_OK(pub); + absl::StatusOr sub = client.CreateSubscriber( + "rel_pub_unrel_sub", subspace::SubscriberOptions().SetReliable(false)); + ASSERT_OK(sub); + + const auto &counters = pub->GetChannelCounters(); + ASSERT_EQ(1, counters.num_reliable_pubs); + ASSERT_EQ(0, counters.num_reliable_subs); + + // An unreliable subscriber may fall behind and drop messages, but it must not + // make reliable publishers wait for every old slot to be observed. + for (int i = 0; i < kNumSlots * 4; i++) { + absl::StatusOr buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + ASSERT_NE(nullptr, *buffer) << "publish " << i; + memcpy(*buffer, "foobar", 6); + absl::StatusOr pub_status = pub->PublishMessage(6); + ASSERT_OK(pub_status); + } +} + TEST_F(ClientTest, ReliablePublisher2) { subspace::Client client; InitClient(client); @@ -2726,6 +3542,61 @@ TEST_F(ClientTest, DroppedMessage) { ASSERT_EQ(4, num_dropped_messages); } +TEST_F(ClientTest, DroppedMessageDetectionCanBeDisabled) { + subspace::Client client; + InitClient(client); + + absl::StatusOr sub = client.CreateSubscriber( + "drop_detection_disabled", + SubOpts().SetKeepActiveMessage(true).SetDetectDroppedMessages(false)); + ASSERT_OK(sub); + + int num_dropped_messages = 0; + ASSERT_OK(sub->RegisterDroppedMessageCallback( + [&num_dropped_messages](Subscriber *, int64_t num_dropped) { + num_dropped_messages += num_dropped; + })); + + absl::StatusOr pub = + client.CreatePublisher("drop_detection_disabled", 32, 5); + ASSERT_OK(pub); + + for (int i = 0; i < 4; i++) { + absl::StatusOr buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + memcpy(*buffer, "foobar", 6); + ASSERT_OK(pub->PublishMessage(6)); + } + + absl::StatusOr msg = sub->ReadMessage(); + ASSERT_OK(msg); + ASSERT_EQ(6, msg->length); + + for (int i = 0; i < 4; i++) { + absl::StatusOr buffer = pub->GetMessageBuffer(); + ASSERT_OK(buffer); + memcpy(*buffer, "foobar", 6); + ASSERT_OK(pub->PublishMessage(6)); + } + + for (;;) { + msg = sub->ReadMessage(); + ASSERT_OK(msg); + if (msg->length == 0) { + break; + } + } + ASSERT_EQ(0, num_dropped_messages); + + uint64_t total_bytes = 0; + uint64_t total_messages = 0; + uint32_t max_message_size = 0; + uint32_t total_drops = 0; + pub->GetStatsCounters(total_bytes, total_messages, max_message_size, + total_drops); + ASSERT_EQ(0u, total_drops); +} + TEST_F(ClientTest, PublishSingleMessageAndReadSharedPtr) { subspace::Client pub_client; subspace::Client sub_client; @@ -5403,6 +6274,43 @@ TEST_F(ClientTest, MaxActiveMessagesTooSmall) { ::testing::HasSubstr("MaxActiveMessages")); } +TEST_F(ClientTest, CapacityErrorIdentifiesExistingClients) { + auto publisher_client_a = EVAL_AND_ASSERT_OK( + subspace::Client::Create(Socket(), "capacity-publisher-a")); + auto publisher_client_b = EVAL_AND_ASSERT_OK( + subspace::Client::Create(Socket(), "capacity-publisher-b")); + auto subscriber_client = EVAL_AND_ASSERT_OK( + subspace::Client::Create(Socket(), "capacity-subscriber")); + auto rejected_client = EVAL_AND_ASSERT_OK( + subspace::Client::Create(Socket(), "capacity-rejected")); + + [[maybe_unused]] auto publisher_a = EVAL_AND_ASSERT_OK( + publisher_client_a->CreatePublisher("capacity_clients", PubOpts(64, 6))); + [[maybe_unused]] auto subscriber = + EVAL_AND_ASSERT_OK(subscriber_client->CreateSubscriber( + "capacity_clients", SubOpts().SetMaxActiveMessages(2))); + [[maybe_unused]] auto publisher_b = EVAL_AND_ASSERT_OK( + publisher_client_b->CreatePublisher("capacity_clients", PubOpts(64, 6))); + + auto rejected = rejected_client->CreateSubscriber( + "capacity_clients", SubOpts().SetMaxActiveMessages(2)); + ASSERT_FALSE(rejected.ok()); + const std::string error(rejected.status().message()); + EXPECT_THAT(error, ::testing::HasSubstr("publishers=[")); + EXPECT_THAT(error, + ::testing::HasSubstr("client=\"capacity-publisher-a\"")); + EXPECT_THAT(error, + ::testing::HasSubstr("client=\"capacity-publisher-b\"")); + EXPECT_THAT(error, ::testing::HasSubstr("subscribers=[")); + EXPECT_THAT(error, ::testing::HasSubstr("client=\"capacity-subscriber\"")); + EXPECT_THAT(error, ::testing::HasSubstr("max_active_messages=2")); + EXPECT_THAT(error, + ::testing::HasSubstr(absl::StrFormat( + "pid=%llu", static_cast(getpid())))); + EXPECT_EQ(std::string::npos, error.find("{id=")); + EXPECT_EQ(std::string::npos, error.find("reliable=")); +} + TEST_F(ClientTest, OnReceiveCallbackSuccess) { subspace::Client pub_client; subspace::Client sub_client; @@ -5455,6 +6363,36 @@ TEST_F(ClientTest, OnReceiveCallbackError) { EXPECT_THAT(msg.status().message(), ::testing::HasSubstr("receive callback failed")); sub.ClearOnReceiveCallback(); + + // The callback runs after NextSlot has claimed a shared slot ref. The error + // path must roll that ref back without clearing the subscriber bit so the + // same message remains readable. + Message retried = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(10, retried.length); + EXPECT_EQ(0, memcmp(retried.buffer, "qqqqqqqqqq", 10)); +} + +TEST_F(ClientTest, OnReceiveCallbackZeroSizeReleasesSlot) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + auto pub = EVAL_AND_ASSERT_OK( + client.CreatePublisher("onrecv_zero", PubOpts(64, 4))); + auto sub = + EVAL_AND_ASSERT_OK(client.CreateSubscriber("onrecv_zero")); + sub.SetOnReceiveCallback( + [](void *, int64_t) -> absl::StatusOr { return 0; }); + + void *buffer = EVAL_AND_ASSERT_OK(pub.GetMessageBuffer()); + memcpy(buffer, "retry", 5); + ASSERT_OK(pub.PublishMessage(5)); + Message empty = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + EXPECT_EQ(0, empty.length); + + sub.ClearOnReceiveCallback(); + Message retried = EVAL_AND_ASSERT_OK(sub.ReadMessage()); + ASSERT_EQ(5, retried.length); + EXPECT_EQ(0, memcmp(retried.buffer, "retry", 5)); } TEST_F(ClientTest, ProcessAllMessagesWithoutCallback) { @@ -5686,6 +6624,7 @@ TEST_F(ClientTest, PublisherOptionsChain) { subspace::PublisherOptions opts; opts.SetSlotSize(128) .SetNumSlots(8) + .SetSubscriberQueueArenaSize(32'000) .SetReliable(true) .SetLocal(true) .SetFixedSize(true) @@ -5702,6 +6641,7 @@ TEST_F(ClientTest, PublisherOptionsChain) { ASSERT_EQ(128, opts.SlotSize()); ASSERT_EQ(8, opts.NumSlots()); + ASSERT_EQ(32'000, opts.SubscriberQueueArenaSize()); ASSERT_TRUE(opts.IsReliable()); ASSERT_TRUE(opts.IsLocal()); ASSERT_TRUE(opts.IsFixedSize()); @@ -5717,9 +6657,62 @@ TEST_F(ClientTest, PublisherOptionsChain) { ASSERT_EQ(3, opts.MaxPublishers()); } +TEST_F(ClientTest, PublisherSubscriberQueueArenaSizeOption) { + subspace::Client client; + ASSERT_OK(client.Init(Socket())); + + auto pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "subscriber_queue_size", + subspace::PublisherOptions() + .SetSlotSize(128) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize(32'000))); + EXPECT_EQ(8, pub.NumSlots()); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + pub.SubscriberQueueSize()); + EXPECT_EQ(32'000, pub.SubscriberQueueArenaSize()); + + auto sub = EVAL_AND_ASSERT_OK( + client.CreateSubscriber("subscriber_queue_size")); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + sub.SubscriberQueueSize()); + + auto info = EVAL_AND_ASSERT_OK(client.GetChannelInfo("subscriber_queue_size")); + EXPECT_EQ(subspace::kDefaultSubscriberQueueSize, + info.subscriber_queue_size); + EXPECT_EQ(32'000, info.subscriber_queue_arena_size); + + auto default_pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "subscriber_queue_size_default", + subspace::PublisherOptions().SetSlotSize(128).SetNumSlots(8))); + EXPECT_EQ(8, default_pub.NumSlots()); + EXPECT_EQ(0, default_pub.SubscriberQueueSize()); + EXPECT_EQ(0, default_pub.SubscriberQueueArenaSize()); + auto default_sub = EVAL_AND_ASSERT_OK( + client.CreateSubscriber("subscriber_queue_size_default")); + EXPECT_EQ(0, default_sub.SubscriberQueueSize()); + auto default_info = + EVAL_AND_ASSERT_OK(client.GetChannelInfo("subscriber_queue_size_default")); + EXPECT_EQ(0, default_info.subscriber_queue_size); + EXPECT_EQ(0, default_info.subscriber_queue_arena_size); + + auto disabled_pub = EVAL_AND_ASSERT_OK(client.CreatePublisher( + "subscriber_queue_size_disabled", + subspace::PublisherOptions() + .SetSlotSize(128) + .SetNumSlots(8) + .SetSubscriberQueueArenaSize(0))); + EXPECT_EQ(0, disabled_pub.SubscriberQueueSize()); + EXPECT_EQ(0, disabled_pub.SubscriberQueueArenaSize()); + auto disabled_sub = EVAL_AND_ASSERT_OK( + client.CreateSubscriber("subscriber_queue_size_disabled")); + EXPECT_EQ(0, disabled_sub.SubscriberQueueSize()); +} + TEST_F(ClientTest, SubscriberOptionsChain) { subspace::SubscriberOptions opts; opts.SetReliable(true) + .SetSubscriberQueueSize(12) .SetType("sub_type") .SetMaxActiveMessages(20) .SetBridge(true) @@ -5730,14 +6723,17 @@ TEST_F(ClientTest, SubscriberOptionsChain) { .SetReadWrite(true) .SetChecksum(true) .SetPassChecksumErrors(true) - .SetKeepActiveMessage(true); + .SetKeepActiveMessage(true) + .SetDetectDroppedMessages(false); opts.SetLogDroppedMessages(true); ASSERT_TRUE(opts.IsReliable()); + ASSERT_EQ(12, opts.SubscriberQueueSize()); ASSERT_EQ("sub_type", opts.Type()); ASSERT_EQ(19, opts.MaxSharedPtrs()); ASSERT_EQ(20, opts.MaxActiveMessages()); ASSERT_TRUE(opts.LogDroppedMessages()); + ASSERT_FALSE(opts.DetectDroppedMessages()); ASSERT_TRUE(opts.IsBridge()); ASSERT_TRUE(opts.ForTunnel()); ASSERT_EQ("/submux", opts.Mux()); diff --git a/client/latency_test.cc b/client/latency_test.cc index 7524948a..0e3ac5ba 100644 --- a/client/latency_test.cc +++ b/client/latency_test.cc @@ -426,7 +426,7 @@ TEST_F(LatencyTest, MultithreadedUnreliableLatency) { ASSERT_OK(pub); absl::StatusOr sub = sub_client.CreateSubscriber( - "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); uint64_t start_time = toolbelt::Now(); @@ -512,7 +512,7 @@ TEST_F(LatencyTest, PublisherLatency) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( - "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); uint64_t total_time = 0; @@ -585,7 +585,7 @@ TEST_F(LatencyTest, PublisherLatencyChecksum) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( "publat", - ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetChecksum(true); return opts; }())); + ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetChecksum(true); return opts; }())); ASSERT_OK(sub); uint64_t total_time = 0; @@ -666,7 +666,7 @@ TEST_F(LatencyTest, PublisherLatencyPayload) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( - "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); uint64_t total_time = 0; @@ -756,7 +756,7 @@ TEST_F(LatencyTest, PublisherLatencyPayloadChecksum) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( "publat", - ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetChecksum(true); return opts; }())); + ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetChecksum(true); return opts; }())); ASSERT_OK(sub); uint64_t total_time = 0; @@ -865,7 +865,7 @@ TEST_F(LatencyTest, PublisherLatencyHistogram) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( - "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); std::vector latencies; @@ -962,7 +962,7 @@ TEST_F(LatencyTest, PublisherLatencyHistogramThreadSafe) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( - "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); std::vector latencies; @@ -1043,7 +1043,7 @@ TEST_F(LatencyTest, PublisherLatencyMultiSub) { for (int i = 0; i < num_subs; i++) { absl::StatusOr sub = sub_client.CreateSubscriber( - "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "publat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); subs.push_back(std::move(*sub)); } @@ -1121,7 +1121,7 @@ TEST_F(LatencyTest, VirtualPublisherLatency) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( "publat", - ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); + ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); ASSERT_OK(sub); uint64_t total_time = 0; @@ -1200,7 +1200,7 @@ TEST_F(LatencyTest, VirtualPublisherLatencyMultiSub) { for (int i = 0; i < num_subs; i++) { absl::StatusOr sub = sub_client.CreateSubscriber( "publat", - ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); + ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); ASSERT_OK(sub); subs.push_back(std::move(*sub)); } @@ -1278,12 +1278,12 @@ TEST_F(LatencyTest, VirtualPublisherMuxLatency) { std::cerr << num_slots << ","; absl::StatusOr sub = sub_client.CreateSubscriber( "publat", - ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); + ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetMux("/foo"); return opts; }())); ASSERT_OK(sub); // Mux subscriber. absl::StatusOr mux_sub = sub_client.CreateSubscriber( - "/foo", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "/foo", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(mux_sub); uint64_t total_time = 0; @@ -1337,9 +1337,10 @@ TEST_F(LatencyTest, VirtualPublisherMuxLatency) { } } -// This measures unreliable latency by sending as fast as possible. It will -// drop messages because the publisher will run faster than the subscriber -// most of the time. +// This measures unreliable latency by sending as fast as possible. It uses a +// subscriber queue large enough to track every retained slot (up to the +// supported queue limit), but can still drop messages when the publisher +// overwrites slots faster than the subscriber consumes them. TEST_F(LatencyTest, MultithreadedUnreliableLatencyHistogram) { subspace::Client pub_client; subspace::Client sub_client; @@ -1352,13 +1353,28 @@ TEST_F(LatencyTest, MultithreadedUnreliableLatencyHistogram) { for (int num_slots = 3; num_slots < LatencyValueForSplitBuffers(20000, 4096); num_slots *= 2) { - std::cerr << "num_slots: " << num_slots << "\n"; + const int subscriber_queue_size = std::min(num_slots, 1024); + std::cerr << "num_slots: " << num_slots + << ", subscriber_queue_size: " << subscriber_queue_size << "\n"; absl::StatusOr pub = pub_client.CreatePublisher( - "lustress", 256, num_slots, subspace::PublisherOptions().SetReliable(false)); + "lustress", + subspace::PublisherOptions() + .SetSlotSize(256) + .SetNumSlots(num_slots) + .SetReliable(false) + .SetSubscriberQueueArenaSize( + subscriber_queue_size == 0 + ? 0 + : subspace::SlotQueueBlockSize(subscriber_queue_size))); ASSERT_OK(pub); + subspace::SubscriberOptions subscriber_options; + subscriber_options.SetReliable(false); + subscriber_options.SetLogDroppedMessages(false); + subscriber_options.SetDetectDroppedMessages(false); + subscriber_options.SetSubscriberQueueSize(subscriber_queue_size); absl::StatusOr sub = sub_client.CreateSubscriber( - "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "lustress", subscriber_options); ASSERT_OK(sub); uint64_t start_time = toolbelt::Now(); @@ -1447,7 +1463,7 @@ TEST_F(LatencyTest, MultithreadedUnreliableLatencyPayload) { ASSERT_OK(pub); absl::StatusOr sub = sub_client.CreateSubscriber( - "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); // Create a subscriber thread to read from the channel and write to random @@ -1572,7 +1588,7 @@ TEST_F(LatencyTest, MultithreadedUnreliableLatencyPayloadHistogram) { ASSERT_OK(pub); absl::StatusOr sub = sub_client.CreateSubscriber( - "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "lustress", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); // Create a subscriber thread to read from the channel and write to random @@ -1679,6 +1695,173 @@ TEST_F(LatencyTest, MultithreadedUnreliableLatencyPayloadHistogram) { } } +TEST_F(LatencyTest, FlatOutSubscriberQueueLatency) { + subspace::Client pub_client; + subspace::Client sub_client; + ASSERT_OK(pub_client.Init(Socket())); + ASSERT_OK(sub_client.Init(Socket())); + + const int kNumMessages = std::atoi( + LatencyEnvOrDefault("SUBSPACE_QUEUE_SWEEP_MESSAGES", "50000")); + const int kNumSlots = std::atoi( + LatencyEnvOrDefault("SUBSPACE_QUEUE_SWEEP_SLOTS", "1024")); + const std::vector queue_sizes = {0, 1, 2, 4, 8, 16, + 32, 64, 128, 256, 512, 1024}; + + struct Stats { + int subscriber_queue_size = 0; + int received = 0; + int dropped = 0; + uint64_t min = 0; + uint64_t max = 0; + uint64_t p50 = 0; + uint64_t p99 = 0; + uint64_t avg = 0; + uint64_t publish_avg = 0; + uint64_t elapsed = 0; + }; + + auto publish_timestamp = [](Publisher &pub) { + for (;;) { + absl::StatusOr buffer = pub.GetMessageBuffer(); + ASSERT_OK(buffer); + if (*buffer == nullptr) { + ASSERT_OK(pub.Wait()); + continue; + } + const uint64_t send_time = toolbelt::Now(); + memcpy(*buffer, &send_time, sizeof(send_time)); + absl::StatusOr pub_status = + pub.PublishMessage(sizeof(send_time)); + ASSERT_OK(pub_status); + return; + } + }; + + std::vector stats; + stats.reserve(queue_sizes.size()); + for (int subscriber_queue_size : queue_sizes) { + const std::string channel_name = + absl::StrFormat("flatout_queue_latency_%d", subscriber_queue_size); + absl::StatusOr pub = pub_client.CreatePublisher( + channel_name, + subspace::PublisherOptions() + .SetSlotSize(256) + .SetNumSlots(kNumSlots) + .SetSubscriberQueueArenaSize( + subscriber_queue_size == 0 + ? 0 + : subspace::SlotQueueBlockSize(subscriber_queue_size)) + .SetReliable(false)); + ASSERT_OK(pub); + + subspace::SubscriberOptions subscriber_options; + subscriber_options.SetReliable(false); + subscriber_options.SetSubscriberQueueSize(subscriber_queue_size); + subscriber_options.SetLogDroppedMessages(false); + subscriber_options.SetDetectDroppedMessages(false); + absl::StatusOr sub = + sub_client.CreateSubscriber(channel_name, subscriber_options); + ASSERT_OK(sub); + + Stats result; + result.subscriber_queue_size = subscriber_queue_size; + std::atomic received{0}; + std::atomic dropped{0}; + std::vector latencies; + latencies.reserve(kNumMessages); + + const uint64_t start_time = toolbelt::Now(); + std::thread sub_thread([&sub, &received, &dropped, &latencies, + kNumMessages]() { + uint64_t last_ordinal = 0; + ASSERT_OK(sub->Wait()); + while (last_ordinal < static_cast(kNumMessages)) { + absl::StatusOr msg = sub->ReadMessage(); + ASSERT_OK(msg); + if (msg->length == 0) { + continue; + } + + const uint64_t receive_time = toolbelt::Now(); + const uint64_t ordinal = msg->ordinal; + if (ordinal > last_ordinal + 1) { + const uint64_t last_original_ordinal = + std::min(ordinal - 1, kNumMessages); + dropped += last_original_ordinal - last_ordinal; + } + last_ordinal = ordinal; + + if (ordinal <= static_cast(kNumMessages)) { + const uint64_t send_time = + *reinterpret_cast(msg->buffer); + latencies.push_back(receive_time - send_time); + received++; + } + } + }); + + const uint64_t publish_start = toolbelt::Now(); + for (int i = 0; i < kNumMessages; i++) { + publish_timestamp(*pub); + } + const uint64_t publish_end = toolbelt::Now(); + + // If the subscriber missed the final run of original messages, publish a + // few extra wakeups so it can observe the ordinal gap and terminate. + for (int i = 0; i < 1000; i++) { + publish_timestamp(*pub); + if (received.load() + dropped.load() >= kNumMessages) { + break; + } + } + sub_thread.join(); + result.elapsed = toolbelt::Now() - start_time; + result.publish_avg = (publish_end - publish_start) / kNumMessages; + result.received = received.load(); + result.dropped = dropped.load(); + + if (!latencies.empty()) { + std::sort(latencies.begin(), latencies.end()); + result.min = latencies.front(); + result.max = latencies.back(); + result.p50 = latencies[latencies.size() / 2]; + result.p99 = latencies[latencies.size() * 99 / 100]; + uint64_t sum = 0; + for (uint64_t latency : latencies) { + sum += latency; + } + result.avg = sum / latencies.size(); + } + stats.push_back(result); + + EmitLatencyMetric("FlatOutSubscriberQueueLatency", "receive_latency", + "subscriber_queue_size", subscriber_queue_size, "min", + result.min); + EmitLatencyMetric("FlatOutSubscriberQueueLatency", "receive_latency", + "subscriber_queue_size", subscriber_queue_size, "median", + result.p50); + EmitLatencyMetric("FlatOutSubscriberQueueLatency", "receive_latency", + "subscriber_queue_size", subscriber_queue_size, "p99", + result.p99); + EmitLatencyMetric("FlatOutSubscriberQueueLatency", "receive_latency", + "subscriber_queue_size", subscriber_queue_size, "average", + result.avg); + EmitLatencyMetric("FlatOutSubscriberQueueLatency", "publisher_latency", + "subscriber_queue_size", subscriber_queue_size, "average", + result.publish_avg); + } + + std::cerr << "subscriber_queue_size,received,dropped,min_ns,p50_ns,p99_ns," + "max_ns,avg_ns,publish_avg_ns,elapsed_ns\n"; + for (const Stats &result : stats) { + std::cerr << result.subscriber_queue_size << "," << result.received << "," + << result.dropped << "," << result.min << "," << result.p50 + << "," << result.p99 << "," << result.max << "," << result.avg + << "," << result.publish_avg << "," << result.elapsed << "\n"; + } +} + TEST_F(LatencyTest, ManyChannelsNonMultiplexed) { std::vector pub_clients; subspace::Client sub_client; @@ -1710,7 +1893,7 @@ TEST_F(LatencyTest, ManyChannelsNonMultiplexed) { std::vector subs; for (int i = 0; i < kNumChannels; i++) { absl::StatusOr sub = sub_client.CreateSubscriber( - channels[i], ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); return opts; }())); + channels[i], ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); // std::cerr << "sub status " << sub.status() << "\n"; ASSERT_OK(sub); subs.push_back(std::move(*sub)); @@ -1836,7 +2019,7 @@ TEST_F(LatencyTest, ManyChannelsMultiplexed) { std::vector subs; for (int i = 0; i < kNumChannels; i++) { absl::StatusOr sub = sub_client.CreateSubscriber( - channels[i], ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); opts.SetMux(kMux); return opts; }())); + channels[i], ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); opts.SetMux(kMux); return opts; }())); // std::cerr << "sub status " << sub.status() << "\n"; ASSERT_OK(sub); subs.push_back(std::move(*sub)); @@ -1960,7 +2143,7 @@ TEST_F(LatencyTest, ManyChannelsMultiplexedSubscribedToMux) { // Create subscriber to multiplexer. absl::StatusOr sub = - sub_client.CreateSubscriber(kMux, ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); return opts; }())); + sub_client.CreateSubscriber(kMux, ([] { subspace::SubscriberOptions opts; opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); // std::cerr << "sub status " << sub.status() << "\n"; ASSERT_OK(sub); @@ -2066,7 +2249,7 @@ TEST_F(LatencyTest, SubscriberLatency) { ASSERT_OK(pub); // Create subscriber. absl::StatusOr sub = sub_client.CreateSubscriber( - "sublat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "sublat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); // Fill channel. @@ -2107,7 +2290,7 @@ TEST_F(LatencyTest, PubSubLatency) { ASSERT_OK(pub); // Create subscriber. absl::StatusOr sub = sub_client.CreateSubscriber( - "pubsublat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); return opts; }())); + "pubsublat", ([] { subspace::SubscriberOptions opts; opts.SetReliable(false); opts.SetLogDroppedMessages(false); opts.SetDetectDroppedMessages(false); return opts; }())); ASSERT_OK(sub); // Send and receive messages, measuring time taken. diff --git a/client/options.h b/client/options.h index a183b863..f59179ad 100644 --- a/client/options.h +++ b/client/options.h @@ -38,6 +38,9 @@ class Subscriber; struct PublisherOptions { int32_t SlotSize() const { return slot_size; } int32_t NumSlots() const { return num_slots; } + uint64_t SubscriberQueueArenaSize() const { + return subscriber_queue_arena_size; + } PublisherOptions &SetSlotSize(int32_t size) { slot_size = size; return *this; @@ -46,6 +49,14 @@ struct PublisherOptions { num_slots = num; return *this; } + // Total bytes reserved for packed per-subscriber queues in the CCB. Queues + // are disabled by default. A non-empty arena gives subscribers that do not + // request an override the fixed kDefaultSubscriberQueueSize capacity. All + // publishers on a channel must agree on this value. + PublisherOptions &SetSubscriberQueueArenaSize(uint64_t size) { + subscriber_queue_arena_size = size; + return *this; + } // A public publisher's messages will be seen outside of the // publishing computer. @@ -219,6 +230,7 @@ struct PublisherOptions { // here. int32_t slot_size = 0; int32_t num_slots = 0; + uint64_t subscriber_queue_arena_size = 0; bool local = false; bool reliable = false; @@ -245,6 +257,14 @@ struct PublisherOptions { }; struct SubscriberOptions { + // Capacity of this subscriber's CCB slot queue. Zero uses the publisher's + // channel default. + SubscriberOptions &SetSubscriberQueueSize(int32_t size) { + subscriber_queue_size = size; + return *this; + } + int32_t SubscriberQueueSize() const { return subscriber_queue_size; } + // A reliable subscriber will never miss a message from a reliable // publisher. SubscriberOptions &SetReliable(bool v) { @@ -277,6 +297,11 @@ struct SubscriberOptions { int MaxActiveMessages() const { return max_active_messages; } bool LogDroppedMessages() const { return log_dropped_messages; } void SetLogDroppedMessages(bool v) { log_dropped_messages = v; } + bool DetectDroppedMessages() const { return detect_dropped_messages; } + SubscriberOptions &SetDetectDroppedMessages(bool v) { + detect_dropped_messages = v; + return *this; + } SubscriberOptions &SetBridge(bool v) { bridge = v; @@ -357,11 +382,13 @@ struct SubscriberOptions { } bool reliable = false; + int32_t subscriber_queue_size = 0; bool bridge = false; bool for_tunnel = false; std::string type; int max_active_messages = 1; bool log_dropped_messages = true; + bool detect_dropped_messages = true; bool pass_activation = false; // If true, the subscriber will pass activation // messages to the user. bool read_write = false; diff --git a/client/publisher.cc b/client/publisher.cc index d5ccd2bc..d9947768 100644 --- a/client/publisher.cc +++ b/client/publisher.cc @@ -13,6 +13,18 @@ namespace subspace { namespace details { +class SubscriberQueuePublishGuard { +public: + explicit SubscriberQueuePublishGuard(PublisherImpl &publisher) + : publisher_(publisher) { + publisher_.BeginSubscriberQueuePublish(); + } + ~SubscriberQueuePublishGuard() { publisher_.EndSubscriberQueuePublish(); } + +private: + PublisherImpl &publisher_; +}; + absl::Status PublisherImpl::CreateOrAttachBuffers(uint64_t final_slot_size) { if (final_slot_size == 0) { // If we are being asked for a slot size of 0, we will just use 64 bytes. @@ -174,15 +186,18 @@ void PublisherImpl::SetSlotToBiggestBuffer(MessageSlot *slot) { if (slot == nullptr) { return; } - if (slot->buffer_index != -1) { + const int old_buffer_index = + slot->buffer_index.load(std::memory_order_relaxed); + if (old_buffer_index != -1) { // If the slot has a buffer (it's not in the free list), decrement the // refs for the buffer. - if (bcb_->refs[slot->buffer_index].load(std::memory_order_relaxed) > 0) { - DecrementBufferRefs(slot->buffer_index); + if (bcb_->refs[old_buffer_index].load(std::memory_order_relaxed) > 0) { + DecrementBufferRefs(old_buffer_index); } } - slot->buffer_index = buffers_.size() - 1; // Use biggest buffer. - IncrementBufferRefs(slot->buffer_index); + const int new_buffer_index = buffers_.size() - 1; + slot->buffer_index.store(new_buffer_index, std::memory_order_relaxed); + IncrementBufferRefs(new_buffer_index); } MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) { @@ -275,9 +290,11 @@ MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) { if ((refs & kPubOwned) != 0) { continue; } - if ((refs & kRefsMask) == 0 && s->timestamp < earliest_timestamp) { + const uint64_t timestamp = + s->timestamp.load(std::memory_order_relaxed); + if ((refs & kRefsMask) == 0 && timestamp < earliest_timestamp) { slot = s; - earliest_timestamp = s->timestamp; + earliest_timestamp = timestamp; } } } @@ -294,7 +311,8 @@ MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) { uint64_t old_refs = slot->refs.load(std::memory_order_relaxed); uint64_t ref = kPubOwned | owner; uint64_t expected = BuildRefsBitField( - slot->ordinal, (old_refs >> kVchanIdShift) & kVchanIdMask, + slot->ordinal.load(std::memory_order_relaxed), + (old_refs >> kVchanIdShift) & kVchanIdMask, (old_refs >> kRetiredRefsShift) & kRetiredRefsMask); if (slot->refs.compare_exchange_weak(expected, ref, std::memory_order_acquire, @@ -316,9 +334,9 @@ MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) { std::this_thread::yield(); } } - slot->ordinal = 0; - slot->timestamp = 0; - slot->vchan_id = vchan_id_; + slot->ordinal.store(0, std::memory_order_relaxed); + slot->timestamp.store(0, std::memory_order_relaxed); + slot->vchan_id.store(vchan_id_, std::memory_order_relaxed); SetSlotToBiggestBuffer(slot); MessagePrefix *p = Prefix(slot); @@ -328,7 +346,9 @@ MessageSlot *PublisherImpl::FindFreeSlotUnreliable(int owner) { // We have a slot. Clear it in all the subscriber bitsets. ccb_->subscribers.Traverse([this, slot](int sub_id) { int vid = GetSubVchanId(sub_id); - if (vid != -1 && slot->vchan_id != -1 && vid != slot->vchan_id) { + const int slot_vchan_id = + slot->vchan_id.load(std::memory_order_relaxed); + if (vid != -1 && slot_vchan_id != -1 && vid != slot_vchan_id) { return; } @@ -374,7 +394,9 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { } MessageSlot *s = &ccb_->slots[free_slot]; - ActiveSlot active_slot = {s, s->ordinal, s->timestamp}; + ActiveSlot active_slot = { + s, s->ordinal.load(std::memory_order_relaxed), + s->timestamp.load(std::memory_order_relaxed)}; active_slots_.push_back(active_slot); } else if (!ForTunnel() && (retired_slot = RetiredSlots().FindFirstSet()) != -1) { if (embargoed_slots_.IsSet(retired_slot)) { @@ -386,7 +408,9 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { } MessageSlot *s = &ccb_->slots[retired_slot]; - ActiveSlot active_slot = {s, s->ordinal, s->timestamp}; + ActiveSlot active_slot = { + s, s->ordinal.load(std::memory_order_relaxed), + s->timestamp.load(std::memory_order_relaxed)}; active_slots_.push_back(active_slot); } else { for (int i = 0; i < NumSlots(); i++) { @@ -396,7 +420,9 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { MessageSlot *s = &ccb_->slots[i]; uint64_t refs = s->refs.load(std::memory_order_relaxed); if ((refs & kPubOwned) == 0) { - ActiveSlot active_slot = {s, s->ordinal, s->timestamp}; + ActiveSlot active_slot = { + s, s->ordinal.load(std::memory_order_relaxed), + s->timestamp.load(std::memory_order_relaxed)}; active_slots_.push_back(active_slot); } } @@ -412,13 +438,17 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { // Look for a slot with zero refs but don't go past one with non-zero // reliable ref count. + const bool require_reliable_seen = GetCounters().num_reliable_subs != 0; for (auto &s : active_slots_) { uint64_t refs = s.slot->refs.load(std::memory_order_relaxed); if (((refs >> kReliableRefCountShift) & kRefCountMask) != 0) { break; } - // Don't go past one without the kMessageSeen flag set. - if (s.ordinal != 0 && (s.slot->flags & kMessageSeen) == 0) { + // Don't let unreliable subscribers create reliable-publisher + // backpressure. Only reliable subscribers require ordered visibility. + if (require_reliable_seen && s.ordinal != 0 && + (s.slot->flags.load(std::memory_order_relaxed) & + kMessageSeenByReliable) == 0) { break; } // If the refs have no references we can claim it. @@ -434,7 +464,8 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { uint64_t old_refs = slot->refs.load(std::memory_order_relaxed); uint64_t ref = kPubOwned | owner; uint64_t expected = BuildRefsBitField( - slot->ordinal, (old_refs >> kVchanIdShift) & kVchanIdMask, + slot->ordinal.load(std::memory_order_relaxed), + (old_refs >> kVchanIdShift) & kVchanIdMask, (old_refs >> kRetiredRefsShift) & kRetiredRefsMask); if (slot->refs.compare_exchange_weak(expected, ref, std::memory_order_acquire, @@ -457,9 +488,9 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { std::this_thread::yield(); } } - slot->ordinal = 0; - slot->timestamp = 0; - slot->vchan_id = vchan_id_; + slot->ordinal.store(0, std::memory_order_relaxed); + slot->timestamp.store(0, std::memory_order_relaxed); + slot->vchan_id.store(vchan_id_, std::memory_order_relaxed); SetSlotToBiggestBuffer(slot); MessagePrefix *p = Prefix(slot); @@ -469,7 +500,9 @@ MessageSlot *PublisherImpl::FindFreeSlotReliable(int owner) { // We have a slot. Clear it in all the subscriber bitsets. ccb_->subscribers.Traverse([this, slot](int sub_id) { int vid = GetSubVchanId(sub_id); - if (vid != -1 && slot->vchan_id != -1 && vid != slot->vchan_id) { + const int slot_vchan_id = + slot->vchan_id.load(std::memory_order_relaxed); + if (vid != -1 && slot_vchan_id != -1 && vid != slot_vchan_id) { return; } GetAvailableSlots(sub_id).Clear(slot->id); @@ -490,41 +523,49 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother( void *buffer = GetBufferAddress(slot); MessagePrefix *prefix = Prefix(slot); - slot->ordinal = ccb_->ordinals.Next(slot->vchan_id); - slot->timestamp = toolbelt::Now(); - slot->flags = 0; + const int initial_vchan_id = + slot->vchan_id.load(std::memory_order_relaxed); + slot->ordinal.store(ccb_->ordinals.Next(initial_vchan_id), + std::memory_order_relaxed); + slot->timestamp.store(toolbelt::Now(), std::memory_order_relaxed); + slot->flags.store(0, std::memory_order_relaxed); // Copy message parameters into message prefix in buffer. if (omit_prefix) { if (for_tunnel) { prefix->SetIsCrossMachine(); } - slot->timestamp = prefix->timestamp; - slot->vchan_id = prefix->vchan_id; + slot->timestamp.store(prefix->timestamp, std::memory_order_relaxed); + slot->vchan_id.store(prefix->vchan_id, std::memory_order_relaxed); // The bridged_slot_id is the slot is used for the retirement notification. - slot->bridged_slot_id = use_prefix_slot_id ? prefix->slot_id : slot->id; + slot->bridged_slot_id.store( + use_prefix_slot_id ? prefix->slot_id : slot->id, + std::memory_order_relaxed); } else { - prefix->message_size = slot->message_size; - prefix->ordinal = slot->ordinal; - prefix->timestamp = slot->timestamp; - prefix->vchan_id = slot->vchan_id; + prefix->message_size = + slot->message_size.load(std::memory_order_relaxed); + prefix->ordinal = slot->ordinal.load(std::memory_order_relaxed); + prefix->timestamp = slot->timestamp.load(std::memory_order_relaxed); + prefix->vchan_id = slot->vchan_id.load(std::memory_order_relaxed); prefix->checksum_size = static_cast(ChecksumSize()); prefix->metadata_size = static_cast(MetadataSize()); prefix->flags = 0; prefix->slot_id = slot->id; - slot->bridged_slot_id = slot->id; + slot->bridged_slot_id.store(slot->id, std::memory_order_relaxed); if (is_activation) { prefix->SetIsActivation(); - slot->flags |= kMessageIsActivation; - ccb_->activation_tracker.Activate(slot->vchan_id); + slot->flags.fetch_or(kMessageIsActivation, std::memory_order_relaxed); + ccb_->activation_tracker.Activate( + slot->vchan_id.load(std::memory_order_relaxed)); } if (for_tunnel) { prefix->SetIsCrossMachine(); } if (options_.Checksum()) { prefix->SetHasChecksum(); - auto data = GetMessageChecksumData(prefix, buffer, slot->message_size, - ChecksumSize(), MetadataSize()); + auto data = GetMessageChecksumData( + prefix, buffer, slot->message_size.load(std::memory_order_relaxed), + ChecksumSize(), MetadataSize()); absl::Span cksum = GetChecksumSpan(prefix, ChecksumSize()); if (checksum_callback_ != nullptr) { checksum_callback_(data, cksum); @@ -535,38 +576,61 @@ Channel::PublishedMessage PublisherImpl::ActivateSlotAndGetAnother( } // Set the refs to the ordinal with no refs. - slot->refs.store(BuildRefsBitField(slot->ordinal, vchan_id_, 0), - std::memory_order_release); + 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. SubscriberImpl::NextSlot() uses total_messages as - // a version stamp for its cached active_slots_ snapshot: a subscriber - // that observes a bumped total_messages 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 total_messages++ 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. If we incremented total_messages - // first, a subscriber could read the new total, run - // CollectVisibleSlots() before the bit was visible, cache that - // snapshot under next_slot_cached_total_, and then reuse the stale - // cache forever (no further total bump arrives to invalidate it). - ccb_->subscribers.Traverse([this, slot](int sub_id) { + // 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. + // + // 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. + SubscriberQueuePublishGuard publish_guard(*this); + std::vector 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); + InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(sub_id); + if (queue != nullptr && + !queue->Push(slot->id, + slot->ordinal.load(std::memory_order_relaxed), + /*report_insertion_failure=*/false)) { + failed_queues.push_back(queue); + } }); - // Update counters AFTER setting the available-slot bits (see above). + // Update counters AFTER notifying subscribers (see above). if (!is_activation) { - ccb_->total_bytes += slot->message_size; - if (slot->message_size > ccb_->max_message_size) { - ccb_->max_message_size = slot->message_size; + const uint64_t message_size = + slot->message_size.load(std::memory_order_relaxed); + ccb_->total_bytes += message_size; + if (message_size > ccb_->max_message_size) { + ccb_->max_message_size = message_size; } - ccb_->total_messages++; + } + 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(); } // A reliable publisher doesn't allocate a slot until it is asked for. diff --git a/client/publisher.h b/client/publisher.h index 022428f4..e0975511 100644 --- a/client/publisher.h +++ b/client/publisher.h @@ -14,11 +14,15 @@ namespace details { // messages to be published. class PublisherImpl : public ClientChannel { public: - PublisherImpl(const std::string &name, int num_slots, int channel_id, + PublisherImpl(const std::string &name, int num_slots, + int subscriber_queue_size, + uint64_t subscriber_queue_arena_size, int channel_id, int publisher_id, int vchan_id, uint64_t session_id, - std::string type, const PublisherOptions &options, + std::string type, + const PublisherOptions &options, std::function reload, int user_id, int group_id) - : ClientChannel(name, num_slots, channel_id, vchan_id, + : ClientChannel(name, num_slots, subscriber_queue_size, + subscriber_queue_arena_size, channel_id, vchan_id, std::move(session_id), std::move(type), std::move(reload), user_id, group_id), publisher_id_(publisher_id), options_(options) {} @@ -28,6 +32,17 @@ class PublisherImpl : public ClientChannel { bool IsLocal() const { return options_.IsLocal(); } bool IsFixedSize() const { return options_.IsFixedSize(); } bool UsesSplitBuffers() const { return UseSplitBuffers(); } + void BeginSubscriberQueuePublish() { + active_queue_publish_depth_.fetch_add(1, std::memory_order_seq_cst); + Channel::BeginSubscriberQueuePublish(publisher_id_); + } + void EndSubscriberQueuePublish() { + Channel::EndSubscriberQueuePublish(publisher_id_); + active_queue_publish_depth_.fetch_sub(1, std::memory_order_seq_cst); + } + uint32_t ActiveQueuePublishDepth() const { + return active_queue_publish_depth_.load(std::memory_order_seq_cst); + } // Trigger the publisher's reliable trigger fd, waking anything that is // waiting on the publisher's reliable event fd (e.g. a reliable publisher @@ -148,6 +163,7 @@ class PublisherImpl : public ClientChannel { toolbelt::TriggerFd trigger_; int publisher_id_; + std::atomic active_queue_publish_depth_{0}; std::vector subscribers_; PublisherOptions options_; toolbelt::FileDescriptor retirement_fd_ = {}; diff --git a/client/python/client.cc b/client/python/client.cc index c3a46b00..620bfc54 100644 --- a/client/python/client.cc +++ b/client/python/client.cc @@ -57,6 +57,10 @@ PYBIND11_MODULE(subspace, m) { .def_readonly("type", &ChannelInfo::type) .def_readonly("slot_size", &ChannelInfo::slot_size) .def_readonly("num_slots", &ChannelInfo::num_slots) + .def_readonly("subscriber_queue_size", + &ChannelInfo::subscriber_queue_size) + .def_readonly("subscriber_queue_arena_size", + &ChannelInfo::subscriber_queue_arena_size) .def_readonly("reliable", &ChannelInfo::reliable); // ChannelStats struct. @@ -109,6 +113,13 @@ PYBIND11_MODULE(subspace, m) { "Set the number of slots for the publisher.") .def("num_slots", &PublisherOptions::NumSlots, "Get the number of slots for the publisher.") + .def("set_subscriber_queue_arena_size", + &PublisherOptions::SetSubscriberQueueArenaSize, + "Set the bytes reserved for packed subscriber queues. Queues are " + "disabled by default; a non-zero size enables them.") + .def("subscriber_queue_arena_size", + &PublisherOptions::SubscriberQueueArenaSize, + "Get the configured subscriber queue arena size in bytes.") .def("set_notify_retirement", &PublisherOptions::SetNotifyRetirement, "Set whether the publisher notifies on message retirement.") .def("notify_retirement", &PublisherOptions::NotifyRetirement, @@ -138,6 +149,13 @@ PYBIND11_MODULE(subspace, m) { .def(py::init<>()) .def("set_reliable", &SubscriberOptions::SetReliable, "Set whether the subscriber is reliable.") + .def("set_subscriber_queue_size", + &SubscriberOptions::SetSubscriberQueueSize, + "Set this subscriber's queue capacity; zero uses the publisher " + "default.") + .def("subscriber_queue_size", + &SubscriberOptions::SubscriberQueueSize, + "Get this subscriber's requested queue capacity.") .def("set_pass_activation", &SubscriberOptions::SetPassActivation, "Set whether the subscriber passes activation messages.") .def("is_reliable", &SubscriberOptions::IsReliable, @@ -156,6 +174,12 @@ PYBIND11_MODULE(subspace, m) { "Sets whether the subscriber logs dropped messages.") .def("log_dropped_messages", &SubscriberOptions::LogDroppedMessages, "Get whether the subscriber logs dropped messages.") + .def("set_detect_dropped_messages", + &SubscriberOptions::SetDetectDroppedMessages, + "Sets whether the subscriber detects dropped messages internally.") + .def("detect_dropped_messages", + &SubscriberOptions::DetectDroppedMessages, + "Get whether the subscriber detects dropped messages internally.") .def("set_bridge", &SubscriberOptions::SetBridge, "Set whether the subscriber is a bridge.") .def("is_bridge", &SubscriberOptions::IsBridge, @@ -290,6 +314,13 @@ PYBIND11_MODULE(subspace, m) { publisher_class.def("num_slots", &Publisher::NumSlots, "Get the number of message slots."); + publisher_class.def("subscriber_queue_size", + &Publisher::SubscriberQueueSize, + "Get each subscriber queue's resolved capacity."); + publisher_class.def("subscriber_queue_arena_size", + &Publisher::SubscriberQueueArenaSize, + "Get the subscriber queue arena size in bytes."); + publisher_class.def("virtual_channel_id", &Publisher::VirtualChannelId, "Get the virtual channel ID assigned to this publisher."); @@ -565,6 +596,10 @@ checksum_error). Use as a context manager to auto-release the slot: subscriber_class.def("num_slots", &Subscriber::NumSlots, "Get the number of message slots."); + subscriber_class.def("subscriber_queue_size", + &Subscriber::SubscriberQueueSize, + "Get each subscriber queue's resolved capacity."); + subscriber_class.def("get_current_ordinal", &Subscriber::GetCurrentOrdinal, "Get the most recently received ordinal."); diff --git a/client/python/client_test.py b/client/python/client_test.py index 601b5c26..cdffce2d 100644 --- a/client/python/client_test.py +++ b/client/python/client_test.py @@ -109,12 +109,18 @@ def test_skip_to_newest(self): # ------------------------------------------------------------------ def test_publisher_accessors(self): client = self._make_client("pub_acc") + opts = subspace.PublisherOptions() + opts.set_slot_size(512) + opts.set_num_slots(8) + opts.set_type("my_type") + opts.set_subscriber_queue_arena_size(11_000) pub = client.create_publisher(channel_name="ch_pub_acc", - slot_size=512, num_slots=8, - type="my_type") + options=opts) self.assertEqual(pub.type(), "my_type") self.assertEqual(pub.slot_size(), 512) self.assertEqual(pub.num_slots(), 8) + self.assertEqual(pub.subscriber_queue_size(), 16) + self.assertEqual(pub.subscriber_queue_arena_size(), 11_000) self.assertFalse(pub.is_reliable()) self.assertFalse(pub.is_fixed_size()) self.assertEqual(pub.name(), "ch_pub_acc") @@ -124,11 +130,18 @@ def test_publisher_accessors(self): def test_subscriber_accessors(self): client = self._make_client("sub_acc") + opts = subspace.PublisherOptions() + opts.set_slot_size(256) + opts.set_num_slots(10) + opts.set_type("sub_type") + opts.set_subscriber_queue_arena_size(11_000) pub = client.create_publisher(channel_name="ch_sub_acc", - slot_size=256, num_slots=10, - type="sub_type") + options=opts) + sub_opts = subspace.SubscriberOptions() + sub_opts.set_subscriber_queue_size(7) + sub_opts.set_type("sub_type") sub = client.create_subscriber(channel_name="ch_sub_acc", - type="sub_type") + options=sub_opts) pub.publish_message(b"probe") sub.wait() @@ -138,6 +151,7 @@ def test_subscriber_accessors(self): self.assertFalse(sub.is_reliable()) self.assertEqual(sub.slot_size(), 256) self.assertEqual(sub.num_slots(), 10) + self.assertEqual(sub.subscriber_queue_size(), 7) self.assertEqual(sub.name(), "ch_sub_acc") self.assertIsInstance(sub.get_virtual_memory_usage(), int) self.assertGreater(sub.get_virtual_memory_usage(), 0) @@ -278,6 +292,7 @@ def test_cancel_publish(self): # ------------------------------------------------------------------ def test_publisher_options(self): opts = subspace.PublisherOptions() + self.assertEqual(opts.subscriber_queue_arena_size(), 0) opts.set_slot_size(1024) opts.set_num_slots(4) opts.set_reliable(True) @@ -285,9 +300,11 @@ def test_publisher_options(self): opts.set_local(True) opts.set_fixed_size(True) opts.set_checksum(True) + opts.set_subscriber_queue_arena_size(9_000) self.assertEqual(opts.slot_size(), 1024) self.assertEqual(opts.num_slots(), 4) + self.assertEqual(opts.subscriber_queue_arena_size(), 9_000) self.assertTrue(opts.is_reliable()) self.assertEqual(opts.type(), "opts_type") self.assertTrue(opts.is_local()) @@ -297,6 +314,7 @@ def test_publisher_options(self): def test_subscriber_options(self): opts = subspace.SubscriberOptions() opts.set_reliable(True) + opts.set_subscriber_queue_size(3) opts.set_type("sub_opts_type") opts.set_max_active_messages(5) opts.set_checksum(True) @@ -304,6 +322,7 @@ def test_subscriber_options(self): opts.set_keep_active_message(True) self.assertTrue(opts.is_reliable()) + self.assertEqual(opts.subscriber_queue_size(), 3) self.assertEqual(opts.type(), "sub_opts_type") self.assertEqual(opts.max_active_messages(), 5) self.assertTrue(opts.checksum()) @@ -316,11 +335,14 @@ def test_create_publisher_with_options(self): opts.set_slot_size(128) opts.set_num_slots(6) opts.set_type("opt_chan_type") + opts.set_subscriber_queue_arena_size(13_000) pub = client.create_publisher(channel_name="ch_opts_pub", options=opts) self.assertEqual(pub.slot_size(), 128) self.assertEqual(pub.num_slots(), 6) + self.assertEqual(pub.subscriber_queue_size(), 16) + self.assertEqual(pub.subscriber_queue_arena_size(), 13_000) self.assertEqual(pub.type(), "opt_chan_type") pub = None diff --git a/client/stress_test.cc b/client/stress_test.cc index 91515124..dbd61ca2 100644 --- a/client/stress_test.cc +++ b/client/stress_test.cc @@ -399,6 +399,283 @@ TEST_F(StressTest, ThreadSafety) { signal(SIGQUIT, oldSig); } +TEST_F(StressTest, SubscriberQueuesManyPublishersAndSubscribers) { + const int kNumPublishers = StressValueForSplitBuffers(8, 4); + const int kNumSubscribers = StressValueForSplitBuffers(16, 8); + const int kMessagesPerPublisher = + StressValueForSplitBuffers(10000, 2000); + const int kNumSlots = StressValueForSplitBuffers(256, 128); + constexpr int kDefaultQueueSize = 64; + constexpr uint64_t kMagic = 0x5155455545535452; + constexpr char kChannel[] = "/subscriber_queue_stress"; + + struct Payload { + uint64_t magic; + uint32_t publisher; + uint32_t sequence; + uint64_t checksum; + }; + + std::vector> publisher_clients; + std::vector publishers; + publisher_clients.reserve(kNumPublishers); + publishers.reserve(kNumPublishers); + for (int i = 0; i < kNumPublishers; ++i) { + publisher_clients.push_back( + EVAL_AND_ASSERT_OK(subspace::Client::Create( + Socket(), absl::StrFormat("queue_publisher_%d", i)))); + publishers.push_back( + EVAL_AND_ASSERT_OK(publisher_clients.back()->CreatePublisher( + kChannel, + subspace::PublisherOptions() + .SetSlotSize(sizeof(Payload)) + .SetNumSlots(kNumSlots) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize)))); + } + + std::vector> subscriber_clients; + std::vector subscribers; + subscriber_clients.reserve(kNumSubscribers); + subscribers.reserve(kNumSubscribers); + for (int i = 0; i < kNumSubscribers; ++i) { + const int queue_size = 1 << (i % 7); + subscriber_clients.push_back( + EVAL_AND_ASSERT_OK(subspace::Client::Create( + Socket(), absl::StrFormat("queue_subscriber_%d", i)))); + subspace::SubscriberOptions options; + options.SetSubscriberQueueSize(queue_size); + options.SetLogDroppedMessages(false); + subscribers.push_back( + EVAL_AND_ASSERT_OK(subscriber_clients.back()->CreateSubscriber( + kChannel, options))); + ASSERT_EQ(queue_size, subscribers.back().SubscriberQueueSize()); + } + + std::atomic start = false; + std::atomic publishers_done = false; + std::atomic failures = 0; + std::vector received(kNumSubscribers, 0); + std::vector subscriber_threads; + subscriber_threads.reserve(kNumSubscribers); + for (int sub_id = 0; sub_id < kNumSubscribers; ++sub_id) { + subscriber_threads.emplace_back([&, sub_id]() { + std::vector last_sequence(kNumPublishers, -1); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (;;) { + absl::StatusOr message = subscribers[sub_id].ReadMessage(); + if (!message.ok()) { + ++failures; + return; + } + if (message->length == 0) { + if (publishers_done.load(std::memory_order_acquire)) { + return; + } + std::this_thread::yield(); + continue; + } + if (message->length != sizeof(Payload)) { + ++failures; + continue; + } + + Payload payload; + memcpy(&payload, message->buffer, sizeof(payload)); + const uint64_t checksum = + payload.magic ^ + (static_cast(payload.publisher) << 32) ^ + payload.sequence; + if (payload.magic != kMagic || + payload.publisher >= static_cast(kNumPublishers) || + payload.sequence >= + static_cast(kMessagesPerPublisher) || + payload.checksum != checksum) { + ++failures; + continue; + } + if (static_cast(payload.sequence) <= + last_sequence[payload.publisher]) { + ++failures; + continue; + } + last_sequence[payload.publisher] = payload.sequence; + ++received[sub_id]; + } + }); + } + + std::vector publisher_threads; + publisher_threads.reserve(kNumPublishers); + for (int pub_id = 0; pub_id < kNumPublishers; ++pub_id) { + publisher_threads.emplace_back([&, pub_id]() { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (int sequence = 0; sequence < kMessagesPerPublisher; ++sequence) { + Payload payload = { + kMagic, + static_cast(pub_id), + static_cast(sequence), + kMagic ^ (static_cast(pub_id) << 32) ^ + static_cast(sequence), + }; + absl::StatusOr buffer = + publishers[pub_id].GetMessageBuffer(sizeof(payload)); + if (!buffer.ok()) { + ++failures; + return; + } + memcpy(*buffer, &payload, sizeof(payload)); + if (!publishers[pub_id].PublishMessage(sizeof(payload)).ok()) { + ++failures; + return; + } + } + }); + } + + start.store(true, std::memory_order_release); + for (auto &thread : publisher_threads) { + thread.join(); + } + publishers_done.store(true, std::memory_order_release); + for (auto &thread : subscriber_threads) { + thread.join(); + } + + EXPECT_EQ(0, failures.load()); + for (int sub_id = 0; sub_id < kNumSubscribers; ++sub_id) { + EXPECT_GT(received[sub_id], 0) << "subscriber " << sub_id; + } +} + +TEST_F(StressTest, SubscriberQueueChurnDuringConcurrentPublishing) { + const int kNumPublishers = StressValueForSplitBuffers(4, 2); + const int kNumSubscriberThreads = StressValueForSplitBuffers(8, 4); + const int kCyclesPerThread = StressValueForSplitBuffers(800, 1400); + constexpr int kDefaultQueueSize = 64; + constexpr int kNumSlots = 128; + constexpr char kChannel[] = "/subscriber_queue_churn_stress"; + + std::vector> publisher_clients; + std::vector publishers; + publisher_clients.reserve(kNumPublishers); + publishers.reserve(kNumPublishers); + for (int i = 0; i < kNumPublishers; ++i) { + publisher_clients.push_back( + EVAL_AND_ASSERT_OK(subspace::Client::Create( + Socket(), absl::StrFormat("queue_churn_publisher_%d", i)))); + publishers.push_back( + EVAL_AND_ASSERT_OK(publisher_clients.back()->CreatePublisher( + kChannel, + subspace::PublisherOptions() + .SetSlotSize(sizeof(uint64_t)) + .SetNumSlots(kNumSlots) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize)))); + } + + std::vector> subscriber_clients; + subscriber_clients.reserve(kNumSubscriberThreads); + for (int i = 0; i < kNumSubscriberThreads; ++i) { + subscriber_clients.push_back( + EVAL_AND_ASSERT_OK(subspace::Client::Create( + Socket(), absl::StrFormat("queue_churn_subscriber_%d", i)))); + } + + std::atomic start = false; + std::atomic stop_publishers = false; + std::atomic failures = 0; + std::atomic messages_published = 0; + std::atomic messages_received = 0; + + std::vector publisher_threads; + publisher_threads.reserve(kNumPublishers); + for (int pub_id = 0; pub_id < kNumPublishers; ++pub_id) { + publisher_threads.emplace_back([&, pub_id]() { + uint64_t sequence = static_cast(pub_id) << 56; + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + while (!stop_publishers.load(std::memory_order_acquire)) { + absl::StatusOr buffer = + publishers[pub_id].GetMessageBuffer(sizeof(sequence)); + if (!buffer.ok()) { + ++failures; + return; + } + memcpy(*buffer, &sequence, sizeof(sequence)); + if (!publishers[pub_id].PublishMessage(sizeof(sequence)).ok()) { + ++failures; + return; + } + ++sequence; + ++messages_published; + } + }); + } + + std::vector subscriber_threads; + subscriber_threads.reserve(kNumSubscriberThreads); + for (int thread_id = 0; thread_id < kNumSubscriberThreads; ++thread_id) { + subscriber_threads.emplace_back([&, thread_id]() { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (int cycle = 0; cycle < kCyclesPerThread; ++cycle) { + const int queue_size = 1 << ((thread_id + cycle) % 6); + subspace::SubscriberOptions options; + options.SetSubscriberQueueSize(queue_size); + options.SetLogDroppedMessages(false); + absl::StatusOr subscriber = + subscriber_clients[thread_id]->CreateSubscriber( + kChannel, options); + if (!subscriber.ok()) { + ++failures; + return; + } + if (subscriber->SubscriberQueueSize() != queue_size) { + ++failures; + return; + } + + for (int attempt = 0; attempt < 32; ++attempt) { + const subspace::ReadMode mode = + (cycle + attempt) % 2 == 0 + ? subspace::ReadMode::kReadNext + : subspace::ReadMode::kReadNewest; + absl::StatusOr message = subscriber->ReadMessage(mode); + if (!message.ok()) { + ++failures; + return; + } + if (message->length == sizeof(uint64_t)) { + ++messages_received; + break; + } + std::this_thread::yield(); + } + } + }); + } + + start.store(true, std::memory_order_release); + for (auto &thread : subscriber_threads) { + thread.join(); + } + stop_publishers.store(true, std::memory_order_release); + for (auto &thread : publisher_threads) { + thread.join(); + } + + EXPECT_EQ(0, failures.load()); + EXPECT_GT(messages_published.load(), 0); + EXPECT_GT(messages_received.load(), 0); +} + TEST_F(StressTest, ActiveMessages) { auto oldSig = signal(SIGQUIT, SigQuitHandler); diff --git a/client/subscriber.cc b/client/subscriber.cc index a5ba621b..da7451cc 100644 --- a/client/subscriber.cc +++ b/client/subscriber.cc @@ -4,6 +4,8 @@ #include "client/subscriber.h" +#include + namespace subspace { namespace details { @@ -23,6 +25,25 @@ void SubscriberImpl::InitActiveMessages() { } } +void SubscriberImpl::ResetDeliveryState() { + ClearActiveMessage(); + SetSlot(nullptr); + active_slots_.clear(); + search_buffer_.clear(); + newest_snapshot_.clear(); + embargoed_slots_.ClearAll(); + ordinal_trackers_.clear(); + (void)GetOrdinalTracker(vchan_id_); + next_slot_cached_total_ = 0; + next_slot_cursor_ = 0; + next_slot_cache_valid_ = false; + poll_drain_exhausted_ = false; + queue_bitset_fallback_ = false; + pending_queue_drops_ = 0; + queue_drain_tail_ = 0; + queue_drain_tail_valid_ = false; +} + // For non-virtual channels both the slots' vchan_id and the subsriber's // are -1. This is the common case. // For virtual subscribers, if the slot's vchan_id is -1 it means that @@ -32,7 +53,10 @@ void SubscriberImpl::InitActiveMessages() { // If the subscriber's vchan_id is -1 it means that the subscriber is on the // multiplexer and should see all messages, regardless of the vchan_id static inline bool VirtualChannelIdMatch(MessageSlot *slot, int vchan_id) { - return vchan_id == -1 || slot->vchan_id == -1 || slot->vchan_id == vchan_id; + const int slot_vchan_id = + slot->vchan_id.load(std::memory_order_relaxed); + return vchan_id == -1 || slot_vchan_id == -1 || + slot_vchan_id == vchan_id; } bool SubscriberImpl::AddActiveMessage([[maybe_unused]] MessageSlot *slot) { @@ -51,7 +75,9 @@ void SubscriberImpl::RemoveActiveMessage(MessageSlot *slot) { // << slot->ordinal << " refs " << std::hex << slot->refs.load() << // std::dec << "\n"; slot->sub_owners.Clear(subscriber_id_); - AtomicIncRefCount(slot, IsReliable(), -1, slot->ordinal, slot->vchan_id, true, + 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 @@ -69,7 +95,8 @@ void SubscriberImpl::RemoveActiveMessage(MessageSlot *slot) { // slot->bridged_slot_id, slot->ordinal, // slot->vchan_id); // std::cerr << details; - TriggerRetirement(slot->bridged_slot_id); + TriggerRetirement( + slot->bridged_slot_id.load(std::memory_order_relaxed)); }); if (--num_active_messages_ < options_.MaxActiveMessages()) { Trigger(); @@ -82,18 +109,20 @@ void SubscriberImpl::RemoveActiveMessage(MessageSlot *slot) { void SubscriberImpl::PopulateActiveSlots(InPlaceAtomicBitset &bits) { uint64_t num_messages = 0; do { - num_messages = ccb_->total_messages; + 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_relaxed); - if (VirtualChannelIdMatch(s, vchan_id_) && s->ordinal != 0 && + uint64_t refs = s->refs.load(std::memory_order_acquire); + if (VirtualChannelIdMatch(s, vchan_id_) && + s->ordinal.load(std::memory_order_relaxed) != 0 && (refs & kPubOwned) == 0) { bits.Set(i); } } - } while (num_messages != ccb_->total_messages); + } while (num_messages != + ccb_->total_messages.load(std::memory_order_seq_cst)); } SubscriberImpl::OrdinalTracker & @@ -109,49 +138,39 @@ SubscriberImpl::GetOrdinalTracker(int vchan_id) { } int SubscriberImpl::DetectDrops(int vchan_id) { - std::vector ordinals; - auto &tracker = GetOrdinalTracker(vchan_id_); - ordinals.reserve(tracker.ordinals.Size()); - tracker.ordinals.Traverse( - [&tracker, &ordinals, vchan_id](const OrdinalAndVchanId &o) { - if (vchan_id == o.vchan_id && o.ordinal >= tracker.last_ordinal_seen) { - ordinals.push_back(o); - } - }); - if (ordinals.empty()) { + auto &tracker = GetOrdinalTracker(vchan_id); + const uint64_t ordinal = CurrentOrdinal(); + if (ordinal == 0 || ordinal <= tracker.last_ordinal_seen) { return 0; } - std::sort(ordinals.begin(), ordinals.end()); - tracker.last_ordinal_seen = ordinals.back().ordinal; - - // Look for gaps in the ordinals. - int drops = 0; - for (size_t i = 1; i < ordinals.size(); i++) { - if (ordinals[i].vchan_id != vchan_id) { - // Must be same vchan_id as ordinals are per vchan. - continue; - } - if (ordinals[i].ordinal - ordinals[i - 1].ordinal == 1) { - continue; - } - drops += - static_cast(ordinals[i].ordinal - ordinals[i - 1].ordinal - 1); + const uint64_t last_seen = tracker.last_ordinal_seen; + tracker.last_ordinal_seen = ordinal; + if (last_seen == 0 || ordinal == last_seen + 1) { + return 0; } - return drops; + return static_cast(ordinal - last_seen - 1); } void SubscriberImpl::RememberOrdinal(uint64_t ordinal, int vchan_id) { - auto &tracker = GetOrdinalTracker(vchan_id_); + auto &tracker = GetOrdinalTracker(vchan_id); + if (ordinal > tracker.last_ordinal_seen) { + tracker.last_ordinal_seen = ordinal; + } tracker.ordinals.Insert(OrdinalAndVchanId{ordinal, vchan_id}); } const ActiveSlot *SubscriberImpl::FindUnseenOrdinal() { // Traverse the active slots looking for the first ordinal that is not zero // and has not been seen by a subscriber. - auto &tracker = GetOrdinalTracker(vchan_id_); + int cached_vchan_id = std::numeric_limits::min(); + OrdinalTracker *cached_tracker = nullptr; for (auto &s : active_slots_) { - if (s.ordinal != 0 && - !tracker.ordinals.Contains(OrdinalAndVchanId{s.ordinal, s.vchan_id})) { + if (s.vchan_id != cached_vchan_id) { + cached_vchan_id = s.vchan_id; + cached_tracker = &GetOrdinalTracker(s.vchan_id); + } + if (!cached_tracker->ordinals.Contains( + OrdinalAndVchanId{s.ordinal, s.vchan_id})) { // std::cerr << absl::StrFormat("Found unseen ordinal %d in slot %d\n", s.ordinal, s.slot->id); return &s; } @@ -163,20 +182,48 @@ void SubscriberImpl::ClaimSlot(MessageSlot *slot, int vchan_id, bool was_newest) { slot->sub_owners.Set(subscriber_id_); if (was_newest) { - // We read the newest slot so there can't be any other messages for this - // subscriber. - GetAvailableSlots(subscriber_id_).ClearAll(); + InPlaceAtomicBitset &bits = GetAvailableSlots(subscriber_id_); + for (const ActiveSlot &snapshot : newest_snapshot_) { + bool pinned = snapshot.slot == slot; + if (!pinned) { + pinned = AtomicIncRefCount(snapshot.slot, IsReliable(), 1, + snapshot.ordinal, snapshot.vchan_id, false); + } + if (!pinned) { + // The slot was recycled after the ReadNewest snapshot. Its current bit + // belongs to the new generation and must remain set. + continue; + } + bits.Clear(snapshot.slot->id); + RememberOrdinal(snapshot.ordinal, snapshot.vchan_id); + if (snapshot.slot != slot) { + AtomicIncRefCount(snapshot.slot, IsReliable(), -1, snapshot.ordinal, + snapshot.vchan_id, false); + } + } + newest_snapshot_.clear(); } else { // Clear the bit in the subscriber bitset. GetAvailableSlots(subscriber_id_).Clear(slot->id); } - RememberOrdinal(slot->ordinal, vchan_id); - slot->flags |= kMessageSeen; + RememberOrdinal(slot->ordinal.load(std::memory_order_relaxed), vchan_id); + slot->flags.fetch_or(kMessageSeen, std::memory_order_relaxed); + if (IsReliable()) { + slot->flags.fetch_or(kMessageSeenByReliable, std::memory_order_relaxed); + } } void SubscriberImpl::UnreadSlot(MessageSlot *slot) { - slot->flags &= ~kMessageSeen; 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 + // the ordinal-sorted bitset path until that backlog has been recovered; + // otherwise the next queue entry would be delivered first and permanently + // skip this ordinal. + if (SubscriberQueueSize() > 0) { + queue_bitset_fallback_ = true; + } // NextSlot()'s cache advanced next_slot_cursor_ past this slot when it // returned, on the assumption that ReadMessageInternal would either // ClaimSlot() it (recording the ordinal in the tracker) or accept that it @@ -189,10 +236,10 @@ void SubscriberImpl::UnreadSlot(MessageSlot *slot) { next_slot_cache_valid_ = false; } -void SubscriberImpl::CollectVisibleSlots(InPlaceAtomicBitset &bits) { +uint64_t SubscriberImpl::CollectVisibleSlots(InPlaceAtomicBitset &bits) { uint64_t num_messages = 0; do { - num_messages = ccb_->total_messages; + num_messages = ccb_->total_messages.load(std::memory_order_seq_cst); active_slots_.clear(); // Traverse the bits and add an active slot for each bit set. @@ -204,20 +251,158 @@ void SubscriberImpl::CollectVisibleSlots(InPlaceAtomicBitset &bits) { if (!VirtualChannelIdMatch(s, vchan_id_)) { return; } - if (s->buffer_index == -1) { + if (s->buffer_index.load(std::memory_order_relaxed) == -1) { return; } - ActiveSlot active_slot = {s, s->ordinal, s->timestamp, s->vchan_id}; + ActiveSlot active_slot = { + s, s->ordinal.load(std::memory_order_relaxed), + s->timestamp.load(std::memory_order_relaxed), + s->vchan_id.load(std::memory_order_relaxed)}; active_slots_.push_back(active_slot); }); - } while (num_messages != ccb_->total_messages); + } while (num_messages != + ccb_->total_messages.load(std::memory_order_seq_cst)); + return num_messages; +} + +std::optional +SubscriberImpl::FindNextQueuedSlot(uint64_t max_queue_position) { + InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(subscriber_id_); + if (queue == nullptr || queue->Capacity() == 0) { + return std::nullopt; + } + + int cached_vchan_id = std::numeric_limits::min(); + OrdinalTracker *cached_tracker = nullptr; + + QueuedSlot queued; + for (size_t i = 0; i < queue->Capacity(); i++) { + if (queue->Head() >= max_queue_position) { + return std::nullopt; + } + if (!queue->TryPeek(queued)) { + return std::nullopt; + } + if (queued.slot_id < 0 || queued.slot_id >= NumSlots()) { + queue->DropFront(); + continue; + } + QueuedSlot popped; + if (!queue->TryPop(popped)) { + continue; + } + queued = popped; + if (queued.slot_id < 0 || queued.slot_id >= NumSlots()) { + continue; + } + MessageSlot *s = &ccb_->slots[queued.slot_id]; + const uint64_t refs = s->refs.load(std::memory_order_acquire); + if ((refs & kPubOwned) != 0) { + continue; + } + const uint64_t ref_ordinal = (refs >> kOrdinalShift) & kOrdinalMask; + int ref_vchan_id = (refs >> kVchanIdShift) & kVchanIdMask; + if (ref_vchan_id == kVchanIdMask) { + ref_vchan_id = -1; + } + if (queued.ordinal == 0 || + (queued.ordinal & kOrdinalMask) != ref_ordinal || + (vchan_id_ != -1 && ref_vchan_id != -1 && + vchan_id_ != ref_vchan_id)) { + continue; + } + if (ref_vchan_id != cached_vchan_id) { + cached_vchan_id = ref_vchan_id; + cached_tracker = &GetOrdinalTracker(ref_vchan_id); + } + const OrdinalAndVchanId queued_key{queued.ordinal, ref_vchan_id}; + if (queued.ordinal <= cached_tracker->last_ordinal_seen && + cached_tracker->ordinals.Contains(queued_key)) { + continue; + } + if (options_.SubscriberQueueSize() == 0 && + queued.ordinal > cached_tracker->last_ordinal_seen && + queued.ordinal - cached_tracker->last_ordinal_seen > 1) { + // Queue reservations from concurrent publishers need not follow ordinal + // order. Before accepting a gap, including the first queued ordinal, + // recover any older authoritative bit in ordinal order. If an older + // queue entry arrives after a newer one was already delivered, the exact + // ordinal tracker above still permits that unseen entry instead of + // silently discarding it. + InPlaceAtomicBitset &bits = GetAvailableSlots(subscriber_id_); + if (FindNextVisibleSlot(bits, queued.ordinal - 1) != nullptr) { + queue_bitset_fallback_ = true; + next_slot_cache_valid_ = false; + return std::nullopt; + } + } + if (!AtomicIncRefCount(s, /*reliable=*/false, 1, queued.ordinal, + ref_vchan_id, false)) { + continue; + } + // The CAS above validates only the low kOrdinalBits stored in refs. After + // those bits wrap, a stale queue entry can therefore claim a newer slot + // generation. Verify the full ordinal while holding the reference and roll + // it back if the queue entry was an alias. + if (s->ordinal.load(std::memory_order_relaxed) != queued.ordinal || + s->vchan_id.load(std::memory_order_relaxed) != ref_vchan_id) { + AtomicIncRefCount(s, /*reliable=*/false, -1, queued.ordinal, + ref_vchan_id, false); + continue; + } + return ClaimedQueuedSlot{ + .slot = s, .ordinal = queued.ordinal, .vchan_id = ref_vchan_id}; + } + + return std::nullopt; +} + +MessageSlot *SubscriberImpl::FindNextVisibleSlot(InPlaceAtomicBitset &bits, + uint64_t max_ordinal) { + MessageSlot *best_slot = nullptr; + uint64_t best_ordinal = 0; + int cached_vchan_id = std::numeric_limits::min(); + OrdinalTracker *cached_tracker = nullptr; + + bits.Traverse([this, &best_slot, &best_ordinal, &cached_vchan_id, + &cached_tracker, max_ordinal](size_t i) { + if (embargoed_slots_.IsSet(i)) { + return; + } + MessageSlot *s = &ccb_->slots[i]; + const uint64_t refs = s->refs.load(std::memory_order_acquire); + if ((refs & kPubOwned) != 0) { + return; + } + const uint64_t ordinal = s->ordinal.load(std::memory_order_relaxed); + const int slot_vchan_id = + s->vchan_id.load(std::memory_order_relaxed); + if (ordinal == 0 || ordinal > max_ordinal || + (vchan_id_ != -1 && slot_vchan_id != -1 && + slot_vchan_id != vchan_id_) || + s->buffer_index.load(std::memory_order_relaxed) == -1) { + return; + } + if (slot_vchan_id != cached_vchan_id) { + cached_vchan_id = slot_vchan_id; + cached_tracker = &GetOrdinalTracker(slot_vchan_id); + } + if (cached_tracker->ordinals.Contains( + OrdinalAndVchanId{ordinal, slot_vchan_id})) { + return; + } + if (best_slot == nullptr || ordinal < best_ordinal) { + best_slot = s; + best_ordinal = ordinal; + } + }); + + return best_slot; } MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, int owner) { - InPlaceAtomicBitset &bits = GetAvailableSlots(owner); - embargoed_slots_.ClearAll(); constexpr int kMaxRetries = 1000; @@ -230,11 +415,110 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, const bool print_errors = false; #endif CheckReload(); + InPlaceAtomicBitset &bits = GetAvailableSlots(owner); + const bool stable_poll_drain = PollDrainPending(); + if (stable_poll_drain && poll_drain_exhausted_) { + if (ccb_->total_messages.load(std::memory_order_seq_cst) != + next_slot_cached_total_) { + poll_drain_exhausted_ = false; + queue_drain_tail_valid_ = false; + next_slot_cache_valid_ = false; + } else { + return nullptr; + } + } if (slot == nullptr) { // Prepopulate the active slots. PopulateActiveSlots(bits); } + if (!reliable && SubscriberQueueSize() > 0) { + InPlaceSlotQueue *queue = + GetAvailableSlotQueueAddress(subscriber_id_); + uint32_t queue_overflow_baseline = 0; + if (queue != nullptr) { + const uint32_t queue_drops = queue->ConsumeOverflow(); + const bool insertion_failed = queue->ConsumeInsertionFailure(); + queue_overflow_baseline = queue->OverflowCount(); + const bool recover_overflow = + options_.SubscriberQueueSize() == 0 && + (queue_drops != 0 || queue_overflow_baseline != 0); + if (options_.DetectDroppedMessages() && !recover_overflow) { + pending_queue_drops_ += static_cast(queue_drops); + } + if (recover_overflow || insertion_failed) { + // An evicted hint can refer to an older slot that remains readable in + // the authoritative bitset. For an inherited default queue, preserve + // the legacy no-drop behavior by delivering that backlog in ordinal + // order before accepting newer queue entries. Explicit queue sizes + // retain their requested bounded/drop-oldest semantics. + queue_bitset_fallback_ = true; + next_slot_cache_valid_ = false; + } + } + if (stable_poll_drain && !queue_drain_tail_valid_) { + next_slot_cached_total_ = CollectVisibleSlots(bits); + std::sort(active_slots_.begin(), active_slots_.end(), ActiveSlotLess); + next_slot_cursor_ = 0; + next_slot_cache_valid_ = true; + queue_drain_tail_ = queue == nullptr ? 0 : queue->Tail(); + queue_drain_tail_valid_ = true; + } + const uint64_t max_queue_position = + stable_poll_drain ? queue_drain_tail_ + : std::numeric_limits::max(); + std::optional claimed = + queue_bitset_fallback_ + ? std::nullopt + : FindNextQueuedSlot(max_queue_position); + if (claimed.has_value()) { + MessageSlot *new_slot = claimed->slot; + const uint64_t ordinal = claimed->ordinal; + const int vchan_id = claimed->vchan_id; + if (queue != nullptr && queue->InsertionFailed()) { + AtomicIncRefCount(new_slot, reliable, -1, ordinal, vchan_id, false); + queue->ConsumeInsertionFailure(); + queue_bitset_fallback_ = true; + next_slot_cache_valid_ = false; + continue; + } + if (queue != nullptr && options_.SubscriberQueueSize() == 0 && + queue->OverflowCount() != queue_overflow_baseline) { + AtomicIncRefCount(new_slot, reliable, -1, ordinal, vchan_id, false); + queue->ConsumeOverflow(); + queue_bitset_fallback_ = true; + next_slot_cache_valid_ = false; + continue; + } + if (queue != nullptr && options_.SubscriberQueueSize() != 0) { + const uint32_t concurrent_drops = queue->ConsumeOverflow(); + if (options_.DetectDroppedMessages()) { + pending_queue_drops_ += static_cast(concurrent_drops); + } + } + const int buffer_index = + new_slot->buffer_index.load(std::memory_order_relaxed); + if (!ValidateSlotBuffer(new_slot) || buffer_index == -1) { + if (print_errors) { + std::cerr << "Subscriber for " << Name() + << " detected buffer failure on slot: " << new_slot->id + << " buffer index: " << buffer_index; + new_slot->Dump(std::cerr); + } + embargoed_slots_.Set(new_slot->id); + AtomicIncRefCount(new_slot, reliable, -1, ordinal, vchan_id, false); + continue; + } + if (!stable_poll_drain) { + next_slot_cache_valid_ = false; + } + return new_slot; + } + // Push() may fail after a peer dies or loses a bounded CAS race. The + // publisher always records the slot in the bitset, so continue below and + // recover it through the authoritative path. + } + // Fast path: if the publisher hasn't appended any new messages since the // last successful NextSlot() call, the cached, already-sorted // active_slots_ list is still valid. We just need to scan forward from @@ -260,29 +544,44 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, // total_messages increment is seq_cst, so the relaxed bit write is // happens-before this seq_cst load and visible to the relaxed // bits.Traverse() inside CollectVisibleSlots(). - const uint64_t total = ccb_->total_messages; - const bool stable_poll_drain = PollDrainPending(); + const uint64_t total = + ccb_->total_messages.load(std::memory_order_seq_cst); if (!next_slot_cache_valid_ || (!stable_poll_drain && total != next_slot_cached_total_)) { - CollectVisibleSlots(bits); - std::sort(active_slots_.begin(), active_slots_.end(), ActiveSlotLess); - next_slot_cached_total_ = total; + const uint64_t snapshot_total = CollectVisibleSlots(bits); + if (queue_bitset_fallback_) { + std::sort(active_slots_.begin(), active_slots_.end(), + [](const ActiveSlot &a, const ActiveSlot &b) { + return a.ordinal < b.ordinal; + }); + } else { + std::sort(active_slots_.begin(), active_slots_.end(), ActiveSlotLess); + } + next_slot_cached_total_ = snapshot_total; next_slot_cursor_ = 0; next_slot_cache_valid_ = true; } // Walk forward from the cursor, skipping anything we've embargoed in // this NextSlot() invocation or already delivered to this subscriber. - auto &tracker = GetOrdinalTracker(vchan_id_); const ActiveSlot *new_slot = nullptr; + int cached_vchan_id = std::numeric_limits::min(); + OrdinalTracker *cached_tracker = nullptr; while (next_slot_cursor_ < active_slots_.size()) { const ActiveSlot &s = active_slots_[next_slot_cursor_]; if (embargoed_slots_.IsSet(s.slot->id)) { ++next_slot_cursor_; continue; } - if (s.ordinal != 0 && - !tracker.ordinals.Contains( + if (s.vchan_id != cached_vchan_id) { + cached_vchan_id = s.vchan_id; + cached_tracker = &GetOrdinalTracker(s.vchan_id); + } + // Concurrent publishers can reserve ordinals in one order and publish + // them in another. The bitset is the authoritative unread-message + // record, so an older ordinal remains deliverable until its exact + // ordinal has been claimed. + if (!cached_tracker->ordinals.Contains( OrdinalAndVchanId{s.ordinal, s.vchan_id})) { new_slot = &s; break; @@ -290,7 +589,19 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, ++next_slot_cursor_; } if (new_slot == nullptr) { - next_slot_cache_valid_ = false; + if (queue_bitset_fallback_) { + if (InPlaceSlotQueue *queue = + GetAvailableSlotQueueAddress(subscriber_id_); + queue != nullptr) { + queue->DiscardAll(); + } + queue_bitset_fallback_ = false; + } + if (stable_poll_drain) { + poll_drain_exhausted_ = true; + } else { + next_slot_cache_valid_ = false; + } // If we suppressed newer messages to keep a stable poll-drain snapshot, // re-arm the trigger so the next poll/Wait wakes promptly and // re-snapshots them. ClearPollFd() during the drain may already have @@ -298,7 +609,8 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, // caller that drains until empty and then re-waits (e.g. the bridge // transmitter) can block forever on the final message of a batch. if (stable_poll_drain && - ccb_->total_messages != next_slot_cached_total_) { + ccb_->total_messages.load(std::memory_order_seq_cst) != + next_slot_cached_total_) { Trigger(); } return nullptr; @@ -314,13 +626,14 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, // we just go back and try again. if (AtomicIncRefCount(new_slot->slot, reliable, 1, new_slot->ordinal, new_slot->vchan_id, false)) { - if (!ValidateSlotBuffer(new_slot->slot) || - new_slot->slot->buffer_index == -1) { + const int buffer_index = + new_slot->slot->buffer_index.load(std::memory_order_relaxed); + if (!ValidateSlotBuffer(new_slot->slot) || buffer_index == -1) { if (print_errors) { std::cerr << "Subscriber for " << Name() << " detected buffer failure on slot: " << new_slot->slot->id - << " buffer index: " << new_slot->slot->buffer_index; + << " buffer index: " << buffer_index; new_slot->slot->Dump(std::cerr); } // Failed to get a buffer for the slot. Embargo the slot so we don't @@ -330,7 +643,11 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, embargoed_slots_.Set(new_slot->slot->id); AtomicIncRefCount(new_slot->slot, reliable, -1, new_slot->ordinal, new_slot->vchan_id, false); - next_slot_cache_valid_ = false; + if (stable_poll_drain) { + ++next_slot_cursor_; + } else { + next_slot_cache_valid_ = false; + } continue; } // Successful claim. Advance the cursor so the next NextSlot() call @@ -340,7 +657,11 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, } // CAS failed: another subscriber raced us, or the slot was retired and // overwritten with a new ordinal. Drop the cache and re-snapshot. - next_slot_cache_valid_ = false; + if (stable_poll_drain) { + ++next_slot_cursor_; + } else { + next_slot_cache_valid_ = false; + } } return nullptr; } @@ -348,19 +669,30 @@ MessageSlot *SubscriberImpl::NextSlot(MessageSlot *slot, bool reliable, MessageSlot *SubscriberImpl::LastSlot(MessageSlot *slot, bool reliable, int owner) { - InPlaceAtomicBitset &bits = GetAvailableSlots(owner); - embargoed_slots_.ClearAll(); for (;;) { CheckReload(); + InPlaceAtomicBitset &bits = GetAvailableSlots(owner); + if (!reliable && SubscriberQueueSize() > 0) { + if (InPlaceSlotQueue *queue = + GetAvailableSlotQueueAddress(subscriber_id_); + queue != nullptr) { + const uint32_t queue_drops = queue->ConsumeOverflow(); + if (options_.DetectDroppedMessages()) { + pending_queue_drops_ += static_cast(queue_drops); + } + queue->ConsumeInsertionFailure(); + } + } if (slot == nullptr) { // Prepopulate the active slots. PopulateActiveSlots(bits); } - CollectVisibleSlots(bits); + (void)CollectVisibleSlots(bits); // Sort the active slots by timestamp. std::sort(active_slots_.begin(), active_slots_.end(), ActiveSlotLess); + newest_snapshot_ = active_slots_; ActiveSlot *new_slot = nullptr; if (!active_slots_.empty()) { @@ -372,14 +704,16 @@ MessageSlot *SubscriberImpl::LastSlot(MessageSlot *slot, bool reliable, } } if (new_slot == nullptr) { + newest_snapshot_.clear(); return nullptr; } // Increment the ref count. if (AtomicIncRefCount(new_slot->slot, reliable, 1, new_slot->ordinal, new_slot->vchan_id, false)) { - if (!ValidateSlotBuffer(new_slot->slot) || - new_slot->slot->buffer_index == -1) { + const int buffer_index = + new_slot->slot->buffer_index.load(std::memory_order_relaxed); + if (!ValidateSlotBuffer(new_slot->slot) || buffer_index == -1) { // Failed to get a buffer for the slot. Embargo the slot so we don't // see it again this loop and try again. embargoed_slots_.Set(new_slot->slot->id); @@ -389,6 +723,7 @@ MessageSlot *SubscriberImpl::LastSlot(MessageSlot *slot, bool reliable, } return new_slot->slot; } + newest_snapshot_.clear(); } } @@ -407,9 +742,12 @@ MessageSlot *SubscriberImpl::FindActiveSlotByTimestamp( continue; } MessageSlot *s = &ccb_->slots[i]; - uint64_t refs = s->refs.load(std::memory_order_relaxed); - if (s->ordinal != 0 && (refs & kPubOwned) == 0) { - buffer.push_back({s, s->ordinal, Prefix(s)->timestamp, s->vchan_id}); + uint64_t refs = s->refs.load(std::memory_order_acquire); + const uint64_t ordinal = s->ordinal.load(std::memory_order_relaxed); + if (ordinal != 0 && (refs & kPubOwned) == 0) { + buffer.push_back( + {s, ordinal, Prefix(s)->timestamp, + s->vchan_id.load(std::memory_order_relaxed)}); } } // Sort by timestamp. @@ -436,7 +774,8 @@ MessageSlot *SubscriberImpl::FindActiveSlotByTimestamp( // Try to increment the ref count. if (AtomicIncRefCount(it->slot, reliable, 1, it->ordinal, it->vchan_id, false)) { - if (!ValidateSlotBuffer(it->slot) || it->slot->buffer_index == -1) { + if (!ValidateSlotBuffer(it->slot) || + it->slot->buffer_index.load(std::memory_order_relaxed) == -1) { // Failed to get a buffer for the slot. Embargo the slot so we don't // see it again this loop and try again. embargoed_slots_.Set(it->slot->id); @@ -444,7 +783,11 @@ MessageSlot *SubscriberImpl::FindActiveSlotByTimestamp( false); continue; } - it->slot->flags |= kMessageSeen; + it->slot->flags.fetch_or(kMessageSeen, std::memory_order_relaxed); + if (reliable) { + it->slot->flags.fetch_or(kMessageSeenByReliable, + std::memory_order_relaxed); + } it->slot->sub_owners.Set(owner); return it->slot; } diff --git a/client/subscriber.h b/client/subscriber.h index 9c97a08a..c33fbd3e 100644 --- a/client/subscriber.h +++ b/client/subscriber.h @@ -8,6 +8,8 @@ #include "common/fast_ring_buffer.h" #include #include +#include +#include namespace subspace { namespace details { @@ -32,6 +34,12 @@ struct OrdinalAndVchanId { } }; +struct ClaimedQueuedSlot { + MessageSlot *slot = nullptr; + uint64_t ordinal = 0; + int vchan_id = -1; +}; + template inline H AbslHashValue(H h, const OrdinalAndVchanId &x) { return H::combine(std::move(h), x.ordinal, x.vchan_id); } @@ -40,21 +48,27 @@ template inline H AbslHashValue(H h, const OrdinalAndVchanId &x) { // shared memory. class SubscriberImpl : public ClientChannel { public: - SubscriberImpl(const std::string &name, int num_slots, int channel_id, - int subscriber_id, int vchan_id, uint64_t session_id, - std::string type, const SubscriberOptions &options, + SubscriberImpl(const std::string &name, int num_slots, + int default_subscriber_queue_size, + uint64_t subscriber_queue_arena_size, + int subscriber_queue_size, int channel_id, int subscriber_id, + int vchan_id, uint64_t session_id, std::string type, + const SubscriberOptions &options, std::function reload, int user_id, int group_id) - : ClientChannel(name, num_slots, channel_id, vchan_id, + : ClientChannel(name, num_slots, default_subscriber_queue_size, + subscriber_queue_arena_size, channel_id, vchan_id, std::move(session_id), std::move(type), std::move(reload), user_id, group_id), - subscriber_id_(subscriber_id), options_(options) { + subscriber_id_(subscriber_id), + subscriber_queue_size_(subscriber_queue_size), options_(options) { // Preallocate to avoid malloc later. (void)GetOrdinalTracker(vchan_id_); } ~SubscriberImpl() override { Unmap(); } void InitActiveMessages(); + void ResetDeliveryState(); bool UsesSplitBuffers() const { return UseSplitBuffers(); } std::shared_ptr shared_from_this() { @@ -63,13 +77,23 @@ class SubscriberImpl : public ClientChannel { } int64_t CurrentOrdinal() const { - return CurrentSlot() == nullptr ? -1 : CurrentSlot()->ordinal; + return CurrentSlot() == nullptr + ? -1 + : static_cast( + CurrentSlot()->ordinal.load(std::memory_order_relaxed)); } int64_t Timestamp() const { return Timestamp(CurrentSlot()); } int64_t Timestamp(MessageSlot *slot) const { - return slot == nullptr ? 0 : slot->timestamp; + return slot == nullptr + ? 0 + : static_cast( + slot->timestamp.load(std::memory_order_relaxed)); } bool IsReliable() const { return options_.IsReliable(); } + int SubscriberQueueSize() const override { return subscriber_queue_size_; } + void SetEffectiveSubscriberQueueSize(int size) { + subscriber_queue_size_ = ResolveSubscriberQueueSize(NumSlots(), size); + } int32_t SlotSize() const { return ClientChannel::SlotSize(CurrentSlot()); } @@ -97,12 +121,20 @@ class SubscriberImpl : public ClientChannel { void ClaimSlot(MessageSlot *slot, int vchan_id, bool was_newest); void UnreadSlot(MessageSlot *slot); void RememberOrdinal(uint64_t ordinal, int vchan_id); - void CollectVisibleSlots(InPlaceAtomicBitset &bits); + uint64_t CollectVisibleSlots(InPlaceAtomicBitset &bits); + std::optional + FindNextQueuedSlot(uint64_t max_queue_position); + MessageSlot *FindNextVisibleSlot(InPlaceAtomicBitset &bits, + uint64_t max_ordinal); void IgnoreActivation(MessageSlot *slot) { - RememberOrdinal(slot->ordinal, slot->vchan_id); + RememberOrdinal(slot->ordinal.load(std::memory_order_relaxed), + slot->vchan_id.load(std::memory_order_relaxed)); DecrementSlotRef(slot, true); - slot->flags |= kMessageSeen; + slot->flags.fetch_or(kMessageSeen, std::memory_order_relaxed); + if (IsReliable()) { + slot->flags.fetch_or(kMessageSeenByReliable, std::memory_order_relaxed); + } } // A subscriber wants to find a slot with a message in it. There are // two ways to get this: @@ -141,12 +173,14 @@ class SubscriberImpl : public ClientChannel { } void DecrementSlotRef(MessageSlot *slot, bool retire) { - AtomicIncRefCount(slot, IsReliable(), -1, slot->ordinal & kOrdinalMask, - vchan_id_, retire); + AtomicIncRefCount(slot, IsReliable(), -1, + slot->ordinal.load(std::memory_order_relaxed) & + kOrdinalMask, + slot->vchan_id.load(std::memory_order_relaxed), retire); } bool SlotExpired(MessageSlot *slot, uint32_t ordinal) { - return slot->ordinal != ordinal; + return slot->ordinal.load(std::memory_order_relaxed) != ordinal; } std::shared_ptr LockWeakMessage(MessageSlot *slot, @@ -154,7 +188,7 @@ class SubscriberImpl : public ClientChannel { if (slot == nullptr) { return nullptr; } - if (slot->ordinal != ordinal) { + if (slot->ordinal.load(std::memory_order_relaxed) != ordinal) { return nullptr; } // If we are still holding on to the same active message, return it. @@ -163,8 +197,10 @@ class SubscriberImpl : public ClientChannel { return active_message_; } std::shared_ptr &msg = active_messages_[slot->id]; - msg->Set(slot->message_size, GetBufferAddress(slot), slot->ordinal, - Timestamp(slot), slot->vchan_id, false, false); + msg->Set(slot->message_size.load(std::memory_order_relaxed), + GetBufferAddress(slot), + slot->ordinal.load(std::memory_order_relaxed), Timestamp(slot), + slot->vchan_id.load(std::memory_order_relaxed), false, false); if (msg->length == 0) { // Failed to get an active message, return an empty shared_ptr. return nullptr; @@ -173,6 +209,9 @@ class SubscriberImpl : public ClientChannel { } int DetectDrops(int vchan_id); + int ConsumeQueueDrops() { + return std::exchange(pending_queue_drops_, 0); + } // Search the active list for a message with the given timestamp. If found, // take ownership of the slot found. Return nullptr if nothing found in which @@ -303,6 +342,8 @@ class SubscriberImpl : public ClientChannel { // stable snapshot and cannot be kept chasing concurrently published // messages forever. poll_drain_pending_ = true; + poll_drain_exhausted_ = false; + queue_drain_tail_valid_ = false; next_slot_cache_valid_ = false; return trigger_.GetPollFd(); } @@ -326,6 +367,7 @@ class SubscriberImpl : public ClientChannel { } int subscriber_id_; + int subscriber_queue_size_ = 0; toolbelt::TriggerFd trigger_; std::vector reliable_publishers_; SubscriberOptions options_; @@ -336,6 +378,10 @@ class SubscriberImpl : public ClientChannel { // will keep the memory allocation to the first search on a subscriber. Most // subscribers won't use this. std::vector search_buffer_; + // Generation snapshots skipped by a successful ReadNewest. ClaimSlot + // temporarily pins each entry before clearing its bit so a concurrent slot + // recycle cannot have its newly published bit cleared. + std::vector newest_snapshot_; // We have one active message per slot. These are allocated when the // subscriber is created to avoid memory allocation for every message we @@ -372,6 +418,14 @@ class SubscriberImpl : public ClientChannel { size_t next_slot_cursor_ = 0; bool next_slot_cache_valid_ = false; bool poll_drain_pending_ = false; + bool poll_drain_exhausted_ = false; + // Queue overflow can remove an older hint while its slot is still readable. + // Stay on the ordered bitset path until that retained backlog is exhausted, + // otherwise a newer queue entry would advance last_ordinal_seen past it. + bool queue_bitset_fallback_ = false; + int pending_queue_drops_ = 0; + uint64_t queue_drain_tail_ = 0; + bool queue_drain_tail_valid_ = false; }; } // namespace details } // namespace subspace diff --git a/common/BUILD.bazel b/common/BUILD.bazel index 105e5994..82a41900 100644 --- a/common/BUILD.bazel +++ b/common/BUILD.bazel @@ -50,6 +50,7 @@ cc_test( "fast_ring_buffer.h", ], deps = [ + ":subspace_common", "@abseil-cpp//absl/container:flat_hash_set", "@googletest//:gtest_main", ], diff --git a/common/atomic_bitset.h b/common/atomic_bitset.h index 28dfdc76..e26f2c01 100644 --- a/common/atomic_bitset.h +++ b/common/atomic_bitset.h @@ -49,6 +49,12 @@ template class AtomicBitSet { bits_[word].fetch_and(~(1ULL << offset), std::memory_order_relaxed); } + void ClearSeqCst(size_t bit) { + size_t word = bit / 64; + size_t offset = bit % 64; + bits_[word].fetch_and(~(1ULL << offset), std::memory_order_seq_cst); + } + // Atomically clear bit and return whether it was previously set. // Use this when racing concurrent producers must establish unique // ownership of a bit (e.g. claiming a free slot from a shared pool): @@ -139,6 +145,23 @@ template class AtomicBitSet { } } + void TraverseSeqCst(std::function func) const { + for (size_t i = 0; i < BitsToWords(num_bits_); i++) { + size_t shift = 0; + size_t bit = i * 64; + while (bit < num_bits_ && shift < 64) { + uint64_t word = bits_[i].load(std::memory_order_seq_cst) >> shift; + size_t n = ffsll(word); + if (n == 0) { + break; + } + bit += n; + func(bit - 1); + shift += n; + } + } + } + private: // If SizeInBits is 0, then kNumWords is 0 (allows bits to be stored outside // the object). diff --git a/common/channel.cc b/common/channel.cc index 787e9f4f..3ea2993a 100644 --- a/common/channel.cc +++ b/common/channel.cc @@ -142,8 +142,14 @@ void UnmapMemory(void *p, size_t size, } Channel::Channel(const std::string &name, int num_slots, int channel_id, - std::string type, std::function reload) - : name_(name), num_slots_(num_slots), channel_id_(channel_id), + int subscriber_queue_size, + uint64_t subscriber_queue_arena_size, std::string type, + std::function reload) + : name_(name), num_slots_(num_slots), + subscriber_queue_size_( + ResolveSubscriberQueueSize(num_slots, subscriber_queue_size)), + subscriber_queue_arena_size_(subscriber_queue_arena_size), + channel_id_(channel_id), type_(std::move(type)), reload_callback_(std::move(reload)) {} void Channel::Unmap() { @@ -158,7 +164,7 @@ void Channel::Unmap() { ccb_ = nullptr; bcb_ = nullptr; UnmapMemory(scb, sizeof(SystemControlBlock), "SCB"); - UnmapMemory(ccb, CcbSize(num_slots_), "CCB"); + UnmapMemory(ccb, CcbSize(num_slots_, subscriber_queue_arena_size_), "CCB"); UnmapMemory(bcb, sizeof(BufferControlBlock), "BCB"); } @@ -285,10 +291,13 @@ void MessageSlot::Dump(std::ostream &os) const { os << " refs: " << just_refs << " reliable refs: " << reliable_refs << " ord: " << ref_ord; } - os << " ordinal: " << ordinal << " buffer_index: " << buffer_index - << " vchan_id: " << vchan_id << " timestamp: " << timestamp - << " message size: " << message_size << " raw refs: " << std::hex << refs - << " flags: " << flags << std::dec << "\n"; + os << " ordinal: " << ordinal.load(std::memory_order_relaxed) + << " buffer_index: " << buffer_index.load(std::memory_order_relaxed) + << " vchan_id: " << vchan_id.load(std::memory_order_relaxed) + << " timestamp: " << timestamp.load(std::memory_order_relaxed) + << " message size: " << message_size.load(std::memory_order_relaxed) + << " raw refs: " << std::hex << l_refs + << " flags: " << flags.load(std::memory_order_relaxed) << std::dec << "\n"; } void Channel::DumpSlots(std::ostream &os) const { @@ -307,7 +316,7 @@ void Channel::Dump(std::ostream &os) const { toolbelt::Hexdump(scb_, 64); os << "CCB:\n"; - toolbelt::Hexdump(ccb_, CcbSize(num_slots_)); + toolbelt::Hexdump(ccb_, CcbSize(num_slots_, subscriber_queue_arena_size_)); os << "Slots:\n"; DumpSlots(os); @@ -340,8 +349,10 @@ void Channel::GetStatsCounters(uint64_t &total_bytes, uint64_t &total_messages, } uint64_t Channel::GetVirtualMemoryUsage() const { - uint64_t size = sizeof(SystemControlBlock) + CcbSize(num_slots_) + - sizeof(BufferControlBlock); + uint64_t size = + sizeof(SystemControlBlock) + + CcbSize(num_slots_, subscriber_queue_arena_size_) + + sizeof(BufferControlBlock); for (int i = 0; i < ccb_->num_buffers; i++) { if (bcb_->refs[i] > 0) { size += bcb_->sizes[i]; @@ -360,9 +371,8 @@ void Channel::CleanupSlots(int owner, bool reliable, bool is_pub, // Is the slot owned by this publisher? if (refs == (kPubOwned | uint64_t(owner))) { // Owned by this publisher, clear slot. - slot->ordinal = 0; - slot->refs = - 0; // Sequentially consistent because we've changed the ordinal too. + slot->ordinal.store(0, std::memory_order_relaxed); + slot->refs.store(0, std::memory_order_release); // Clear the slot in all the subscriber bitsets. ccb_->subscribers.Traverse([this, slot](int sub_id) { diff --git a/common/channel.h b/common/channel.h index f2aa352b..0e86b64d 100644 --- a/common/channel.h +++ b/common/channel.h @@ -15,9 +15,11 @@ #include "toolbelt/bitset.h" #include "toolbelt/fd.h" +#include #include #include #include +#include #include #include @@ -110,8 +112,10 @@ static_assert(sizeof(MessagePrefix) == 64, "MessagePrefix size is not 64 bytes"); // Flags for MessageSlot flags. -constexpr int kMessageSeen = 1; // Message has been seen. +constexpr int kMessageSeen = 1; // Message has been seen by any subscriber. constexpr int kMessageIsActivation = 2; // This is an activation message. +constexpr int kMessageSeenByReliable = + 4; // Message has been seen by a reliable subscriber. // We need a max channels number because the size of things in // shared memory needs to be fixed. @@ -121,6 +125,17 @@ constexpr int kMaxChannels = 1024; // and publisher reference. Best if it's a multiple of 64 because // it's used as the size in a toolbelt::BitSet. constexpr int kMaxSlotOwners = 1024; +// Default per-subscriber queue depth used whenever the publisher provisions a +// non-empty arena and the subscriber does not request an override. +constexpr int kDefaultSubscriberQueueSize = 16; +// Standard packed arena size for callers that opt into subscriber queues. This +// fits 100 default-sized (16-entry) queues. Publisher options default to zero, +// which keeps the available-slot bitset path and omits the queue arena. +constexpr uint64_t kDefaultSubscriberQueueArenaSize = 64'000; +constexpr size_t kDefaultMaxAvailableSlotQueueCapacity = 1024; +constexpr size_t kMaxSlotQueueCasAttempts = 64; +constexpr uint32_t kChannelControlBlockVersion = 4; +constexpr size_t kMaxChannelControlBlockSize = 1ULL << 30; // This limits the number of virtual channels. Each virtual channel // needs its own ordinal counter in the CCB (8 bytes each). @@ -205,18 +220,30 @@ struct SystemControlBlock { // This is the meta data for a slot. struct MessageSlot { std::atomic refs; // Number of subscribers referring to this slot. - uint64_t ordinal; // Message ordinal held currently in slot. - uint64_t message_size; // Size of message held in slot. + std::atomic ordinal; // Message ordinal held currently in slot. + std::atomic message_size; // Size of message held in slot. int32_t id; // Unique ID for slot (0...num_slots-1). - int16_t buffer_index; // Index of buffer. - int16_t vchan_id; // Virtual channel ID. + std::atomic buffer_index; // Index of buffer. + std::atomic vchan_id; // Virtual channel ID. AtomicBitSet sub_owners; // One bit per subscriber. - uint64_t timestamp; // Timestamp of message. - uint32_t flags; - int32_t bridged_slot_id; // Slot ID of other side of bridge. + std::atomic timestamp; // Timestamp of message. + std::atomic flags; + std::atomic + bridged_slot_id; // Slot ID of other side of bridge. void Dump(std::ostream &os) const; }; +static_assert(sizeof(MessageSlot) == 184); +static_assert(offsetof(MessageSlot, refs) == 0); +static_assert(offsetof(MessageSlot, ordinal) == 8); +static_assert(offsetof(MessageSlot, message_size) == 16); +static_assert(offsetof(MessageSlot, id) == 24); +static_assert(offsetof(MessageSlot, buffer_index) == 28); +static_assert(offsetof(MessageSlot, vchan_id) == 30); +static_assert(offsetof(MessageSlot, sub_owners) == 32); +static_assert(offsetof(MessageSlot, timestamp) == 168); +static_assert(offsetof(MessageSlot, flags) == 176); +static_assert(offsetof(MessageSlot, bridged_slot_id) == 180); struct ActiveSlot { MessageSlot *slot; @@ -225,6 +252,266 @@ struct ActiveSlot { int vchan_id; }; +struct QueuedSlot { + int32_t slot_id; + uint64_t ordinal; +}; + +struct SlotQueueEntry { + // Sequence number used to publish an entry after its payload is written and + // to mark it reusable after the consumer has popped it. + std::atomic sequence; + std::atomic ordinal; + std::atomic slot_id; +}; +static_assert(sizeof(SlotQueueEntry) == 24); +static_assert(offsetof(SlotQueueEntry, sequence) == 0); +static_assert(offsetof(SlotQueueEntry, ordinal) == 8); +static_assert(offsetof(SlotQueueEntry, slot_id) == 16); + +// A bounded MPSC queue stored in shared memory after the available-slots +// bitsets. Publishers push slot IDs as they publish; the single owning +// subscriber pops them to avoid scanning its bitset. For unreliable subscribers +// this queue is the hot-path source of truth; the bitset is still maintained for +// reliable mode and diagnostics while the queue path is proven out. +class InPlaceSlotQueue { +public: + InPlaceSlotQueue(size_t capacity, bool drop_oldest = true) { + Init(capacity, drop_oldest); + } + + // Initialize queue metadata and mark every ring entry as free. `capacity` + // is the number of SlotQueueEntry objects laid out immediately after this + // header in shared memory. + void Init(size_t capacity, bool drop_oldest = true) { + capacity_ = capacity; + head_.store(0, std::memory_order_relaxed); + tail_.store(0, std::memory_order_relaxed); + overflow_count_.store(0, std::memory_order_relaxed); + insertion_failed_.store(false, std::memory_order_relaxed); + drop_oldest_ = drop_oldest; + for (size_t i = 0; i < capacity_; i++) { + entries_[i].sequence.store(i, std::memory_order_relaxed); + entries_[i].ordinal.store(0, std::memory_order_relaxed); + entries_[i].slot_id.store(-1, std::memory_order_relaxed); + } + } + + size_t Capacity() const { return capacity_; } + uint64_t Head() const { return head_.load(std::memory_order_acquire); } + uint64_t Tail() const { return tail_.load(std::memory_order_acquire); } + + // Push a published slot. Multiple publishers may call this concurrently. + // If an explicit queue is full, evict its oldest hint and enqueue the newest + // one. An inherited queue instead rejects the hint so its subscriber recovers + // every unread ordinal from the authoritative bitset. + bool Push(int32_t slot_id, uint64_t ordinal, + bool report_insertion_failure = true) { + if (capacity_ == 0) { + if (report_insertion_failure) { + MarkInsertionFailure(); + } + return false; + } + + SlotQueueEntry *entry = nullptr; + uint64_t tail = tail_.load(std::memory_order_relaxed); + for (size_t attempt = 0; attempt < kMaxSlotQueueCasAttempts; ++attempt) { + const uint64_t head = head_.load(std::memory_order_acquire); + if (tail - head >= capacity_) { + // Inherited queues preserve legacy no-drop delivery through the + // authoritative bitset. Do not expose a newer queue entry before the + // subscriber observes the fallback signal. + if (!drop_oldest_) { + if (report_insertion_failure) { + MarkInsertionFailure(); + } + return false; + } + if (!DropFront()) { + if (report_insertion_failure) { + MarkInsertionFailure(); + } + return false; + } + overflow_count_.fetch_add(1, std::memory_order_release); + tail = tail_.load(std::memory_order_relaxed); + continue; + } + SlotQueueEntry &candidate = entries_[tail % capacity_]; + // A consumer publishes the reusable sequence after advancing head_. It + // may be paused or terminated between those operations. Do not reserve + // the entry until it is reusable: reserving first would force this + // producer to wait indefinitely for that consumer. + if (candidate.sequence.load(std::memory_order_acquire) != tail) { + if (report_insertion_failure) { + MarkInsertionFailure(); + } + return false; + } + if (tail_.compare_exchange_strong(tail, tail + 1, + std::memory_order_acq_rel, + std::memory_order_relaxed)) { + entry = &candidate; + break; + } + } + if (entry == nullptr) { + if (report_insertion_failure) { + MarkInsertionFailure(); + } + return false; + } + + entry->slot_id.store(slot_id, std::memory_order_relaxed); + entry->ordinal.store(ordinal, std::memory_order_relaxed); + entry->sequence.store(tail + 1, std::memory_order_release); + return true; + } + + void MarkInsertionFailure() { + insertion_failed_.store(true, std::memory_order_release); + } + + // Read the oldest queued slot without consuming it. This lets poll-driven + // subscribers stop at the end of a stable drain snapshot without losing the + // first newer message. + bool TryPeek(QueuedSlot &slot) { + if (capacity_ == 0) { + return false; + } + + const uint64_t head = head_.load(std::memory_order_acquire); + SlotQueueEntry &entry = entries_[head % capacity_]; + if (entry.sequence.load(std::memory_order_acquire) != head + 1) { + return false; + } + + QueuedSlot candidate = { + entry.slot_id.load(std::memory_order_relaxed), + entry.ordinal.load(std::memory_order_relaxed), + }; + if (entry.sequence.load(std::memory_order_acquire) != head + 1) { + return false; + } + slot = candidate; + return true; + } + + // Drop the oldest queued slot. Producers use this on overflow, and the + // subscriber uses it to discard stale entries whose slot has been reused. + bool DropFront() { + if (capacity_ == 0) { + return false; + } + + uint64_t head = head_.load(std::memory_order_relaxed); + for (size_t attempt = 0; attempt < kMaxSlotQueueCasAttempts; ++attempt) { + SlotQueueEntry &entry = entries_[head % capacity_]; + if (entry.sequence.load(std::memory_order_acquire) != head + 1) { + return false; + } + if (head_.compare_exchange_weak(head, head + 1, + std::memory_order_acq_rel, + std::memory_order_relaxed)) { + entry.sequence.store(head + capacity_, std::memory_order_release); + return true; + } + } + return false; + } + + // Discard a bounded snapshot of queued hints. The available-slot bitset + // remains authoritative, so subscribers use this after switching to bitset + // recovery following queue overflow or insertion failure. + void DiscardAll() { + for (size_t i = 0; i < capacity_; ++i) { + if (!DropFront()) { + return; + } + } + } + + // Pop one slot for the owning subscriber. There is exactly one consumer per + // queue, but producers may advance head_ to evict on overflow, so the + // consumer claims the front entry with a CAS. + bool TryPop(QueuedSlot &slot) { + if (capacity_ == 0) { + return false; + } + + uint64_t head = head_.load(std::memory_order_relaxed); + for (size_t attempt = 0; attempt < kMaxSlotQueueCasAttempts; ++attempt) { + SlotQueueEntry &entry = entries_[head % capacity_]; + if (entry.sequence.load(std::memory_order_acquire) != head + 1) { + return false; + } + QueuedSlot candidate = { + entry.slot_id.load(std::memory_order_relaxed), + entry.ordinal.load(std::memory_order_relaxed), + }; + if (head_.compare_exchange_weak(head, head + 1, + std::memory_order_acq_rel, + std::memory_order_relaxed)) { + slot = candidate; + entry.sequence.store(head + capacity_, std::memory_order_release); + return true; + } + } + return false; + } + + // Return and clear the number of older queued slots evicted to preserve the + // newest data. + uint32_t ConsumeOverflow() { + return overflow_count_.exchange(0, std::memory_order_acq_rel); + } + + uint32_t OverflowCount() const { + return overflow_count_.load(std::memory_order_acquire); + } + + bool ConsumeInsertionFailure() { + return insertion_failed_.exchange(false, std::memory_order_acq_rel); + } + + bool InsertionFailed() const { + return insertion_failed_.load(std::memory_order_acquire); + } + +private: + // Fixed ring capacity for this queue, capped independently of the channel's + // slot count to keep shared-memory usage bounded. + size_t capacity_ = 0; + // Next sequence number the single subscriber will try to pop. + std::atomic head_{0}; + // Next sequence number producers will reserve for Push(). + std::atomic tail_{0}; + std::atomic overflow_count_{0}; + std::atomic insertion_failed_{false}; + // Explicit queues drop their oldest hint on overflow. Inherited queues leave + // the queue unchanged and force the subscriber to recover from its bitset. + bool drop_oldest_ = true; + // Flexible array of `capacity_` entries stored immediately after the header. + SlotQueueEntry entries_[0]; +}; +static_assert(sizeof(InPlaceSlotQueue) == 32); + +inline size_t SizeofSlotQueue(size_t capacity) { + return sizeof(InPlaceSlotQueue) + sizeof(SlotQueueEntry) * capacity; +} + +constexpr uint64_t kInvalidSlotQueueOffset = + std::numeric_limits::max(); + +inline int ResolveSubscriberQueueSize(int num_slots, + int subscriber_queue_size) { + if (num_slots <= 0 || subscriber_queue_size <= 0) { + return 0; + } + return subscriber_queue_size; +} + struct BufferControlBlock { std::atomic refs[kMaxBuffers]; // Number of references to this buffer. @@ -341,6 +628,8 @@ struct ChannelControlBlock { // a.k.a CCB char channel_name[kMaxChannelName]; // So that you can see the name in a // debugger or hexdump. int num_slots; + int subscriber_queue_size; // Fixed inherited per-subscriber queue capacity. + uint32_t version; OrdinalAccumulator ordinals; // Ordinal accumulator for virtual channels. ActivationTracker activation_tracker; // Tracks which vchan_ids have been // activated by a publisher. @@ -355,6 +644,8 @@ struct ChannelControlBlock { // a.k.a CCB // Statistics counters. std::atomic total_bytes; + // Number of completed publications, including activation messages. This is + // also the version stamp for subscriber delivery snapshots. std::atomic total_messages; std::atomic max_message_size; std::atomic total_drops; @@ -371,18 +662,94 @@ struct ChannelControlBlock { // a.k.a CCB // AtomicBitSet<0> freeSlots[num_slots]; // Followed by: // AtomicBitSet<0> availableSlots[kMaxSlotOwners]; + // Followed by: + // AvailableSlotQueueIndex availableSlotQueueIndex; + // Followed by: + // A packed arena of variable-capacity InPlaceSlotQueue objects. // }; +static_assert(offsetof(ChannelControlBlock, version) == 72); + +// Locates each subscriber's variable-capacity queue in the packed queue arena. +// Offsets are relative to the start of the arena. +struct AvailableSlotQueueIndex { + std::atomic next_offset; + std::array, kMaxSlotOwners> offsets; + std::array, kMaxSlotOwners> active_publishers; +}; +static_assert(offsetof(AvailableSlotQueueIndex, next_offset) == 0); +static_assert(offsetof(AvailableSlotQueueIndex, offsets) == 8); +static_assert(offsetof(AvailableSlotQueueIndex, active_publishers) == 8200); +static_assert(sizeof(AvailableSlotQueueIndex) == 12296); inline size_t AvailableSlotsSize(int num_slots) { return SizeofAtomicBitSet(num_slots) * kMaxSlotOwners; } -inline size_t CcbSize(int num_slots) { +inline size_t AvailableSlotQueueIndexSize() { + return Aligned(sizeof(AvailableSlotQueueIndex)); +} + +enum class SlotQueueBlockState : uint32_t { + kAllocated = 0, + kRetired = 1, + kFree = 2, +}; + +struct alignas(64) SlotQueueBlockHeader { + uint64_t block_size = 0; + std::atomic state{ + static_cast(SlotQueueBlockState::kFree)}; + uint32_t reserved = 0; + AtomicBitSet waiting_publishers; +}; +static_assert(sizeof(SlotQueueBlockHeader) == 192); +static_assert(offsetof(SlotQueueBlockHeader, waiting_publishers) == 16); + +inline size_t SlotQueueBlockHeaderSize() { + return Aligned(sizeof(SlotQueueBlockHeader)); +} + +inline size_t SlotQueueBlockSize(size_t capacity) { + return SlotQueueBlockHeaderSize() + Aligned(SizeofSlotQueue(capacity)); +} + +inline size_t CcbSize(int num_slots, uint64_t subscriber_queue_arena_size) { return Aligned(sizeof(ChannelControlBlock) + num_slots * sizeof(MessageSlot)) + Aligned(SizeofAtomicBitSet(num_slots)) * 2 + - AvailableSlotsSize(num_slots); + AvailableSlotsSize(num_slots) + + AvailableSlotQueueIndexSize() + + static_cast(subscriber_queue_arena_size); +} + +inline size_t CcbSize(int num_slots) { + return CcbSize(num_slots, /*subscriber_queue_arena_size=*/0); +} + +inline absl::StatusOr +CheckedCcbSize(int num_slots, uint64_t subscriber_queue_arena_size) { + if (num_slots < 0) { + return absl::InvalidArgumentError("num_slots must be non-negative"); + } + const size_t slots = static_cast(num_slots); + if (slots > kMaxChannelControlBlockSize / sizeof(MessageSlot)) { + return absl::ResourceExhaustedError( + "num_slots exceeds the channel control block limit"); + } + if (slots > (std::numeric_limits::max() - + sizeof(ChannelControlBlock)) / + sizeof(MessageSlot)) { + return absl::ResourceExhaustedError("channel control block size overflow"); + } + const size_t base_size = CcbSize(num_slots, 0); + if (base_size > kMaxChannelControlBlockSize || + subscriber_queue_arena_size > + kMaxChannelControlBlockSize - base_size) { + return absl::ResourceExhaustedError( + "channel control block exceeds the 1 GiB limit"); + } + return base_size + static_cast(subscriber_queue_arena_size); } struct SlotBuffer { @@ -442,7 +809,9 @@ class Channel : public std::enable_shared_from_this { }; Channel(const std::string &name, int num_slots, int channel_id, - std::string type, std::function reload = nullptr); + int subscriber_queue_size, uint64_t subscriber_queue_arena_size, + std::string type, + std::function reload = nullptr); virtual ~Channel() { Unmap(); } virtual void Unmap(); @@ -466,9 +835,15 @@ class Channel : public std::enable_shared_from_this { std::string BufferSharedMemoryName(uint64_t session_id, int buffer_index) const; - void RegisterSubscriber(int sub_id, int vchan_id, bool /*is_new*/) { - ccb_->subscribers.Set(sub_id); + 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); + } SubscriberCounter num_subs; ccb_->subscribers.Traverse([this, &num_subs](size_t id) { num_subs.AddSubscriber(ccb_->sub_vchan_ids[id]); @@ -478,6 +853,59 @@ class Channel : public std::enable_shared_from_this { int GetSubVchanId(int32_t i) const { return ccb_->sub_vchan_ids[i]; } + void SeedAvailableSlotQueue(int sub_id, int vchan_id) { + InPlaceAtomicBitset &bits = GetAvailableSlots(sub_id); + InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(sub_id); + 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); + if (ordinal == 0 || buffer_index == -1) { + return false; + } + const int slot_vchan_id = + slot.vchan_id.load(std::memory_order_relaxed); + if (vchan_id != -1 && slot_vchan_id != -1 && + vchan_id != slot_vchan_id) { + return false; + } + return true; + }; + + uint64_t last_ordinal = 0; + for (;;) { + MessageSlot *best = nullptr; + for (int i = 0; i < NumSlots(); i++) { + MessageSlot &slot = ccb_->slots[i]; + if (!visible(slot)) { + continue; + } + const uint64_t ordinal = + slot.ordinal.load(std::memory_order_relaxed); + if (ordinal <= last_ordinal) { + continue; + } + if (best == nullptr || + ordinal < best->ordinal.load(std::memory_order_relaxed)) { + best = &slot; + } + } + if (best == nullptr) { + return; + } + bits.Set(best->id); + if (queue != nullptr) { + queue->Push(best->id, + best->ordinal.load(std::memory_order_relaxed)); + } + last_ordinal = best->ordinal.load(std::memory_order_relaxed); + } + } + void DumpSlots(std::ostream &os) const; virtual void Dump(std::ostream &os) const; @@ -521,6 +949,16 @@ class Channel : public std::enable_shared_from_this { // Get the number of slots in the channel (can't be changed) int NumSlots() const { return num_slots_; } virtual void SetNumSlots(int n) { num_slots_ = n; } + virtual int SubscriberQueueSize() const { return subscriber_queue_size_; } + virtual void SetSubscriberQueueSize(int n) { + subscriber_queue_size_ = ResolveSubscriberQueueSize(num_slots_, n); + } + virtual uint64_t SubscriberQueueArenaSize() const { + return subscriber_queue_arena_size_; + } + virtual void SetSubscriberQueueArenaSize(uint64_t size) { + subscriber_queue_arena_size_ = size; + } std::string SlotType() const { return type_; } void CleanupSlots(int owner, bool reliable, bool is_pub, int vchan_id); @@ -567,6 +1005,12 @@ class Channel : public std::enable_shared_from_this { char *EndOfFreeSlots() const { return EndOfRetiredSlots() + Aligned(SizeofAtomicBitSet(num_slots_)); } + char *EndOfAvailableSlots() const { + return EndOfFreeSlots() + AvailableSlotsSize(num_slots_); + } + char *EndOfAvailableSlotQueueIndex() const { + return EndOfAvailableSlots() + AvailableSlotQueueIndexSize(); + } InPlaceAtomicBitset *RetiredSlotsAddr() { return reinterpret_cast(EndOfSlots()); @@ -601,6 +1045,61 @@ class Channel : public std::enable_shared_from_this { EndOfFreeSlots() + SizeofAtomicBitSet(num_slots_) * sub_id); } + AvailableSlotQueueIndex *GetAvailableSlotQueueIndexAddress() { + return reinterpret_cast(EndOfAvailableSlots()); + } + + const AvailableSlotQueueIndex *GetAvailableSlotQueueIndexAddress() const { + return reinterpret_cast( + EndOfAvailableSlots()); + } + + InPlaceSlotQueue *GetAvailableSlotQueueAddress(int sub_id) { + uint64_t offset = GetAvailableSlotQueueIndexAddress() + ->offsets[sub_id] + .load(std::memory_order_acquire); + if (offset == kInvalidSlotQueueOffset) { + return nullptr; + } + return reinterpret_cast( + EndOfAvailableSlotQueueIndex() + offset); + } + + const InPlaceSlotQueue *GetAvailableSlotQueueAddress(int sub_id) const { + uint64_t offset = GetAvailableSlotQueueIndexAddress() + ->offsets[sub_id] + .load(std::memory_order_acquire); + if (offset == kInvalidSlotQueueOffset) { + return nullptr; + } + return reinterpret_cast( + EndOfAvailableSlotQueueIndex() + offset); + } + + virtual int SubscriberQueueSize(int sub_id) const { + const InPlaceSlotQueue *queue = GetAvailableSlotQueueAddress(sub_id); + return queue == nullptr ? 0 : static_cast(queue->Capacity()); + } + + void BeginSubscriberQueuePublish(int pub_id) { + GetAvailableSlotQueueIndexAddress() + ->active_publishers[pub_id] + .fetch_add(1, std::memory_order_seq_cst); + } + + void EndSubscriberQueuePublish(int pub_id) { + auto &counter = + GetAvailableSlotQueueIndexAddress()->active_publishers[pub_id]; + uint32_t active = counter.load(std::memory_order_seq_cst); + while (active != 0) { + if (counter.compare_exchange_strong(active, active - 1, + std::memory_order_seq_cst, + std::memory_order_seq_cst)) { + return; + } + } + } + bool IsActivated(int vchan_id) const { return ccb_->activation_tracker.IsActivated(vchan_id); } @@ -633,6 +1132,8 @@ class Channel : public std::enable_shared_from_this { std::string name_; int num_slots_; + int subscriber_queue_size_; + uint64_t subscriber_queue_arena_size_; int channel_id_; // ID allocated from server. std::string type_; diff --git a/common/common_test.cc b/common/common_test.cc index 05b55f95..5fa9e54c 100644 --- a/common/common_test.cc +++ b/common/common_test.cc @@ -1,8 +1,24 @@ #include "common/atomic_bitset.h" +#include "common/channel.h" #include "common/fast_ring_buffer.h" +#include +#include +#include + #include +TEST(CommonTest, SubscriberQueueArenaSizeIsExplicitBytes) { + constexpr int kNumSlots = 8; + constexpr uint64_t kArenaSize = 12'345; + EXPECT_EQ(kArenaSize, + subspace::CcbSize(kNumSlots, kArenaSize) - + subspace::CcbSize(kNumSlots, 0)); + EXPECT_EQ(100 * subspace::SlotQueueBlockSize( + subspace::kDefaultSubscriberQueueSize), + subspace::kDefaultSubscriberQueueArenaSize); +} + TEST(CommonTest, AtomicBitset) { subspace::AtomicBitSet<6144> bitset; bitset.Set(0); @@ -50,6 +66,100 @@ TEST(CommonTest, FastRingBuffer) { EXPECT_TRUE(buffer.Contains(4)); } +TEST(CommonTest, InPlaceSlotQueueEvictsOldestOnOverflow) { + constexpr size_t kCapacity = 2; + std::unique_ptr storage( + new std::byte[subspace::SizeofSlotQueue(kCapacity)]); + auto *queue = new (storage.get()) subspace::InPlaceSlotQueue(kCapacity); + + EXPECT_TRUE(queue->Push(1, 10)); + EXPECT_TRUE(queue->Push(2, 20)); + EXPECT_TRUE(queue->Push(3, 30)); + EXPECT_EQ(1, queue->OverflowCount()); + EXPECT_EQ(1, queue->OverflowCount()); + EXPECT_EQ(1, queue->ConsumeOverflow()); + EXPECT_EQ(0, queue->OverflowCount()); + + subspace::QueuedSlot slot; + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 2); + EXPECT_EQ(slot.ordinal, 20); + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 3); + EXPECT_EQ(slot.ordinal, 30); + EXPECT_FALSE(queue->TryPop(slot)); +} + +TEST(CommonTest, InPlaceSlotQueueCanRejectOverflowWithoutEviction) { + constexpr size_t kCapacity = 2; + std::unique_ptr storage( + new std::byte[subspace::SizeofSlotQueue(kCapacity)]); + auto *queue = + new (storage.get()) subspace::InPlaceSlotQueue(kCapacity, false); + + EXPECT_TRUE(queue->Push(1, 10)); + EXPECT_TRUE(queue->Push(2, 20)); + EXPECT_FALSE(queue->Push(3, 30)); + EXPECT_TRUE(queue->ConsumeInsertionFailure()); + EXPECT_EQ(0, queue->ConsumeOverflow()); + + subspace::QueuedSlot slot; + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 1); + EXPECT_EQ(slot.ordinal, 10); + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 2); + EXPECT_EQ(slot.ordinal, 20); + EXPECT_FALSE(queue->TryPop(slot)); +} + +TEST(CommonTest, InPlaceSlotQueueDoesNotWaitForUnreleasedEntry) { + constexpr size_t kCapacity = 2; + std::unique_ptr storage( + new std::byte[subspace::SizeofSlotQueue(kCapacity)]); + auto *queue = new (storage.get()) subspace::InPlaceSlotQueue(kCapacity); + auto *entries = reinterpret_cast( + storage.get() + sizeof(subspace::InPlaceSlotQueue)); + + ASSERT_TRUE(queue->Push(1, 10)); + ASSERT_TRUE(queue->Push(2, 20)); + subspace::QueuedSlot slot; + ASSERT_TRUE(queue->TryPop(slot)); + + // Model a consumer that advanced head but did not mark the entry reusable. + entries[0].sequence.store(1, std::memory_order_release); + EXPECT_FALSE(queue->Push(3, 30)); + EXPECT_TRUE(queue->ConsumeInsertionFailure()); + EXPECT_EQ(0, queue->ConsumeOverflow()); + + // Once the consumer releases the entry, producers can use it again. + entries[0].sequence.store(2, std::memory_order_release); + EXPECT_TRUE(queue->Push(3, 30)); + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 2); + ASSERT_TRUE(queue->TryPop(slot)); + EXPECT_EQ(slot.slot_id, 3); +} + +TEST(CommonTest, InPlaceSlotQueueDoesNotWaitAfterProducerReservationDeath) { + constexpr size_t kCapacity = 2; + std::unique_ptr storage( + new std::byte[subspace::SizeofSlotQueue(kCapacity)]); + auto *queue = new (storage.get()) subspace::InPlaceSlotQueue(kCapacity); + + // Model a producer that advanced tail from 0 to 1 and died before publishing + // entry 0's sequence. No operation may wait indefinitely behind the hole. + auto *tail = reinterpret_cast *>( + storage.get() + sizeof(size_t) + sizeof(std::atomic)); + tail->store(1, std::memory_order_release); + + EXPECT_TRUE(queue->Push(2, 20)); + subspace::QueuedSlot slot; + EXPECT_FALSE(queue->TryPop(slot)); + EXPECT_FALSE(queue->Push(3, 30)); + EXPECT_TRUE(queue->ConsumeInsertionFailure()); +} + TEST(CommonTest, BitsetTraverse1) { subspace::AtomicBitSet<10000> bitset; for (int i = 0; i < 10000; i++) { diff --git a/docs/client_design.md b/docs/client_design.md index 35c2590c..f3e1329b 100644 --- a/docs/client_design.md +++ b/docs/client_design.md @@ -158,7 +158,7 @@ The CCB contains: - **OrdinalAccumulator** — per-virtual-channel atomic ordinal counters. - **ActivationTracker** — bitset of activated virtual channels. - **Subscriber tracking** — bitset of active subscribers, per-subscriber vchan_id array, subscriber counter per vchan. -- **Statistics** — `total_bytes`, `total_messages`, `max_message_size`, `total_drops` (atomics). +- **Statistics** — `total_bytes`, `total_messages`, `max_message_size`, `total_drops` (atomics). `total_messages` includes activations and also versions subscriber snapshots. - **free_slots_exhausted** — atomic bool, optimization to skip scanning the free-slots bitset. Following the slot array (with 64-byte alignment): @@ -259,9 +259,15 @@ Message msg = subscriber.ReadMessage(ReadMode::kReadNext); 2. If reliable publisher triggers need refreshing (detected via SCB counters), reload them. 3. Clear the subscriber's poll trigger. 4. **Slot selection:** - - `kReadNext`: Scans `AvailableSlots` for this subscriber, collects all slots with non-zero ordinal that are not publisher-owned and match the vchan_id filter. Sorts by timestamp. Returns the first slot whose ordinal has not been seen. - - `kReadNewest`: Same scan, but returns only the most recent slot. -5. **Claim the slot:** `AtomicIncRefCount(slot, +1)` increments the ref count via CAS. If the CAS fails (slot was recycled), retries from scratch. + - `kReadNext`: Unreliable subscribers normally pop their per-subscriber + queue, carrying the queued `(slot_id, ordinal, vchan_id)` generation + through the ref-count CAS. The `AvailableSlots` bitset remains + authoritative and is scanned in ordinal order after queue overflow or + insertion failure. + - `kReadNewest`: Selects the most recent slot from an authoritative bitset + snapshot. Snapshot entries are temporarily pinned before their bits are + cleared, so a concurrently recycled generation is not erased. +5. **Claim the slot:** `AtomicIncRefCount(slot, +1)` increments the ref count via CAS using the frozen ordinal and vchan. If the CAS fails (slot was recycled), retries from scratch. 6. **Dropped message detection:** Compares the new message's ordinal against the ordinal tracker. Gaps indicate dropped messages; the dropped-message callback is invoked with the count. 7. **Checksum verification:** If the prefix has the `kMessageHasChecksum` flag: - Computes the checksum over the same three data regions used by the publisher. @@ -485,7 +491,7 @@ Publishers and subscribers can specify a `type` string. The server enforces: ### 13.1 Subscriber Polling -Each subscriber has a trigger file descriptor (pipe or eventfd). Publishers write to this fd when a new message is published. Subscribers use `GetPollFd()` to get a `struct pollfd` for use with `poll()` or `epoll()`, or call `Wait()` to block until a message is available. +Each subscriber has a trigger file descriptor (pipe or eventfd). Publishers write to this fd when a new message is published. Subscribers use `GetPollFd()` to get a `struct pollfd` for use with `poll()` or `epoll()`, or call `Wait()` to block until a message is available. A poll-driven drain uses a bounded queue-tail/bitset snapshot; `total_messages`, which includes activations, re-arms the next poll burst when a publication arrives after that snapshot. **Important:** After `Wait()` returns, the subscriber should read **all** available messages before waiting again. The trigger fd may not be re-armed until all messages are consumed. @@ -751,7 +757,7 @@ monitoring tools to distinguish between local, bridged, and tunneled users. ### Channel Statistics (from CCB) - `total_bytes`: Total bytes published. -- `total_messages`: Total messages published. +- `total_messages`: Total publications, including activations. - `max_message_size`: Largest message seen. - `total_drops`: Total messages dropped by unreliable publishers. diff --git a/docs/server-architecture.md b/docs/server-architecture.md index 75824db6..44e873cb 100644 --- a/docs/server-architecture.md +++ b/docs/server-architecture.md @@ -74,8 +74,25 @@ Each channel requires three shared memory regions, created via `shm_open()` (POS - One per channel. - Contains: channel name, num_slots, ordinals, activation tracker. -- Variable-length: `MessageSlot` array + bitsets for retired/free/available slots. -- Size: `CcbSize(num_slots)` = base + slots + bitsets. +- CCB version 4 uses atomic slot metadata. `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 + subscriber queue index, and a packed subscriber queue arena. +- Size: `CcbSize(num_slots, subscriber_queue_arena_size)`. Publisher client + APIs explicitly configure the packed arena in bytes and default to zero, + omitting subscriber queues and using the available-slot bitset path. Opting + into the standard 64,000-byte arena supports 100 default-sized queues. A + subscriber that does not request an override then gets the fixed 16-entry + default. Subscriber IDs still support the full 1024 owner limit, but queue + allocation fails once the packed arena is full. +- Per-subscriber queues are acceleration hints. The available-slot bitset is + authoritative, and consumers fall back to an ordinal-ordered bitset snapshot + if queue overflow or insertion failure races a claim. +- Queue blocks are retired before reuse while publisher traversal hazards are + active. Shadow recovery reconciles subscriber offsets with allocated blocks, + conservatively retires orphan blocks, and only reclaims them after their + recorded publisher hazards have quiesced. ### Buffer Control Block (BCB) diff --git a/proto/subspace.proto b/proto/subspace.proto index e87d3571..7392239a 100644 --- a/proto/subspace.proto +++ b/proto/subspace.proto @@ -42,6 +42,12 @@ message CreatePublisherRequest { int32 max_publishers = 17; // 0 means no explicit publisher limit. bool split_buffers_over_bridge = 18; // Remote bridge publisher uses split buffers. uint64 process_id = 19; // Client process id for introspection. + // Bytes reserved for packed per-subscriber queues in the CCB. + uint64 subscriber_queue_arena_size = 20; + // Local number of subscriber-queue traversals in progress when reclaiming + // after server failover. Reclaim runs under the client lock, so a zero value + // proves that a stale shared-memory hazard counter can be cleared. + uint32 active_queue_publish_depth = 21; } message CreatePublisherResponse { @@ -58,6 +64,8 @@ message CreatePublisherResponse { int32 vchan_id = 11; int32 retirement_fd_index = 12; // My retirement fd index (read end) repeated int32 retirement_fd_indexes = 13; // Write end of all retirement fds. + int32 subscriber_queue_size = 14; // Resolved capacity; 0 means disabled. + uint64 subscriber_queue_arena_size = 15; } // This is used both to create a new subscriber and to reload @@ -73,7 +81,9 @@ message CreateSubscriberRequest { bool for_tunnel = 7; string mux = 8; int32 vchan_id = 9; + // Requested queue capacity. 0 uses the publisher's channel default. uint64 process_id = 10; // Client process id for introspection. + int32 subscriber_queue_size = 11; } message CreateSubscriberResponse { @@ -94,6 +104,10 @@ message CreateSubscriberResponse { int32 checksum_size = 15; // Bytes reserved for checksum (from publisher). int32 metadata_size = 16; // Bytes reserved for user metadata (from publisher). bool use_split_buffers = 17; + int32 subscriber_queue_size = 18; // This subscriber's resolved capacity. + // Fixed publisher default used when subscriber_queue_size is zero. + int32 default_subscriber_queue_size = 19; + uint64 subscriber_queue_arena_size = 20; } message GetTriggersRequest { string channel_name = 1; } @@ -255,6 +269,8 @@ message ChannelInfoProto { string mux = 14; int32 channel_id = 15; repeated ChannelParticipantInfoProto participants = 16; + int32 subscriber_queue_size = 17; + uint64 subscriber_queue_arena_size = 18; } // This is published to the /subspace/ChannelDirectory channel. @@ -315,6 +331,7 @@ message Subscribed { int32 metadata_size = 8; // Bytes reserved for user metadata. bool split_buffers = 9; // Bridge messages are sent as prefix and payload chunks. bool split_buffers_over_bridge = 10; // Receiving bridge publisher uses split buffers. + uint64 subscriber_queue_arena_size = 11; } // This is sent over a TCP connection from the peer server when the @@ -488,6 +505,7 @@ message ShadowCreateChannel { bool has_max_publishers = 15; int32 max_publishers = 16; bool split_buffers_over_bridge = 17; + uint64 subscriber_queue_arena_size = 18; // FDs sent via SCM_RIGHTS: [ccb_fd, bcb_fd] } @@ -505,6 +523,7 @@ message ShadowAddPublisher { bool is_fixed_size = 6; bool notify_retirement = 7; bool for_tunnel = 8; + uint64 process_id = 9; // FDs sent via SCM_RIGHTS: [poll_fd, trigger_fd] // If notify_retirement: also [retirement_read_fd, retirement_write_fd] } @@ -521,6 +540,9 @@ message ShadowAddSubscriber { bool is_bridge = 4; int32 max_active_messages = 5; bool for_tunnel = 6; + // Requested capacity. 0 uses the publisher's channel default. + int32 subscriber_queue_size = 7; + uint64 process_id = 8; // FDs sent via SCM_RIGHTS: [trigger_fd, poll_fd] } diff --git a/rust_client/src/bitset.rs b/rust_client/src/bitset.rs index 4fe65889..098d6eaa 100644 --- a/rust_client/src/bitset.rs +++ b/rust_client/src/bitset.rs @@ -82,6 +82,24 @@ impl AtomicBitSet { } } } + + pub fn traverse_seq_cst(&self, mut func: F) { + let num_bits = self.num_bits; + for i in 0..WORDS { + let mut shift = 0usize; + let mut bit = i * 64; + while bit < num_bits && shift < 64 { + let word = self.bits[i].load(Ordering::SeqCst) >> shift; + let n = ffs64(word); + if n == 0 { + break; + } + bit += n; + func(bit - 1); + shift += n; + } + } + } } /// In-place atomic bitset accessor for shared memory. diff --git a/rust_client/src/channel.rs b/rust_client/src/channel.rs index 2cdd3ea0..ee09acc5 100644 --- a/rust_client/src/channel.rs +++ b/rust_client/src/channel.rs @@ -14,7 +14,9 @@ use std::num::NonZeroUsize; use std::os::fd::BorrowedFd; use std::os::unix::io::RawFd; use std::ptr::NonNull; -use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, AtomicU64, Ordering}; +use std::sync::atomic::{ + AtomicBool, AtomicI16, AtomicI32, AtomicU32, AtomicU64, Ordering, +}; // ── Flag constants ────────────────────────────────────────────────────────── @@ -24,9 +26,13 @@ pub const MESSAGE_HAS_CHECKSUM: i64 = 4; pub const MESSAGE_SEEN: u32 = 1; pub const MESSAGE_IS_ACTIVATION: u32 = 2; +pub const MESSAGE_SEEN_BY_RELIABLE: u32 = 4; 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 MAX_VCHAN_ID: usize = 1023; pub const MAX_CHANNEL_NAME: usize = 64; pub const MAX_BUFFERS: usize = 1024; @@ -119,15 +125,96 @@ const SLOT_OWNER_WORDS: usize = bits_to_words(MAX_SLOT_OWNERS); #[repr(C)] pub struct MessageSlot { pub refs: AtomicU64, - pub ordinal: u64, - pub message_size: u64, + pub ordinal: AtomicU64, + pub message_size: AtomicU64, pub id: i32, - pub buffer_index: i16, - pub vchan_id: i16, + pub buffer_index: AtomicI16, + pub vchan_id: AtomicI16, pub sub_owners: AtomicBitSet, - pub timestamp: u64, - pub flags: u32, - pub bridged_slot_id: i32, + pub timestamp: AtomicU64, + pub flags: AtomicU32, + pub bridged_slot_id: AtomicI32, +} +const _: () = assert!(std::mem::size_of::() == 184); +const _: () = assert!(std::mem::offset_of!(MessageSlot, refs) == 0); +const _: () = assert!(std::mem::offset_of!(MessageSlot, ordinal) == 8); +const _: () = assert!(std::mem::offset_of!(MessageSlot, message_size) == 16); +const _: () = assert!(std::mem::offset_of!(MessageSlot, id) == 24); +const _: () = assert!(std::mem::offset_of!(MessageSlot, buffer_index) == 28); +const _: () = assert!(std::mem::offset_of!(MessageSlot, vchan_id) == 30); +const _: () = assert!(std::mem::offset_of!(MessageSlot, sub_owners) == 32); +const _: () = assert!( + std::mem::offset_of!(MessageSlot, timestamp) + == std::mem::offset_of!(MessageSlot, sub_owners) + + std::mem::size_of::>() +); +const _: () = assert!(std::mem::offset_of!(MessageSlot, flags) == 176); +const _: () = assert!(std::mem::offset_of!(MessageSlot, bridged_slot_id) == 180); + +impl MessageSlot { + pub fn ordinal(&self) -> u64 { + self.ordinal.load(Ordering::Relaxed) + } + + pub fn set_ordinal(&self, v: u64) { + self.ordinal.store(v, Ordering::Relaxed); + } + + pub fn message_size(&self) -> u64 { + self.message_size.load(Ordering::Relaxed) + } + + pub fn set_message_size(&self, v: u64) { + self.message_size.store(v, Ordering::Relaxed); + } + + pub fn buffer_index(&self) -> i16 { + self.buffer_index.load(Ordering::Relaxed) + } + + pub fn set_buffer_index(&self, v: i16) { + self.buffer_index.store(v, Ordering::Relaxed); + } + + pub fn vchan_id(&self) -> i16 { + self.vchan_id.load(Ordering::Relaxed) + } + + pub fn set_vchan_id(&self, v: i16) { + self.vchan_id.store(v, Ordering::Relaxed); + } + + pub fn timestamp(&self) -> u64 { + self.timestamp.load(Ordering::Relaxed) + } + + pub fn set_timestamp(&self, v: u64) { + self.timestamp.store(v, Ordering::Relaxed); + } + + pub fn flags(&self) -> u32 { + self.flags.load(Ordering::Relaxed) + } + + pub fn set_flags(&self, v: u32) { + self.flags.store(v, Ordering::Relaxed); + } + + pub fn set_flag(&self, flag: u32) { + self.flags.fetch_or(flag, Ordering::Relaxed); + } + + pub fn clear_flags(&self, mask: u32) { + self.flags.fetch_and(!mask, Ordering::Relaxed); + } + + pub fn bridged_slot_id(&self) -> i32 { + self.bridged_slot_id.load(Ordering::Relaxed) + } + + pub fn set_bridged_slot_id(&self, v: i32) { + self.bridged_slot_id.store(v, Ordering::Relaxed); + } } #[derive(Clone)] @@ -138,6 +225,213 @@ pub struct ActiveSlot { pub vchan_id: i32, } +#[repr(C)] +pub struct SlotQueueEntry { + sequence: AtomicU64, + ordinal: AtomicU64, + slot_id: AtomicI32, +} +const _: () = assert!(std::mem::size_of::() == 24); +const _: () = assert!(std::mem::offset_of!(SlotQueueEntry, sequence) == 0); +const _: () = assert!(std::mem::offset_of!(SlotQueueEntry, ordinal) == 8); +const _: () = assert!(std::mem::offset_of!(SlotQueueEntry, slot_id) == 16); + +#[repr(C)] +pub struct SlotQueueHeader { + capacity: usize, + head: AtomicU64, + tail: AtomicU64, + overflow_count: AtomicU32, + insertion_failed: AtomicBool, + drop_oldest: bool, +} +const _: () = assert!(std::mem::size_of::() == 32); + +pub fn sizeof_slot_queue(capacity: usize) -> usize { + std::mem::size_of::() + + std::mem::size_of::() * capacity +} + +impl SlotQueueHeader { + fn entries(&self) -> *mut SlotQueueEntry { + unsafe { (self as *const Self as *mut u8).add(std::mem::size_of::()) as *mut SlotQueueEntry } + } + + pub fn head(&self) -> u64 { + self.head.load(Ordering::Acquire) + } + + pub fn tail(&self) -> u64 { + self.tail.load(Ordering::Acquire) + } + + fn drop_front(&self) -> bool { + if self.capacity == 0 { + return false; + } + let mut head = self.head.load(Ordering::Relaxed); + for _ in 0..MAX_SLOT_QUEUE_CAS_ATTEMPTS { + let entry = unsafe { &*self.entries().add((head % self.capacity as u64) as usize) }; + if entry.sequence.load(Ordering::Acquire) != head + 1 { + return false; + } + match self.head.compare_exchange_weak( + head, + head + 1, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + entry + .sequence + .store(head + self.capacity as u64, Ordering::Release); + return true; + } + Err(v) => head = v, + } + } + false + } + + pub fn discard_all(&self) { + for _ in 0..self.capacity { + if !self.drop_front() { + return; + } + } + } + + pub fn push( + &self, + slot_id: i32, + ordinal: u64, + report_insertion_failure: bool, + ) -> bool { + if self.capacity == 0 { + if report_insertion_failure { + self.mark_insertion_failure(); + } + return false; + } + + let mut tail = self.tail.load(Ordering::Relaxed); + let mut reserved_entry = None; + for _ in 0..MAX_SLOT_QUEUE_CAS_ATTEMPTS { + let head = self.head.load(Ordering::Acquire); + if tail - head >= self.capacity as u64 { + if !self.drop_oldest { + if report_insertion_failure { + self.mark_insertion_failure(); + } + return false; + } + if !self.drop_front() { + if report_insertion_failure { + self.mark_insertion_failure(); + } + return false; + } + self.overflow_count.fetch_add(1, Ordering::Release); + tail = self.tail.load(Ordering::Relaxed); + continue; + } + let candidate = + unsafe { &*self.entries().add((tail % self.capacity as u64) as usize) }; + // The consumer may stop after advancing head but before publishing + // the reusable sequence. Reserve only entries that are already + // reusable so a dead consumer cannot make this producer wait. + if candidate.sequence.load(Ordering::Acquire) != tail { + if report_insertion_failure { + self.mark_insertion_failure(); + } + return false; + } + match self.tail.compare_exchange( + tail, + tail + 1, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + reserved_entry = Some(candidate); + break; + } + Err(v) => tail = v, + } + } + let Some(entry) = reserved_entry else { + if report_insertion_failure { + self.mark_insertion_failure(); + } + return false; + }; + + entry.slot_id.store(slot_id, Ordering::Relaxed); + entry.ordinal.store(ordinal, Ordering::Relaxed); + entry.sequence.store(tail + 1, Ordering::Release); + true + } + + pub fn mark_insertion_failure(&self) { + self.insertion_failed.store(true, Ordering::Release); + } + + pub fn try_pop(&self) -> Option<(i32, u64)> { + if self.capacity == 0 { + return None; + } + + let mut head = self.head.load(Ordering::Relaxed); + for _ in 0..MAX_SLOT_QUEUE_CAS_ATTEMPTS { + let entry = unsafe { &*self.entries().add((head % self.capacity as u64) as usize) }; + if entry.sequence.load(Ordering::Acquire) != head + 1 { + return None; + } + let candidate = ( + entry.slot_id.load(Ordering::Relaxed), + entry.ordinal.load(Ordering::Relaxed), + ); + match self.head.compare_exchange_weak( + head, + head + 1, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => { + entry + .sequence + .store(head + self.capacity as u64, Ordering::Release); + return Some(candidate); + } + Err(v) => head = v, + } + } + None + } + + pub fn consume_overflow(&self) -> u32 { + self.overflow_count.swap(0, Ordering::AcqRel) + } + + pub fn overflow_count(&self) -> u32 { + self.overflow_count.load(Ordering::Acquire) + } + + pub fn consume_insertion_failure(&self) -> bool { + self.insertion_failed.swap(false, Ordering::AcqRel) + } + + pub fn insertion_failed(&self) -> bool { + self.insertion_failed.load(Ordering::Acquire) + } +} + +pub fn available_slot_queue_capacity(num_slots: usize) -> usize { + resolve_subscriber_queue_size(num_slots as i32, 0) as usize +} + +pub const INVALID_SLOT_QUEUE_OFFSET: u64 = u64::MAX; + // ── ChannelCounters ───────────────────────────────────────────────────────── #[repr(C)] @@ -234,6 +528,8 @@ impl SubscriberCounter { pub struct ChannelControlBlock { pub channel_name: [u8; MAX_CHANNEL_NAME], pub num_slots: i32, + pub subscriber_queue_size: i32, + pub version: u32, pub ordinals: OrdinalAccumulator, pub activation_tracker: ActivationTracker, pub buffer_index: i32, @@ -252,8 +548,42 @@ pub struct ChannelControlBlock { // Followed by: slots[num_slots], then trailing bitsets. // Accessed via unsafe pointer arithmetic. } +const _: () = assert!(std::mem::offset_of!(ChannelControlBlock, version) == 72); + +#[repr(C)] +pub struct AvailableSlotQueueIndex { + pub next_offset: AtomicU64, + pub offsets: [AtomicU64; MAX_SLOT_OWNERS], + pub active_publishers: [AtomicU32; MAX_SLOT_OWNERS], +} +const _: () = assert!(std::mem::size_of::() == 12296); +const _: () = assert!(std::mem::offset_of!(AvailableSlotQueueIndex, offsets) == 8); +const _: () = + assert!(std::mem::offset_of!(AvailableSlotQueueIndex, active_publishers) == 8200); -pub fn ccb_size(num_slots: i32) -> usize { +fn available_slot_queue_index_size() -> usize { + aligned64(std::mem::size_of::() as i64) as usize +} + +#[repr(C, align(64))] +pub struct SlotQueueBlockHeader { + pub block_size: u64, + pub state: AtomicU32, + pub reserved: u32, + pub waiting_publishers: AtomicBitSet, +} +const _: () = assert!(std::mem::size_of::() == 192); +const _: () = assert!(std::mem::offset_of!(SlotQueueBlockHeader, waiting_publishers) == 16); + +pub fn resolve_subscriber_queue_size(num_slots: i32, subscriber_queue_size: i32) -> i32 { + if num_slots <= 0 || subscriber_queue_size <= 0 { + 0 + } else { + subscriber_queue_size + } +} + +pub fn ccb_size(num_slots: i32, subscriber_queue_arena_size: u64) -> usize { let ns = num_slots as usize; let base = aligned64( (std::mem::size_of::() + ns * std::mem::size_of::()) @@ -261,6 +591,8 @@ pub fn ccb_size(num_slots: i32) -> usize { ) as usize; base + aligned64(sizeof_atomic_bitset(ns) as i64) as usize * 2 + sizeof_atomic_bitset(ns) * MAX_SLOT_OWNERS + + available_slot_queue_index_size() + + subscriber_queue_arena_size as usize } // ── Channel: shared memory accessor ───────────────────────────────────────── @@ -269,6 +601,8 @@ pub fn ccb_size(num_slots: i32) -> usize { pub struct Channel { pub name: String, pub num_slots: i32, + pub subscriber_queue_size: i32, + pub subscriber_queue_arena_size: u64, pub channel_id: i32, pub channel_type: String, pub vchan_id: i32, @@ -444,6 +778,8 @@ impl Channel { pub fn new( name: String, num_slots: i32, + subscriber_queue_size: i32, + subscriber_queue_arena_size: u64, channel_id: i32, channel_type: String, vchan_id: i32, @@ -453,6 +789,8 @@ impl Channel { Self { name, num_slots, + subscriber_queue_size: resolve_subscriber_queue_size(num_slots, subscriber_queue_size), + subscriber_queue_arena_size, channel_id, channel_type, vchan_id, @@ -485,12 +823,24 @@ impl Channel { prot: ProtFlags, ) -> crate::error::Result<()> { let scb_sz = std::mem::size_of::(); - let ccb_sz = ccb_size(self.num_slots); + let ccb_sz = ccb_size(self.num_slots, self.subscriber_queue_arena_size); let bcb_sz = std::mem::size_of::(); self.scb = map_memory(scb_fd, scb_sz, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? as *mut SystemControlBlock; self.ccb = map_memory(ccb_fd, ccb_sz, prot)? as *mut ChannelControlBlock; + if self.ccb().version != CHANNEL_CONTROL_BLOCK_VERSION { + let found = self.ccb().version; + unsafe { + let _ = shim_munmap(NonNull::new_unchecked(self.ccb as *mut _), ccb_sz); + let _ = shim_munmap(NonNull::new_unchecked(self.scb as *mut _), scb_sz); + } + self.ccb = std::ptr::null_mut(); + self.scb = std::ptr::null_mut(); + return Err(crate::error::SubspaceError::Internal(format!( + "unsupported channel control block version {found} (expected {CHANNEL_CONTROL_BLOCK_VERSION})" + ))); + } self.bcb = map_memory(bcb_fd, bcb_sz, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? as *mut BufferControlBlock; self.scb_size = scb_sz; @@ -591,6 +941,62 @@ impl Channel { } } + fn end_of_available_slots(&self) -> *mut u8 { + unsafe { + self.end_of_free_slots().add( + sizeof_atomic_bitset(self.num_slots as usize) * MAX_SLOT_OWNERS, + ) + } + } + + fn available_slot_queue_index(&self) -> &AvailableSlotQueueIndex { + unsafe { + &*(self.end_of_available_slots() as *const AvailableSlotQueueIndex) + } + } + + fn end_of_available_slot_queue_index(&self) -> *mut u8 { + unsafe { + self.end_of_available_slots() + .add(available_slot_queue_index_size()) + } + } + + pub fn get_available_slot_queue(&self, sub_id: usize) -> Option<&SlotQueueHeader> { + let offset = self.available_slot_queue_index().offsets[sub_id].load(Ordering::Acquire); + if offset == INVALID_SLOT_QUEUE_OFFSET { + return None; + } + unsafe { + Some( + &*(self + .end_of_available_slot_queue_index() + .add(offset as usize) as *const SlotQueueHeader), + ) + } + } + + pub fn begin_subscriber_queue_publish(&self, pub_id: usize) { + self.available_slot_queue_index().active_publishers[pub_id] + .fetch_add(1, Ordering::SeqCst); + } + + pub fn end_subscriber_queue_publish(&self, pub_id: usize) { + let counter = &self.available_slot_queue_index().active_publishers[pub_id]; + let mut active = counter.load(Ordering::SeqCst); + while active != 0 { + match counter.compare_exchange( + active, + active - 1, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => return, + Err(value) => active = value, + } + } + } + pub fn num_subscribers(&self, vchan_id: i32) -> i32 { self.ccb().num_subs.num_subscribers(vchan_id) } @@ -696,7 +1102,7 @@ impl Channel { /// 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); - let buf_idx = slot.buffer_index; + let buf_idx = slot.buffer_index(); if buf_idx < 0 || buf_idx as usize >= self.buffers.len() { return std::ptr::null_mut(); } @@ -715,7 +1121,7 @@ impl Channel { pub fn get_prefix(&self, slot_idx: usize) -> *mut MessagePrefix { let slot = self.slot_ref(slot_idx); - let buf_idx = slot.buffer_index; + let buf_idx = slot.buffer_index(); if buf_idx < 0 || buf_idx as usize >= self.buffers.len() { return std::ptr::null_mut(); } @@ -744,7 +1150,7 @@ impl Channel { pub fn slot_size_for_slot(&self, slot_idx: usize) -> u64 { let slot = self.slot_ref(slot_idx); - let buf_idx = slot.buffer_index; + let buf_idx = slot.buffer_index(); if buf_idx < 0 || buf_idx as usize >= self.buffers.len() { return 0; } @@ -779,7 +1185,7 @@ impl Channel { pub fn validate_slot_buffer(&self, slot_idx: usize) -> bool { let slot = self.slot_ref(slot_idx); - let buf_idx = slot.buffer_index; + let buf_idx = slot.buffer_index(); if buf_idx < 0 { return true; } @@ -796,12 +1202,13 @@ impl Channel { } pub fn set_slot_to_biggest_buffer(&mut self, slot_idx: usize) { - let slot = self.slot_mut(slot_idx); - if slot.buffer_index != -1 { - self.decrement_buffer_refs(slot.buffer_index as usize); + let slot = self.slot_ref(slot_idx); + if slot.buffer_index() != -1 { + self.decrement_buffer_refs(slot.buffer_index() as usize); } - slot.buffer_index = (self.buffers.len() - 1) as i16; - self.increment_buffer_refs(slot.buffer_index as usize); + let new_index = (self.buffers.len() - 1) as i16; + slot.set_buffer_index(new_index); + self.increment_buffer_refs(new_index as usize); } pub fn decrement_buffer_refs(&self, buffer_index: usize) { @@ -822,8 +1229,8 @@ impl Channel { let slot = self.slot_ref(i); let refs = slot.refs.load(Ordering::Relaxed); if refs == (PUB_OWNED | owner as u64) { - self.slot_mut(i).ordinal = 0; - slot.refs.store(0, Ordering::SeqCst); + self.slot_ref(i).set_ordinal(0); + slot.refs.store(0, Ordering::Release); let ccb = self.ccb(); ccb.subscribers.traverse(|sub_id| { diff --git a/rust_client/src/client.rs b/rust_client/src/client.rs index 329d9616..ed9690ee 100644 --- a/rust_client/src/client.rs +++ b/rust_client/src/client.rs @@ -71,6 +71,8 @@ pub struct ChannelInfo { pub channel_type: String, pub slot_size: u64, pub num_slots: i32, + pub subscriber_queue_size: i32, + pub subscriber_queue_arena_size: u64, pub reliable: bool, } @@ -129,6 +131,18 @@ impl Publisher { self.imp.lock().unwrap().channel.num_slots } + pub fn subscriber_queue_size(&self) -> i32 { + self.imp.lock().unwrap().channel.subscriber_queue_size + } + + pub fn subscriber_queue_arena_size(&self) -> u64 { + self.imp + .lock() + .unwrap() + .channel + .subscriber_queue_arena_size + } + /// Get a mutable pointer to the message buffer for writing. /// Returns None if no slot is available (reliable publisher). /// @@ -234,7 +248,9 @@ impl Publisher { } let slot_idx = pub_impl.channel.slot.unwrap(); - pub_impl.channel.slot_mut(slot_idx).message_size = message_size as u64; + pub_impl.channel + .slot_ref(slot_idx) + .set_message_size(message_size as u64); let owner = pub_impl.publisher_id; let reliable = pub_impl.options.reliable; @@ -487,10 +503,14 @@ impl Subscriber { self.imp.lock().unwrap().channel.num_slots } + pub fn subscriber_queue_size(&self) -> i32 { + self.imp.lock().unwrap().subscriber_queue_size + } + pub fn current_ordinal(&self) -> i64 { let sub = self.imp.lock().unwrap(); match sub.channel.slot { - Some(si) => sub.channel.slot_ref(si).ordinal as i64, + Some(si) => sub.channel.slot_ref(si).ordinal() as i64, None => -1, } } @@ -498,7 +518,7 @@ impl Subscriber { pub fn timestamp(&self) -> u64 { let sub = self.imp.lock().unwrap(); match sub.channel.slot { - Some(si) => sub.channel.slot_ref(si).timestamp, + Some(si) => sub.channel.slot_ref(si).timestamp(), None => 0, } } @@ -540,7 +560,12 @@ impl Subscriber { } pub fn get_poll_fd(&self) -> RawFd { - self.imp.lock().unwrap().poll_fd + let mut imp = self.imp.lock().unwrap(); + imp.poll_drain_pending = true; + imp.poll_drain_exhausted = false; + imp.queue_drain_tail = None; + imp.poll_snapshot_valid = false; + imp.poll_fd } pub fn trigger(&self) { @@ -702,15 +727,15 @@ impl Subscriber { sub_impl.channel.slot = Some(slot_idx); let slot = sub_impl.channel.slot_ref(slot_idx); - if slot.message_size == 0 { + if slot.message_size() == 0 { return Ok(Message::default()); } let buffer = sub_impl.channel.get_buffer_address(slot_idx); - let msg_size = slot.message_size as usize; - let ordinal = slot.ordinal; - let timestamp = slot.timestamp; - let vchan_id = slot.vchan_id as i32; + let msg_size = slot.message_size() as usize; + let ordinal = slot.ordinal(); + let timestamp = slot.timestamp(); + let vchan_id = slot.vchan_id() as i32; let slot_id = slot.id; Ok(Message { @@ -868,9 +893,11 @@ impl Client { metadata_size: opts.metadata_size, use_split_buffers: opts.use_split_buffers, split_buffers_over_bridge: opts.split_buffers_over_bridge, + subscriber_queue_arena_size: opts.subscriber_queue_arena_size, max_publishers: 0, publisher_id: -1, process_id: std::process::id() as u64, + active_queue_publish_depth: 0, }, )), }; @@ -893,6 +920,8 @@ impl Client { let mut pub_impl = PublisherImpl::new( channel_name.to_string(), opts.num_slots, + pub_resp.subscriber_queue_size, + pub_resp.subscriber_queue_arena_size, pub_resp.channel_id, pub_resp.publisher_id, pub_resp.vchan_id, @@ -1000,6 +1029,7 @@ impl Client { max_active_messages: opts.max_active_messages, mux: opts.mux.clone(), vchan_id: opts.vchan_id, + subscriber_queue_size: opts.subscriber_queue_size, process_id: std::process::id() as u64, }, )), @@ -1023,6 +1053,9 @@ impl Client { let mut sub_impl = SubscriberImpl::new( channel_name.to_string(), sub_resp.num_slots, + sub_resp.default_subscriber_queue_size, + sub_resp.subscriber_queue_arena_size, + sub_resp.subscriber_queue_size, sub_resp.channel_id, sub_resp.subscriber_id, sub_resp.vchan_id, @@ -1040,6 +1073,10 @@ impl Client { }; sub_impl.channel.num_slots = sub_resp.num_slots; + sub_impl.channel.subscriber_queue_size = sub_resp.default_subscriber_queue_size; + sub_impl.channel.subscriber_queue_arena_size = + sub_resp.subscriber_queue_arena_size; + sub_impl.subscriber_queue_size = sub_resp.subscriber_queue_size; sub_impl .channel .embargoed_slots @@ -1140,6 +1177,8 @@ impl Client { channel_type: String::from_utf8_lossy(&info.r#type).to_string(), slot_size: info.slot_size as u64, num_slots: info.num_slots, + subscriber_queue_size: info.subscriber_queue_size, + subscriber_queue_arena_size: info.subscriber_queue_arena_size, reliable: info.is_reliable, }) } @@ -1178,6 +1217,8 @@ impl Client { channel_type: String::from_utf8_lossy(&info.r#type).to_string(), slot_size: info.slot_size as u64, num_slots: info.num_slots, + subscriber_queue_size: info.subscriber_queue_size, + subscriber_queue_arena_size: info.subscriber_queue_arena_size, reliable: info.is_reliable, }) .collect()) @@ -1329,10 +1370,9 @@ fn read_message_internal( let old_slot = sub.channel.slot; let last_ordinal: i64 = match old_slot { - Some(si) => sub.channel.slot_ref(si).ordinal as i64, + Some(si) => sub.channel.slot_ref(si).ordinal() as i64, None => -1, }; - let new_slot_idx = match mode { ReadMode::ReadNext => sub.next_slot(), ReadMode::ReadNewest => sub.last_slot(), @@ -1348,29 +1388,11 @@ fn read_message_internal( sub.channel.slot = Some(new_idx); - if mode == ReadMode::ReadNext && last_ordinal != -1 { - let new_vchan_id = sub.channel.slot_ref(new_idx).vchan_id as i32; - let drops = sub.detect_drops(new_vchan_id); - if drops > 0 { - if let Some(ref cb) = sub.dropped_message_callback { - cb(drops as i64); - } - if sub.options.log_dropped_messages { - log::warn!( - "Dropped {} message{} on channel {}", - drops, - if drops == 1 { "" } else { "s" }, - sub.channel.name - ); - } - sub.channel - .ccb() - .total_drops - .fetch_add(drops as u32, Ordering::Relaxed); - } - } - let prefix = sub.channel.get_prefix(new_idx); + let slot = sub.channel.slot_ref(new_idx); + let frozen_ordinal = slot.ordinal(); + let frozen_vchan_id = slot.vchan_id() as i32; + let mut delivered_message_size = slot.message_size() as i64; let mut is_activation = false; let mut checksum_error = false; @@ -1379,13 +1401,12 @@ fn read_message_internal( let p = &*prefix; if p.has_checksum() && sub.options.checksum { let buffer = sub.channel.get_buffer_address(new_idx); - let slot = sub.channel.slot_ref(new_idx); let cs = sub.channel.checksum_size; let ms = sub.channel.metadata_size; let data = checksum::get_message_checksum_data( prefix, buffer, - slot.message_size as usize, + delivered_message_size as usize, cs, ms, ); @@ -1403,6 +1424,7 @@ fn read_message_internal( is_activation = true; if !pass_activation { sub.ignore_activation(new_idx); + sub.channel.slot = old_slot; if sub.options.reliable { sub.trigger_reliable_publishers(); } @@ -1414,31 +1436,63 @@ fn read_message_internal( if let Some(ref cb) = sub.on_receive_callback { let buffer = sub.channel.get_buffer_address(new_idx); - let slot = sub.channel.slot_ref(new_idx); - let new_size = cb(buffer as *mut u8, slot.message_size as i64)?; - sub.channel.slot_mut(new_idx).message_size = new_size as u64; + delivered_message_size = match cb(buffer as *mut u8, delivered_message_size) { + Ok(size) => size, + Err(e) => { + sub.release_unclaimed_slot(new_idx, frozen_ordinal, frozen_vchan_id); + sub.channel.slot = old_slot; + return Err(e); + } + }; } - let slot = sub.channel.slot_ref(new_idx); - if slot.message_size == 0 { + if delivered_message_size <= 0 { + sub.release_unclaimed_slot(new_idx, frozen_ordinal, frozen_vchan_id); + sub.channel.slot = old_slot; return Ok(Message::default()); } let buffer = sub.channel.get_buffer_address(new_idx); - let msg_size = slot.message_size as usize; - let ordinal = slot.ordinal; - let timestamp = slot.timestamp; - let vchan_id = slot.vchan_id as i32; + let slot = sub.channel.slot_ref(new_idx); + let msg_size = delivered_message_size as usize; + let ordinal = frozen_ordinal; + let timestamp = slot.timestamp(); + let vchan_id = frozen_vchan_id; let slot_id = slot.id; sub.clear_active_message(); if !sub.add_active_message() { - sub.unread_slot(new_idx); + sub.unread_slot(new_idx, frozen_ordinal, frozen_vchan_id); + sub.channel.slot = old_slot; return Ok(Message::default()); } - sub.claim_slot(new_idx, sub.channel.vchan_id, mode == ReadMode::ReadNewest); + if mode == ReadMode::ReadNext && sub.options.detect_dropped_messages { + let mut drops = std::mem::take(&mut sub.pending_queue_drops); + if last_ordinal != -1 { + drops = drops.max(sub.detect_drops(vchan_id)); + } + if drops > 0 { + if let Some(ref cb) = sub.dropped_message_callback { + cb(drops as i64); + } + if sub.options.log_dropped_messages { + log::warn!( + "Dropped {} message{} on channel {}", + drops, + if drops == 1 { "" } else { "s" }, + sub.channel.name + ); + } + sub.channel + .ccb() + .total_drops + .fetch_add(drops as u32, Ordering::Relaxed); + } + } + + sub.claim_slot(new_idx, vchan_id, mode == ReadMode::ReadNewest); if checksum_error && !sub.options.pass_checksum_errors { return Err(SubspaceError::ChecksumError); @@ -1472,14 +1526,13 @@ fn reload_subscriber(client: &mut ClientInner, sub: &mut SubscriberImpl) -> Resu if sub.channel.num_updates == updates { return Ok(()); } - sub.channel.num_updates = updates; - let req = proto::Request { request: Some(proto::request::Request::CreateSubscriber( proto::CreateSubscriberRequest { channel_name: sub.channel.name.clone(), subscriber_id: sub.subscriber_id, mux: sub.options.mux.clone(), + subscriber_queue_size: sub.options.subscriber_queue_size, process_id: std::process::id() as u64, ..Default::default() }, @@ -1495,11 +1548,22 @@ fn reload_subscriber(client: &mut ClientInner, sub: &mut SubscriberImpl) -> Resu return Err(SubspaceError::ServerError(sub_resp.error)); } - sub.channel.unmap(); + // A subscriber-created placeholder is the only case where the server + // replaces the CCB. Established channels retain their CCB across + // publisher updates. + let remap_ccb = sub.channel.num_slots == 0; + if remap_ccb { + sub.reset_delivery_state(); + sub.channel.unmap(); + } if !sub_resp.r#type.is_empty() { sub.channel.channel_type = String::from_utf8_lossy(&sub_resp.r#type).to_string(); } sub.channel.num_slots = sub_resp.num_slots; + sub.channel.subscriber_queue_size = sub_resp.default_subscriber_queue_size; + sub.channel.subscriber_queue_arena_size = + sub_resp.subscriber_queue_arena_size; + sub.subscriber_queue_size = sub_resp.subscriber_queue_size; sub.channel .embargoed_slots .resize(sub_resp.num_slots as usize); @@ -1522,13 +1586,15 @@ fn reload_subscriber(client: &mut ClientInner, sub: &mut SubscriberImpl) -> Resu sub.checksum_tmp = vec![0u8; cs as usize]; } - let prot = ProtFlags::PROT_READ | ProtFlags::PROT_WRITE; - sub.channel.map( - client.scb_fd, - fds[sub_resp.ccb_fd_index as usize], - fds[sub_resp.bcb_fd_index as usize], - prot, - )?; + if remap_ccb { + let prot = ProtFlags::PROT_READ | ProtFlags::PROT_WRITE; + sub.channel.map( + client.scb_fd, + fds[sub_resp.ccb_fd_index as usize], + fds[sub_resp.bcb_fd_index as usize], + prot, + )?; + } sub.attach_buffers()?; @@ -1548,7 +1614,10 @@ fn reload_subscriber(client: &mut ClientInner, sub: &mut SubscriberImpl) -> Resu sub.retirement_trigger_fds.push(fds[idx as usize]); } - sub.init_active_messages(); + if remap_ccb { + sub.init_active_messages(); + } + sub.channel.num_updates = updates; Ok(()) } @@ -1652,7 +1721,7 @@ fn activate_reliable_channel(publisher: &mut PublisherImpl) -> Result<()> { publisher.channel.name ))); } - publisher.channel.slot_mut(si).message_size = 1; + publisher.channel.slot_ref(si).set_message_size(1); let owner = publisher.publisher_id; publisher.activate_slot_and_get_another(si, true, true, owner, false, false); @@ -1675,7 +1744,7 @@ fn activate_channel(publisher: &mut PublisherImpl) -> Result<()> { publisher.channel.name ))); } - publisher.channel.slot_mut(si).message_size = 1; + publisher.channel.slot_ref(si).set_message_size(1); let owner = publisher.publisher_id; let published = publisher.activate_slot_and_get_another(si, false, true, owner, false, false); @@ -1757,7 +1826,7 @@ fn expand_slot_size(slot_size: u64) -> u64 { fn get_virtual_memory_usage(channel: &Channel) -> u64 { let mut size = std::mem::size_of::() as u64 - + ccb_size(channel.num_slots) as u64 + + ccb_size(channel.num_slots, channel.subscriber_queue_arena_size) as u64 + std::mem::size_of::() as u64; if !channel.bcb.is_null() { let bcb = unsafe { &*channel.bcb }; diff --git a/rust_client/src/options.rs b/rust_client/src/options.rs index 6db0551b..48dbbd40 100644 --- a/rust_client/src/options.rs +++ b/rust_client/src/options.rs @@ -9,10 +9,16 @@ use crate::split_buffer::{ }; use std::sync::Arc; +/// Fixed queue depth inherited by subscribers when an arena is provisioned. +pub const DEFAULT_SUBSCRIBER_QUEUE_SIZE: i32 = 16; +/// Standard packed subscriber queue arena size for callers that opt in. +pub const DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE: u64 = 64_000; + #[derive(Debug, Clone)] pub struct PublisherOptions { pub slot_size: i32, pub num_slots: i32, + pub subscriber_queue_arena_size: u64, pub local: bool, pub reliable: bool, pub bridge: bool, @@ -36,6 +42,7 @@ impl Default for PublisherOptions { Self { slot_size: 0, num_slots: 0, + subscriber_queue_arena_size: 0, local: false, reliable: false, bridge: false, @@ -71,6 +78,13 @@ impl PublisherOptions { self } + /// Set the bytes reserved for packed per-subscriber queues in the CCB. + /// Subscriber queues are disabled by default; a non-zero size opts in. + pub fn set_subscriber_queue_arena_size(mut self, size: u64) -> Self { + self.subscriber_queue_arena_size = size; + self + } + pub fn set_local(mut self, v: bool) -> Self { self.local = v; self @@ -183,11 +197,13 @@ impl PublisherOptions { #[derive(Debug, Clone)] pub struct SubscriberOptions { pub reliable: bool, + pub subscriber_queue_size: i32, pub bridge: bool, pub for_tunnel: bool, pub channel_type: String, pub max_active_messages: i32, pub log_dropped_messages: bool, + pub detect_dropped_messages: bool, pub pass_activation: bool, pub read_write: bool, pub mux: String, @@ -202,11 +218,13 @@ impl Default for SubscriberOptions { fn default() -> Self { Self { reliable: false, + subscriber_queue_size: 0, bridge: false, for_tunnel: false, channel_type: String::new(), max_active_messages: 1, log_dropped_messages: true, + detect_dropped_messages: true, pass_activation: false, read_write: false, mux: String::new(), @@ -229,6 +247,11 @@ impl SubscriberOptions { self } + pub fn set_subscriber_queue_size(mut self, size: i32) -> Self { + self.subscriber_queue_size = size; + self + } + pub fn set_type(mut self, t: String) -> Self { self.channel_type = t; self @@ -249,6 +272,11 @@ impl SubscriberOptions { self } + pub fn set_detect_dropped_messages(mut self, v: bool) -> Self { + self.detect_dropped_messages = v; + self + } + pub fn set_bridge(mut self, v: bool) -> Self { self.bridge = v; self diff --git a/rust_client/src/publisher.rs b/rust_client/src/publisher.rs index f69b1d3a..2d401801 100644 --- a/rust_client/src/publisher.rs +++ b/rust_client/src/publisher.rs @@ -12,14 +12,14 @@ use crate::split_buffer::{ read_split_buffer_metadata_file, split_buffer_object_name, write_split_buffer_metadata_file, SplitBufferMetadata, }; -use crate::syscall_shim::{ - shim_close, shim_fstat, shim_ftruncate, shim_open, shim_read, shim_write, -}; +#[cfg(target_os = "linux")] +use crate::syscall_shim::shim_fstat; +use crate::syscall_shim::{shim_close, shim_ftruncate, shim_open, shim_read, shim_write}; use nix::fcntl::OFlag; use nix::sys::mman::ProtFlags; use nix::sys::stat::Mode; use std::os::unix::io::RawFd; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicU32, Ordering}; pub type OnSendCallback = Box Result + Send + Sync>; pub type ResizeCallback = Box Result<()> + Send + Sync>; @@ -30,10 +30,37 @@ pub struct PublishedMessage { pub timestamp: u64, } +struct SubscriberQueuePublishGuard<'a> { + channel: &'a Channel, + publisher_id: usize, + local_depth: &'a AtomicU32, +} + +impl<'a> SubscriberQueuePublishGuard<'a> { + fn new(channel: &'a Channel, publisher_id: usize, local_depth: &'a AtomicU32) -> Self { + local_depth.fetch_add(1, Ordering::SeqCst); + channel.begin_subscriber_queue_publish(publisher_id); + Self { + channel, + publisher_id, + local_depth, + } + } +} + +impl Drop for SubscriberQueuePublishGuard<'_> { + fn drop(&mut self) { + self.channel + .end_subscriber_queue_publish(self.publisher_id); + self.local_depth.fetch_sub(1, Ordering::SeqCst); + } +} + pub struct PublisherImpl { pub channel: Channel, pub publisher_id: i32, pub options: PublisherOptions, + pub active_queue_publish_depth: AtomicU32, pub subscriber_trigger_fds: Vec, pub poll_fd: RawFd, @@ -50,6 +77,8 @@ impl PublisherImpl { pub fn new( name: String, num_slots: i32, + subscriber_queue_size: i32, + subscriber_queue_arena_size: u64, channel_id: i32, publisher_id: i32, vchan_id: i32, @@ -61,6 +90,8 @@ impl PublisherImpl { channel: Channel::new( name, num_slots, + subscriber_queue_size, + subscriber_queue_arena_size, channel_id, channel_type, vchan_id, @@ -68,6 +99,7 @@ impl PublisherImpl { ), publisher_id, options, + active_queue_publish_depth: AtomicU32::new(0), subscriber_trigger_fds: Vec::new(), poll_fd: -1, trigger_fd: -1, @@ -180,9 +212,9 @@ impl PublisherImpl { if (refs & PUB_OWNED) != 0 { continue; } - if (refs & REFS_MASK) == 0 && s.timestamp < earliest_timestamp { + if (refs & REFS_MASK) == 0 && s.timestamp() < earliest_timestamp { slot_idx = Some(i); - earliest_timestamp = s.timestamp; + earliest_timestamp = s.timestamp(); } } } @@ -201,7 +233,7 @@ impl PublisherImpl { let old_refs = (*slot_ptr).refs.load(Ordering::Relaxed); let ref_val = PUB_OWNED | owner as u64; let expected = build_refs_bit_field( - (*slot_ptr).ordinal, + (*slot_ptr).ordinal(), ((old_refs >> VCHAN_ID_SHIFT) & VCHAN_ID_MASK) as i32, ((old_refs >> RETIRED_REFS_SHIFT) & RETIRED_REFS_MASK) as i32, ); @@ -228,10 +260,10 @@ impl PublisherImpl { } let si = slot_idx.unwrap(); - let slot = self.channel.slot_mut(si); - slot.ordinal = 0; - slot.timestamp = 0; - slot.vchan_id = self.channel.vchan_id as i16; + let slot = self.channel.slot_ref(si); + slot.set_ordinal(0); + slot.set_timestamp(0); + slot.set_vchan_id(self.channel.vchan_id as i16); self.channel.set_slot_to_biggest_buffer(si); let prefix = self.channel.get_prefix(si); @@ -291,9 +323,9 @@ impl PublisherImpl { let s = self.channel.slot_ref(fs); self.channel.active_slots.push(ActiveSlot { slot_index: fs, - ordinal: s.ordinal, - timestamp: s.timestamp, - vchan_id: s.vchan_id as i32, + ordinal: s.ordinal(), + timestamp: s.timestamp(), + vchan_id: s.vchan_id() as i32, }); } } @@ -306,9 +338,9 @@ impl PublisherImpl { let s = self.channel.slot_ref(rs); self.channel.active_slots.push(ActiveSlot { slot_index: rs, - ordinal: s.ordinal, - timestamp: s.timestamp, - vchan_id: s.vchan_id as i32, + ordinal: s.ordinal(), + timestamp: s.timestamp(), + vchan_id: s.vchan_id() as i32, }); } else { continue; @@ -326,9 +358,9 @@ impl PublisherImpl { if (refs & PUB_OWNED) == 0 { self.channel.active_slots.push(ActiveSlot { slot_index: i, - ordinal: s.ordinal, - timestamp: s.timestamp, - vchan_id: s.vchan_id as i32, + ordinal: s.ordinal(), + timestamp: s.timestamp(), + vchan_id: s.vchan_id() as i32, }); } } @@ -337,13 +369,19 @@ impl PublisherImpl { self.channel.active_slots.sort_by_key(|s| s.timestamp); slot_idx = None; + let require_reliable_seen = + self.channel.scb().counters[self.channel.channel_id as usize].num_reliable_subs + != 0; for active in &self.channel.active_slots { let s = self.channel.slot_ref(active.slot_index); let refs = s.refs.load(Ordering::Relaxed); if ((refs >> RELIABLE_REF_COUNT_SHIFT) & REF_COUNT_MASK) != 0 { break; } - if active.ordinal != 0 && (s.flags & MESSAGE_SEEN) == 0 { + if require_reliable_seen + && active.ordinal != 0 + && (s.flags() & MESSAGE_SEEN_BY_RELIABLE) == 0 + { break; } if (refs & REFS_MASK) == 0 { @@ -362,7 +400,7 @@ impl PublisherImpl { let old_refs = (*slot_ptr).refs.load(Ordering::Relaxed); let ref_val = PUB_OWNED | owner as u64; let expected = build_refs_bit_field( - (*slot_ptr).ordinal, + (*slot_ptr).ordinal(), ((old_refs >> VCHAN_ID_SHIFT) & VCHAN_ID_MASK) as i32, ((old_refs >> RETIRED_REFS_SHIFT) & RETIRED_REFS_MASK) as i32, ); @@ -384,10 +422,10 @@ impl PublisherImpl { } let si = slot_idx.unwrap(); - let slot = self.channel.slot_mut(si); - slot.ordinal = 0; - slot.timestamp = 0; - slot.vchan_id = self.channel.vchan_id as i16; + let slot = self.channel.slot_ref(si); + slot.set_ordinal(0); + slot.set_timestamp(0); + slot.set_vchan_id(self.channel.vchan_id as i16); self.channel.set_slot_to_biggest_buffer(si); let prefix = self.channel.get_prefix(si); @@ -425,42 +463,48 @@ impl PublisherImpl { omit_prefix: bool, use_prefix_slot_id: bool, ) -> PublishedMessage { - let slot = self.channel.slot_mut(slot_idx); + let slot = self.channel.slot_ref(slot_idx); let vchan_id = self.channel.vchan_id; + let slot_vchan_id = slot.vchan_id(); - slot.ordinal = self.channel.ccb().ordinals.next(slot.vchan_id as i32); - slot.timestamp = now_ns(); - slot.flags = 0; + let ordinal = self.channel.ccb().ordinals.next(slot_vchan_id as i32); + slot.set_ordinal(ordinal); + slot.set_timestamp(now_ns()); + slot.set_flags(0); let prefix = self.channel.get_prefix(slot_idx); if !prefix.is_null() { unsafe { let p = &mut *prefix; if omit_prefix { - let slot = self.channel.slot_mut(slot_idx); - slot.timestamp = p.timestamp; - slot.vchan_id = p.vchan_id as i16; - slot.bridged_slot_id = if use_prefix_slot_id { + slot.set_timestamp(p.timestamp); + slot.set_vchan_id(p.vchan_id as i16); + slot.set_bridged_slot_id(if use_prefix_slot_id { p.slot_id } else { slot.id - }; + }); } else { - let slot = self.channel.slot_ref(slot_idx); - p.message_size = slot.message_size; - p.ordinal = slot.ordinal; - p.timestamp = slot.timestamp; - p.vchan_id = slot.vchan_id as i32; + let message_size = slot.message_size(); + let ordinal = slot.ordinal(); + let timestamp = slot.timestamp(); + let vchan_id_i16 = slot.vchan_id(); + p.message_size = message_size; + p.ordinal = ordinal; + p.timestamp = timestamp; + p.vchan_id = vchan_id_i16 as i32; p.checksum_size = self.channel.checksum_size as u16; p.metadata_size = self.channel.metadata_size as u16; p.flags = 0; p.slot_id = slot.id; - let slot = self.channel.slot_mut(slot_idx); - slot.bridged_slot_id = slot.id; + slot.set_bridged_slot_id(slot.id); if is_activation { p.set_is_activation(); - slot.flags |= MESSAGE_IS_ACTIVATION; - self.channel.ccb().activation_tracker.activate(vchan_id); + slot.set_flag(MESSAGE_IS_ACTIVATION); + self.channel + .ccb() + .activation_tracker + .activate(vchan_id); } if self.options.checksum { p.set_has_checksum(); @@ -470,7 +514,7 @@ impl PublisherImpl { let data = checksum::get_message_checksum_data( prefix, buffer, - slot.message_size as usize, + message_size as usize, cs, ms, ); @@ -486,18 +530,55 @@ impl PublisherImpl { } } - let slot = self.channel.slot_ref(slot_idx); + // 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, + ); + + // Tell all subscribers the slot is available. + let ccb = self.channel.ccb(); + { + let _publish_guard = + SubscriberQueuePublishGuard::new( + &self.channel, + owner as usize, + &self.active_queue_publish_depth, + ); + let mut failed_queues: Vec<*const SlotQueueHeader> = Vec::new(); + ccb.subscribers.traverse_seq_cst(|sub_id| { + if vchan_id != -1 + && self.channel.get_sub_vchan_id(sub_id) != -1 + && vchan_id != self.channel.get_sub_vchan_id(sub_id) + { + return; + } + self.channel.get_available_slots(sub_id).set(slot_idx); + let queue = self.channel.get_available_slot_queue(sub_id); + if let Some(queue) = queue { + if !queue.push( + slot.id, + ordinal, + /* report_insertion_failure= */ false, + ) { + failed_queues.push(queue as *const SlotQueueHeader); + } + } + }); + ccb.total_messages.fetch_add(1, Ordering::SeqCst); + for queue in failed_queues { + unsafe { (&*queue).mark_insertion_failure() }; + } + } + if !is_activation { - self.channel - .ccb() - .total_messages - .fetch_add(1, Ordering::Relaxed); + let message_size = slot.message_size(); self.channel .ccb() .total_bytes - .fetch_add(slot.message_size, Ordering::Relaxed); - - let msg_size = slot.message_size as u32; + .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( @@ -511,26 +592,6 @@ impl PublisherImpl { } } } - - // Release the slot: store refs with ordinal, no PUB_OWNED. - let slot = self.channel.slot_ref(slot_idx); - slot.refs.store( - build_refs_bit_field(slot.ordinal, vchan_id, 0), - Ordering::Release, - ); - - // Tell all subscribers the slot is available. - let ccb = self.channel.ccb(); - ccb.subscribers.traverse(|sub_id| { - if vchan_id != -1 - && self.channel.get_sub_vchan_id(sub_id) != -1 - && vchan_id != self.channel.get_sub_vchan_id(sub_id) - { - return; - } - self.channel.get_available_slots(sub_id).set(slot_idx); - }); - if reliable { return PublishedMessage { new_slot: None, @@ -961,19 +1022,26 @@ pub fn clear_trigger(fd: RawFd) { } } -pub fn attach_buffers(channel: &mut Channel, read_write: bool) -> crate::error::Result<()> { +pub fn attach_buffers( + channel: &mut Channel, + resolved_name: &str, + read_write: bool, +) -> crate::error::Result<()> { if channel.use_split_buffers { - return attach_split_buffers(channel, read_write); + return attach_split_buffers(channel, resolved_name, read_write); } - attach_shm_buffers(channel, read_write) + attach_shm_buffers(channel, resolved_name, read_write) } -fn attach_shm_buffers(channel: &mut Channel, read_write: bool) -> crate::error::Result<()> { +fn attach_shm_buffers( + channel: &mut Channel, + resolved_name: &str, + read_write: bool, +) -> crate::error::Result<()> { let num_buffers = channel.ccb().num_buffers.load(Ordering::Acquire) as usize; - let resolved_name = channel.name.clone(); while channel.buffers.len() < num_buffers { let buffer_index = channel.buffers.len(); - let shm_name = channel.buffer_shared_memory_name(&resolved_name, buffer_index); + let shm_name = channel.buffer_shared_memory_name(resolved_name, buffer_index); let fd = open_shm(&shm_name)?; let size = get_shm_size(fd, &shm_name)?; @@ -996,16 +1064,19 @@ fn attach_shm_buffers(channel: &mut Channel, read_write: bool) -> crate::error:: Ok(()) } -fn attach_split_buffers(channel: &mut Channel, read_write: bool) -> crate::error::Result<()> { +fn attach_split_buffers( + channel: &mut Channel, + resolved_name: &str, + read_write: bool, +) -> crate::error::Result<()> { let num_buffers = channel.ccb().num_buffers.load(Ordering::Acquire) as usize; - let resolved_name = channel.name.clone(); while channel.buffers.len() < num_buffers { let buffer_index = channel.buffers.len(); let full_size = channel.bcb().sizes[buffer_index].load(Ordering::Acquire); let slot_size = channel.buffer_size_to_slot_size(full_size); let buffer = open_split_buffer_set( channel, - &resolved_name, + resolved_name, buffer_index, full_size, slot_size, diff --git a/rust_client/src/subscriber.rs b/rust_client/src/subscriber.rs index b9c9af1d..71f91211 100644 --- a/rust_client/src/subscriber.rs +++ b/rust_client/src/subscriber.rs @@ -23,6 +23,7 @@ pub type OnReceiveCallback = Box Result + Send + Sy pub struct SubscriberImpl { pub channel: Channel, pub subscriber_id: i32, + pub subscriber_queue_size: i32, pub options: SubscriberOptions, pub poll_fd: RawFd, @@ -42,6 +43,15 @@ pub struct SubscriberImpl { pub(crate) on_receive_callback: Option, pub(crate) checksum_callback: Option, pub checksum_tmp: Vec, + pub poll_drain_pending: bool, + pub(crate) poll_drain_exhausted: bool, + pub(crate) queue_drain_tail: Option, + queue_bitset_fallback: bool, + pub(crate) poll_snapshot_valid: bool, + poll_snapshot_total: u64, + poll_snapshot: Vec, + newest_snapshot: Option>, + pub(crate) pending_queue_drops: i32, } struct OrdinalTracker { @@ -94,15 +104,6 @@ impl FastRingBuffer { self.set.contains(value) } - fn traverse(&self, mut func: F) { - for v in &self.buffer { - func(v); - } - } - - fn size(&self) -> usize { - self.buffer.len() - } } fn virtual_channel_id_match(slot_vchan_id: i16, subscriber_vchan_id: i32) -> bool { @@ -115,6 +116,9 @@ impl SubscriberImpl { pub fn new( name: String, num_slots: i32, + default_subscriber_queue_size: i32, + subscriber_queue_arena_size: u64, + subscriber_queue_size: i32, channel_id: i32, subscriber_id: i32, vchan_id: i32, @@ -126,12 +130,15 @@ impl SubscriberImpl { channel: Channel::new( name, num_slots, + default_subscriber_queue_size, + subscriber_queue_arena_size, channel_id, channel_type, vchan_id, session_id, ), subscriber_id, + subscriber_queue_size, options, poll_fd: -1, trigger_fd: -1, @@ -147,6 +154,15 @@ impl SubscriberImpl { on_receive_callback: None, checksum_callback: None, checksum_tmp: vec![0u8; 4], + poll_drain_pending: false, + poll_drain_exhausted: false, + queue_drain_tail: None, + queue_bitset_fallback: false, + poll_snapshot_valid: false, + poll_snapshot_total: 0, + poll_snapshot: Vec::new(), + newest_snapshot: None, + pending_queue_drops: 0, }; s.get_or_create_tracker(vchan_id); s @@ -173,6 +189,26 @@ impl SubscriberImpl { } } + pub fn reset_delivery_state(&mut self) { + self.channel.active_slots.clear(); + self.channel.embargoed_slots.clear_all(); + self.channel.slot = None; + self.poll_snapshot_valid = false; + self.poll_snapshot_total = 0; + self.poll_snapshot.clear(); + self.poll_drain_exhausted = false; + self.queue_drain_tail = None; + self.queue_bitset_fallback = false; + self.pending_queue_drops = 0; + self.newest_snapshot = None; + self.ordinal_trackers.clear(); + self.get_or_create_tracker(self.channel.vchan_id); + } + + pub fn total_messages(&self) -> u64 { + self.channel.ccb().total_messages.load(Ordering::SeqCst) + } + pub fn resolved_name(&self) -> &str { if !self.options.mux.is_empty() && self.channel.vchan_id != -1 { &self.options.mux @@ -223,7 +259,10 @@ impl SubscriberImpl { } pub fn remember_ordinal(&mut self, ordinal: u64, vchan_id: i32) { - let tracker = self.get_or_create_tracker(self.channel.vchan_id); + let tracker = self.get_or_create_tracker(vchan_id); + if ordinal > tracker.last_ordinal_seen { + tracker.last_ordinal_seen = ordinal; + } tracker.ring.insert(OrdinalAndVchanId { ordinal, vchan_id, @@ -242,9 +281,9 @@ 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 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( @@ -273,43 +312,32 @@ impl SubscriberImpl { pub fn populate_active_slots(&self, bits: &crate::bitset::InPlaceAtomicBitSet) { loop { - let num_messages = self - .channel - .ccb() - .total_messages - .load(Ordering::Relaxed); + 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::Relaxed); - if virtual_channel_id_match(s.vchan_id, self.channel.vchan_id) - && s.ordinal != 0 + let refs = s.refs.load(Ordering::Acquire); + if virtual_channel_id_match(s.vchan_id(), self.channel.vchan_id) + && s.ordinal() != 0 && (refs & PUB_OWNED) == 0 { bits.set(i); } } - if num_messages - == self - .channel - .ccb() - .total_messages - .load(Ordering::Relaxed) - { + if total == self.total_messages() { break; } } } - pub fn collect_visible_slots(&mut self, bits: &crate::bitset::InPlaceAtomicBitSet) { + pub fn collect_visible_slots( + &mut self, + bits: &crate::bitset::InPlaceAtomicBitSet, + ) -> u64 { loop { - let num_messages = self - .channel - .ccb() - .total_messages - .load(Ordering::Relaxed); + let total = self.total_messages(); self.channel.active_slots.clear(); bits.traverse(|i| { @@ -317,47 +345,261 @@ impl SubscriberImpl { return; } let s = self.channel.slot_ref(i); - if !virtual_channel_id_match(s.vchan_id, self.channel.vchan_id) { + if !virtual_channel_id_match(s.vchan_id(), self.channel.vchan_id) { return; } - if s.buffer_index == -1 { + if s.buffer_index() == -1 { return; } self.channel.active_slots.push(ActiveSlot { slot_index: i, - ordinal: s.ordinal, - timestamp: s.timestamp, - vchan_id: s.vchan_id as i32, + ordinal: s.ordinal(), + timestamp: s.timestamp(), + vchan_id: s.vchan_id() as i32, }); }); - if num_messages - == self - .channel - .ccb() - .total_messages - .load(Ordering::Relaxed) - { - break; + if total == self.total_messages() { + return total; } } } fn find_unseen_ordinal(&self) -> Option { - let tracker = self.ordinal_trackers.get(&self.channel.vchan_id)?; for (i, active) in self.channel.active_slots.iter().enumerate() { - if active.ordinal != 0 - && !tracker.ring.contains(&OrdinalAndVchanId { - ordinal: active.ordinal, - vchan_id: active.vchan_id, - }) - { + let seen = self.ordinal_trackers.get(&active.vchan_id).is_some_and( + |tracker| { + active.ordinal <= tracker.last_ordinal_seen + || tracker.ring.contains(&OrdinalAndVchanId { + ordinal: active.ordinal, + vchan_id: active.vchan_id, + }) + }, + ); + if active.ordinal != 0 && !seen { return Some(i); } } None } + fn has_visible_ordinal_before( + &self, + vchan_id: i32, + last_ordinal_seen: u64, + max_ordinal: u64, + ) -> bool { + let bits = self + .channel + .get_available_slots(self.subscriber_id as usize); + let mut found = false; + bits.traverse(|slot_idx| { + if found { + return; + } + let slot = self.channel.slot_ref(slot_idx); + let ordinal = slot.ordinal(); + if ordinal > last_ordinal_seen + && ordinal <= max_ordinal + && (slot.refs.load(Ordering::Acquire) & PUB_OWNED) == 0 + && slot.vchan_id() as i32 == vchan_id + && slot.buffer_index() != -1 + { + found = true; + } + }); + found + } + + fn next_queued_slot( + &mut self, + max_queue_position: u64, + overflow_baseline: u32, + ) -> Option { + if self.options.reliable { + return None; + } + + if self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + .is_none() + { + return None; + } + loop { + let queue_at_boundary = match self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + { + Some(queue) => queue.head() >= max_queue_position, + None => true, + }; + if queue_at_boundary { + break; + } + let Some((slot_id, ordinal)) = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + .and_then(|queue| queue.try_pop()) + else { + break; + }; + if slot_id < 0 || slot_id as usize >= self.channel.num_slots as usize { + continue; + } + let slot_idx = slot_id as usize; + let slot = self.channel.slot_ref(slot_idx); + let refs = slot.refs.load(Ordering::Acquire); + if (refs & PUB_OWNED) != 0 { + continue; + } + let slot_ordinal = slot.ordinal(); + if slot_ordinal != ordinal || slot_ordinal == 0 { + continue; + } + if !virtual_channel_id_match(slot.vchan_id(), self.channel.vchan_id) { + continue; + } + + let vchan_id = slot.vchan_id() as i32; + let last_ordinal_seen = self + .ordinal_trackers + .get(&vchan_id) + .map_or(0, |tracker| tracker.last_ordinal_seen); + let queued_key = OrdinalAndVchanId { ordinal, vchan_id }; + let already_seen = self + .ordinal_trackers + .get(&vchan_id) + .is_some_and(|tracker| tracker.ring.contains(&queued_key)); + // Concurrent publishers can reserve queue positions out of ordinal + // order. Only discard a lower ordinal when this subscriber has + // actually delivered that exact generation. + if ordinal <= last_ordinal_seen && already_seen { + continue; + } + if self.options.subscriber_queue_size == 0 + && ordinal > last_ordinal_seen + && ordinal - last_ordinal_seen > 1 + && self.has_visible_ordinal_before( + vchan_id, + last_ordinal_seen, + ordinal - 1, + ) + { + self.queue_bitset_fallback = true; + self.poll_snapshot_valid = false; + return None; + } + if self.channel.atomic_inc_ref_count::( + slot_idx, + false, + 1, + ordinal, + vchan_id, + false, + None, + ) { + if self.channel.slot_ref(slot_idx).ordinal() != ordinal + || self.channel.slot_ref(slot_idx).vchan_id() as i32 != vchan_id + { + self.channel.atomic_inc_ref_count::( + slot_idx, + false, + -1, + ordinal, + vchan_id, + false, + None, + ); + continue; + } + if let Some(queue) = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + { + if queue.insertion_failed() { + self.channel.atomic_inc_ref_count::( + slot_idx, + false, + -1, + ordinal, + vchan_id, + false, + None, + ); + queue.consume_insertion_failure(); + self.queue_bitset_fallback = true; + self.poll_snapshot_valid = false; + return None; + } + } + if self.options.subscriber_queue_size == 0 { + if let Some(queue) = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + { + if queue.overflow_count() != overflow_baseline { + self.channel.atomic_inc_ref_count::( + slot_idx, + false, + -1, + ordinal, + vchan_id, + false, + None, + ); + self.queue_bitset_fallback = true; + self.poll_snapshot_valid = false; + return None; + } + } + } + if !self.channel.validate_slot_buffer(slot_idx) + || self.channel.slot_ref(slot_idx).buffer_index() == -1 + { + if self.channel.buffers_changed() { + self.channel.atomic_inc_ref_count::( + slot_idx, + false, + -1, + ordinal, + vchan_id, + false, + None, + ); + self.reload_buffers_if_necessary(); + continue; + } + self.channel.atomic_inc_ref_count::( + slot_idx, + false, + -1, + ordinal, + vchan_id, + false, + None, + ); + continue; + } + if self.options.subscriber_queue_size != 0 { + let concurrent_drops = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + .map(|queue| queue.consume_overflow()) + .unwrap_or(0); + if self.options.detect_dropped_messages { + self.pending_queue_drops = self + .pending_queue_drops + .saturating_add(concurrent_drops as i32); + } + } + return Some(slot_idx); + } + } + None + } + pub fn next_slot(&mut self) -> Option { let bits = self .channel @@ -372,19 +614,103 @@ impl SubscriberImpl { self.reload_buffers_if_necessary(); + if self.poll_drain_pending && self.poll_drain_exhausted { + if self.total_messages() != self.poll_snapshot_total { + self.poll_drain_exhausted = false; + self.queue_drain_tail = None; + self.poll_snapshot_valid = false; + } else { + return None; + } + } + + if self.poll_drain_pending && !self.poll_snapshot_valid { + self.poll_snapshot_total = self.collect_visible_slots(&bits); + self.channel + .active_slots + .sort_by_key(|slot| (slot.timestamp, slot.ordinal)); + self.poll_snapshot.clone_from(&self.channel.active_slots); + self.poll_snapshot_valid = true; + self.queue_drain_tail = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + .map(|queue| queue.tail()); + } + let mut queue_overflow_baseline = 0; + let queue_status = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + .map(|queue| { + let queue_drops = queue.consume_overflow(); + let insertion_failed = queue.consume_insertion_failure(); + let overflow_after_consume = queue.overflow_count(); + ( + queue_drops, + insertion_failed, + overflow_after_consume, + ) + }); + if let Some(( + queue_drops, + insertion_failed, + overflow_after_consume, + )) = queue_status + { + queue_overflow_baseline = overflow_after_consume; + let recover_overflow = self.options.subscriber_queue_size == 0 + && (queue_drops != 0 || overflow_after_consume != 0); + if self.options.detect_dropped_messages && !recover_overflow { + self.pending_queue_drops = self + .pending_queue_drops + .saturating_add(queue_drops as i32); + } + if recover_overflow || insertion_failed { + self.queue_bitset_fallback = true; + } + } + let max_queue_position = self.queue_drain_tail.unwrap_or(u64::MAX); + if !self.queue_bitset_fallback { + if let Some(slot_idx) = + self.next_queued_slot(max_queue_position, queue_overflow_baseline) + { + return Some(slot_idx); + } + } + if self.channel.slot.is_none() { self.populate_active_slots(&bits); } - self.collect_visible_slots(&bits); + if self.poll_drain_pending && self.poll_snapshot_valid { + self.channel + .active_slots + .clone_from(&self.poll_snapshot); + } else { + self.collect_visible_slots(&bits); + } - self.channel - .active_slots - .sort_by_key(|s| s.timestamp); + if self.queue_bitset_fallback { + self.channel.active_slots.sort_by_key(|s| s.ordinal); + } else { + self.channel + .active_slots + .sort_by_key(|s| (s.timestamp, s.ordinal)); + } let unseen_idx = match self.find_unseen_ordinal() { Some(idx) => idx, - None => return None, + None => { + if self.queue_bitset_fallback { + if let Some(queue) = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + { + queue.discard_all(); + } + self.queue_bitset_fallback = false; + } + break; + } }; let active = self.channel.active_slots[unseen_idx].clone(); @@ -400,7 +726,7 @@ impl SubscriberImpl { None, ) { if !self.channel.validate_slot_buffer(active.slot_index) - || self.channel.slot_ref(active.slot_index).buffer_index == -1 + || self.channel.slot_ref(active.slot_index).buffer_index() == -1 { if self.channel.buffers_changed() { self.channel.atomic_inc_ref_count::( @@ -430,6 +756,12 @@ impl SubscriberImpl { return Some(active.slot_index); } } + if self.poll_drain_pending && self.total_messages() != self.poll_snapshot_total { + self.trigger(); + } + if self.poll_drain_pending { + self.poll_drain_exhausted = true; + } None } @@ -438,6 +770,19 @@ impl SubscriberImpl { .channel .get_available_slots(self.subscriber_id as usize); self.channel.embargoed_slots.clear_all(); + self.newest_snapshot = None; + if let Some(queue) = self + .channel + .get_available_slot_queue(self.subscriber_id as usize) + { + let queue_drops = queue.consume_overflow(); + if self.options.detect_dropped_messages { + self.pending_queue_drops = self + .pending_queue_drops + .saturating_add(queue_drops as i32); + } + queue.consume_insertion_failure(); + } loop { self.reload_buffers_if_necessary(); @@ -450,7 +795,7 @@ impl SubscriberImpl { self.channel .active_slots - .sort_by_key(|s| s.timestamp); + .sort_by_key(|s| (s.timestamp, s.ordinal)); let new_active = if let Some(last) = self.channel.active_slots.last() { if let Some(current_idx) = self.channel.slot { @@ -471,6 +816,8 @@ impl SubscriberImpl { None => return None, }; + self.newest_snapshot = Some(self.channel.active_slots.clone()); + let reliable = self.options.reliable; if self.channel.atomic_inc_ref_count::( active.slot_index, @@ -482,7 +829,7 @@ impl SubscriberImpl { None, ) { if !self.channel.validate_slot_buffer(active.slot_index) - || self.channel.slot_ref(active.slot_index).buffer_index == -1 + || self.channel.slot_ref(active.slot_index).buffer_index() == -1 { if self.channel.buffers_changed() { self.channel.atomic_inc_ref_count::( @@ -515,88 +862,131 @@ impl SubscriberImpl { } pub fn claim_slot(&mut self, slot_idx: usize, vchan_id: i32, was_newest: bool) { - let slot = self.channel.slot_ref(slot_idx); - slot.sub_owners.set(self.subscriber_id as usize); + let ordinal = self.channel.slot_ref(slot_idx).ordinal(); + self.channel + .slot_ref(slot_idx) + .sub_owners + .set(self.subscriber_id as usize); + let bits = self + .channel + .get_available_slots(self.subscriber_id as usize); if was_newest { - self.channel - .get_available_slots(self.subscriber_id as usize) - .clear_all(); + let mut skipped = Vec::new(); + if let Some(snapshot) = self.newest_snapshot.take() { + for active in snapshot { + let pinned = active.slot_index == slot_idx + || self.channel.atomic_inc_ref_count::( + active.slot_index, + self.options.reliable, + 1, + active.ordinal, + active.vchan_id, + false, + None, + ); + if pinned { + bits.clear(active.slot_index); + skipped.push((active.ordinal, active.vchan_id)); + if active.slot_index != slot_idx { + self.channel.atomic_inc_ref_count::( + active.slot_index, + self.options.reliable, + -1, + active.ordinal, + active.vchan_id, + false, + None, + ); + } + } + } + } else { + bits.clear(slot_idx); + } + for (skipped_ordinal, skipped_vchan_id) in skipped { + self.remember_ordinal(skipped_ordinal, skipped_vchan_id); + } } else { - self.channel - .get_available_slots(self.subscriber_id as usize) - .clear(slot_idx); + bits.clear(slot_idx); } - let ordinal = slot.ordinal; self.remember_ordinal(ordinal, vchan_id); - self.channel.slot_mut(slot_idx).flags |= MESSAGE_SEEN; + let slot = self.channel.slot_ref(slot_idx); + slot.set_flag(MESSAGE_SEEN); + if self.options.reliable { + slot.set_flag(MESSAGE_SEEN_BY_RELIABLE); + } } - pub fn unread_slot(&self, slot_idx: usize) { - self.channel.slot_mut(slot_idx).flags &= !MESSAGE_SEEN; - self.decrement_slot_ref(slot_idx, false); + 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 + .get_available_slot_queue(self.subscriber_id as usize) + .is_some() + { + // The queue entry was already popped. Recover the rejected ordinal + // from the authoritative bitset before accepting newer hints. + self.queue_bitset_fallback = true; + } + self.poll_snapshot_valid = false; + self.newest_snapshot = None; } pub fn ignore_activation(&mut self, slot_idx: usize) { - let slot = self.channel.slot_ref(slot_idx); - let ordinal = slot.ordinal; - let vchan_id = slot.vchan_id as i32; + let ordinal = self.channel.slot_ref(slot_idx).ordinal(); + let vchan_id = self.channel.slot_ref(slot_idx).vchan_id() as i32; self.remember_ordinal(ordinal, vchan_id); - self.decrement_slot_ref(slot_idx, true); - self.channel.slot_mut(slot_idx).flags |= MESSAGE_SEEN; + self.decrement_slot_ref(slot_idx, ordinal, vchan_id, true); + self.channel.slot_ref(slot_idx).set_flag(MESSAGE_SEEN); + if self.options.reliable { + self.channel.slot_ref(slot_idx).set_flag(MESSAGE_SEEN_BY_RELIABLE); + } } - pub fn decrement_slot_ref(&self, slot_idx: usize, retire: bool) { - let slot = self.channel.slot_ref(slot_idx); - let ordinal = slot.ordinal & ORDINAL_MASK; - let vchan_id = self.channel.vchan_id; + pub fn decrement_slot_ref( + &self, + slot_idx: usize, + ordinal: u64, + vchan_id: i32, + retire: bool, + ) { let reliable = self.options.reliable; - self.channel - .atomic_inc_ref_count::(slot_idx, reliable, -1, ordinal, vchan_id, retire, None); + self.channel.atomic_inc_ref_count::( + slot_idx, + reliable, + -1, + ordinal & ORDINAL_MASK, + vchan_id, + retire, + None, + ); + } + + pub fn release_unclaimed_slot( + &self, + slot_idx: usize, + ordinal: u64, + vchan_id: i32, + ) { + self.decrement_slot_ref(slot_idx, ordinal, vchan_id, false); } pub fn detect_drops(&mut self, vchan_id: i32) -> i32 { - let tracker_vchan = self.channel.vchan_id; - let tracker = match self.ordinal_trackers.get(&tracker_vchan) { - Some(t) => t, - None => return 0, + let Some(slot_idx) = self.channel.slot else { + return 0; }; - - let mut ordinals: Vec = Vec::with_capacity(tracker.ring.size()); - let last_seen = tracker.last_ordinal_seen; - tracker.ring.traverse(|o| { - if o.vchan_id == vchan_id && o.ordinal >= last_seen { - ordinals.push(*o); - } - }); - - if ordinals.is_empty() { + let ordinal = self.channel.slot_ref(slot_idx).ordinal(); + let tracker = self.get_or_create_tracker(vchan_id); + if ordinal == 0 || ordinal <= tracker.last_ordinal_seen { return 0; } - - ordinals.sort_by(|a, b| { - a.vchan_id - .cmp(&b.vchan_id) - .then(a.ordinal.cmp(&b.ordinal)) - }); - - let last_ordinal = ordinals.last().unwrap().ordinal; - - let mut drops: i32 = 0; - for i in 1..ordinals.len() { - if ordinals[i].vchan_id != vchan_id { - continue; - } - let gap = ordinals[i].ordinal.wrapping_sub(ordinals[i - 1].ordinal); - if gap > 1 { - drops += (gap - 1) as i32; - } + let last_seen = tracker.last_ordinal_seen; + tracker.last_ordinal_seen = ordinal; + if last_seen == 0 || ordinal == last_seen + 1 { + 0 + } else { + (ordinal - last_seen - 1) as i32 } - - // Update last_ordinal_seen. - let tracker = self.ordinal_trackers.get_mut(&tracker_vchan).unwrap(); - tracker.last_ordinal_seen = last_ordinal; - - drops } pub fn find_active_slot_by_timestamp( @@ -614,19 +1004,20 @@ impl SubscriberImpl { continue; } let s = self.channel.slot_ref(i); - let refs = s.refs.load(Ordering::Relaxed); - if s.ordinal != 0 && (refs & PUB_OWNED) == 0 { + let refs = s.refs.load(Ordering::Acquire); + let ordinal = s.ordinal(); + if ordinal != 0 && (refs & PUB_OWNED) == 0 { let prefix = self.channel.get_prefix(i); let ts = if !prefix.is_null() { unsafe { (*prefix).timestamp } } else { - s.timestamp + s.timestamp() }; buffer.push(ActiveSlot { slot_index: i, - ordinal: 0, + ordinal, timestamp: ts, - vchan_id: s.vchan_id as i32, + vchan_id: s.vchan_id() as i32, }); } } @@ -654,7 +1045,7 @@ impl SubscriberImpl { None, ) { if !self.channel.validate_slot_buffer(active.slot_index) - || self.channel.slot_ref(active.slot_index).buffer_index == -1 + || self.channel.slot_ref(active.slot_index).buffer_index() == -1 { if self.channel.buffers_changed() { self.channel.atomic_inc_ref_count::( @@ -681,8 +1072,11 @@ impl SubscriberImpl { ); continue; } - let slot = self.channel.slot_mut(active.slot_index); - slot.flags |= MESSAGE_SEEN; + let slot = self.channel.slot_ref(active.slot_index); + slot.set_flag(MESSAGE_SEEN); + if reliable { + slot.set_flag(MESSAGE_SEEN_BY_RELIABLE); + } slot.sub_owners.set(self.subscriber_id as usize); return Some(active.slot_index); } @@ -736,6 +1130,7 @@ impl SubscriberImpl { pub fn attach_buffers(&mut self) -> crate::error::Result<()> { let read_write = self.options.bridge || self.options.read_write; - attach_buffers(&mut self.channel, read_write) + let resolved_name = self.resolved_name().to_string(); + attach_buffers(&mut self.channel, &resolved_name, read_write) } } diff --git a/rust_client/tests/client_test.rs b/rust_client/tests/client_test.rs index f896728e..f80df702 100644 --- a/rust_client/tests/client_test.rs +++ b/rust_client/tests/client_test.rs @@ -22,7 +22,10 @@ fn calculate_checksum(spans: &[&[u8]]) -> u32 { fn verify_checksum(spans: &[&[u8]], checksum: u32) -> bool { verify_crc32_checksum(spans, &checksum.to_ne_bytes()) } -use subspace_client::options::{PublisherOptions, SubscriberOptions}; +use subspace_client::options::{ + PublisherOptions, SubscriberOptions, DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE, + DEFAULT_SUBSCRIBER_QUEUE_SIZE, +}; use subspace_client::{Client, ReadMode, SubspaceError}; fn unique_socket_path() -> String { @@ -54,6 +57,7 @@ fn publisher_options_defaults() { let opts = PublisherOptions::new(); assert_eq!(opts.slot_size, 0); assert_eq!(opts.num_slots, 0); + assert_eq!(opts.subscriber_queue_arena_size, 0); assert!(!opts.local); assert!(!opts.reliable); assert!(!opts.bridge); @@ -72,6 +76,7 @@ fn publisher_options_builder_chain() { let opts = PublisherOptions::new() .set_slot_size(4096) .set_num_slots(16) + .set_subscriber_queue_arena_size(32_000) .set_reliable(true) .set_local(true) .set_fixed_size(true) @@ -85,6 +90,7 @@ fn publisher_options_builder_chain() { assert_eq!(opts.slot_size, 4096); assert_eq!(opts.num_slots, 16); + assert_eq!(opts.subscriber_queue_arena_size, 32_000); assert!(opts.reliable); assert!(opts.local); assert!(opts.fixed_size); @@ -101,9 +107,11 @@ fn publisher_options_builder_chain() { fn subscriber_options_defaults() { let opts = SubscriberOptions::new(); assert!(!opts.reliable); + assert_eq!(opts.subscriber_queue_size, 0); assert!(!opts.bridge); assert_eq!(opts.max_active_messages, 1); assert!(opts.log_dropped_messages); + assert!(opts.detect_dropped_messages); assert!(!opts.pass_activation); assert!(!opts.read_write); assert!(!opts.checksum); @@ -116,8 +124,10 @@ fn subscriber_options_defaults() { fn subscriber_options_builder_chain() { let opts = SubscriberOptions::new() .set_reliable(true) + .set_subscriber_queue_size(3) .set_max_active_messages(8) .set_log_dropped_messages(false) + .set_detect_dropped_messages(false) .set_pass_activation(true) .set_checksum(true) .set_pass_checksum_errors(true) @@ -126,8 +136,10 @@ fn subscriber_options_builder_chain() { .set_type("image".into()); assert!(opts.reliable); + assert_eq!(opts.subscriber_queue_size, 3); assert_eq!(opts.max_active_messages, 8); assert!(!opts.log_dropped_messages); + assert!(!opts.detect_dropped_messages); assert!(opts.pass_activation); assert!(opts.checksum); assert!(opts.pass_checksum_errors); @@ -865,6 +877,119 @@ fn integration_publish_multiple_messages() { } } +#[test] +fn integration_subscriber_queue_overflow_preserves_newest() { + let client = new_client("rust_queue_overflow"); + let pub_opts = PublisherOptions::new() + .set_slot_size(64) + .set_num_slots(8) + .set_subscriber_queue_arena_size(DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE); + let publisher = client + .create_publisher("rust_queue_overflow_ch", &pub_opts) + .unwrap(); + let sub_opts = SubscriberOptions::new().set_subscriber_queue_size(2); + let subscriber = client + .create_subscriber("rust_queue_overflow_ch", &sub_opts) + .unwrap(); + let reported_drops = + std::sync::Arc::new(std::sync::atomic::AtomicI64::new(0)); + let callback_drops = reported_drops.clone(); + subscriber.register_dropped_message_callback(move |drops| { + callback_drops.fetch_add(drops, std::sync::atomic::Ordering::Relaxed); + }); + + for value in 1u8..=4 { + let (buffer, _) = publisher.get_message_buffer(1).unwrap().unwrap(); + unsafe { + *buffer = value; + } + publisher.publish_message(1).unwrap(); + } + + let first = subscriber.read_message(ReadMode::ReadNext).unwrap(); + assert_eq!(unsafe { *first.buffer }, 3); + assert_eq!( + reported_drops.load(std::sync::atomic::Ordering::Relaxed), + 2 + ); + drop(first); + let second = subscriber.read_message(ReadMode::ReadNext).unwrap(); + assert_eq!(unsafe { *second.buffer }, 4); + drop(second); + assert!(subscriber + .read_message(ReadMode::ReadNext) + .unwrap() + .is_empty()); +} + +#[test] +fn integration_default_subscriber_queue_overflow_recovers_from_bitset() { + let client = new_client("rust_default_queue_overflow"); + let publisher = client + .create_publisher( + "rust_default_queue_overflow_ch", + &PublisherOptions::new() + .set_slot_size(64) + .set_num_slots(64), + ) + .unwrap(); + let subscriber = client + .create_subscriber( + "rust_default_queue_overflow_ch", + &SubscriberOptions::new(), + ) + .unwrap(); + + for value in 1u8..=32 { + let (buffer, _) = publisher.get_message_buffer(1).unwrap().unwrap(); + unsafe { + *buffer = value; + } + publisher.publish_message(1).unwrap(); + } + + for expected in 1u8..=32 { + let message = subscriber.read_message(ReadMode::ReadNext).unwrap(); + assert_eq!(unsafe { *message.buffer }, expected); + } + assert!(subscriber + .read_message(ReadMode::ReadNext) + .unwrap() + .is_empty()); +} + +#[test] +fn integration_subscriber_queue_read_newest_does_not_redeliver_old_entries() { + let client = new_client("rust_queue_newest"); + let pub_opts = PublisherOptions::new() + .set_slot_size(64) + .set_num_slots(8) + .set_subscriber_queue_arena_size(DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE); + let publisher = client + .create_publisher("rust_queue_newest_ch", &pub_opts) + .unwrap(); + let sub_opts = SubscriberOptions::new().set_subscriber_queue_size(4); + let subscriber = client + .create_subscriber("rust_queue_newest_ch", &sub_opts) + .unwrap(); + + for value in 1u8..=3 { + let (buffer, _) = publisher.get_message_buffer(1).unwrap().unwrap(); + unsafe { + *buffer = value; + } + publisher.publish_message(1).unwrap(); + } + + let newest = subscriber.read_message(ReadMode::ReadNewest).unwrap(); + assert_eq!(unsafe { *newest.buffer }, 3); + drop(newest); + assert!(subscriber + .read_message(ReadMode::ReadNext) + .unwrap() + .is_empty()); +} + // ── Read newest skips intermediate messages ────────────────────────────────── #[test] @@ -3201,6 +3326,7 @@ fn coverage_publisher_accessors() { let opts = PublisherOptions::new() .set_slot_size(128) .set_num_slots(8) + .set_subscriber_queue_arena_size(5_000) .set_type("pub_type".to_string()) .set_fixed_size(true); let pub_handle = client.create_publisher("cov_pub_acc_ch", &opts).unwrap(); @@ -3209,6 +3335,11 @@ fn coverage_publisher_accessors() { assert!(!pub_handle.is_reliable()); assert!(pub_handle.is_fixed_size()); assert_eq!(pub_handle.num_slots(), 8); + assert_eq!( + pub_handle.subscriber_queue_size(), + DEFAULT_SUBSCRIBER_QUEUE_SIZE + ); + assert_eq!(pub_handle.subscriber_queue_arena_size(), 5_000); assert!(pub_handle.slot_size() > 0); assert!(pub_handle.get_poll_fd() >= 0); assert!(pub_handle.prefix_size() > 0); @@ -3218,9 +3349,12 @@ fn coverage_publisher_accessors() { #[test] fn coverage_subscriber_accessors() { let client = new_client("cov_sub_acc"); - let opts = PublisherOptions::new().set_slot_size(128).set_num_slots(16); + let opts = PublisherOptions::new() + .set_slot_size(128) + .set_num_slots(16) + .set_subscriber_queue_arena_size(DEFAULT_SUBSCRIBER_QUEUE_ARENA_SIZE); let _pub = client.create_publisher("cov_sub_acc_ch", &opts).unwrap(); - let sub_opts = SubscriberOptions::new(); + let sub_opts = SubscriberOptions::new().set_subscriber_queue_size(3); let sub = client .create_subscriber("cov_sub_acc_ch", &sub_opts) .unwrap(); @@ -3229,6 +3363,7 @@ fn coverage_subscriber_accessors() { assert!(!sub.is_reliable()); assert!(!sub.is_placeholder()); assert!(sub.num_slots() > 0); + assert_eq!(sub.subscriber_queue_size(), 3); assert!(sub.get_poll_fd() >= 0); assert!(sub.prefix_size() > 0); assert!(sub.checksum_size() > 0); diff --git a/rust_client/tests/latency_test.rs b/rust_client/tests/latency_test.rs index 5b8d3301..115699fe 100644 --- a/rust_client/tests/latency_test.rs +++ b/rust_client/tests/latency_test.rs @@ -19,6 +19,12 @@ use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; use std::sync::Arc; use std::time::Instant; +fn latency_subscriber_options() -> SubscriberOptions { + SubscriberOptions::new() + .set_log_dropped_messages(false) + .set_detect_dropped_messages(false) +} + fn unique_socket_path() -> String { let mut template = b"/tmp/ss_lat_XXXXXX\0".to_vec(); let fd = unsafe { libc::mkstemp(template.as_mut_ptr() as *mut libc::c_char) }; @@ -272,7 +278,7 @@ fn stress_multithreaded_unreliable() { .create_publisher("stress_unrel", &pub_opts) .unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("stress_unrel", &sub_opts) .unwrap(); @@ -502,7 +508,7 @@ fn latency_unreliable_round_trip() { .create_publisher("lat_unrel_rt", &pub_opts) .unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("lat_unrel_rt", &sub_opts) .unwrap(); @@ -609,7 +615,7 @@ fn latency_publisher_with_retirement() { .create_publisher("lat_pub_ret", &pub_opts) .unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("lat_pub_ret", &sub_opts) .unwrap(); @@ -674,7 +680,8 @@ fn latency_publisher_checksum() { let sub_opts = SubscriberOptions::new() .set_checksum(true) - .set_log_dropped_messages(false); + .set_log_dropped_messages(false) + .set_detect_dropped_messages(false); let subscriber = sub_client .create_subscriber("lat_pub_csum", &sub_opts) .unwrap(); @@ -714,7 +721,7 @@ fn latency_pub_sub_single_thread() { let pub_opts = PublisherOptions::new().set_slot_size(256).set_num_slots(10); let publisher = pub_client.create_publisher("lat_ps_st", &pub_opts).unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("lat_ps_st", &sub_opts) .unwrap(); @@ -756,7 +763,7 @@ fn latency_subscriber_drain() { .create_publisher("lat_sub_drain", &pub_opts) .unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("lat_sub_drain", &sub_opts) .unwrap(); @@ -885,7 +892,7 @@ fn latency_publisher_multi_subscriber() { let mut subscribers = Vec::new(); for _ in 0..num_subs { - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); subscribers.push( sub_client .create_subscriber("lat_pub_msub", &sub_opts) @@ -932,7 +939,7 @@ fn latency_publisher_histogram() { .create_publisher("lat_pub_hist", &pub_opts) .unwrap(); - let sub_opts = SubscriberOptions::new().set_log_dropped_messages(false); + let sub_opts = latency_subscriber_options(); let subscriber = sub_client .create_subscriber("lat_pub_hist", &sub_opts) .unwrap(); diff --git a/server/client_handler.cc b/server/client_handler.cc index 7c464075..76343def 100644 --- a/server/client_handler.cc +++ b/server/client_handler.cc @@ -350,6 +350,16 @@ void ClientHandler::HandleCreatePublisher( const subspace::CreatePublisherRequest &req, subspace::CreatePublisherResponse *response, std::vector &fds) { + if (req.num_slots() <= 0 || req.slot_size() <= 0) { + response->set_error("num_slots and slot_size must be greater than 0"); + return; + } + absl::StatusOr checked_ccb_size = + CheckedCcbSize(req.num_slots(), req.subscriber_queue_arena_size()); + if (!checked_ccb_size.ok()) { + response->set_error(checked_ccb_size.status().ToString()); + return; + } ServerChannel *channel = server_->FindChannel(req.channel_name()); if (channel == nullptr) { server_->logger_.Log(toolbelt::LogLevel::kDebug, @@ -359,8 +369,9 @@ void ClientHandler::HandleCreatePublisher( req.slot_size(), req.num_slots(), req.type().size(), server_->GetNumChannels()); absl::StatusOr ch = server_->CreateChannel( - req.channel_name(), req.slot_size(), req.num_slots(), req.mux(), - req.vchan_id(), req.type()); + req.channel_name(), req.slot_size(), req.num_slots(), + req.subscriber_queue_arena_size(), req.mux(), req.vchan_id(), + req.type()); if (!ch.ok()) { response->set_error(ch.status().ToString()); return; @@ -375,8 +386,9 @@ void ClientHandler::HandleCreatePublisher( req.num_slots(), req.type().size(), server_->GetNumChannels()); // Channel exists, but it's just a placeholder. Remap the memory now // that we know the slots. - absl::Status status = - server_->RemapChannel(channel, req.slot_size(), req.num_slots()); + absl::Status status = server_->RemapChannel( + channel, req.slot_size(), req.num_slots(), + req.subscriber_queue_arena_size()); if (!status.ok()) { response->set_error(status.ToString()); return; @@ -460,6 +472,21 @@ void ClientHandler::HandleCreatePublisher( int num_tunnel_pubs, num_tunnel_subs; channel->CountUsers(num_pubs, num_subs, num_bridge_pubs, num_bridge_subs, num_tunnel_pubs, num_tunnel_subs); + // The subscriber queue arena size defines the physical CCB layout and must + // remain fixed even when this channel currently has no publishers. Virtual + // channels delegate SubscriberQueueArenaSize() to their shared multiplexer, + // so + // this also enforces consistency across all vchans on a mux. + if (req.subscriber_queue_arena_size() != + channel->SubscriberQueueArenaSize()) { + response->set_error(absl::StrFormat( + "Inconsistent publisher parameters for channel %s: subscriber queue " + "arena size is %llu, not %llu", + req.channel_name(), + static_cast(channel->SubscriberQueueArenaSize()), + static_cast(req.subscriber_queue_arena_size()))); + return; + } // Check consistency of publisher parameters. if (num_pubs > 0) { if (req.is_fixed_size() != channel->IsFixedSize()) { @@ -482,7 +509,6 @@ void ClientHandler::HandleCreatePublisher( req.channel_name(), req.num_slots(), current_num_slots)); return; } - if (slot_size_changed) { if (slot_size_changed) { if (channel->IsFixedSize()) { @@ -556,6 +582,9 @@ void ClientHandler::HandleCreatePublisher( pub = static_cast(*user); pub->SetHandler(this); pub->SetProcessId(req.process_id()); + split_channel->GetAvailableSlotQueueIndexAddress() + ->active_publishers[req.publisher_id()] + .store(req.active_queue_publish_depth(), std::memory_order_seq_cst); reclaimed = true; server_->logger_.Log(toolbelt::LogLevel::kDebug, "Client %s reclaiming publisher %d on channel %s", @@ -647,6 +676,9 @@ void ClientHandler::HandleCreatePublisher( response->set_type(channel->Type()); response->set_vchan_id(channel->GetVirtualChannelId()); response->set_publisher_id(pub->GetId()); + response->set_subscriber_queue_size(channel->SubscriberQueueSize()); + response->set_subscriber_queue_arena_size( + channel->SubscriberQueueArenaSize()); const SharedMemoryFds &channel_fds = channel->GetFds(); response->set_ccb_fd_index(0); @@ -706,6 +738,17 @@ void ClientHandler::HandleCreateSubscriber( const subspace::CreateSubscriberRequest &req, subspace::CreateSubscriberResponse *response, std::vector &fds) { + if (req.subscriber_queue_size() < 0) { + response->set_error("subscriber_queue_size must be >= 0"); + return; + } + if (static_cast(req.subscriber_queue_size()) > + kDefaultMaxAvailableSlotQueueCapacity) { + response->set_error(absl::StrFormat( + "subscriber_queue_size must be <= %zu", + kDefaultMaxAvailableSlotQueueCapacity)); + return; + } ServerChannel *channel = server_->FindChannel(req.channel_name()); if (channel == nullptr) { // No channel exists, map an empty channel. @@ -715,7 +758,7 @@ void ClientHandler::HandleCreateSubscriber( client_name_.c_str(), req.channel_name().c_str(), req.type().size(), server_->GetNumChannels()); absl::StatusOr ch = server_->CreateChannel( - req.channel_name(), 0, 0, req.mux(), req.vchan_id(), req.type()); + req.channel_name(), 0, 0, 0, req.mux(), req.vchan_id(), req.type()); if (!ch.ok()) { response->set_error(ch.status().ToString()); return; @@ -801,7 +844,7 @@ void ClientHandler::HandleCreateSubscriber( absl::StatusOr subscriber = channel->AddSubscriber(this, req.is_reliable(), req.is_bridge(), req.for_tunnel(), req.max_active_messages(), - req.process_id()); + req.subscriber_queue_size(), req.process_id()); if (!subscriber.ok()) { response->set_error(subscriber.status().ToString()); return; @@ -851,6 +894,11 @@ void ClientHandler::HandleCreateSubscriber( response->set_slot_size(channel->SlotSize()); response->set_num_slots(channel->NumSlots()); + response->set_subscriber_queue_size( + channel->SubscriberQueueSize(sub->GetId())); + response->set_default_subscriber_queue_size(channel->SubscriberQueueSize()); + response->set_subscriber_queue_arena_size( + channel->SubscriberQueueArenaSize()); response->set_checksum_size(channel->ChecksumSize()); response->set_metadata_size(channel->MetadataSize()); ServerChannel *split_response_channel = diff --git a/server/server.cc b/server/server.cc index f3ea5b6d..35fb1ce2 100644 --- a/server/server.cc +++ b/server/server.cc @@ -1187,7 +1187,11 @@ Server::HandleIncomingConnection(async::Context ctx, absl::StatusOr Server::CreateMultiplexer(const std::string &channel_name, int slot_size, - int num_slots, std::string type) { + int num_slots, + uint64_t subscriber_queue_arena_size, + std::string type) { + const int subscriber_queue_size = + subscriber_queue_arena_size == 0 ? 0 : kDefaultSubscriberQueueSize; absl::StatusOr channel_id = channel_ids_.Allocate("mux"); if (!channel_id.ok()) { return channel_id.status(); @@ -1196,29 +1200,38 @@ Server::CreateMultiplexer(const std::string &channel_name, int slot_size, "Creating multiplexer %s with %d slots", channel_name.c_str(), num_slots); ServerChannel *channel = new ChannelMultiplexer( - *channel_id, channel_name, num_slots, std::move(type), session_id_); + *channel_id, channel_name, num_slots, subscriber_queue_size, + subscriber_queue_arena_size, std::move(type), session_id_); channel->SetDebug(logger_.GetLogLevel() <= toolbelt::LogLevel::kVerboseDebug); absl::StatusOr fds = - channel->Allocate(scb_fd_, slot_size, num_slots, initial_ordinal_); + channel->Allocate(scb_fd_, slot_size, num_slots, + subscriber_queue_arena_size, initial_ordinal_); if (!fds.ok()) { return fds.status(); } channel->SetSharedMemoryFds(std::move(*fds)); channels_.emplace(std::make_pair(channel_name, channel)); + OnNewChannel(channel_name); + ForEachShadow([channel](const std::unique_ptr &s) { + s->SendCreateChannel(channel); + }); return channel; } absl::StatusOr Server::CreateChannel(const std::string &channel_name, int slot_size, - int num_slots, const std::string &mux, int vchan_id, - std::string type) { + int num_slots, uint64_t subscriber_queue_arena_size, + const std::string &mux, int vchan_id, std::string type) { + const int subscriber_queue_size = + subscriber_queue_arena_size == 0 ? 0 : kDefaultSubscriberQueueSize; if (!mux.empty()) { ServerChannel *mux_channel = FindChannel(mux); if (mux_channel == nullptr) { // No mux found, create one. absl::StatusOr m = - CreateMultiplexer(mux, slot_size, num_slots, type); + CreateMultiplexer(mux, slot_size, num_slots, + subscriber_queue_arena_size, type); if (!m.ok()) { return m.status(); } @@ -1228,9 +1241,22 @@ Server::CreateChannel(const std::string &channel_name, int slot_size, return absl::InternalError( absl::StrFormat("Channel %s is not a multiplexer", mux)); } + if (!mux_channel->IsPlaceholder() && num_slots > 0 && + subscriber_queue_arena_size != + mux_channel->SubscriberQueueArenaSize()) { + return absl::InternalError(absl::StrFormat( + "Inconsistent publisher parameters for mux %s: subscriber queue " + "arena size is %llu, not %llu", + mux, + static_cast( + mux_channel->SubscriberQueueArenaSize()), + static_cast(subscriber_queue_arena_size))); + } if (mux_channel->IsPlaceholder()) { // Remap the memory now that we know the slots. - absl::Status status = RemapChannel(mux_channel, slot_size, num_slots); + absl::Status status = + RemapChannel(mux_channel, slot_size, num_slots, + subscriber_queue_arena_size); if (!status.ok()) { return status; } @@ -1261,13 +1287,15 @@ Server::CreateChannel(const std::string &channel_name, int slot_size, return channel_id.status(); } ServerChannel *channel = - new ServerChannel(*channel_id, channel_name, num_slots, std::move(type), - false, session_id_); + new ServerChannel(*channel_id, channel_name, num_slots, + subscriber_queue_size, subscriber_queue_arena_size, + std::move(type), false, session_id_); channel->SetDebug(logger_.GetLogLevel() <= toolbelt::LogLevel::kVerboseDebug); channel->SetLastKnownSlotSize(slot_size); absl::StatusOr fds = - channel->Allocate(scb_fd_, slot_size, num_slots, initial_ordinal_); + channel->Allocate(scb_fd_, slot_size, num_slots, + subscriber_queue_arena_size, initial_ordinal_); if (!fds.ok()) { return fds.status(); } @@ -1290,22 +1318,27 @@ uint64_t Server::GetVirtualMemoryUsage() const { } absl::Status Server::RemapChannel(ServerChannel *channel, int slot_size, - int num_slots) { + int num_slots, + uint64_t subscriber_queue_arena_size) { if (channel->IsVirtual()) { ChannelMultiplexer *mux = static_cast(channel)->GetMux(); logger_.Log(toolbelt::LogLevel::kDebug, "Remapping multiplexer %s with %d slots", channel->Name().c_str(), num_slots); - return RemapChannel(mux, slot_size, num_slots); + return RemapChannel(mux, slot_size, num_slots, + subscriber_queue_arena_size); } absl::StatusOr fds = - channel->Allocate(scb_fd_, slot_size, num_slots, initial_ordinal_); + channel->Allocate(scb_fd_, slot_size, num_slots, + subscriber_queue_arena_size, initial_ordinal_); if (!fds.ok()) { return fds.status(); } channel->SetLastKnownSlotSize(slot_size); channel->SetSharedMemoryFds(std::move(*fds)); - channel->RegisterExistingSubscribers(); + for (const std::string &warning : channel->RegisterExistingSubscribers()) { + logger_.Log(toolbelt::LogLevel::kWarning, "%s", warning.c_str()); + } // Remapping replaces the CCB/BCB FDs; shadow recovery must receive the // refreshed descriptors instead of retaining the placeholder mappings. ForEachShadow([channel](const std::unique_ptr &s) { @@ -1323,11 +1356,9 @@ ServerChannel *Server::FindChannel(const std::string &channel_name) { } absl::Status Server::RecoverFromShadow(RecoveredState &state) { - for (auto &rch : state.channels) { - channel_ids_.Set(rch.channel_id); - - auto *channel = new ServerChannel(rch.channel_id, rch.name, rch.num_slots, - rch.type, false, session_id_); + auto configure_channel = [this](ServerChannel *channel, + RecoveredChannel &rch, + bool map_storage) -> absl::Status { channel->SetDebug(logger_.GetLogLevel() <= toolbelt::LogLevel::kVerboseDebug); channel->SetLastKnownSlotSize(rch.slot_size); @@ -1350,20 +1381,26 @@ absl::Status Server::RecoverFromShadow(RecoveredState &state) { } } - if (absl::Status s = channel->MapExisting(scb_fd_, std::move(rch.ccb_fd), - std::move(rch.bcb_fd)); - !s.ok()) { - return s; - } - for (RegisteredClientBuffer &buffer : rch.client_buffers) { - channel->RegisterClientBuffer(std::move(buffer.metadata), - std::move(buffer.fd)); + if (map_storage) { + if (absl::Status s = channel->MapExisting( + scb_fd_, std::move(rch.ccb_fd), std::move(rch.bcb_fd)); + !s.ok()) { + return s; + } + for (RegisteredClientBuffer &buffer : rch.client_buffers) { + channel->RegisterClientBuffer(std::move(buffer.metadata), + std::move(buffer.fd)); + } } + return absl::OkStatus(); + }; + auto restore_users = [this](ServerChannel *channel, RecoveredChannel &rch) { for (auto &rpub : rch.publishers) { auto pub = std::make_unique( nullptr, rpub.id, rpub.is_reliable, rpub.is_local, rpub.is_bridge, rpub.for_tunnel, rpub.is_fixed_size); + pub->SetProcessId(rpub.process_id); toolbelt::TriggerFd tfd(rpub.poll_fd, rpub.trigger_fd); pub->SetTriggerFd(std::move(tfd)); @@ -1376,20 +1413,72 @@ absl::Status Server::RecoverFromShadow(RecoveredState &state) { } channel->AddUser(rpub.id, std::move(pub)); + channel->ClearPublisherQueueHazardIfDead(rpub.id, rpub.process_id); } for (auto &rsub : rch.subscribers) { auto sub = std::make_unique( nullptr, rsub.id, rsub.is_reliable, rsub.is_bridge, rsub.for_tunnel, - rsub.max_active_messages); + rsub.max_active_messages, rsub.subscriber_queue_size); + sub->SetProcessId(rsub.process_id); toolbelt::TriggerFd tfd(rsub.trigger_fd, rsub.poll_fd); sub->SetTriggerFd(std::move(tfd)); channel->AddUser(rsub.id, std::move(sub)); } - channel->RegisterExistingSubscribers(); + for (const std::string &warning : + channel->RegisterExistingSubscribers()) { + logger_.Log(toolbelt::LogLevel::kWarning, "%s", warning.c_str()); + } + }; + absl::flat_hash_set mux_names; + for (const RecoveredChannel &rch : state.channels) { + if (!rch.mux.empty()) { + mux_names.insert(rch.mux); + } + } + + absl::flat_hash_map recovered_muxes; + absl::flat_hash_map physical_channel_ids; + + // Recover physical channels first so virtual channels can attach to their + // single shared CCB/BCB mapping. + for (RecoveredChannel &rch : state.channels) { + if (!rch.mux.empty()) { + continue; + } + if (channels_.contains(rch.name) || + physical_channel_ids.contains(rch.channel_id)) { + return absl::FailedPreconditionError(absl::StrFormat( + "duplicate recovered physical channel mapping for %s (id=%d)", + rch.name, rch.channel_id)); + } + physical_channel_ids.emplace(rch.channel_id, rch.name); + channel_ids_.Set(rch.channel_id); + ServerChannel *channel = nullptr; + if (mux_names.contains(rch.name)) { + auto *mux = new ChannelMultiplexer( + rch.channel_id, rch.name, rch.num_slots, + rch.subscriber_queue_arena_size == 0 ? 0 + : kDefaultSubscriberQueueSize, + rch.subscriber_queue_arena_size, rch.type, session_id_); + channel = mux; + recovered_muxes.emplace(rch.name, mux); + } else { + channel = new ServerChannel( + rch.channel_id, rch.name, rch.num_slots, + rch.subscriber_queue_arena_size == 0 ? 0 + : kDefaultSubscriberQueueSize, + rch.subscriber_queue_arena_size, rch.type, false, session_id_); + } + if (absl::Status status = configure_channel(channel, rch, true); + !status.ok()) { + delete channel; + return status; + } + restore_users(channel, rch); channels_.emplace(rch.name, channel); logger_.Log(toolbelt::LogLevel::kInfo, "Recovered channel '%s' (id=%d, %d pubs, %d subs)", @@ -1397,6 +1486,83 @@ absl::Status Server::RecoverFromShadow(RecoveredState &state) { static_cast(rch.publishers.size()), static_cast(rch.subscribers.size())); } + + // Older shadow state may contain only virtual-channel records. In that case + // infer and create the physical mux from the first virtual record. + for (RecoveredChannel &rch : state.channels) { + if (rch.mux.empty()) { + continue; + } + ChannelMultiplexer *mux = nullptr; + auto mux_it = recovered_muxes.find(rch.mux); + if (mux_it == recovered_muxes.end()) { + if (channels_.contains(rch.mux) || + physical_channel_ids.contains(rch.channel_id)) { + return absl::FailedPreconditionError(absl::StrFormat( + "duplicate recovered mux mapping for %s (id=%d)", rch.mux, + rch.channel_id)); + } + physical_channel_ids.emplace(rch.channel_id, rch.mux); + channel_ids_.Set(rch.channel_id); + mux = new ChannelMultiplexer( + rch.channel_id, rch.mux, rch.num_slots, + rch.subscriber_queue_arena_size == 0 ? 0 + : kDefaultSubscriberQueueSize, + rch.subscriber_queue_arena_size, rch.type, session_id_); + if (absl::Status status = configure_channel(mux, rch, true); + !status.ok()) { + delete mux; + return status; + } + channels_.emplace(rch.mux, mux); + recovered_muxes.emplace(rch.mux, mux); + } else { + mux = mux_it->second; + if (rch.channel_id != mux->GetChannelId() || + rch.num_slots != mux->NumSlots() || + rch.subscriber_queue_arena_size != + mux->SubscriberQueueArenaSize()) { + return absl::FailedPreconditionError(absl::StrFormat( + "inconsistent recovered virtual channel %s for mux %s", rch.name, + rch.mux)); + } + } + + if (channels_.contains(rch.name)) { + return absl::FailedPreconditionError( + absl::StrFormat("duplicate recovered virtual channel %s", rch.name)); + } + absl::StatusOr> recovered_vchan = + mux->CreateVirtualChannel(*this, rch.name, rch.vchan_id); + if (!recovered_vchan.ok()) { + return recovered_vchan.status(); + } + VirtualChannel *vchan = recovered_vchan->get(); + if (absl::Status status = configure_channel(vchan, rch, false); + !status.ok()) { + return status; + } + restore_users(vchan, rch); + for (const auto &entry : vchan->GetUsers()) { + mux->AddUserId(entry.first); + } + channels_.emplace(rch.name, std::move(*recovered_vchan)); + logger_.Log(toolbelt::LogLevel::kInfo, + "Recovered virtual channel '%s' on mux '%s' " + "(%d pubs, %d subs)", + rch.name.c_str(), rch.mux.c_str(), + static_cast(rch.publishers.size()), + static_cast(rch.subscribers.size())); + } + for (auto &[name, channel] : channels_) { + (void)name; + if (!channel->IsVirtual()) { + if (absl::Status status = channel->ReconcileSubscriberQueueArena(); + !status.ok()) { + return status; + } + } + } return absl::OkStatus(); } @@ -1932,6 +2098,8 @@ void Server::BridgeTransmitterCoroutine(async::Context ctx, subscribed.set_channel_name(channel_name); subscribed.set_slot_size(info.slot_size); subscribed.set_num_slots(info.num_slots); + subscribed.set_subscriber_queue_arena_size( + info.subscriber_queue_arena_size); subscribed.set_reliable(pub_reliable); subscribed.set_checksum_size(info.checksum_size); subscribed.set_metadata_size(info.metadata_size); @@ -2419,6 +2587,8 @@ void Server::BridgeReceiverCoroutine(async::Context ctx, absl::StatusOr pub = client.CreatePublisher( channel_name, subscribed.slot_size(), subscribed.num_slots(), PublisherOptions() + .SetSubscriberQueueArenaSize( + subscribed.subscriber_queue_arena_size()) .SetReliable(subscribed.reliable()) .SetBridge(true) .SetNotifyRetirement(subscribed.notify_retirement()) @@ -2750,6 +2920,7 @@ void Server::IncomingSubscribe(const Discovery::Subscribe &subscribe, .channel_name = ch->Name(), .slot_size = ch->SlotSize(), .num_slots = ch->NumSlots(), + .subscriber_queue_arena_size = ch->SubscriberQueueArenaSize(), .checksum_size = ch->ChecksumSize(), .metadata_size = ch->MetadataSize(), .wire_split_buffers = ChannelUsesSplitBuffers(ch), diff --git a/server/server.h b/server/server.h index 6341fbb6..a1c965ec 100644 --- a/server/server.h +++ b/server/server.h @@ -214,13 +214,16 @@ class Server { // num_slots will be zero. absl::StatusOr CreateChannel(const std::string &channel_name, int slot_size, int num_slots, + uint64_t subscriber_queue_arena_size, const std::string &mux, int vchan_id, std::string type); absl::StatusOr CreateMultiplexer(const std::string &channel_name, int slot_size, - int num_slots, std::string type); + int num_slots, uint64_t subscriber_queue_arena_size, + std::string type); absl::Status RemapChannel(ServerChannel *channel, int slot_size, - int num_slots); + int num_slots, + uint64_t subscriber_queue_arena_size); ServerChannel *FindChannel(const std::string &channel_name); void RemoveChannel(ServerChannel *channel); @@ -298,6 +301,7 @@ class Server { std::string channel_name; int slot_size = 0; int num_slots = 0; + uint64_t subscriber_queue_arena_size = 0; int32_t checksum_size = 0; int32_t metadata_size = 0; bool wire_split_buffers = false; diff --git a/server/server_channel.cc b/server/server_channel.cc index 3b119b42..bcb834fd 100644 --- a/server/server_channel.cc +++ b/server/server_channel.cc @@ -6,6 +6,8 @@ #include "absl/strings/str_format.h" #include "server/client_handler.h" #include "server/server.h" +#include +#include #include #include #if SUBSPACE_SHMEM_MODE == SUBSPACE_SHMEM_MODE_MEMFD @@ -19,6 +21,18 @@ #endif namespace subspace { +namespace { + +bool ProcessDefinitelyDead(uint64_t process_id) { + if (process_id == 0) { + return false; + } + errno = 0; + return kill(static_cast(process_id), 0) == -1 && errno == ESRCH; +} + +} // namespace + ServerChannel::~ServerChannel() { if (is_virtual_ || skip_cleanup_) { return; @@ -247,13 +261,15 @@ uint64_t ServerChannel::GetVirtualMemoryUsage() const { if (split_buffer_size == 0) { return Channel::GetVirtualMemoryUsage(); } - return sizeof(SystemControlBlock) + CcbSize(num_slots_) + + return sizeof(SystemControlBlock) + + CcbSize(num_slots_, subscriber_queue_arena_size_) + sizeof(BufferControlBlock) + split_buffer_size; } absl::StatusOr ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, [[maybe_unused]] int slot_size, int num_slots, + uint64_t subscriber_queue_arena_size, int initial_ordinal) { // Unmap existing memory. Unmap(); @@ -268,6 +284,10 @@ ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, } else { num_slots_ = num_slots; } + SetSubscriberQueueArenaSize(subscriber_queue_arena_size); + SetSubscriberQueueSize(subscriber_queue_arena_size == 0 + ? 0 + : kDefaultSubscriberQueueSize); // Map SCB into process memory. scb_ = reinterpret_cast(MapMemory( @@ -280,9 +300,15 @@ ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, SharedMemoryFds fds; // Create CCB in shared memory and map into process memory. - absl::StatusOr p = - CreateSharedMemory(channel_id_, "ccb", CcbSize(num_slots_), /*map=*/true, - fds.ccb, session_id_); + absl::StatusOr checked_ccb_size = + CheckedCcbSize(num_slots_, subscriber_queue_arena_size_); + if (!checked_ccb_size.ok()) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + return checked_ccb_size.status(); + } + absl::StatusOr p = CreateSharedMemory( + channel_id_, "ccb", *checked_ccb_size, + /*map=*/true, fds.ccb, session_id_); if (!p.ok()) { UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); return p.status(); @@ -295,7 +321,7 @@ ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, /*map=*/true, fds.bcb, session_id_); if (!p.ok()) { UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); - UnmapMemory(ccb_, CcbSize(num_slots_), "CCB"); + UnmapMemory(ccb_, CcbSize(num_slots_, subscriber_queue_arena_size_), "CCB"); return p.status(); } bcb_ = reinterpret_cast(*p); @@ -306,19 +332,36 @@ ServerChannel::Allocate(const toolbelt::FileDescriptor &scb_fd, // of debugging (you can see it in all processes). strncpy(ccb_->channel_name, name_.c_str(), kMaxChannelName - 1); ccb_->num_slots = num_slots_; + ccb_->subscriber_queue_size = subscriber_queue_size_; + ccb_->version = kChannelControlBlockVersion; // Initialize all ordinals. ccb_->ordinals.Init(initial_ordinal); new (&ccb_->subscribers) AtomicBitSet(); + auto *queue_index = + new (GetAvailableSlotQueueIndexAddress()) AvailableSlotQueueIndex; + queue_index->next_offset.store(0, std::memory_order_relaxed); + for (auto &offset : queue_index->offsets) { + offset.store(kInvalidSlotQueueOffset, std::memory_order_relaxed); + } + for (auto &active : queue_index->active_publishers) { + active.store(0, std::memory_order_relaxed); + } // Initialize all slots for (int32_t i = 0; i < num_slots_; i++) { MessageSlot *slot = &ccb_->slots[i]; slot->id = i; - slot->refs = 0; - slot->vchan_id = -1; - slot->buffer_index = -1; // No buffer in the free list. + slot->refs.store(0, std::memory_order_relaxed); + slot->ordinal.store(0, std::memory_order_relaxed); + slot->message_size.store(0, std::memory_order_relaxed); + slot->vchan_id.store(-1, std::memory_order_relaxed); + slot->buffer_index.store(-1, + std::memory_order_relaxed); // No buffer in the free list. + slot->timestamp.store(0, std::memory_order_relaxed); + slot->flags.store(0, std::memory_order_relaxed); + slot->bridged_slot_id.store(-1, std::memory_order_relaxed); new (&slot->sub_owners) AtomicBitSet(); } @@ -354,19 +397,82 @@ ServerChannel::MapExisting(const toolbelt::FileDescriptor &scb_fd, "Failed to map recovered SCB: %s", strerror(errno))); } + absl::StatusOr checked_ccb_size = + CheckedCcbSize(num_slots_, subscriber_queue_arena_size_); + if (!checked_ccb_size.ok()) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + return checked_ccb_size.status(); + } ccb_ = reinterpret_cast(MapMemory( - ccb_fd.Fd(), CcbSize(num_slots_), PROT_READ | PROT_WRITE, "CCB")); + ccb_fd.Fd(), *checked_ccb_size, PROT_READ | PROT_WRITE, "CCB")); if (ccb_ == MAP_FAILED) { UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); return absl::InternalError(absl::StrFormat( "Failed to map recovered CCB: %s", strerror(errno))); } + if (ccb_->version != kChannelControlBlockVersion) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + return absl::FailedPreconditionError(absl::StrFormat( + "unsupported channel control block version %u (expected %u)", + ccb_->version, kChannelControlBlockVersion)); + } + AvailableSlotQueueIndex *queue_index = GetAvailableSlotQueueIndexAddress(); + const uint64_t arena_size = SubscriberQueueArenaSize(); + const uint64_t next_offset = + queue_index->next_offset.load(std::memory_order_acquire); + if (next_offset > arena_size) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + return absl::FailedPreconditionError( + "recovered subscriber queue arena high-water mark is out of range"); + } + char *arena = EndOfAvailableSlotQueueIndex(); + for (uint64_t offset = 0; offset < next_offset;) { + auto *block = reinterpret_cast(arena + offset); + const uint32_t state = block->state.load(std::memory_order_acquire); + if (block->block_size < SlotQueueBlockSize(0) || + block->block_size > next_offset - offset || + state > static_cast(SlotQueueBlockState::kFree)) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + return absl::FailedPreconditionError( + "recovered subscriber queue arena contains a corrupt block"); + } + offset += block->block_size; + } + for (int sub_id = 0; sub_id < kMaxSlotOwners; ++sub_id) { + const uint64_t offset = + queue_index->offsets[sub_id].load(std::memory_order_acquire); + if (offset == kInvalidSlotQueueOffset) { + continue; + } + if (offset < SlotQueueBlockHeaderSize() || offset >= next_offset) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + return absl::FailedPreconditionError( + "recovered subscriber queue offset is out of range"); + } + auto *block = reinterpret_cast( + arena + offset - SlotQueueBlockHeaderSize()); + auto *queue = reinterpret_cast(arena + offset); + if (block->state.load(std::memory_order_acquire) != + static_cast(SlotQueueBlockState::kAllocated) || + queue->Capacity() > kDefaultMaxAvailableSlotQueueCapacity || + Aligned(SizeofSlotQueue(queue->Capacity())) > + block->block_size - SlotQueueBlockHeaderSize()) { + UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); + return absl::FailedPreconditionError( + "recovered subscriber queue metadata is inconsistent"); + } + } bcb_ = reinterpret_cast(MapMemory( bcb_fd.Fd(), sizeof(BufferControlBlock), PROT_READ | PROT_WRITE, "BCB")); if (bcb_ == MAP_FAILED) { UnmapMemory(scb_, sizeof(SystemControlBlock), "SCB"); - UnmapMemory(ccb_, CcbSize(num_slots_), "CCB"); + UnmapMemory(ccb_, *checked_ccb_size, "CCB"); return absl::InternalError(absl::StrFormat( "Failed to map recovered BCB: %s", strerror(errno))); } @@ -453,17 +559,26 @@ ServerChannel::AddPublisher(ClientHandler *handler, bool is_reliable, absl::StatusOr ServerChannel::AddSubscriber(ClientHandler *handler, bool is_reliable, bool is_bridge, bool for_tunnel, - int max_active_messages, uint64_t process_id) { + int max_active_messages, + int subscriber_queue_size, uint64_t process_id) { absl::StatusOr user_id = AllocateUserId("subscriber"); if (!user_id.ok()) { return user_id.status(); } + if (absl::Status status = + AllocateSubscriberQueue(*user_id, subscriber_queue_size); + !status.ok()) { + RemoveUserId(*user_id); + return status; + } std::unique_ptr sub = std::make_unique( handler, *user_id, is_reliable, is_bridge, for_tunnel, - max_active_messages); + max_active_messages, subscriber_queue_size); sub->SetProcessId(process_id); absl::Status status = sub->Init(); if (!status.ok()) { + RetireSubscriberQueue(*user_id); + RemoveUserId(*user_id); return status; } SubscriberUser *result = sub.get(); @@ -471,20 +586,330 @@ ServerChannel::AddSubscriber(ClientHandler *handler, bool is_reliable, return result; } -void ServerChannel::RegisterExistingSubscribers() { +absl::Status +ServerChannel::AllocateSubscriberQueue(int sub_id, + int subscriber_queue_size) { + if (IsVirtual()) { + return static_cast(this) + ->GetMux() + ->AllocateSubscriberQueue(sub_id, subscriber_queue_size); + } + if (IsPlaceholder()) { + return absl::OkStatus(); + } + const int capacity = + ResolveSubscriberQueueSize(NumSlots(), subscriber_queue_size == 0 + ? SubscriberQueueSize() + : subscriber_queue_size); + AvailableSlotQueueIndex *index = GetAvailableSlotQueueIndexAddress(); + if (capacity == 0) { + index->offsets[sub_id].store(kInvalidSlotQueueOffset, + std::memory_order_release); + return absl::OkStatus(); + } + + const size_t allocation_size = + SlotQueueBlockSize(static_cast(capacity)); + const size_t arena_size = SubscriberQueueArenaSize(); + uint64_t next_offset = + index->next_offset.load(std::memory_order_relaxed); + char *arena = EndOfAvailableSlotQueueIndex(); + auto block_at = [arena](uint64_t offset) { + return reinterpret_cast(arena + offset); + }; + auto state_of = [](SlotQueueBlockHeader *block) { + return static_cast( + block->state.load(std::memory_order_acquire)); + }; + + // Publishers that started after subscriber retirement cannot observe the + // retired subscriber bit. Once every publisher that was active at retirement + // has left its traversal, the block is safe to reuse. + for (uint64_t offset = 0; offset < next_offset;) { + SlotQueueBlockHeader *block = block_at(offset); + if (block->block_size < SlotQueueBlockSize(0) || + block->block_size > next_offset - offset) { + return absl::InternalError( + absl::StrFormat("corrupt subscriber queue arena for channel %s", + Name())); + } + if (state_of(block) == SlotQueueBlockState::kRetired) { + block->waiting_publishers.Traverse([block, index](int pub_id) { + if (index->active_publishers[pub_id].load( + std::memory_order_seq_cst) == 0) { + block->waiting_publishers.Clear(pub_id); + } + }); + if (block->waiting_publishers.IsEmpty()) { + block->state.store(static_cast(SlotQueueBlockState::kFree), + std::memory_order_release); + } + } + offset += block->block_size; + } + + // Coalesce adjacent safe free blocks to avoid permanent fragmentation when + // subscribers churn between different queue capacities. + for (uint64_t offset = 0; offset < next_offset;) { + SlotQueueBlockHeader *block = block_at(offset); + if (state_of(block) == SlotQueueBlockState::kFree) { + while (offset + block->block_size < next_offset) { + SlotQueueBlockHeader *next = block_at(offset + block->block_size); + if (state_of(next) != SlotQueueBlockState::kFree) { + break; + } + block->block_size += next->block_size; + } + } + offset += block->block_size; + } + + uint64_t block_offset = kInvalidSlotQueueOffset; + uint64_t best_size = std::numeric_limits::max(); + for (uint64_t offset = 0; offset < next_offset;) { + SlotQueueBlockHeader *block = block_at(offset); + if (state_of(block) == SlotQueueBlockState::kFree && + block->block_size >= allocation_size && block->block_size < best_size) { + block_offset = offset; + best_size = block->block_size; + } + offset += block->block_size; + } + + if (block_offset == kInvalidSlotQueueOffset) { + if (allocation_size > arena_size || + next_offset > arena_size - allocation_size) { + return absl::ResourceExhaustedError(absl::StrFormat( + "subscriber queue capacity %d does not fit in channel %s queue arena " + "(%zu of %zu bytes remain)", + capacity, Name(), + arena_size - std::min(next_offset, arena_size), arena_size)); + } + block_offset = next_offset; + SlotQueueBlockHeader *block = + new (arena + block_offset) SlotQueueBlockHeader; + block->block_size = allocation_size; + next_offset += allocation_size; + index->next_offset.store(next_offset, std::memory_order_relaxed); + } else { + SlotQueueBlockHeader *block = block_at(block_offset); + const uint64_t remainder = block->block_size - allocation_size; + if (remainder >= SlotQueueBlockSize(0)) { + block->block_size = allocation_size; + SlotQueueBlockHeader *split = + new (arena + block_offset + allocation_size) SlotQueueBlockHeader; + split->block_size = remainder; + split->state.store( + static_cast(SlotQueueBlockState::kFree), + std::memory_order_relaxed); + } + } + + SlotQueueBlockHeader *block = block_at(block_offset); + block->waiting_publishers.ClearAll(); + block->state.store( + static_cast(SlotQueueBlockState::kAllocated), + std::memory_order_relaxed); + const uint64_t queue_offset = block_offset + SlotQueueBlockHeaderSize(); + new (arena + queue_offset) + InPlaceSlotQueue(static_cast(capacity), + /*drop_oldest=*/subscriber_queue_size != 0); + index->offsets[sub_id].store(queue_offset, std::memory_order_release); + return absl::OkStatus(); +} + +void ServerChannel::RetireSubscriberQueue(int sub_id) { + if (IsVirtual()) { + static_cast(this)->GetMux()->RetireSubscriberQueue(sub_id); + return; + } + if (IsPlaceholder()) { + return; + } + AvailableSlotQueueIndex *index = GetAvailableSlotQueueIndexAddress(); + const uint64_t queue_offset = + index->offsets[sub_id].exchange(kInvalidSlotQueueOffset, + std::memory_order_seq_cst); + if (queue_offset == kInvalidSlotQueueOffset || + queue_offset < SlotQueueBlockHeaderSize()) { + return; + } + const uint64_t block_offset = queue_offset - SlotQueueBlockHeaderSize(); + if (block_offset >= + index->next_offset.load(std::memory_order_acquire)) { + return; + } + auto *block = reinterpret_cast( + EndOfAvailableSlotQueueIndex() + block_offset); + block->waiting_publishers.ClearAll(); + for (int pub_id = 0; pub_id < kMaxSlotOwners; ++pub_id) { + if (index->active_publishers[pub_id].load(std::memory_order_seq_cst) != 0) { + block->waiting_publishers.Set(pub_id); + } + } + block->state.store( + static_cast(block->waiting_publishers.IsEmpty() + ? SlotQueueBlockState::kFree + : SlotQueueBlockState::kRetired), + std::memory_order_release); +} + +absl::Status ServerChannel::ReconcileSubscriberQueueArena() { + if (IsVirtual()) { + return static_cast(this) + ->GetMux() + ->ReconcileSubscriberQueueArena(); + } + if (IsPlaceholder()) { + return absl::OkStatus(); + } + + AvailableSlotQueueIndex *index = GetAvailableSlotQueueIndexAddress(); + const uint64_t next_offset = + index->next_offset.load(std::memory_order_acquire); + char *arena = EndOfAvailableSlotQueueIndex(); + auto block_at = [arena](uint64_t offset) { + return reinterpret_cast(arena + offset); + }; + + absl::flat_hash_map owners; + for (int sub_id = 0; sub_id < kMaxSlotOwners; ++sub_id) { + const uint64_t queue_offset = + index->offsets[sub_id].load(std::memory_order_acquire); + if (queue_offset == kInvalidSlotQueueOffset) { + continue; + } + if (!ccb_->subscribers.IsSet(sub_id)) { + RetireSubscriberQueue(sub_id); + continue; + } + const uint64_t block_offset = queue_offset - SlotQueueBlockHeaderSize(); + if (!owners.emplace(block_offset, sub_id).second) { + return absl::FailedPreconditionError(absl::StrFormat( + "subscriber queue arena for channel %s has duplicate ownership of " + "block offset %llu", + Name(), static_cast(block_offset))); + } + } + + for (uint64_t offset = 0; offset < next_offset;) { + SlotQueueBlockHeader *block = block_at(offset); + SlotQueueBlockState state = static_cast( + block->state.load(std::memory_order_acquire)); + if (state == SlotQueueBlockState::kAllocated && + !owners.contains(offset)) { + // A crash may occur after constructing a block but before publishing its + // subscriber offset, or after clearing the offset but before retirement. + // Retire conservatively against every publisher that may still hold an + // arena pointer. + block->waiting_publishers.ClearAll(); + for (int pub_id = 0; pub_id < kMaxSlotOwners; ++pub_id) { + if (index->active_publishers[pub_id].load( + std::memory_order_seq_cst) != 0) { + block->waiting_publishers.Set(pub_id); + } + } + state = block->waiting_publishers.IsEmpty() + ? SlotQueueBlockState::kFree + : SlotQueueBlockState::kRetired; + block->state.store(static_cast(state), + std::memory_order_release); + } + if (state == SlotQueueBlockState::kRetired) { + block->waiting_publishers.Traverse([block, index](int pub_id) { + if (index->active_publishers[pub_id].load( + std::memory_order_seq_cst) == 0) { + block->waiting_publishers.Clear(pub_id); + } + }); + if (block->waiting_publishers.IsEmpty()) { + block->state.store(static_cast(SlotQueueBlockState::kFree), + std::memory_order_release); + } + } + offset += block->block_size; + } + + for (uint64_t offset = 0; offset < next_offset;) { + SlotQueueBlockHeader *block = block_at(offset); + if (static_cast( + block->state.load(std::memory_order_acquire)) == + SlotQueueBlockState::kFree) { + while (offset + block->block_size < next_offset) { + SlotQueueBlockHeader *next = block_at(offset + block->block_size); + if (static_cast( + next->state.load(std::memory_order_acquire)) != + SlotQueueBlockState::kFree) { + break; + } + block->block_size += next->block_size; + } + } + offset += block->block_size; + } + return absl::OkStatus(); +} + +void ServerChannel::ClearPublisherQueueHazardIfDead(int publisher_id, + uint64_t process_id) { + if (!ProcessDefinitelyDead(process_id) || IsPlaceholder()) { + return; + } + ServerChannel *storage_channel = + IsVirtual() + ? static_cast( + static_cast(this)->GetMux()) + : this; + storage_channel->GetAvailableSlotQueueIndexAddress() + ->active_publishers[publisher_id] + .store(0, std::memory_order_seq_cst); +} + +void ServerChannel::CleanupSlots(int owner, bool reliable, bool is_pub, + int vchan_id) { + if (!is_pub) { + ccb_->subscribers.ClearSeqCst(owner); + } + Channel::CleanupSlots(owner, reliable, is_pub, vchan_id); + if (!is_pub) { + RetireSubscriberQueue(owner); + } +} + +std::vector ServerChannel::RegisterExistingSubscribers() { + std::vector warnings; for (auto &[id, user] : users_) { if (user == nullptr || !user->IsSubscriber()) { continue; } + auto *sub = static_cast(user.get()); + if (absl::Status status = + AllocateSubscriberQueue(id, sub->SubscriberQueueSize()); + !status.ok()) { + // The arena was provisioned from the publisher default and should fit + // every default-sized subscriber. A pre-publisher override can still + // exceed that budget, so leave that subscriber on the bitset path. + RetireSubscriberQueue(id); + warnings.push_back(absl::StrFormat( + "Subscriber %d on channel %s requested queue capacity %d but the " + "publisher-provisioned arena cannot fit it; using the bitset path: %s", + id, Name(), sub->SubscriberQueueSize(), status.ToString())); + } RegisterSubscriber(id, GetVirtualChannelId(), /*is_new=*/true); } + return warnings; } -void ChannelMultiplexer::RegisterExistingSubscribers() { - ServerChannel::RegisterExistingSubscribers(); +std::vector ChannelMultiplexer::RegisterExistingSubscribers() { + std::vector warnings = + ServerChannel::RegisterExistingSubscribers(); for (VirtualChannel *vchan : virtual_channels_) { - vchan->RegisterExistingSubscribers(); + std::vector vchan_warnings = + vchan->RegisterExistingSubscribers(); + warnings.insert(warnings.end(), vchan_warnings.begin(), + vchan_warnings.end()); } + return warnings; } void ServerChannel::TriggerAllSubscribers() { @@ -528,6 +953,18 @@ void ServerChannel::RemoveUser(Server *server, int user_id) { } CleanupSlots(user->GetId(), user->IsReliable(), user->IsPublisher(), GetVirtualChannelId()); + if (user->IsPublisher() && !IsPlaceholder()) { + ServerChannel *storage_channel = + IsVirtual() + ? static_cast( + static_cast(this)->GetMux()) + : this; + // RemoveUser is an explicit client request serialized with publication, so + // no local SubscriberQueuePublishGuard can still be live. + storage_channel->GetAvailableSlotQueueIndexAddress() + ->active_publishers[user->GetId()] + .store(0, std::memory_order_seq_cst); + } RemoveUserId(user->GetId()); RecordUpdate(user->IsPublisher(), /*add=*/false, user->IsReliable()); if (user->IsPublisher()) { @@ -548,6 +985,17 @@ void ServerChannel::RemoveAllUsersFor(ClientHandler *handler) { if (user->GetHandler() == handler) { CleanupSlots(user->GetId(), user->IsReliable(), user->IsPublisher(), GetVirtualChannelId()); + if (user->IsPublisher() && !IsPlaceholder() && + ProcessDefinitelyDead(user->ProcessId())) { + ServerChannel *storage_channel = + IsVirtual() + ? static_cast( + static_cast(this)->GetMux()) + : this; + storage_channel->GetAvailableSlotQueueIndexAddress() + ->active_publishers[user->GetId()] + .store(0, std::memory_order_seq_cst); + } RemoveUserId(user->GetId()); RecordUpdate(user->IsPublisher(), /*add=*/false, user->IsReliable()); if (user->IsPublisher()) { @@ -710,19 +1158,52 @@ ServerChannel::HasSufficientCapacity(int new_max_active_messages) const { } absl::Status ServerChannel::CapacityError(const CapacityInfo &info) const { - return absl::InternalError(absl::StrFormat( + std::string message = absl::StrFormat( "there are %d slots with %d publisher%s and %d " "subscriber%s with %d additional active message%s; you " "need at least %d slots", NumSlots(), info.num_pubs, (info.num_pubs == 1 ? "" : "s"), info.num_subs, (info.num_subs == 1 ? "" : "s"), info.max_active_messages, - (info.max_active_messages == 1 ? "" : "s"), info.slots_needed + 1)); + (info.max_active_messages == 1 ? "" : "s"), info.slots_needed + 1); + + auto append_users = [this, &message](bool publishers) { + message += publishers ? "; publishers=[" : "; subscribers=["; + bool first = true; + for (int id = 0; id < kMaxUsers; ++id) { + auto it = users_.find(id); + if (it == users_.end() || it->second == nullptr || + it->second->IsPublisher() != publishers) { + continue; + } + const User &user = *it->second; + const ClientHandler *handler = user.GetHandler(); + const std::string client_name = + handler == nullptr ? std::string("") + : handler->ClientName(); + message += absl::StrFormat( + "%s{pid=%llu, client=\"%s\"", first ? "" : ", ", + static_cast(user.ProcessId()), client_name); + if (user.IsSubscriber()) { + const auto &subscriber = static_cast(user); + message += absl::StrFormat( + ", max_active_messages=%d", subscriber.MaxActiveMessages()); + } + message += "}"; + first = false; + } + message += "]"; + }; + append_users(/*publishers=*/true); + append_users(/*publishers=*/false); + return absl::InternalError(message); } void ServerChannel::GetChannelInfo(subspace::ChannelInfoProto *info) { info->set_name(Name()); info->set_slot_size(SlotSize()); info->set_num_slots(NumSlots()); + info->set_subscriber_queue_size(SubscriberQueueSize()); + info->set_subscriber_queue_arena_size(SubscriberQueueArenaSize()); info->set_type(Type()); info->set_channel_id(GetChannelId()); diff --git a/server/server_channel.h b/server/server_channel.h index e9fdbf59..1484ecd5 100644 --- a/server/server_channel.h +++ b/server/server_channel.h @@ -86,14 +86,19 @@ class User { class SubscriberUser : public User { public: SubscriberUser(ClientHandler *handler, int id, bool is_reliable, - bool is_bridge, bool for_tunnel, int max_active_messages) + bool is_bridge, bool for_tunnel, int max_active_messages, + int subscriber_queue_size) : User(handler, id, is_reliable, is_bridge, for_tunnel), - max_active_messages_(max_active_messages) {} + max_active_messages_(max_active_messages), + subscriber_queue_size_(subscriber_queue_size) {} bool IsSubscriber() const override { return true; } int MaxActiveMessages() const { return max_active_messages_; } + int SubscriberQueueSize() const { return subscriber_queue_size_; } private: int max_active_messages_; + // Requested capacity. Zero means use the publisher's channel default. + int subscriber_queue_size_; }; class PublisherUser : public User { @@ -205,9 +210,12 @@ struct ClientBufferSlotKey { class ServerChannel : public Channel { public: ServerChannel(int id, const std::string &name, int num_slots, - std::string type, bool is_virtual, int session_id) - : Channel(name, num_slots, id, std::move(type)), is_virtual_(is_virtual), - session_id_(session_id) {} + int subscriber_queue_size, + uint64_t subscriber_queue_arena_size, std::string type, + bool is_virtual, int session_id) + : Channel(name, num_slots, id, subscriber_queue_size, + subscriber_queue_arena_size, std::move(type)), + is_virtual_(is_virtual), session_id_(session_id) {} virtual ~ServerChannel(); @@ -220,8 +228,12 @@ class ServerChannel : public Channel { uint64_t process_id); absl::StatusOr AddSubscriber(ClientHandler *handler, bool is_reliable, bool is_bridge, - bool for_tunnel, int max_active_messages, uint64_t process_id); - virtual void RegisterExistingSubscribers(); + bool for_tunnel, int max_active_messages, + int subscriber_queue_size, uint64_t process_id); + virtual std::vector RegisterExistingSubscribers(); + absl::Status ReconcileSubscriberQueueArena(); + void ClearPublisherQueueHazardIfDead(int publisher_id, + uint64_t process_id); virtual std::string Type() const { return Channel::Type(); } virtual void SetType(const std::string &type) { Channel::SetType(type); } @@ -305,9 +317,7 @@ class ServerChannel : public Channel { virtual int NumSlots() const { return Channel::NumSlots(); } virtual void CleanupSlots(int owner, bool reliable, bool is_pub, - int vchan_id) { - Channel::CleanupSlots(owner, reliable, is_pub, vchan_id); - } + int vchan_id); virtual void RemoveBuffer(uint64_t session_id, Server *server = nullptr); @@ -437,7 +447,7 @@ class ServerChannel : public Channel { // this channel. This is only used in the server. virtual absl::StatusOr Allocate(const toolbelt::FileDescriptor &scb_fd, int slot_size, int num_slots, - int initial_ordinal); + uint64_t subscriber_queue_arena_size, int initial_ordinal); // Map existing shared memory from recovered FDs (after a server crash). // Does not initialize CCB/BCB -- they already contain valid data. @@ -464,6 +474,10 @@ class ServerChannel : public Channel { } protected: + absl::Status AllocateSubscriberQueue(int sub_id, + int subscriber_queue_size); + void RetireSubscriberQueue(int sub_id); + absl::flat_hash_map> users_; toolbelt::BitSet user_ids_; absl::flat_hash_map bridged_publishers_; @@ -487,8 +501,11 @@ class VirtualChannel; class ChannelMultiplexer : public ServerChannel { public: ChannelMultiplexer(int id, const std::string &name, int num_slots, - std::string type, int session_id) - : ServerChannel(id, name, num_slots, type, false, session_id) {} + int subscriber_queue_size, + uint64_t subscriber_queue_arena_size, std::string type, + int session_id) + : ServerChannel(id, name, num_slots, subscriber_queue_size, + subscriber_queue_arena_size, type, false, session_id) {} absl::StatusOr> CreateVirtualChannel(Server &server, const std::string &name, int vchan_id); @@ -497,7 +514,7 @@ class ChannelMultiplexer : public ServerChannel { bool IsMux() const override { return true; } bool HasPublisherOwnedBy(const ClientHandler *handler) const override; - void RegisterExistingSubscribers() override; + std::vector RegisterExistingSubscribers() override; bool IsEmpty() const override { return virtual_channels_.empty() && ServerChannel::IsEmpty(); } @@ -529,8 +546,9 @@ class VirtualChannel : public ServerChannel { public: VirtualChannel(ChannelMultiplexer *mux, int vchan_id, const std::string &name, int num_slots, std::string type, int session_id) - : ServerChannel(mux->GetChannelId(), name, num_slots, type, true, - session_id), + : ServerChannel(mux->GetChannelId(), name, num_slots, + mux->SubscriberQueueSize(), + mux->SubscriberQueueArenaSize(), type, true, session_id), mux_(mux), vchan_id_(vchan_id) {} std::string Type() const override { return mux_->Type(); } @@ -559,6 +577,19 @@ class VirtualChannel : public ServerChannel { int GetVirtualChannelId() const override { return vchan_id_; } bool IsPlaceholder() const override { return mux_->IsPlaceholder(); } + int SubscriberQueueSize() const override { return mux_->SubscriberQueueSize(); } + int SubscriberQueueSize(int sub_id) const override { + return mux_->SubscriberQueueSize(sub_id); + } + void SetSubscriberQueueSize(int n) override { + mux_->SetSubscriberQueueSize(n); + } + uint64_t SubscriberQueueArenaSize() const override { + return mux_->SubscriberQueueArenaSize(); + } + void SetSubscriberQueueArenaSize(uint64_t size) override { + mux_->SetSubscriberQueueArenaSize(size); + } const SharedMemoryFds &GetFds() override { return mux_->GetFds(); } diff --git a/server/server_test.cc b/server/server_test.cc index 3977d6cf..21055635 100644 --- a/server/server_test.cc +++ b/server/server_test.cc @@ -14,6 +14,7 @@ #include "toolbelt/fd.h" #include "toolbelt/sockets.h" #include +#include #include #include @@ -155,7 +156,8 @@ class RawConnection { bool fixed_size = false, const std::string &mux = "", int vchan_id = 0, bool for_tunnel = false, bool notify_retirement = false, int checksum_size = 0, - int metadata_size = 0, int max_publishers = 0) { + int metadata_size = 0, int max_publishers = 0, + uint64_t subscriber_queue_arena_size = 0) { subspace::Request req; auto *cmd = req.mutable_create_publisher(); cmd->set_channel_name(channel); @@ -172,6 +174,7 @@ class RawConnection { cmd->set_checksum_size(checksum_size); cmd->set_metadata_size(metadata_size); cmd->set_max_publishers(max_publishers); + cmd->set_subscriber_queue_arena_size(subscriber_queue_arena_size); cmd->set_publisher_id(-1); auto result = Send(req); return std::move(*result); @@ -182,7 +185,8 @@ class RawConnection { CreateSubscriber(const std::string &channel, const std::string &type = "", bool reliable = false, int max_active_messages = 4, const std::string &mux = "", - int vchan_id = 0, bool for_tunnel = false) { + int vchan_id = 0, bool for_tunnel = false, + int subscriber_queue_size = 0) { subspace::Request req; auto *cmd = req.mutable_create_subscriber(); cmd->set_channel_name(channel); @@ -193,6 +197,7 @@ class RawConnection { cmd->set_mux(mux); cmd->set_vchan_id(vchan_id); cmd->set_for_tunnel(for_tunnel); + cmd->set_subscriber_queue_size(subscriber_queue_size); auto result = Send(req); return std::move(*result); } @@ -325,6 +330,74 @@ TEST_F(ServerTest, PubNumSlotsIncrease) { ::testing::HasSubstr("more slots")); } +TEST_F(ServerTest, PubSubscriberQueueArenaSizeMismatchFromDisabled) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + conn.CreatePublisher("queue_size_disabled_ch", 64, 4); + auto [resp, fds] = conn.CreatePublisher( + "queue_size_disabled_ch", 64, 4, "", false, true, false, "", 0, false, + false, 0, 0, 0, /*subscriber_queue_arena_size=*/8000); + EXPECT_THAT(resp.create_publisher().error(), + ::testing::HasSubstr("subscriber queue arena size")); +} + +TEST_F(ServerTest, PubSubscriberQueueArenaSizeTooLarge) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + auto [resp, fds] = conn.CreatePublisher( + "queue_size_too_large", 64, 4, "", false, true, false, "", 0, false, + false, 0, 0, 0, + /*subscriber_queue_arena_size=*/ + subspace::kMaxChannelControlBlockSize + 1); + EXPECT_THAT(resp.create_publisher().error(), + ::testing::HasSubstr("channel control block exceeds")); +} + +TEST_F(ServerTest, PubCcbSizeLimitIsEnforced) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + auto [resp, fds] = + conn.CreatePublisher("ccb_too_large", 64, + std::numeric_limits::max()); + EXPECT_THAT(resp.create_publisher().error(), + ::testing::HasSubstr("channel control block limit")); +} + +TEST_F(ServerTest, PubSubscriberQueueArenaSizeMismatchToDisabled) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + conn.CreatePublisher("queue_size_enabled_ch", 64, 4, "", false, true, false, + "", 0, false, false, 0, 0, 0, + /*subscriber_queue_arena_size=*/8000); + auto [resp, fds] = conn.CreatePublisher("queue_size_enabled_ch", 64, 4); + EXPECT_THAT(resp.create_publisher().error(), + ::testing::HasSubstr("subscriber queue arena size")); +} + +TEST_F(ServerTest, PubSubscriberQueueArenaSizeMismatchForMux) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + conn.CreatePublisher("queue_size_vchan1", 64, 4, "", false, true, false, + "/queue_size_mux", 0, false, false, 0, 0, 0, + /*subscriber_queue_arena_size=*/8000); + auto [resp, fds] = conn.CreatePublisher( + "queue_size_vchan2", 64, 4, "", false, true, false, "/queue_size_mux", + 1, false, false, 0, 0, 0, + /*subscriber_queue_arena_size=*/16000); + EXPECT_THAT(resp.create_publisher().error(), + ::testing::HasSubstr("subscriber queue arena size")); +} + TEST_F(ServerTest, PubSlotSizeIncreaseOnFixedSize) { RawConnection conn; ASSERT_OK(conn.Connect(Socket())); @@ -509,6 +582,32 @@ TEST_F(ServerTest, PubVirtualRetirementNotSupported) { // CreateSubscriber error paths // --------------------------------------------------------------------------- +TEST_F(ServerTest, SubNegativeSubscriberQueueSize) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + conn.CreatePublisher("sub_negative_queue_size", 64, 4); + auto [resp, fds] = conn.CreateSubscriber( + "sub_negative_queue_size", "", false, 4, "", 0, false, + /*subscriber_queue_size=*/-1); + EXPECT_THAT(resp.create_subscriber().error(), + ::testing::HasSubstr("subscriber_queue_size must be >= 0")); +} + +TEST_F(ServerTest, SubQueueSizeTooLarge) { + RawConnection conn; + ASSERT_OK(conn.Connect(Socket())); + ASSERT_OK(conn.Init()); + + conn.CreatePublisher("sub_queue_size_too_large", 64, 4); + auto [resp, fds] = conn.CreateSubscriber( + "sub_queue_size_too_large", "", false, 4, "", 0, false, + /*subscriber_queue_size=*/1025); + EXPECT_THAT(resp.create_subscriber().error(), + ::testing::HasSubstr("subscriber_queue_size must be <= 1024")); +} + TEST_F(ServerTest, SubTypeMismatch) { RawConnection conn; ASSERT_OK(conn.Connect(Socket())); diff --git a/server/shadow_replicator.cc b/server/shadow_replicator.cc index 5eaa35df..b4185fc4 100644 --- a/server/shadow_replicator.cc +++ b/server/shadow_replicator.cc @@ -151,6 +151,8 @@ void ShadowReplicator::SendCreateChannel(ServerChannel *channel) { msg->set_channel_id(channel->GetChannelId()); msg->set_slot_size(channel->SlotSize()); msg->set_num_slots(channel->NumSlots()); + msg->set_subscriber_queue_arena_size( + channel->SubscriberQueueArenaSize()); msg->set_type(channel->Type()); msg->set_is_local(channel->IsLocal()); msg->set_is_reliable(channel->IsReliable()); @@ -202,6 +204,7 @@ void ShadowReplicator::SendAddPublisher(const std::string &channel_name, msg->set_is_bridge(pub->IsBridge()); msg->set_for_tunnel(pub->ForTunnel()); msg->set_is_fixed_size(pub->IsFixedSize()); + msg->set_process_id(pub->ProcessId()); std::vector fds; fds.push_back(const_cast(pub)->GetPollFd()); @@ -237,6 +240,8 @@ void ShadowReplicator::SendAddSubscriber(const std::string &channel_name, msg->set_is_bridge(sub->IsBridge()); msg->set_for_tunnel(sub->ForTunnel()); msg->set_max_active_messages(sub->MaxActiveMessages()); + msg->set_subscriber_queue_size(sub->SubscriberQueueSize()); + msg->set_process_id(sub->ProcessId()); std::vector fds; fds.push_back(const_cast(sub)->GetTriggerFd()); @@ -399,6 +404,8 @@ absl::StatusOr ShadowReplicator::ReceiveStateDump() { .channel_id = msg.channel_id(), .slot_size = msg.slot_size(), .num_slots = msg.num_slots(), + .subscriber_queue_arena_size = + msg.subscriber_queue_arena_size(), .type = msg.type(), .is_local = msg.is_local(), .is_reliable = msg.is_reliable(), @@ -430,9 +437,8 @@ absl::StatusOr ShadowReplicator::ReceiveStateDump() { if (msg.has_fd() && static_cast(msg.fd_index()) < fds.size()) { fd = std::move(fds[size_t(msg.fd_index())]); } - (*ch)->client_buffers.push_back( - RegisteredClientBuffer{.metadata = std::move(metadata), - .fd = std::move(fd)}); + (*ch)->client_buffers.push_back(RegisteredClientBuffer{ + .metadata = std::move(metadata), .fd = std::move(fd)}); continue; } @@ -454,6 +460,7 @@ absl::StatusOr ShadowReplicator::ReceiveStateDump() { .for_tunnel = msg.for_tunnel(), .is_fixed_size = msg.is_fixed_size(), .notify_retirement = msg.notify_retirement(), + .process_id = msg.process_id(), .poll_fd = std::move(fds[0]), .trigger_fd = std::move(fds[1]), .retirement_read_fd = msg.notify_retirement() @@ -481,6 +488,8 @@ absl::StatusOr ShadowReplicator::ReceiveStateDump() { .is_bridge = msg.is_bridge(), .for_tunnel = msg.for_tunnel(), .max_active_messages = msg.max_active_messages(), + .subscriber_queue_size = msg.subscriber_queue_size(), + .process_id = msg.process_id(), .trigger_fd = std::move(fds[0]), .poll_fd = std::move(fds[1]), }); diff --git a/server/shadow_replicator.h b/server/shadow_replicator.h index 85b94cba..bc055ffe 100644 --- a/server/shadow_replicator.h +++ b/server/shadow_replicator.h @@ -30,6 +30,7 @@ struct RecoveredPublisher { bool for_tunnel = false; bool is_fixed_size = false; bool notify_retirement = false; + uint64_t process_id = 0; toolbelt::FileDescriptor poll_fd; toolbelt::FileDescriptor trigger_fd; toolbelt::FileDescriptor retirement_read_fd; @@ -42,6 +43,8 @@ struct RecoveredSubscriber { bool is_bridge = false; bool for_tunnel = false; int max_active_messages = 0; + int subscriber_queue_size = 0; + uint64_t process_id = 0; toolbelt::FileDescriptor trigger_fd; toolbelt::FileDescriptor poll_fd; }; @@ -51,6 +54,7 @@ struct RecoveredChannel { int channel_id = 0; int slot_size = 0; int num_slots = 0; + uint64_t subscriber_queue_arena_size = 0; std::string type; bool is_local = false; bool is_reliable = false; diff --git a/shadow/shadow.cc b/shadow/shadow.cc index 1c5f1012..2e52e3af 100644 --- a/shadow/shadow.cc +++ b/shadow/shadow.cc @@ -251,6 +251,7 @@ Shadow::HandleCreateChannel(const ShadowCreateChannel &msg, ch.channel_id = msg.channel_id(); ch.slot_size = msg.slot_size(); ch.num_slots = msg.num_slots(); + ch.subscriber_queue_arena_size = msg.subscriber_queue_arena_size(); ch.type = msg.type(); ch.is_local = msg.is_local(); ch.is_reliable = msg.is_reliable(); @@ -317,6 +318,7 @@ Shadow::HandleAddPublisher(const ShadowAddPublisher &msg, .for_tunnel = msg.for_tunnel(), .is_fixed_size = msg.is_fixed_size(), .notify_retirement = msg.notify_retirement(), + .process_id = msg.process_id(), .poll_fd = std::move(fds[0]), .trigger_fd = std::move(fds[1]), .retirement_read_fd = msg.notify_retirement() @@ -370,6 +372,8 @@ Shadow::HandleAddSubscriber(const ShadowAddSubscriber &msg, .is_bridge = msg.is_bridge(), .for_tunnel = msg.for_tunnel(), .max_active_messages = msg.max_active_messages(), + .subscriber_queue_size = msg.subscriber_queue_size(), + .process_id = msg.process_id(), .trigger_fd = std::move(fds[0]), .poll_fd = std::move(fds[1]), }; @@ -528,6 +532,8 @@ absl::Status Shadow::SendStateDump(toolbelt::UnixSocket &socket) { msg->set_channel_id(ch.channel_id); msg->set_slot_size(ch.slot_size); msg->set_num_slots(ch.num_slots); + msg->set_subscriber_queue_arena_size( + ch.subscriber_queue_arena_size); msg->set_type(ch.type); msg->set_is_local(ch.is_local); msg->set_is_reliable(ch.is_reliable); @@ -578,6 +584,7 @@ absl::Status Shadow::SendStateDump(toolbelt::UnixSocket &socket) { msg->set_for_tunnel(pub.for_tunnel); msg->set_is_fixed_size(pub.is_fixed_size); msg->set_notify_retirement(pub.notify_retirement); + msg->set_process_id(pub.process_id); std::vector fds; fds.push_back(pub.poll_fd); @@ -601,6 +608,8 @@ absl::Status Shadow::SendStateDump(toolbelt::UnixSocket &socket) { msg->set_is_bridge(sub.is_bridge); msg->set_for_tunnel(sub.for_tunnel); msg->set_max_active_messages(sub.max_active_messages); + msg->set_subscriber_queue_size(sub.subscriber_queue_size); + msg->set_process_id(sub.process_id); std::vector fds; fds.push_back(sub.trigger_fd); diff --git a/shadow/shadow.h b/shadow/shadow.h index 09dbb835..d5f458ad 100644 --- a/shadow/shadow.h +++ b/shadow/shadow.h @@ -27,6 +27,7 @@ struct ShadowPublisher { bool for_tunnel = false; bool is_fixed_size = false; bool notify_retirement = false; + uint64_t process_id = 0; toolbelt::FileDescriptor poll_fd; toolbelt::FileDescriptor trigger_fd; toolbelt::FileDescriptor retirement_read_fd; @@ -39,6 +40,8 @@ struct ShadowSubscriber { bool is_bridge = false; bool for_tunnel = false; int max_active_messages = 0; + int subscriber_queue_size = 0; + uint64_t process_id = 0; toolbelt::FileDescriptor trigger_fd; toolbelt::FileDescriptor poll_fd; }; @@ -48,6 +51,7 @@ struct ShadowChannel { int channel_id = 0; int slot_size = 0; int num_slots = 0; + uint64_t subscriber_queue_arena_size = 0; std::string type; bool is_local = false; bool is_reliable = false; diff --git a/shadow/shadow_test.cc b/shadow/shadow_test.cc index 7b7ef9d5..fbec1727 100644 --- a/shadow/shadow_test.cc +++ b/shadow/shadow_test.cc @@ -910,6 +910,74 @@ TEST_F(ShadowRecoveryTest, ServerFunctionalAfterRecovery) { StopShadow(); } +TEST_F(ShadowRecoveryTest, RecoversMuxSubscriberQueueTopology) { + signal(SIGPIPE, SIG_IGN); + + StartShadow(); + StartServer(); + + constexpr char kMux[] = "/queue_recovery/*"; + constexpr char kVchan[] = "/queue_recovery/0"; + subspace::Client pre_client; + pre_client.SetThreadSafe(true); + ASSERT_THAT(pre_client.Init(RecoveryServerSocket()), IsOk()); + + subspace::PublisherOptions pub_options; + pub_options.SetSlotSize(64) + .SetNumSlots(32) + .SetSubscriberQueueArenaSize( + subspace::kDefaultSubscriberQueueArenaSize) + .SetMux(kMux); + auto pre_pub = pre_client.CreatePublisher(kVchan, pub_options); + ASSERT_THAT(pre_pub, IsOk()); + subspace::SubscriberOptions sub_options; + sub_options.SetSubscriberQueueSize(4); + auto pre_sub = pre_client.CreateSubscriber(kMux, sub_options); + ASSERT_THAT(pre_sub, IsOk()); + + ASSERT_TRUE(WaitForShadowState([this]() { + return shadow_->WithChannels([](auto &channels) { + return channels.contains("/queue_recovery/*") && + channels.contains("/queue_recovery/0"); + }); + })); + + server_->ForEachShadow( + [](const std::unique_ptr &shadow) { + shadow->Close(); + }); + StopServer(); + StartServer(); + + subspace::ServerChannel *mux = server_->FindChannel(kMux); + subspace::ServerChannel *vchan = server_->FindChannel(kVchan); + ASSERT_NE(nullptr, mux); + ASSERT_NE(nullptr, vchan); + EXPECT_TRUE(mux->IsMux()); + EXPECT_TRUE(vchan->IsVirtual()); + + subspace::Client post_client; + post_client.SetThreadSafe(true); + ASSERT_THAT(post_client.Init(RecoveryServerSocket()), IsOk()); + auto post_pub = post_client.CreatePublisher(kVchan, pub_options); + ASSERT_THAT(post_pub, IsOk()); + auto post_sub = post_client.CreateSubscriber(kMux, sub_options); + ASSERT_THAT(post_sub, IsOk()); + EXPECT_EQ(4, post_sub->SubscriberQueueSize()); + + auto buffer = post_pub->GetMessageBuffer(); + ASSERT_THAT(buffer, IsOk()); + memcpy(*buffer, "recovered_queue", 15); + ASSERT_THAT(post_pub->PublishMessage(15), IsOk()); + auto message = post_sub->ReadMessage(subspace::ReadMode::kReadNewest); + ASSERT_THAT(message, IsOk()); + ASSERT_EQ(15, message->length); + EXPECT_EQ(0, memcmp(message->buffer, "recovered_queue", 15)); + + StopServer(); + StopShadow(); +} + TEST_F(ShadowRecoveryTest, ClientReconnectsAfterServerRestart) { signal(SIGPIPE, SIG_IGN);