fix(json_family): Fix memory tracking for JSON

fixes dragonflydb#4725

Signed-off-by: Stepan Bagritsevich <[email protected]>

Author: BagritsevichStepan

src/core/compact_object.cc

@@ -55,6 +55,29 @@ size_t QlMAllocSize(quicklist* ql, bool slow) {
   return node_size + ql->count * 16;  // we account for each member 16 bytes.
 }
 
+size_t UpdateSize(size_t size, int64_t update) {
+  if (update >= 0) {
+    return size + update;
+  }
+
+  // update < 0, so we need to convert it to unsigned
+  uint64_t unsigned_update = 0;
+  if (update > std::numeric_limits<int64_t>::min()) {
+    unsigned_update = -update;
+  } else {
+    static constexpr uint64_t kMaxUpdate = std::numeric_limits<int64_t>::max();
+    unsigned_update = kMaxUpdate + 1;
+  }
+
+  if (unsigned_update > std::numeric_limits<size_t>::max()) {
+    LOG(DFATAL) << "Overflow. Update size is too large. Size: " << size << ", update: " << update;
+    return 0;
+  }
+
+  DCHECK(static_cast<uint64_t>(size) >= unsigned_update);
+  return size - unsigned_update;
+}
+
 inline void FreeObjSet(unsigned encoding, void* ptr, MemoryResource* mr) {
   switch (encoding) {
     case kEncodingStrMap2: {
@@ -925,10 +948,7 @@ void CompactObj::SetJson(JsonType&& j) {
 void CompactObj::SetJsonSize(int64_t size) {
   if (taglen_ == JSON_TAG && JsonEnconding() == kEncodingJsonCons) {
     // JSON.SET or if mem hasn't changed from a JSON op then we just update.
-    if (size < 0) {
-      DCHECK(static_cast<int64_t>(u_.json_obj.cons.bytes_used) >= size);
-    }
-    u_.json_obj.cons.bytes_used += size;
+    u_.json_obj.cons.bytes_used = UpdateSize(u_.json_obj.cons.bytes_used, size);
   }
 }
 

src/server/json_family.cc

@@ -339,6 +339,7 @@ OpResult<DbSlice::ItAndUpdater> SetJson(const OpArgs& op_args, string_view key,
 
   op_args.shard->search_indices()->RemoveDoc(key, op_args.db_cntx, res.it->second);
 
+  JsonMemTracker tracker;
   std::optional<JsonType> parsed_json = ShardJsonFromString(json_str);
   if (!parsed_json) {
     VLOG(1) << "got invalid JSON string '" << json_str << "' cannot be saved";
@@ -354,6 +355,9 @@ OpResult<DbSlice::ItAndUpdater> SetJson(const OpArgs& op_args, string_view key,
   } else {
     res.it->second.SetJson(std::move(*parsed_json));
   }
+
+  tracker.SetJsonSize(res.it->second, res.is_new);
+
   op_args.shard->search_indices()->AddDoc(key, op_args.db_cntx, res.it->second);
   return std::move(res);
 }
@@ -1329,10 +1333,8 @@ OpResult<bool> OpSet(const OpArgs& op_args, string_view key, string_view path,
       }
     }
 
-    JsonMemTracker mem_tracker;
     auto st = SetJson(op_args, key, json_str);
     RETURN_ON_BAD_STATUS(st);
-    mem_tracker.SetJsonSize(st->it->second, st->is_new);
     return true;
   }
 

tests/dragonfly/replication_test.py

@@ -2951,3 +2951,35 @@ async def test_preempt_in_atomic_section_of_heartbeat(df_factory: DflyInstanceFa
         await wait_for_replicas_state(*c_replicas)
 
     await fill_task
+
+
+async def test_bug_in_json_memory_tracking(df_factory: DflyInstanceFactory):
+    """
+    This test reproduces a bug in the JSON memory tracking.
+    """
+    master = df_factory.create(proactor_threads=2, serialization_max_chunk_size=1)
+    replicas = [df_factory.create(proactor_threads=2) for i in range(2)]
+
+    # Start instances and connect clients
+    df_factory.start_all([master] + replicas)
+    c_master = master.client()
+    c_replicas = [replica.client() for replica in replicas]
+
+    total = 100000
+    await c_master.execute_command(f"DEBUG POPULATE {total} tmp 1000 TYPE SET ELEMENTS 100")
+
+    thresehold = 25000
+    for i in range(thresehold):
+        rand = random.randint(1, 4)
+        await c_master.execute_command(f"EXPIRE tmp:{i} {rand} NX")
+
+    seeder = StaticSeeder(key_target=100_000)
+    fill_task = asyncio.create_task(seeder.run(master.client()))
+
+    for replica in c_replicas:
+        await replica.execute_command(f"REPLICAOF LOCALHOST {master.port}")
+
+    async with async_timeout.timeout(240):
+        await wait_for_replicas_state(*c_replicas)
+
+    await fill_task