feat(server): support client side tracking (resp3 protocol only)

src/facade/dragonfly_connection.cc

@@ -150,6 +150,7 @@ struct Connection::DispatchOperations {
   void operator()(const AclUpdateMessage& msg);
   void operator()(const MigrationRequestMessage& msg);
   void operator()(CheckpointMessage msg);
+  void operator()(const InvalidationMessage& msg);
 
   template <typename T, typename D> void operator()(unique_ptr<T, D>& ptr) {
     operator()(*ptr.get());
@@ -215,6 +216,9 @@ size_t Connection::MessageHandle::UsedMemory() const {
     size_t operator()(const CheckpointMessage& msg) {
       return 0;  // no access to internal type, memory usage negligible
     }
+    size_t operator()(const InvalidationMessage& msg) {
+      return 0;
+    }
   };
 
   return sizeof(MessageHandle) + visit(MessageSize{}, this->handle);
@@ -282,6 +286,19 @@ void Connection::DispatchOperations::operator()(CheckpointMessage msg) {
   msg.bc.Dec();
 }
 
+void Connection::DispatchOperations::operator()(const InvalidationMessage& msg) {
+  RedisReplyBuilder* rbuilder = (RedisReplyBuilder*)builder;
+  rbuilder->SetResp3(true);
+  rbuilder->StartCollection(2, facade::RedisReplyBuilder::CollectionType::PUSH);
+  rbuilder->SendBulkString("invalidate");
+  if (msg.invalidate_due_to_flush) {
+    rbuilder->SendNull();
+  } else {
+    std::vector<string_view> keys{msg.key};
+    rbuilder->SendStringArr(keys);
+  }
+}
+
 Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener, SSL_CTX* ctx,
                        ServiceInterface* service)
     : io_buf_(kMinReadSize), http_listener_(http_listener), ctx_(ctx), service_(service), name_{} {
@@ -313,6 +330,10 @@ Connection::Connection(Protocol protocol, util::HttpListenerBase* http_listener,
 
   migration_enabled_ = absl::GetFlag(FLAGS_migrate_connections);
 
+  // Create shared_ptr with empty value and associate it with `this` pointer (aliasing constructor).
+  // We use it for reference counting and accessing `this` (without managing it).
+  self_ = {std::make_shared<std::monostate>(std::monostate{}), this};
+
 #ifdef DFLY_USE_SSL
   // Increment reference counter so Listener won't free the context while we're
   // still using it.
@@ -1168,11 +1189,19 @@ void Connection::Migrate(util::fb2::ProactorBase* dest) {
   // connections
   CHECK(!cc_->async_dispatch);
   CHECK_EQ(cc_->subscriptions, 0);    // are bound to thread local caches
+  CHECK_EQ(self_.use_count(), 1u);    // references cache our thread and backpressure
   CHECK(!dispatch_fb_.IsJoinable());  // can't move once it started
 
   listener()->Migrate(this, dest);
 }
 
+Connection::WeakRef Connection::Borrow() {
+  DCHECK(self_);
+  DCHECK_GT(cc_->subscriptions, 0);
+
+  return WeakRef(self_, queue_backpressure_, socket_->proactor()->GetPoolIndex());
+}
+
 void Connection::ShutdownThreadLocal() {
   pipeline_req_pool_.clear();
 }
@@ -1189,6 +1218,10 @@ void Connection::SendPubMessageAsync(PubMessage msg) {
   SendAsync({PubMessagePtr{new (ptr) PubMessage{move(msg)}, MessageDeleter{}}});
 }
 
+void Connection::SendInvalidationMessageAsync(InvalidationMessage msg) {
+  SendAsync({std::move(msg)});
+}
+
 void Connection::SendMonitorMessageAsync(string msg) {
   SendAsync({MonitorMessage{move(msg)}});
 }
@@ -1277,10 +1310,6 @@ void Connection::RecycleMessage(MessageHandle msg) {
   }
 }
 
-void Connection::EnsureAsyncMemoryBudget() {
-  queue_backpressure_->EnsureBelowLimit();
-}
-
 std::string Connection::LocalBindStr() const {
   if (socket_->IsUDS())
     return "unix-domain-socket";
@@ -1327,6 +1356,15 @@ void Connection::RequestAsyncMigration(util::fb2::ProactorBase* dest) {
   migration_request_ = dest;
 }
 
+void Connection::EnableTracking() {
+  tracking_enabled_ = true;
+  cc_->subscriptions++;
+}
+
+void Connection::DisableTracking() {
+  tracking_enabled_ = false;
+}
+
 Connection::MemoryUsage Connection::GetMemoryUsage() const {
   size_t mem = sizeof(*this) + dfly::HeapSize(dispatch_q_) + dfly::HeapSize(name_) +
                dfly::HeapSize(tmp_parse_args_) + dfly::HeapSize(tmp_cmd_vec_) +
@@ -1344,6 +1382,35 @@ Connection::MemoryUsage Connection::GetMemoryUsage() const {
   };
 }
 
+Connection::WeakRef::WeakRef(std::shared_ptr<Connection> ptr, QueueBackpressure* backpressure,
+                             unsigned thread)
+    : ptr_{ptr}, backpressure_{backpressure}, thread_{thread}, client_id_{ptr->GetClientId()} {
+}
+
+unsigned Connection::WeakRef::Thread() const {
+  return thread_;
+}
+
+Connection* Connection::WeakRef::Get() const {
+  // DCHECK_EQ(ProactorBase::me()->GetPoolIndex(), int(thread_));
+  return ptr_.lock().get();
+}
+
+uint32_t Connection::WeakRef::GetClientId() const {
+  return client_id_;
+}
+
+bool Connection::WeakRef::EnsureMemoryBudget() const {
+  // Simple optimization: If a connection was closed, don't check memory budget.
+  if (!ptr_.expired()) {
+    // We don't rely on the connection ptr staying valid because we only access
+    // the threads backpressure
+    backpressure_->EnsureBelowLimit();
+    return true;
+  }
+  return false;
+}
+
 void Connection::DecreaseStatsOnClose() {
   stats_->read_buf_capacity -= io_buf_.Capacity();
 

src/facade/dragonfly_connection.h

@@ -53,6 +53,8 @@ class SinkReplyBuilder;
 // For pipelined requests, monitor and pubsub messages it uses
 // a separate dispatch queue that is processed on a separate fiber.
 class Connection : public util::Connection {
+  struct QueueBackpressure;
+
  public:
   Connection(Protocol protocol, util::HttpListenerBase* http_listener, SSL_CTX* ctx,
              ServiceInterface* service);
@@ -75,7 +77,12 @@ class Connection : public util::Connection {
                size_t message_len);
   };
 
-  // Pipeline message, accumulated command to be executed.
+  struct InvalidationMessage {
+    std::string_view key;
+    bool invalidate_due_to_flush = false;
+  };
+
+  // Pipeline message, accumulated command to be execute.d
   struct PipelineMessage {
     PipelineMessage(size_t nargs, size_t capacity) : args(nargs), storage(capacity) {
     }
@@ -134,7 +141,7 @@ class Connection : public util::Connection {
     bool IsPubMsg() const;
 
     std::variant<MonitorMessage, PubMessagePtr, PipelineMessagePtr, AclUpdateMessagePtr,
-                 MigrationRequestMessage, CheckpointMessage>
+                 MigrationRequestMessage, CheckpointMessage, InvalidationMessage>
         handle;
   };
 
@@ -143,11 +150,47 @@ class Connection : public util::Connection {
 
   enum Phase { SETUP, READ_SOCKET, PROCESS, SHUTTING_DOWN, PRECLOSE, NUM_PHASES };
 
+  // Weak reference to a connection, invalidated upon connection close.
+  // Used to dispatch async operations for the connection without worrying about pointer lifetime.
+  struct WeakRef {
+   public:
+    // Get residing thread of connection. Thread-safe.
+    unsigned Thread() const;
+
+    // Get pointer to connection if still valid, nullptr if expired.
+    // Can only be called from connection's thread.
+    Connection* Get() const;
+
+    uint32_t GetClientId() const;
+
+    // Ensure owner thread's memory budget. If expired, skips and returns false. Thread-safe.
+    bool EnsureMemoryBudget() const;
+
+    bool operator==(const WeakRef rhs) const {
+      auto rhs_ptr = rhs.ptr_.lock();
+      auto lhs_ptr = ptr_.lock();
+      return (rhs_ptr == lhs_ptr);
+    };
+
+   private:
+    friend class Connection;
+
+    WeakRef(std::shared_ptr<Connection> ptr, QueueBackpressure* backpressure, unsigned thread);
+
+    std::weak_ptr<Connection> ptr_;
+    QueueBackpressure* backpressure_;
+    unsigned thread_;
+    uint32_t client_id_;
+  };
+
  public:
   // Add PubMessage to dispatch queue.
   // Virtual because behavior is overridden in test_utils.
   virtual void SendPubMessageAsync(PubMessage);
 
+  // Add InvalidationMessage to dispatch queue.
+  void SendInvalidationMessageAsync(InvalidationMessage);
+
   // Add monitor message to dispatch queue.
   void SendMonitorMessageAsync(std::string);
 
@@ -176,6 +219,9 @@ class Connection : public util::Connection {
   // Migrate this connecton to a different thread.
   void Migrate(util::fb2::ProactorBase* dest);
 
+  // Borrow weak reference to connection. Can be called from any thread.
+  WeakRef Borrow();
+
   static void ShutdownThreadLocal();
 
   bool IsCurrentlyDispatching() const;
@@ -220,6 +266,16 @@ class Connection : public util::Connection {
   // Connections will migrate at most once, and only when the flag --migrate_connections is true.
   void RequestAsyncMigration(util::fb2::ProactorBase* dest);
 
+  // Set the flag to enable client side tracking
+  void EnableTracking();
+
+  // Set the flag to disable client side tracking
+  void DisableTracking();
+
+  bool IsTrackingOn() const {
+    return tracking_enabled_;
+  }
+
  protected:
   void OnShutdown() override;
   void OnPreMigrateThread() override;
@@ -340,6 +396,9 @@ class Connection : public util::Connection {
   RespVec tmp_parse_args_;
   CmdArgVec tmp_cmd_vec_;
 
+  // Used to keep track of borrowed references. Does not really own itself
+  std::shared_ptr<Connection> self_;
+
   // Pointer to corresponding queue backpressure struct.
   // Needed for access from different threads by EnsureAsyncMemoryBudget().
   QueueBackpressure* queue_backpressure_;
@@ -354,6 +413,9 @@ class Connection : public util::Connection {
 
   // Per-thread queue backpressure structs.
   static thread_local QueueBackpressure tl_queue_backpressure_;
+
+  // whether client tracking is enabled
+  bool tracking_enabled_ = false;
 };
 
 }  // namespace facade

src/facade/reply_builder.cc

@@ -30,6 +30,8 @@ constexpr char kCRLF[] = "\r\n";
 constexpr char kErrPref[] = "-ERR ";
 constexpr char kSimplePref[] = "+";
 
+constexpr char kRET[] = "$1\r\n\r\r\n";
+
 constexpr unsigned kConvFlags =
     DoubleToStringConverter::UNIQUE_ZERO | DoubleToStringConverter::EMIT_POSITIVE_EXPONENT_SIGN;
 
@@ -261,6 +263,10 @@ char* RedisReplyBuilder::FormatDouble(double val, char* dest, unsigned dest_len)
 RedisReplyBuilder::RedisReplyBuilder(::io::Sink* sink) : SinkReplyBuilder(sink) {
 }
 
+bool RedisReplyBuilder::IsResp3() const {
+  return is_resp3_;
+}
+
 void RedisReplyBuilder::SetResp3(bool is_resp3) {
   is_resp3_ = is_resp3;
 }
@@ -291,7 +297,6 @@ void RedisReplyBuilder::SendProtocolError(std::string_view str) {
 
 void RedisReplyBuilder::SendSimpleString(std::string_view str) {
   iovec v[3] = {IoVec(kSimplePref), IoVec(str), IoVec(kCRLF)};
-
   Send(v, ABSL_ARRAYSIZE(v));
 }
 
@@ -331,6 +336,32 @@ void RedisReplyBuilder::SendBulkString(std::string_view str) {
   return Send(v, ABSL_ARRAYSIZE(v));
 }
 
+void RedisReplyBuilder::SendVerbatimString(std::string_view str, VerbatimFormat format) {
+  if (!is_resp3_)
+    return SendBulkString(str);
+
+  char tmp[absl::numbers_internal::kFastToBufferSize + 3];
+  tmp[0] = '=';
+  // + 4 because format is three byte, and need to be followed by a ":"
+  char* next = absl::numbers_internal::FastIntToBuffer(uint32_t(str.size() + 4), tmp + 1);
+  *next++ = '\r';
+  *next++ = '\n';
+
+  std::string_view lenpref{tmp, size_t(next - tmp)};
+
+  std::string_view format_str;
+  if (format == VerbatimFormat::TXT)
+    format_str = "txt:";
+  else if (format == VerbatimFormat::MARKDOWN)
+    format_str = "mkd:";
+  else {
+    DVLOG(1) << "Unknown verbatim reply format: " << format;
+    return;
+  }
+  iovec v[4] = {IoVec(lenpref), IoVec(format_str), IoVec(str), IoVec(kCRLF)};
+  return Send(v, ABSL_ARRAYSIZE(v));
+}
+
 void RedisReplyBuilder::SendLong(long num) {
   string str = absl::StrCat(":", num, kCRLF);
   SendRaw(str);

src/facade/reply_builder.h

@@ -212,10 +212,14 @@ class RedisReplyBuilder : public SinkReplyBuilder {
  public:
   enum CollectionType { ARRAY, SET, MAP, PUSH };
 
+  enum VerbatimFormat { TXT, MARKDOWN };
+
   using StrSpan = std::variant<absl::Span<const std::string>, absl::Span<const std::string_view>>;
 
   RedisReplyBuilder(::io::Sink* stream);
 
+  bool IsResp3() const;
+
   void SetResp3(bool is_resp3);
 
   void SendError(std::string_view str, std::string_view type = {}) override;
@@ -247,6 +251,8 @@ class RedisReplyBuilder : public SinkReplyBuilder {
 
   static char* FormatDouble(double val, char* dest, unsigned dest_len);
 
+  void SendVerbatimString(std::string_view str, VerbatimFormat format = TXT);
+
  protected:
   struct WrappedStrSpan : public StrSpan {
     size_t Size() const;

src/server/channel_store.cc

@@ -25,18 +25,12 @@ bool Matches(string_view pattern, string_view channel) {
 
 }  // namespace
 
-ChannelStore::Subscriber::Subscriber(ConnectionContext* cntx, uint32_t tid)
-    : conn_cntx(cntx), borrow_token(cntx->conn_state.subscribe_info->borrow_token), thread_id(tid) {
-}
-
-ChannelStore::Subscriber::Subscriber(uint32_t tid)
-    : conn_cntx(nullptr), borrow_token(0), thread_id(tid) {
+bool ChannelStore::Subscriber::ByThread(const Subscriber& lhs, const Subscriber& rhs) {
+  return ByThreadId(lhs, rhs.Thread());
 }
 
-bool ChannelStore::Subscriber::ByThread(const Subscriber& lhs, const Subscriber& rhs) {
-  if (lhs.thread_id == rhs.thread_id)
-    return (lhs.conn_cntx != nullptr) < (rhs.conn_cntx != nullptr);
-  return lhs.thread_id < rhs.thread_id;
+bool ChannelStore::Subscriber::ByThreadId(const Subscriber& lhs, const unsigned thread) {
+  return lhs.Thread() < thread;
 }
 
 ChannelStore::UpdatablePointer::UpdatablePointer(const UpdatablePointer& other) {
@@ -120,12 +114,8 @@ void ChannelStore::Fill(const SubscribeMap& src, const string& pattern, vector<S
   out->reserve(out->size() + src.size());
   for (const auto [cntx, thread_id] : src) {
     CHECK(cntx->conn_state.subscribe_info);
-
-    Subscriber s(cntx, thread_id);
-    s.pattern = pattern;
-    s.borrow_token.Inc();
-
-    out->push_back(std::move(s));
+    Subscriber sub{cntx->conn()->Borrow(), pattern};
+    out->push_back(std::move(sub));
   }
 }