feat(core): Add defrag support for json objects (#6023)

* core: Add utility to reallocate a JSON object

Taking an existing JSON object as input, it is serialized and then
deserialized, allocating a new object which is a copy of the old one,
but now exists in the page reserved for malloc.

It can be used to defragment a JSON object.

* core: Add defrag support for json objects
* core: Add tests for json defrag

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/core/compact_object.cc

@@ -723,7 +723,7 @@ void RobjWrapper::MakeInnerRoom(size_t current_cap, size_t desired, MemoryResour
 }  // namespace detail
 
 uint32_t JsonEnconding() {
-  static thread_local uint32_t json_enc =
+  thread_local uint32_t json_enc =
       absl::GetFlag(FLAGS_experimental_flat_json) ? kEncodingJsonFlat : kEncodingJsonCons;
   return json_enc;
 }
@@ -1094,6 +1094,8 @@ bool CompactObj::DefragIfNeeded(PageUsage* page_usage) {
       return false;
     case SMALL_TAG:
       return u_.small_str.DefragIfNeeded(page_usage);
+    case JSON_TAG:
+      return u_.json_obj.DefragIfNeeded(page_usage);
     case INT_TAG:
       page_usage->RecordNotRequired();
       // this is not relevant in this case
@@ -1620,6 +1622,44 @@ MemoryResource* CompactObj::memory_resource() {
   return tl.local_mr;
 }
 
+bool CompactObj::JsonConsT::DefragIfNeeded(PageUsage* page_usage) {
+  if (JsonType* old = json_ptr; ShouldDefragment(page_usage)) {
+    json_ptr = AllocateMR<JsonType>(DeepCopyJSON(old, memory_resource()));
+    DeleteMR<JsonType>(old);
+    return true;
+  }
+  return false;
+}
+
+bool CompactObj::JsonConsT::ShouldDefragment(PageUsage* page_usage) const {
+  bool should_defragment = false;
+  json_ptr->compute_memory_size([&page_usage, &should_defragment](const void* p) {
+    should_defragment |= page_usage->IsPageForObjectUnderUtilized(const_cast<void*>(p));
+    return 0;
+  });
+  return should_defragment;
+}
+
+bool CompactObj::FlatJsonT::DefragIfNeeded(PageUsage* page_usage) {
+  if (uint8_t* old = flat_ptr; page_usage->IsPageForObjectUnderUtilized(old)) {
+    const uint32_t size = json_len;
+    flat_ptr = static_cast<uint8_t*>(tl.local_mr->allocate(size, kAlignSize));
+    memcpy(flat_ptr, old, size);
+    tl.local_mr->deallocate(old, size, kAlignSize);
+    return true;
+  }
+
+  return false;
+}
+
+bool CompactObj::JsonWrapper::DefragIfNeeded(PageUsage* page_usage) {
+  if (JsonEnconding() == kEncodingJsonCons) {
+    return cons.DefragIfNeeded(page_usage);
+  }
+
+  return flat.DefragIfNeeded(page_usage);
+}
+
 constexpr std::pair<CompactObjType, std::string_view> kObjTypeToString[8] = {
     {OBJ_STRING, "string"sv},  {OBJ_LIST, "list"sv},    {OBJ_SET, "set"sv},
     {OBJ_ZSET, "zset"sv},      {OBJ_HASH, "hash"sv},    {OBJ_STREAM, "stream"sv},

src/core/compact_object.h

@@ -492,18 +492,28 @@ class CompactObj {
   struct JsonConsT {
     JsonType* json_ptr;
     size_t bytes_used;
+
+    bool DefragIfNeeded(PageUsage* page_usage);
+
+    // Computes if the contained object should be defragmented, by examining pointers within it and
+    // returning true if any of them reside in an underutilized page.
+    bool ShouldDefragment(PageUsage* page_usage) const;
   };
 
   struct FlatJsonT {
     uint32_t json_len;
     uint8_t* flat_ptr;
+
+    bool DefragIfNeeded(PageUsage* page_usage);
   };
 
   struct JsonWrapper {
     union {
       JsonConsT cons;
       FlatJsonT flat;
     };
+
+    bool DefragIfNeeded(PageUsage* page_usage);
   };
 
   // My main data structure. Union of representations.

src/core/json/json_object.cc

@@ -43,4 +43,12 @@ optional<JsonType> JsonFromString(string_view input, PMR_NS::memory_resource* mr
   return nullopt;
 }
 
+JsonType DeepCopyJSON(const JsonType* j, PMR_NS::memory_resource* mr) {
+  std::string serialized;
+  j->dump(serialized);
+  auto deserialized = JsonFromString(serialized, mr);
+  DCHECK(deserialized.has_value());
+  return std::move(deserialized.value());
+}
+
 }  // namespace dfly

src/core/json/json_object.h

@@ -26,6 +26,11 @@ using JsonType = jsoncons::pmr::json;
 // Build a json object from string. If the string is not legal json, will return nullopt
 std::optional<JsonType> JsonFromString(std::string_view input, PMR_NS::memory_resource* mr);
 
+// Deep copy a JSON object, by first serializing it to a string and then deserializing the string.
+// The operation is intended to help during defragmentation, by copying into a page reserved for
+// malloc.
+JsonType DeepCopyJSON(const JsonType* j, PMR_NS::memory_resource* mr);
+
 inline auto MakeJsonPathExpr(std::string_view path, std::error_code& ec)
     -> jsoncons::jsonpath::jsonpath_expression<JsonType> {
   return jsoncons::jsonpath::make_expression<JsonType, std::allocator<char>>(

src/core/page_usage_stats_test.cc

@@ -4,23 +4,26 @@
 
 #include "core/page_usage_stats.h"
 
+#include <absl/flags/reflection.h>
 #include <gmock/gmock-matchers.h>
 
 #include "base/gtest.h"
 #include "base/logging.h"
 #include "core/compact_object.h"
-#include "core/mi_memory_resource.h"
 #include "core/qlist.h"
 #include "core/score_map.h"
 #include "core/small_string.h"
 #include "core/sorted_map.h"
 #include "core/string_map.h"
 #include "core/string_set.h"
+#include "redis/redis_aux.h"
 
 extern "C" {
 #include "redis/zmalloc.h"
 }
 
+ABSL_DECLARE_FLAG(bool, experimental_flat_json);
+
 using namespace dfly;
 
 class PageUsageStatsTest : public ::testing::Test {
@@ -45,6 +48,8 @@ class PageUsageStatsTest : public ::testing::Test {
   }
 
   void SetUp() override {
+    CompactObj::InitThreadLocal(&m_);
+
     score_map_ = std::make_unique<ScoreMap>(&m_);
     sorted_map_ = std::make_unique<detail::SortedMap>(&m_);
     string_set_ = std::make_unique<StringSet>(&m_);
@@ -176,3 +181,34 @@ TEST_F(PageUsageStatsTest, StatCollection) {
   const CollectedPageStats::ShardUsageSummary expected{{50, 65}, {90, 85}, {99, 89}};
   EXPECT_EQ(usage.at(1), expected);
 }
+
+TEST_F(PageUsageStatsTest, JSONCons) {
+  // Because of the static encoding it is not possible to easily test the flat encoding. Once the
+  // encoding flag is set, it is not re-read. If friend class is used to access the compact object
+  // inner fields and call `DefragIfNeeded` directly on the flat variant of the union, the test will
+  // still fail. This is because freeing the compact object code path takes the wrong branch based
+  // on encoding. The flat encoding was tested manually adjusting this same test with changed
+  // encoding.
+  std::string_view data{R"#({"data": "some", "count": 1, "checked": false})#"};
+
+  auto parsed = JsonFromString(data, &m_);
+  EXPECT_TRUE(parsed.has_value());
+
+  c_obj_.SetJson(std::move(parsed.value()));
+
+  PageUsage p{CollectPageStats::YES, 0.1};
+  p.SetForceReallocate(true);
+
+  c_obj_.DefragIfNeeded(&p);
+
+  const auto stats = p.CollectedStats();
+  EXPECT_GT(stats.pages_scanned, 0);
+  EXPECT_EQ(stats.objects_skipped_not_required, 0);
+
+  EXPECT_EQ(c_obj_.ObjType(), OBJ_JSON);
+
+  auto json_obj = c_obj_.GetJson();
+  EXPECT_EQ(json_obj->at("data").as_string_view(), "some");
+  EXPECT_EQ(json_obj->at("count").as_integer<uint8_t>(), 1);
+  EXPECT_EQ(json_obj->at("checked").as_bool(), false);
+}