fix(json_family): Fix memory tracking for the JSON.DEL command

fixes dragonflydb#5055

Signed-off-by: Stepan Bagritsevich <stefan@dragonflydb.io>

src/server/json_family.cc

@@ -54,34 +54,6 @@ using CI = CommandId;
 
 namespace {
 
-class JsonMemTracker {
- public:
-  JsonMemTracker() {
-    start_size_ = static_cast<MiMemoryResource*>(CompactObj::memory_resource())->used();
-  }
-
-  void SetJsonSize(PrimeValue& pv, bool is_op_set) {
-    const size_t current = static_cast<MiMemoryResource*>(CompactObj::memory_resource())->used();
-    int64_t diff = static_cast<int64_t>(current) - static_cast<int64_t>(start_size_);
-    // If the diff is 0 it means the object use the same memory as before. No action needed.
-    if (diff == 0) {
-      return;
-    }
-    // If op_set_ it means we JSON.SET or JSON.MSET was called. This is a blind update,
-    // and because the operation sets the size to 0 we also need to include the size of
-    // the pointer.
-    if (is_op_set) {
-      diff += static_cast<int64_t>(mi_usable_size(pv.GetJson()));
-    }
-    pv.SetJsonSize(diff);
-    // Under any flow we must not end up with this special value.
-    DCHECK(pv.MallocUsed() != 0);
-  }
-
- private:
-  size_t start_size_{0};
-};
-
 /* Helper class which must be initialized before any mutate operations on json.
   It will track the memory usage of the json object and update the size in the CompactObj.
   It also contains indexes updates, post update operations on the iterator. */
@@ -107,6 +79,8 @@ class JsonAutoUpdater {
   void SetJsonSize() {
     set_size_was_called_ = true;
 
+    ShrinkJsonIfNeeded();
+
     const size_t current = GetMemoryUsage();
     int64_t diff = static_cast<int64_t>(current) - static_cast<int64_t>(start_size_);
 
@@ -145,6 +119,16 @@ class JsonAutoUpdater {
     return static_cast<MiMemoryResource*>(CompactObj::memory_resource())->used();
   }
 
+  /* Shrinks the json object to fit its current size.
+     Sometimes after mutating the json object, it may have more capacity than needed.
+     This method will reduce the capacity to fit the current size. */
+  void ShrinkJsonIfNeeded() {
+    auto json = GetJson();
+    if (json->size() * 2 < json->capacity()) {
+      json->shrink_to_fit();
+    }
+  }
+
  private:
   const OpArgs& op_args_;
   string_view key_;
@@ -1060,29 +1044,24 @@ OpResult<long> OpDel(const OpArgs& op_args, string_view key, string_view path,
     return 0;
   }
 
-  PrimeValue& pv = it_res->it->second;
-  JsonType* json_val = pv.GetJson();
-
-  JsonMemTracker tracker;
-  absl::Cleanup update_size_on_exit([tracker, &pv]() mutable { tracker.SetJsonSize(pv, false); });
-
   if (json_path.HoldsJsonPath()) {
+    JsonAutoUpdater updater(op_args, key, *std::move(it_res), true);
     const json::Path& path = json_path.AsJsonPath();
     long deletions = json::MutatePath(
-        path, [](optional<string_view>, JsonType* val) { return true; }, json_val);
+        path, [](optional<string_view>, JsonType* val) { return true; }, updater.GetJson());
     return deletions;
   }
 
+  // Allocates memory for the deletion_items.
+  // So we need to initialize JsonAutoUpdater after this callback
   vector<string> deletion_items;
-  auto cb = [&](std::optional<std::string_view> path, JsonType* val) -> MutateCallbackResult<> {
+  auto cb = [&deletion_items](string_view path, const JsonType& val) -> Nothing {
-    deletion_items.emplace_back(*path);
+    deletion_items.emplace_back(path);
     return {};
   };
 
-  auto res = json_path.ExecuteMutateCallback<Nothing>(
+  auto res = json_path.ExecuteReadOnlyCallback<Nothing>(
-      json_val, std::move(cb), CallbackResultOptions::DefaultMutateOptions());
+      it_res->it->second.GetJson(), cb, CallbackResultOptions::DefaultReadOnlyOptions());
-  RETURN_ON_BAD_STATUS(res);
-
   if (deletion_items.empty()) {
     return 0;
   }
@@ -1097,13 +1076,17 @@ OpResult<long> OpDel(const OpArgs& op_args, string_view key, string_view path,
     patch.emplace_back(patch_item);
   }
 
+  JsonAutoUpdater updater(op_args, key, *std::move(it_res));
+
   std::error_code ec;
-  jsoncons::jsonpatch::apply_patch(*json_val, patch, ec);
+  jsoncons::jsonpatch::apply_patch(*updater.GetJson(), patch, ec);
   if (ec) {
     VLOG(1) << "Failed to apply patch on json with error: " << ec.message();
     return 0;
   }
 
+  updater.SetJsonSize();
+
   // SetString(op_args, key, j.as_string());
   return total_deletions;
 }

src/server/json_family_memory_test.cc

@@ -105,4 +105,108 @@ TEST_F(JsonFamilyMemoryTest, PartialSet) {
   EXPECT_THAT(resp, IntArg(start_size));
 }
 
+/* Tests how works memory usage after deleting json object in jsoncons */
+TEST_F(JsonFamilyMemoryTest, JsonConsDelTest) {
+  std::string_view start_json = R"({"a":"some text", "b":" "})";
+
+  size_t start = GetMemoryUsage();
+
+  auto json = dfly::JsonFromString(start_json, JsonFamilyMemoryTest::GetMemoryResource());
+  void* ptr =
+      JsonFamilyMemoryTest::GetMemoryResource()->allocate(sizeof(JsonType), alignof(JsonType));
+  JsonType* json_on_heap = new (ptr) JsonType(std::move(json).value());
+
+  size_t memory_usage_before_erase = GetMemoryUsage() - start;
+
+  json_on_heap->erase("a");
+  /* To deallocate memory we should use shrink_to_fit */
+  json_on_heap->shrink_to_fit();
+
+  size_t memory_usage_after_erase = GetMemoryUsage() - start;
+
+  EXPECT_GT(memory_usage_before_erase, memory_usage_after_erase);
+  EXPECT_EQ(memory_usage_after_erase, GetJsonMemoryUsageFromString(R"({"b":" "})"));
+}
+
+TEST_F(JsonFamilyMemoryTest, SimpleDel) {
+  std::string_view start_json = R"({"a":"some text", "b":" "})";
+  size_t start_size = GetJsonMemoryUsageFromString(start_json);
+
+  auto resp = Run({"JSON.SET", "j1", "$", start_json});
+  EXPECT_EQ(resp, "OK");
+  resp = GetJsonMemoryUsageFromDb("j1");
+  EXPECT_THAT(resp, IntArg(start_size));
+
+  std::string_view json_after_del = R"({"b":" "})";
+  size_t size_after_del = GetJsonMemoryUsageFromString(json_after_del);
+
+  // Test that raw memory usage is correct
+  resp = Run({"JSON.SET", "j2", "$", json_after_del});
+  EXPECT_EQ(resp, "OK");
+  resp = GetJsonMemoryUsageFromDb("j2");
+  EXPECT_THAT(resp, IntArg(size_after_del));
+
+  // Test that after deletion memory usage is correct
+  resp = Run({"JSON.DEL", "j1", "$.a"});
+  EXPECT_THAT(resp, IntArg(1));
+  resp = Run({"JSON.GET", "j1"});
+  EXPECT_EQ(resp, json_after_del);
+  resp = GetJsonMemoryUsageFromDb("j1");
+
+  /* We still expect the initial size here, because after deletion we do not call shrink_to_fit on
+     the JSON object. As a result, the memory will not be deallocated. Check
+     JsonFamilyMemoryTest::JsonConsDelTest for example. */
+  EXPECT_THAT(resp, IntArg(start_size));
+
+  // Again set start json
+  resp = Run({"JSON.SET", "j1", "$.a", "\"some text\""});
+  EXPECT_EQ(resp, "OK");
+  resp = GetJsonMemoryUsageFromDb("j1");
+  EXPECT_THAT(resp, IntArg(start_size));
+}
+
+TEST_F(JsonFamilyMemoryTest, JsonShrinking) {
+  std::string_view start_json = R"({"a":"some text","b":"some another text","c":" "})";
+  size_t start_size = GetJsonMemoryUsageFromString(start_json);
+
+  auto resp = Run({"JSON.SET", "j1", "$", start_json});
+  EXPECT_EQ(resp, "OK");
+  resp = GetJsonMemoryUsageFromDb("j1");
+  EXPECT_THAT(resp, IntArg(start_size));
+
+  std::string_view json_after_del = R"({"c":" "})";
+  size_t size_after_del = GetJsonMemoryUsageFromString(json_after_del);
+
+  // Test that raw memory usage is correct
+  resp = Run({"JSON.SET", "j2", "$", json_after_del});
+  EXPECT_EQ(resp, "OK");
+  resp = GetJsonMemoryUsageFromDb("j2");
+  EXPECT_THAT(resp, IntArg(size_after_del));
+
+  // Test that after deletion memory usage decreases
+  resp = Run({"JSON.DEL", "j1", "$.a"});
+  EXPECT_THAT(resp, IntArg(1));
+  resp = Run({"JSON.DEL", "j1", "$.b"});
+  EXPECT_THAT(resp, IntArg(1));
+  resp = Run({"JSON.GET", "j1"});
+  EXPECT_EQ(resp, json_after_del);
+  resp = GetJsonMemoryUsageFromDb("j1");
+  // Now we expect the size to be smaller, because shrink_to_fit was called
+  EXPECT_THAT(resp, IntArg(size_after_del));
+
+  // Again set start json
+  resp = Run({"JSON.SET", "j1", "$.a", "\"some text\""});
+  EXPECT_EQ(resp, "OK");
+  resp = Run({"JSON.SET", "j1", "$.b", "\"some another text\""});
+  EXPECT_EQ(resp, "OK");
+  resp = Run({"JSON.GET", "j1"});
+  EXPECT_EQ(resp, start_json);
+  resp = GetJsonMemoryUsageFromDb("j1");
+
+  // Jsoncons will allocate more memory for the new json that needed.
+  // This is totally fine, because we will not call shrink_to_fit.
+  EXPECT_THAT(resp, IntArg(368));
+  EXPECT_GT(368, start_size);
+}
+
 }  // namespace dfly

src/server/search/search_family_test.cc

@@ -2767,4 +2767,19 @@ TEST_F(SearchFamilyTest, JsonSetIndexesBug) {
   resp = Run({"FT.AGGREGATE", "index", "*", "GROUPBY", "1", "@text"});
   EXPECT_THAT(resp, IsUnordArrayWithSize(IsMap("text", "some text")));
 }
+
+TEST_F(SearchFamilyTest, JsonDelIndexesBug) {
+  auto resp = Run({"JSON.SET", "j1", "$", R"({"text":"some text"})"});
+  EXPECT_THAT(resp, "OK");
+
+  resp = Run(
+      {"FT.CREATE", "index", "ON", "json", "SCHEMA", "$.text", "AS", "text", "TEXT", "SORTABLE"});
+  EXPECT_THAT(resp, "OK");
+
+  resp = Run({"JSON.DEL", "j1", "$.text"});
+  EXPECT_THAT(resp, IntArg(1));
+
+  resp = Run({"FT.AGGREGATE", "index", "*", "GROUPBY", "1", "@text"});
+  EXPECT_THAT(resp, IsUnordArrayWithSize(IsMap("text", ArgType(RespExpr::NIL))));
+}
 }  // namespace dfly