feat: add noloop subcommand in client tracking (#3164)

* add noloop subcommand
dragonflydb · Jun 11, 2024 · 3f8c817 · 3f8c817
1 parent 2089760
commit 3f8c817
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 6 deletions.
diff --git a/src/server/conn_context.cc b/src/server/conn_context.cc
@@ -270,6 +270,12 @@ bool ConnectionState::ClientTracking::ShouldTrackKeys() const {
     return false;
   }
 
+  if (noloop_ == true) {
+    // Once we implement REDIRECT this should return true since noloop
+    // without it only affects the current connection
+    return false;
+  }
+
   if (option_ == NONE) {
     return true;
   }

diff --git a/src/server/conn_context.h b/src/server/conn_context.h
@@ -210,6 +210,10 @@ struct ConnectionState {
       option_ = option;
     }
 
+    void SetNoLoop(bool noloop) {
+      noloop_ = noloop;
+    }
+
     // Check if the keys should be tracked. Result adheres to the state machine described above.
     bool ShouldTrackKeys() const;
 
@@ -235,6 +239,7 @@ struct ConnectionState {
    private:
     // a flag indicating whether the client has turned on client tracking.
     bool tracking_enabled_ = false;
+    bool noloop_ = false;
     Options option_ = NONE;
     // sequence number
     size_t seq_num_ = 0;

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
@@ -1413,10 +1413,6 @@ size_t Service::DispatchManyCommands(absl::Span<CmdArgList> args_list,
   for (auto args : args_list) {
     ToUpper(&args[0]);
     const auto [cid, tail_args] = FindCmd(args);
-    // disable squashing for client commands
-    if (cid && cid->name() == "CLIENT") {
-      break;
-    }
 
     // MULTI...EXEC commands need to be collected into a single context, so squashing is not
     // possible

diff --git a/src/server/multi_command_squasher.cc b/src/server/multi_command_squasher.cc
@@ -80,7 +80,7 @@ MultiCommandSquasher::SquashResult MultiCommandSquasher::TrySquash(StoredCmd* cm
       (cmd->Cid()->opt_mask() & CO::GLOBAL_TRANS))
     return SquashResult::NOT_SQUASHED;
 
-  if (cmd->Cid()->name() == "CLIENT") {
+  if (cmd->Cid()->name() == "CLIENT" || cntx_->conn_state.tracking_info_.IsTrackingOn()) {
     return SquashResult::NOT_SQUASHED;
   }
 

diff --git a/src/server/server_family.cc b/src/server/server_family.cc
@@ -462,7 +462,7 @@ void ClientTracking(CmdArgList args, ConnectionContext* cntx) {
         "Client tracking is currently not supported for RESP2. Please use RESP3.");
 
   CmdArgParser parser{args};
-  if (!parser.HasAtLeast(1) || args.size() > 2)
+  if (!parser.HasAtLeast(1) || args.size() > 3)
     return cntx->SendError(kSyntaxErr);
 
   bool is_on = false;
@@ -474,11 +474,23 @@ void ClientTracking(CmdArgList args, ConnectionContext* cntx) {
     return cntx->SendError(kSyntaxErr);
   }
 
+  bool noloop = false;
+
   if (parser.HasNext()) {
     if (parser.Check("OPTIN").IgnoreCase()) {
       option = Tracking::OPTIN;
     } else if (parser.Check("OPTOUT").IgnoreCase()) {
       option = Tracking::OPTOUT;
+    } else if (parser.Check("NOLOOP").IgnoreCase()) {
+      noloop = true;
+    } else {
+      return cntx->SendError(kSyntaxErr);
+    }
+  }
+
+  if (parser.HasNext()) {
+    if (!noloop && parser.Check("NOLOOP").IgnoreCase()) {
+      noloop = true;
     } else {
       return cntx->SendError(kSyntaxErr);
     }
@@ -487,8 +499,10 @@ void ClientTracking(CmdArgList args, ConnectionContext* cntx) {
   if (is_on) {
     ++cntx->subscriptions;
   }
+
   cntx->conn_state.tracking_info_.SetClientTracking(is_on);
   cntx->conn_state.tracking_info_.SetOption(option);
+  cntx->conn_state.tracking_info_.SetNoLoop(noloop);
   return cntx->SendOk();
 }