refactor client tracking, fix atomicity, squashing and multi/exec

src/facade/dragonfly_connection.cc

@@ -883,7 +883,6 @@ void Connection::ConnectionFlow(FiberSocketBase* peer) {
   cc_->conn_closing = true;  // Signal dispatch to close.
   evc_.notify();
   phase_ = SHUTTING_DOWN;
-
   VLOG(2) << "Before dispatch_fb.join()";
   dispatch_fb_.JoinIfNeeded();
   VLOG(2) << "After dispatch_fb.join()";
@@ -1119,6 +1118,10 @@ void Connection::HandleMigrateRequest() {
     }
   }
 
+  // This triggers when a pub/sub connection both publish and subscribe to the
+  // same channel. See #3035 on github for details.
+  // DCHECK(dispatch_q_.empty());
+
   // In case we Yield()ed in Migrate() above, dispatch_fb_ might have been started.
   LaunchDispatchFiberIfNeeded();
 }
@@ -1648,16 +1651,6 @@ void Connection::RequestAsyncMigration(util::fb2::ProactorBase* dest) {
   migration_request_ = dest;
 }
 
-void Connection::SetClientTrackingSwitch(bool is_on) {
-  tracking_enabled_ = is_on;
-  if (tracking_enabled_)
-    cc_->subscriptions++;
-}
-
-bool Connection::IsTrackingOn() const {
-  return tracking_enabled_;
-}
-
 void Connection::StartTrafficLogging(string_view path) {
   OpenTrafficLogger(path);
 }

src/facade/dragonfly_connection.h

@@ -298,10 +298,6 @@ 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);
 
-  void SetClientTrackingSwitch(bool is_on);
-
-  bool IsTrackingOn() const;
-
   // Starts traffic logging in the calling thread. Must be a proactor thread.
   // Each thread creates its own log file combining requests from all the connections in
   // that thread. A noop if the thread is already logging.
@@ -450,8 +446,6 @@ class Connection : public util::Connection {
   // Per-thread queue backpressure structs.
   static thread_local QueueBackpressure tl_queue_backpressure_;
 
-  // a flag indicating whether the client has turned on client tracking.
-  bool tracking_enabled_ = false;
   bool skip_next_squashing_ = false;  // Forcefully skip next squashing
 
   // Connection migration vars, see RequestAsyncMigration() above.

src/server/CMakeLists.txt

@@ -33,8 +33,7 @@ add_library(dfly_transaction db_slice.cc malloc_stats.cc blocking_controller.cc
             common.cc journal/journal.cc journal/types.cc journal/journal_slice.cc
             server_state.cc table.cc  top_keys.cc transaction.cc tx_base.cc
             serializer_commons.cc journal/serializer.cc journal/executor.cc journal/streamer.cc
-            ${TX_LINUX_SRCS} acl/acl_log.cc slowlog.cc
-            )
+            ${TX_LINUX_SRCS} acl/acl_log.cc slowlog.cc)
 
 SET(SEARCH_FILES search/search_family.cc search/doc_index.cc search/doc_accessors.cc
     search/aggregator.cc)

src/server/conn_context.cc

@@ -265,4 +265,12 @@ void ConnectionState::ExecInfo::ClearWatched() {
   watched_existed = 0;
 }
 
+bool ConnectionState::ClientTracking::ShouldTrackKeys() const {
+  if (!IsTrackingOn()) {
+    return false;
+  }
+
+  return !optin_ || (seq_num_ == (1 + caching_seq_num_));
+}
+
 }  // namespace dfly

src/server/conn_context.h

@@ -147,6 +147,87 @@ struct ConnectionState {
 
   size_t UsedMemory() const;
 
+  // Client tracking is a per-connection state machine that adheres to the requirements
+  // of the CLIENT TRACKING command. Note that the semantics described below are enforced
+  // by the tests in server_family_test. The rules are:
+  // 1. If CLIENT TRACKING is ON then each READ command must be tracked. Invalidation
+  //    messages are sent `only once`. Subsequent changes of the same key require the
+  //    client to re-read the key in order to receive the next invalidation message.
+  // 2. CLIENT TRACKING ON OPTIN turns on optional tracking. Read commands are not
+  //    tracked unless the client issues a CLIENT CACHING YES command which conditionally
+  //    allows the tracking of the command that follows CACHING YES). For example:
+  //    >> CLIENT TRACKING ON
+  //    >> CLIENT CACHING YES
+  //    >> GET foo  <--------------------- From now foo is being tracked
+  //    However:
+  //    >> CLIENT TRACKING ON
+  //    >> CLIENT CACHING YES
+  //    >> SET foo bar
+  //    >> GET foo <--------------------- is *NOT* tracked since GET does not succeed CACHING
+  //    Also, in the context of multi transactions, CLIENT CACHING YES is *STICKY*:
+  //    >> CLIENT TRACKING ON
+  //    >> CLIENT CACHING YES
+  //    >> MULTI
+  //    >>   GET foo
+  //    >>   SET foo bar
+  //    >>   GET brother_foo
+  //    >> EXEC
+  //    From this point onwards `foo` and `get` keys are tracked. Same aplies if CACHING YES
+  //    is used within the MULTI/EXEC block.
+  //
+  // The state machine implements the above rules. We need to track:
+  // 1. If TRACKING is ON and OPTIN
+  // 2. Stickiness of CACHING as described above
+  //
+  // We introduce a monotonic counter called sequence number which we increment only:
+  // * On InvokeCmd when we are not Collecting (multi)
+  // We introduce another counter called caching_seq_num which is set to seq_num
+  // when the users sends a CLIENT CACHING YES command
+  // If seq_num == caching_seq_num + 1 then we know that we should Track().
+  class ClientTracking {
+   public:
+    // Sets to true when CLIENT TRACKING is ON
+    void SetClientTracking(bool is_on) {
+      tracking_enabled_ = is_on;
+    }
+
+    // Increment current sequence number
+    void IncrementSequenceNumber() {
+      ++seq_num_;
+    }
+
+    // Set if OPTIN subcommand is used in CLIENT TRACKING
+    void SetOptin(bool optin) {
+      optin_ = optin;
+    }
+
+    // Check if the keys should be tracked. Result adheres to the state machine described above.
+    bool ShouldTrackKeys() const;
+
+    // Check only if CLIENT TRACKING is ON
+    bool IsTrackingOn() const {
+      return tracking_enabled_;
+    }
+
+    // Called by CLIENT CACHING YES and caches the current seq_num_
+    void SetCachingSequenceNumber(bool is_multi) {
+      // We need -1 when we are in multi
+      caching_seq_num_ = is_multi && seq_num_ != 0 ? seq_num_ - 1 : seq_num_;
+    }
+
+    void ResetCachingSequenceNumber() {
+      caching_seq_num_ = 0;
+    }
+
+   private:
+    // a flag indicating whether the client has turned on client tracking.
+    bool tracking_enabled_ = false;
+    bool optin_ = false;
+    // sequence number
+    size_t seq_num_ = 0;
+    size_t caching_seq_num_ = 0;
+  };
+
  public:
   DbIndex db_index = 0;
 
@@ -161,6 +242,7 @@ struct ConnectionState {
   std::optional<SquashingInfo> squashing_info;
   std::unique_ptr<ScriptInfo> script_info;
   std::unique_ptr<SubscribeInfo> subscribe_info;
+  ClientTracking tracking_info_;
 };
 
 class ConnectionContext : public facade::ConnectionContext {
@@ -183,6 +265,7 @@ class ConnectionContext : public facade::ConnectionContext {
   // TODO: to introduce proper accessors.
   Transaction* transaction = nullptr;
   const CommandId* cid = nullptr;
+
   ConnectionState conn_state;
 
   DbIndex db_index() const {

src/server/db_slice.cc

@@ -1423,24 +1423,29 @@ void DbSlice::SendInvalidationTrackingMessage(std::string_view key) {
     return;
 
   auto it = client_tracking_map_.find(key);
-  if (it != client_tracking_map_.end()) {
-    // notify all the clients.
-    auto& client_set = it->second;
-    auto cb = [key, client_set = std::move(client_set)](unsigned idx, util::ProactorBase*) {
-      for (auto it = client_set.begin(); it != client_set.end(); ++it) {
-        if ((unsigned int)it->Thread() != idx)
-          continue;
-        facade::Connection* conn = it->Get();
-        if ((conn != nullptr) && conn->IsTrackingOn()) {
-          std::string key_str = {key.begin(), key.end()};
-          conn->SendInvalidationMessageAsync({key_str});
-        }
-      }
-    };
-    shard_set->pool()->DispatchBrief(std::move(cb));
-    // remove this key from the tracking table as the key no longer exists
-    client_tracking_map_.erase(key);
+  if (it == client_tracking_map_.end()) {
+    return;
   }
+  auto& client_set = it->second;
+  // Notify all the clients. We copy key because we dispatch briefly below and
+  // we need to preserve its lifetime
+  // TODO this key is further copied within DispatchFiber. Fix this.
+  auto cb = [key = std::string(key), client_set = std::move(client_set)](unsigned idx,
+                                                                         util::ProactorBase*) {
+    for (auto& client : client_set) {
+      if (client.IsExpired() || (client.Thread() != idx)) {
+        continue;
+      }
+      auto* conn = client.Get();
+      auto* cntx = static_cast<ConnectionContext*>(conn->cntx());
+      if (cntx && cntx->conn_state.tracking_info_.IsTrackingOn()) {
+        conn->SendInvalidationMessageAsync({key});
+      }
+    }
+  };
+  shard_set->pool()->DispatchBrief(std::move(cb));
+  // remove this key from the tracking table as the key no longer exists
+  client_tracking_map_.erase(key);
 }
 
 void DbSlice::PerformDeletion(PrimeIterator del_it, DbTable* table) {