From 0d5f8514b2f09f6d512fd17ae8d5a3ec57935eda Mon Sep 17 00:00:00 2001
From: lappely | Kirill Gnapovsky <82707867+poipoiPIO@users.noreply.github.com>
Date: Sat, 5 Oct 2024 17:37:22 +0300
Subject: [PATCH] fix(json): JSON.MSET handles different paths for the same key
 incorrectly (#2579)

Co-authored-by: hulk <hulk.website@gmail.com>
Co-authored-by: Twice <twice@apache.org>
---
 src/types/redis_json.cc                  | 37 +++++++++++++++++-------
 tests/gocase/unit/type/json/json_test.go |  8 ++++-
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/src/types/redis_json.cc b/src/types/redis_json.cc
index b058a2dfd6a..5120573e7ce 100644
--- a/src/types/redis_json.cc
+++ b/src/types/redis_json.cc
@@ -20,6 +20,8 @@
 
 #include "redis_json.h"
 
+#include <unordered_map>
+
 #include "json.h"
 #include "lock_manager.h"
 #include "storage/redis_metadata.h"
@@ -573,6 +575,11 @@ rocksdb::Status Json::MSet(engine::Context &ctx, const std::vector<std::string>
 
   auto batch = storage_->GetWriteBatchBase();
   WriteBatchLogData log_data(kRedisJson);
+
+  // A single JSON key may be modified multiple times in the MSET command,
+  // so we need to record them temporarily to avoid reading old values from DB.
+  std::unordered_map<std::string, std::pair<JsonValue, JsonMetadata>> dirty_keys{};
+
   auto s = batch->PutLogData(log_data.Encode());
   if (!s.ok()) return s;
 
@@ -583,18 +590,28 @@ rocksdb::Status Json::MSet(engine::Context &ctx, const std::vector<std::string>
     JsonMetadata metadata;
     JsonValue value;
 
-    if (auto s = read(ctx, ns_keys[i], &metadata, &value); s.IsNotFound()) {
-      if (paths[i] != "$") return rocksdb::Status::InvalidArgument("new objects must be created at the root");
-
-      value = *std::move(json_res);
-    } else {
-      if (!s.ok()) return s;
-
-      JsonValue new_val = *std::move(json_res);
-      auto set_res = value.Set(paths[i], std::move(new_val));
+    // If a key has been modified before, just read from memory to find the modified value.
+    if (dirty_keys.count(ns_keys[i])) {
+      value = dirty_keys[ns_keys[i]].first;
+      auto set_res = value.Set(paths[i], *std::move(json_res));
       if (!set_res) return rocksdb::Status::InvalidArgument(set_res.Msg());
+    } else {
+      if (auto s = read(ctx, ns_keys[i], &metadata, &value); s.IsNotFound()) {
+        if (paths[i] != "$") return rocksdb::Status::InvalidArgument("new objects must be created at the root");
+        value = *std::move(json_res);
+      } else {
+        if (!s.ok()) return s;
+
+        auto set_res = value.Set(paths[i], *std::move(json_res));
+        if (!set_res) return rocksdb::Status::InvalidArgument(set_res.Msg());
+      }
     }
 
+    dirty_keys[ns_keys[i]] = std::make_pair(value, metadata);
+  }
+
+  for (auto &[ns_key, updated_object] : dirty_keys) {
+    auto &[value, metadata] = updated_object;
     auto format = storage_->GetConfig()->json_storage_format;
     metadata.format = format;
 
@@ -613,7 +630,7 @@ rocksdb::Status Json::MSet(engine::Context &ctx, const std::vector<std::string>
       return rocksdb::Status::InvalidArgument("Failed to encode JSON into storage: " + res.Msg());
     }
 
-    s = batch->Put(metadata_cf_handle_, ns_keys[i], val);
+    s = batch->Put(metadata_cf_handle_, ns_key, val);
     if (!s.ok()) return s;
   }
 
diff --git a/tests/gocase/unit/type/json/json_test.go b/tests/gocase/unit/type/json/json_test.go
index b943a32a714..1b810fb73d8 100644
--- a/tests/gocase/unit/type/json/json_test.go
+++ b/tests/gocase/unit/type/json/json_test.go
@@ -651,6 +651,13 @@ func testJSON(t *testing.T, configs util.KvrocksServerConfigs) {
 		EqualJSON(t, `[4]`, rdb.Do(ctx, "JSON.GET", "a1", "$.a").Val())
 	})
 
+	t.Run("JSON.MSET multi-command", func(t *testing.T) {
+		require.NoError(t, rdb.Do(ctx, "JSON.SET", "doc", "$", `{"f1":{"a1":0},"f2":{"a2":0}}`).Err())
+		require.NoError(t, rdb.Do(ctx, "JSON.MSET", "doc", "$.f1.a1", "1", "doc", "$.f2.a2", "2").Err())
+
+		EqualJSON(t, `{"f1":{"a1":1},"f2":{"a2":2}}`, rdb.Do(ctx, "JSON.GET", "doc").Val())
+	})
+
 	t.Run("JSON.DEBUG MEMORY basics", func(t *testing.T) {
 		require.NoError(t, rdb.Do(ctx, "JSON.SET", "a", "$", `{"b":true,"x":1, "y":1.2, "z": {"x":[1,2,3], "y": null}, "v":{"x":"y"},"f":{"x":[]}}`).Err())
 		//object
@@ -712,7 +719,6 @@ func testJSON(t *testing.T, configs util.KvrocksServerConfigs) {
 		require.Equal(t, make([]interface{}, 0), rdb.Do(ctx, "JSON.RESP", "item:2", "$.a").Val())
 
 	})
-
 }
 
 func EqualJSON(t *testing.T, expected string, actual interface{}) {