core,server: Use new allocator for json type

The new stateless allocator is used for JSON type. The backing memory
resource is initialized in engine shard.

The function which parses JSON is refactored to no longer accept a
memory resource, as the allocator is bound to its thread local memory
resource.

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

Author: abhijat

src/core/compact_object.cc

@@ -1624,7 +1624,7 @@ MemoryResource* CompactObj::memory_resource() {
 
 bool CompactObj::JsonConsT::DefragIfNeeded(PageUsage* page_usage) {
   if (JsonType* old = json_ptr; ShouldDefragment(page_usage)) {
-    json_ptr = AllocateMR<JsonType>(DeepCopyJSON(old, memory_resource()));
+    json_ptr = AllocateMR<JsonType>(DeepCopyJSON(old));
     DeleteMR<JsonType>(old);
     return true;
   }

src/core/compact_object_test.cc

@@ -455,14 +455,14 @@ TEST_F(CompactObjectTest, JsonTypeTest) {
     "children":[],"spouse":null}
   )";
   std::optional<JsonType> json_option2 =
-      JsonFromString(R"({"a":{}, "b":{"a":1}, "c":{"a":1, "b":2}})", CompactObj::memory_resource());
+      ParseJsonUsingShardHeap(R"({"a":{}, "b":{"a":1}, "c":{"a":1, "b":2}})");
 
   cobj_.SetString(json_str, false);
   ASSERT_TRUE(cobj_.ObjType() == OBJ_STRING);  // we set this as a string
   JsonType* failed_json = cobj_.GetJson();
   ASSERT_TRUE(failed_json == nullptr);
   ASSERT_TRUE(cobj_.ObjType() == OBJ_STRING);
-  std::optional<JsonType> json_option = JsonFromString(json_str, CompactObj::memory_resource());
+  std::optional<JsonType> json_option = ParseJsonUsingShardHeap(json_str);
   ASSERT_TRUE(json_option.has_value());
   cobj_.SetJson(std::move(json_option.value()));
   ASSERT_TRUE(cobj_.ObjType() == OBJ_JSON);  // and now this is a JSON type
@@ -477,7 +477,7 @@ TEST_F(CompactObjectTest, JsonTypeTest) {
   ASSERT_TRUE(json != nullptr);
   ASSERT_TRUE(json->contains("b"));
   ASSERT_FALSE(json->contains("firstName"));
-  std::optional<JsonType> set_array = JsonFromString("", CompactObj::memory_resource());
+  std::optional<JsonType> set_array = ParseJsonUsingShardHeap("");
   // now set it to string again
   cobj_.SetString(R"({"a":{}, "b":{"a":1}, "c":{"a":1, "b":2}})", false);
   ASSERT_TRUE(cobj_.ObjType() == OBJ_STRING);  // we set this as a string
@@ -504,7 +504,7 @@ TEST_F(CompactObjectTest, JsonTypeWithPathTest) {
             "title" : "The Night Watch",
             "author" : "Phillips, David Atlee"
         }]})";
-  std::optional<JsonType> json_array = JsonFromString(books_json, CompactObj::memory_resource());
+  std::optional<JsonType> json_array = ParseJsonUsingShardHeap(books_json);
   ASSERT_TRUE(json_array.has_value());
   cobj_.SetJson(std::move(json_array.value()));
   ASSERT_TRUE(cobj_.ObjType() == OBJ_JSON);  // and now this is a JSON type

src/core/json/json_object.cc

@@ -54,14 +54,14 @@ std::optional<ShortLivedJSON> JsonFromString(std::string_view input) {
   return ParseWithDecoder(input, json_decoder<ShortLivedJSON>{});
 }
 
-optional<JsonType> JsonFromString(string_view input, PMR_NS::memory_resource* mr) {
-  return ParseWithDecoder(input, json_decoder<JsonType>{PMR_NS::polymorphic_allocator<char>{mr}});
+optional<JsonType> ParseJsonUsingShardHeap(string_view input) {
+  return ParseWithDecoder(input, json_decoder<JsonType>{detail::StatelessJsonAllocator<char>{}});
 }
 
-JsonType DeepCopyJSON(const JsonType* j, PMR_NS::memory_resource* mr) {
+JsonType DeepCopyJSON(const JsonType* j) {
   std::string serialized;
   j->dump(serialized);
-  auto deserialized = JsonFromString(serialized, mr);
+  auto deserialized = ParseJsonUsingShardHeap(serialized);
   DCHECK(deserialized.has_value());
   return std::move(deserialized.value());
 }

src/core/json/json_object.h

@@ -68,7 +68,8 @@ bool operator!=(const StatelessJsonAllocator<T>&, const StatelessJsonAllocator<U
 void InitTLJsonHeap(PMR_NS::memory_resource* mr);
 
 using ShortLivedJSON = jsoncons::json;
-using JsonType = jsoncons::pmr::json;
+using JsonType =
+    jsoncons::basic_json<char, jsoncons::sorted_policy, detail::StatelessJsonAllocator<char>>;
 
 // A helper type to use in template functions which are expected to work with both ShortLivedJSON
 // and JsonType
@@ -82,12 +83,12 @@ std::optional<ShortLivedJSON> JsonFromString(std::string_view input);
 
 // Parses string into JSON, using mimalloc heap for allocations. This method should only be used on
 // shards where mimalloc heap is initialized.
-std::optional<JsonType> JsonFromString(std::string_view input, PMR_NS::memory_resource* mr);
+std::optional<JsonType> ParseJsonUsingShardHeap(std::string_view input);
 
 // 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);
+JsonType DeepCopyJSON(const JsonType* j);
 
 inline auto MakeJsonPathExpr(std::string_view path, std::error_code& ec)
     -> jsoncons::jsonpath::jsonpath_expression<JsonType> {

src/core/json/jsonpath_test.cc

@@ -39,7 +39,7 @@ class TestDriver : public Driver {
 template <typename JSON> JSON ValidJson(string_view str);
 
 template <> JsonType ValidJson<JsonType>(string_view str) {
-  auto res = ::dfly::JsonFromString(str, pmr::get_default_resource());
+  auto res = ParseJsonUsingShardHeap(str);
   CHECK(res) << "Failed to parse json: " << str;
   return *res;
 }

src/core/page_usage_stats_test.cc

@@ -191,7 +191,7 @@ TEST_F(PageUsageStatsTest, JSONCons) {
   // encoding.
   std::string_view data{R"#({"data": "some", "count": 1, "checked": false})#"};
 
-  auto parsed = JsonFromString(data, &m_);
+  auto parsed = ParseJsonUsingShardHeap(data);
   EXPECT_TRUE(parsed.has_value());
 
   c_obj_.SetJson(std::move(parsed.value()));

src/server/engine_shard.cc

@@ -422,6 +422,7 @@ void EngineShard::InitThreadLocal(ProactorBase* pb) {
 
   CompactObj::InitThreadLocal(shard_->memory_resource());
   SmallString::InitThreadLocal(data_heap);
+  InitTLJsonHeap(shard_->memory_resource());
 
   shard_->shard_search_indices_ = std::make_unique<ShardDocIndices>();
 }

src/server/json_family.cc

@@ -475,7 +475,7 @@ std::optional<std::string> ConvertJsonPathToJsonPointer(string_view json_path) {
    result before invoking SetJsonSize. Note that even after calling std::move on an optional, it may
    still hold the JSON value, which can lead to incorrect memory tracking. */
 std::optional<JsonType> ShardJsonFromString(std::string_view input) {
-  return JsonFromString(input, CompactObj::memory_resource());
+  return ParseJsonUsingShardHeap(input);
 }
 
 OpStatus SetFullJson(const OpArgs& op_args, string_view key, string_view json_str) {
@@ -556,8 +556,7 @@ OpResult<bool> SetPartialJson(const OpArgs& op_args, string_view key,
     path_exists = true;
     if (!is_nx_condition) {
       value_was_set = true;
-      *val = JsonType(parsed_json.value(),
-                      std::pmr::polymorphic_allocator<char>{CompactObj::memory_resource()});
+      *val = JsonType(parsed_json.value(), detail::StatelessJsonAllocator<char>{});
     }
     return {};
   };
@@ -1526,10 +1525,18 @@ auto OpMemory(const OpArgs& op_args, string_view key, const WrappedJsonPath& jso
       ReadOnlyOperationOptions{false, CallbackResultOptions::DefaultReadOnlyOptions()});
 }
 
-// Returns json vector that represents the result of the json query.
-auto OpResp(const OpArgs& op_args, string_view key, const WrappedJsonPath& json_path) {
-  auto cb = [](const string_view&, const JsonType& val) { return val; };
-  return JsonReadOnlyOperation<JsonType>(op_args, key, json_path, std::move(cb));
+// Returns json vector that represents the result of the json query. A shard local
+// heap allocated JSON cannot be copied and then destroyed on another shard because we use stateless
+// allocators which forward all requests to thread local memory resource. So the value is first
+// copied to the std allocator-backed type ShortLivedJSON.
+OpResult<JsonCallbackResult<ShortLivedJSON>> OpResp(const OpArgs& op_args, string_view key,
+                                                    const WrappedJsonPath& json_path) {
+  auto cb = [](const string_view&, const JsonType& val) {
+    string s;
+    val.dump(s);
+    return JsonFromString(s);
+  };
+  return JsonReadOnlyOperation<ShortLivedJSON>(op_args, key, json_path, std::move(cb));
 }
 
 // Returns boolean that represents the result of the operation.

src/server/json_family_memory_test.cc

@@ -111,9 +111,8 @@ TEST_F(JsonFamilyMemoryTest, JsonConsDelTest) {
 
   size_t start = GetMemoryUsage();
 
-  auto json = dfly::JsonFromString(start_json, JsonFamilyMemoryTest::GetMemoryResource());
-  void* ptr =
-      JsonFamilyMemoryTest::GetMemoryResource()->allocate(sizeof(JsonType), alignof(JsonType));
+  auto json = ParseJsonUsingShardHeap(start_json);
+  void* ptr = 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;

src/server/rdb_load.cc

@@ -984,8 +984,7 @@ void RdbLoaderBase::OpaqueObjLoader::HandleBlob(string_view blob) {
   } else if (rdb_type_ == RDB_TYPE_JSON) {
     size_t start_size = static_cast<MiMemoryResource*>(CompactObj::memory_resource())->used();
     {
-      auto json = JsonFromString(blob, CompactObj::memory_resource());
-      if (json) {
+      if (auto json = ParseJsonUsingShardHeap(blob)) {
         pv_->SetJson(std::move(*json));
       } else {
         LOG(INFO) << "Invalid JSON string during rdb load of JSON object: " << blob;