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

src/facade/dragonfly_connection.cc

@@ -254,8 +254,9 @@ thread_local vector<Connection::PipelineMessagePtr> Connection::pipeline_req_poo
 thread_local Connection::QueueBackpressure Connection::tl_queue_backpressure_;
 
 void Connection::QueueBackpressure::EnsureBelowLimit() {
- ec.await(
- [this] { return subscriber_bytes.load(memory_order_relaxed) <= subscriber_thread_limit; });
+ ec.await([this] {
+ return done || subscriber_bytes.load(memory_order_relaxed) <= subscriber_thread_limit;
+ });
 }
 
 struct Connection::Shutdown {
@@ -885,6 +886,8 @@ void Connection::ConnectionFlow(FiberSocketBase* peer) {
  // After the client disconnected.
  cc_->conn_closing = true; // Signal dispatch to close.
  evc_.notify();
+ queue_backpressure_->done = true;
+ queue_backpressure_->ec.notify();
  phase_ = SHUTTING_DOWN;
 
  VLOG(2) << "Before dispatch_fb.join()";
@@ -1114,7 +1117,7 @@ void Connection::HandleMigrateRequest() {
  this->Migrate(dest);
  }
 
- DCHECK(dispatch_q_.empty());
+ // DCHECK(dispatch_q_.empty());
 
  // In case we Yield()ed in Migrate() above, dispatch_fb_ might have been started.
  LaunchDispatchFiberIfNeeded();
@@ -1641,16 +1644,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

@@ -292,10 +292,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.
@@ -333,6 +329,8 @@ class Connection : public util::Connection {
 
  size_t subscriber_thread_limit = 0; // cached flag subscriber_thread_limit
  size_t pipeline_cache_limit = 0; // cached flag pipeline_cache_limit
+ // cancelation flag
+ bool done = false;
  };
 
  private:
@@ -444,8 +442,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/conn_context.cc

@@ -88,9 +88,9 @@ ConnectionContext::ConnectionContext(::io::Sink* stream, facade::Connection* own
  acl_commands = std::vector<uint64_t>(acl::NumberOfFamilies(), acl::ALL_COMMANDS);
 }
 
-ConnectionContext::ConnectionContext(const ConnectionContext* owner, Transaction* tx,
+ConnectionContext::ConnectionContext(ConnectionContext* owner, Transaction* tx,
  facade::CapturingReplyBuilder* crb)
- : facade::ConnectionContext(nullptr, nullptr), transaction{tx} {
+ : facade::ConnectionContext(nullptr, nullptr), transaction{tx}, parent_cntx_(owner) {
  acl_commands = std::vector<uint64_t>(acl::NumberOfFamilies(), acl::ALL_COMMANDS);
  if (tx) { // If we have a carrier transaction, this context is used for squashing
  DCHECK(owner);
@@ -119,6 +119,13 @@ void ConnectionContext::ChangeMonitor(bool start) {
  EnableMonitoring(start);
 }
 
+ConnectionState::ClientTracking& ConnectionContext::ClientTrackingInfo() {
+ if (parent_cntx_) {
+ return parent_cntx_->conn_state.tracking_info_;
+ }
+ return conn_state.tracking_info_;
+}
+
 vector<unsigned> ChangeSubscriptions(bool pattern, CmdArgList args, bool to_add, bool to_reply,
  ConnectionContext* conn) {
  vector<unsigned> result(to_reply ? args.size() : 0, 0);
@@ -265,4 +272,78 @@ void ConnectionState::ExecInfo::ClearWatched() {
  watched_existed = 0;
 }
 
+void ConnectionState::ClientTracking::SetClientTracking(bool is_on) {
+ tracking_enabled_ = is_on;
+}
+
+void ConnectionState::ClientTracking::TrackClientCaching() {
+ executing_command_ = true;
+}
+
+void ConnectionState::ClientTracking::UpdatePrevAndLastCommand() {
+ if (prev_command_ && multi_) {
+ return;
+ }
+ prev_command_ = std::exchange(executing_command_, false);
+}
+
+void ConnectionState::ClientTracking::SetOptin(bool optin) {
+ optin_ = optin;
+}
+
+void ConnectionState::ClientTracking::SetMulti(bool multi) {
+ multi_ = multi;
+}
+
+bool ConnectionState::ClientTracking::IsTrackingOn() const {
+ return tracking_enabled_;
+}
+
+bool ConnectionState::ClientTracking::ShouldTrackKeys() const {
+ if (!IsTrackingOn()) {
+ return false;
+ }
+
+ return !optin_ || prev_command_;
+}
+
+OpResult<void> OpTrackKeys(const OpArgs slice_args, const facade::Connection::WeakRef& conn_ref,
+ const ShardArgs& args) {
+ if (conn_ref.IsExpired()) {
+ DVLOG(2) << "Connection expired, exiting TrackKey function.";
+ return OpStatus::OK;
+ }
+
+ DVLOG(2) << "Start tracking keys for client ID: " << conn_ref.GetClientId()
+ << " with thread ID: " << conn_ref.Thread();
+
+ auto& db_slice = slice_args.shard->db_slice();
+ // TODO: There is a bug here that we track all arguments instead of tracking only keys.
+ for (auto key : args) {
+ DVLOG(2) << "Inserting client ID " << conn_ref.GetClientId()
+ << " into the tracking client set of key " << key;
+ db_slice.TrackKey(conn_ref, key);
+ }
+
+ return OpStatus::OK;
+}
+
+void ConnectionState::ClientTracking::Track(ConnectionContext* cntx, const CommandId* cid) {
+ auto& info = cntx->ClientTrackingInfo();
+ auto shards = cntx->transaction->GetActiveShards();
+ if ((cid->opt_mask() & CO::READONLY) && cid->IsTransactional() && info.ShouldTrackKeys()) {
+ if (cntx->parent_cntx_) {
+ }
+ auto conn = cntx->parent_cntx_ ? cntx->parent_cntx_->conn()->Borrow() : cntx->conn()->Borrow();
+ auto cb = [&, conn](unsigned i, auto* pb) {
+ if (shards.find(i) != shards.end()) {
+ auto* t = cntx->transaction;
+ CHECK(t);
+ auto* shard = EngineShard::tlocal();
+ OpTrackKeys(t->GetOpArgs(shard), conn, t->GetShardArgs(shard->shard_id()));
+ }
+ };
+ shard_set->pool()->AwaitFiberOnAll(std::move(cb));
+ }
+}
 } // namespace dfly

src/server/conn_context.h

@@ -147,6 +147,78 @@ 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 two commands at each time:
+ // 1. The command invoked previously.
+ // 2. The command that is invoked now (via InvokeCmd).
+ // Which is tracked by current_command_ and prev_command_ respectively. When CACHING YES
+ // is invoked the current_command_ is set to true which is later moved to the prev_command_
+ // when the next command is invoked. This is needed to keep track of the different rules
+ // described above. Stickiness is covered similarly by the multi/exec/discard command which
+ // when called sets the corresponding multi_ variable to true.
+ class ClientTracking {
+ public:
+ // Sets to true when CLIENT TRACKING is ON
+ void SetClientTracking(bool is_on);
+ // Enable tracking on the client
+ void TrackClientCaching();
+
+ void UpdatePrevAndLastCommand();
+ // Set if OPTIN subcommand is used in CLIENT TRACKING
+ void SetOptin(bool optin);
+ // When Multi command is invoked, it calls this to broadcast that we are on a multi
+ // transaction.
+ void SetMulti(bool multi);
+
+ // 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;
+
+ // Iterates over the active shards of the transaction. If a key satisfies
+ // the tracking requirements, is is set for tracking.
+ void Track(ConnectionContext* cntx, const CommandId* cid);
+
+ private:
+ // a flag indicating whether the client has turned on client tracking.
+ bool tracking_enabled_ = false;
+ bool optin_ = false;
+ // remember if CLIENT CACHING TRUE was the last command
+ // true if the previous command invoked is CLIENT CACHING TRUE
+ bool prev_command_ = false;
+ // true if the currently executing command is CLIENT CACHING TRUE
+ bool executing_command_ = false;
+ // true if we are in a multi transaction
+ bool multi_ = false;
+ };
+
  public:
  DbIndex db_index = 0;
 
@@ -161,14 +233,14 @@ 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 {
  public:
  ConnectionContext(::io::Sink* stream, facade::Connection* owner);
 
- ConnectionContext(const ConnectionContext* owner, Transaction* tx,
- facade::CapturingReplyBuilder* crb);
+ ConnectionContext(ConnectionContext* owner, Transaction* tx, facade::CapturingReplyBuilder* crb);
 
  struct DebugInfo {
  uint32_t shards_count = 0;
@@ -183,6 +255,10 @@ class ConnectionContext : public facade::ConnectionContext {
  // TODO: to introduce proper accessors.
  Transaction* transaction = nullptr;
  const CommandId* cid = nullptr;
+ ConnectionContext* parent_cntx_ = nullptr;
+
+ ConnectionState::ClientTracking& ClientTrackingInfo();
+
  ConnectionState conn_state;
 
  DbIndex db_index() const {

src/server/db_slice.cc

@@ -1391,24 +1391,27 @@ 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.
+ 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->ClientTrackingInfo().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) {