fix: serialize subscribe message together with publish message

Fixes #924

Signed-off-by: Roman Gershman <roman@dragonflydb.io>
This commit is contained in:
Roman Gershman 2023-03-11 22:09:57 +02:00 committed by Roman Gershman
parent 2ec3d48534
commit b88fc329a9
5 changed files with 32 additions and 29 deletions

View file

@ -275,22 +275,27 @@ void Connection::DispatchOperations::operator()(const PubMsgRecord& msg) {
++stats->async_writes_cnt;
const PubMessage& pub_msg = msg.pub_msg;
string_view arr[4];
DCHECK(!rbuilder->is_sending);
rbuilder->is_sending = true;
if (pub_msg.pattern.empty()) {
DVLOG(1) << "Sending message, from channel: " << *pub_msg.channel << " " << *pub_msg.message;
arr[0] = "message";
arr[1] = *pub_msg.channel;
arr[2] = *pub_msg.message;
rbuilder->SendStringArr(absl::Span<string_view>{arr, 3});
if (pub_msg.type == PubMessage::kPublish) {
if (pub_msg.pattern.empty()) {
DVLOG(1) << "Sending message, from channel: " << *pub_msg.channel << " " << *pub_msg.message;
arr[0] = "message";
arr[1] = *pub_msg.channel;
arr[2] = *pub_msg.message;
rbuilder->SendStringArr(absl::Span<string_view>{arr, 3});
} else {
arr[0] = "pmessage";
arr[1] = pub_msg.pattern;
arr[2] = *pub_msg.channel;
arr[3] = *pub_msg.message;
rbuilder->SendStringArr(absl::Span<string_view>{arr, 4});
}
} else {
arr[0] = "pmessage";
arr[1] = pub_msg.pattern;
arr[2] = *pub_msg.channel;
arr[3] = *pub_msg.message;
rbuilder->SendStringArr(absl::Span<string_view>{arr, 4});
const char* action[2] = {"unsubscribe", "subscribe"};
rbuilder->StartArray(3);
rbuilder->SendBulkString(action[pub_msg.type == PubMessage::kSubscribe]);
rbuilder->SendBulkString(*pub_msg.channel);
rbuilder->SendLong(pub_msg.channel_cnt);
}
rbuilder->is_sending = false;
}
void Connection::DispatchOperations::operator()(Request::PipelineMsg& msg) {

View file

@ -58,6 +58,8 @@ class Connection : public util::Connection {
std::shared_ptr<std::string> channel;
std::shared_ptr<std::string> message; // ensure that this message would out live passing
// between different threads/fibers
enum Type { kSubscribe, kUnsubscribe, kPublish } type;
uint32_t channel_cnt; // relevant only for kSubscribe and kUnsubscribe
PubMessage() = default;
PubMessage(const PubMessage&) = delete;

View file

@ -82,8 +82,6 @@ class SinkReplyBuilder {
virtual void SendStored() = 0;
virtual void SendSetSkipped() = 0;
bool is_sending = false;
protected:
void Send(const iovec* v, uint32_t len);

View file

@ -135,20 +135,14 @@ void ConnectionContext::ChangeSubscription(bool to_add, bool to_reply, CmdArgLis
}
if (to_reply) {
const char* action[2] = {"unsubscribe", "subscribe"};
facade::RedisReplyBuilder* rbuilder = this->operator->();
DCHECK(!rbuilder->is_sending);
rbuilder->is_sending = true;
using PubMessage = facade::Connection::PubMessage;
for (size_t i = 0; i < result.size(); ++i) {
rbuilder->StartArray(3);
rbuilder->SendBulkString(action[to_add]);
rbuilder->SendBulkString(ArgS(args, i)); // channel
// number of subscribed channels for this connection *right after*
// we subscribe.
rbuilder->SendLong(result[i]);
PubMessage msg;
msg.type = to_add ? PubMessage::kSubscribe : PubMessage::kUnsubscribe;
msg.channel = make_shared<string>(ArgS(args, i));
msg.channel_cnt = result[i];
owner()->SendMsgVecAsync(move(msg));
}
rbuilder->is_sending = false;
}
}

View file

@ -1391,6 +1391,7 @@ void Service::Publish(CmdArgList args, ConnectionContext* cntx) {
// by DispatchOperations.
shared_ptr<string> msg_ptr = make_shared<string>(ArgS(args, 2));
shared_ptr<string> channel_ptr = make_shared<string>(channel);
using PubMessage = facade::Connection::PubMessage;
// We run publish_cb in each subscriber's thread.
auto publish_cb = [&](unsigned idx, util::ProactorBase*) mutable {
@ -1405,10 +1406,13 @@ void Service::Publish(CmdArgList args, ConnectionContext* cntx) {
facade::Connection* conn = subscriber_arr[i].conn_cntx->owner();
DCHECK(conn);
facade::Connection::PubMessage pmsg;
PubMessage pmsg;
pmsg.channel = channel_ptr;
pmsg.message = msg_ptr;
pmsg.pattern = move(subscriber.pattern);
pmsg.type = PubMessage::kPublish;
conn->SendMsgVecAsync(move(pmsg));
}
};