server(JSON): Remove parse method which used default heap

There were two parse methods, one which used the default heap and the
other which used a thread local mi heap.

The former one is removed here. The reason is to prepare for using a
stateless allocator.

The allocator is part of the type of the json object. A stateless
allocator must not hold state of its own and simply forward requests to
a memory resource pointer accessed from a static storage.

In such a model, having two memory resources does not work because two
allocators using different heaps are not interchangeable.

It is possible to use two different types, but in that case the json
object types become different, and it becomes impossible to work with
them transparently.

The switch to mi heap everywhere can cause one problem during
accounting. The JSONAutoUpdater class uses memory usage before and after
parsing to figure out how many bytes are used by a given object. If the
objects parsed by the method which is removed in this PR are
interspersed with JSONAutoUpdater, then this can cause a problem with
accounting.

At present there are no such cases in the code.

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

@@ -88,6 +88,7 @@ static void InitThreadStructs() {
   SmallString::InitThreadLocal(tlh);
   thread_local MiMemoryResource mi_resource(tlh);
   CompactObj::InitThreadLocal(&mi_resource);
+  InitTLJsonHeap(&mi_resource);
 };
 
 static void CheckEverythingDeallocated() {
@@ -455,14 +456,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 +478,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 +505,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

@@ -46,18 +46,22 @@ std::optional<T> ParseWithDecoder(std::string_view input, json_decoder<T>&& deco
 
 namespace dfly {
 
+void InitTLJsonHeap(PMR_NS::memory_resource* mr) {
+  detail::tl_mr = mr;
+}
+
 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

@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <glog/logging.h>
+
 #include <version>  // for __cpp_lib_to_chars macro.
 
 // std::from_chars is available in C++17 if __cpp_lib_to_chars is defined.
@@ -21,8 +23,55 @@
 
 namespace dfly {
 
+namespace detail {
+
+inline thread_local PMR_NS::memory_resource* tl_mr = nullptr;
+
+template <typename T> class StatelessJsonAllocator {
+ public:
+  using value_type = T;
+  using size_type = std::size_t;
+  using difference_type = std::ptrdiff_t;
+  using is_always_equal = std::true_type;
+
+  template <typename U> StatelessJsonAllocator(const StatelessJsonAllocator<U>&) noexcept {
+  }
+
+  StatelessJsonAllocator() noexcept {
+    DCHECK_NE(tl_mr, nullptr) << "JSON allocator created without backing mimalloc heap";
+  }
+
+  static value_type* allocate(size_type n) {
+    void* ptr = tl_mr->allocate(n * sizeof(value_type), alignof(value_type));
+    return static_cast<value_type*>(ptr);
+  }
+
+  static void deallocate(value_type* ptr, size_type n) noexcept {
+    tl_mr->deallocate(ptr, n * sizeof(value_type), alignof(value_type));
+  }
+
+  static PMR_NS::memory_resource* resource() {
+    return tl_mr;
+  }
+};
+
+template <typename T, typename U>
+bool operator==(const StatelessJsonAllocator<T>&, const StatelessJsonAllocator<U>&) noexcept {
+  return true;
+}
+
+template <typename T, typename U>
+bool operator!=(const StatelessJsonAllocator<T>&, const StatelessJsonAllocator<U>&) noexcept {
+  return false;
+}
+
+}  // namespace detail
+
+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
@@ -36,12 +85,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;
 }
@@ -91,6 +91,7 @@ class ScannerTest : public ::testing::Test {
  protected:
   ScannerTest() {
     driver_.lexer()->set_debug(1);
+    InitTLJsonHeap(PMR_NS::get_default_resource());
   }
 
   void SetInput(const std::string& str) {

src/core/page_usage_stats_test.cc

@@ -45,6 +45,7 @@ class PageUsageStatsTest : public ::testing::Test {
   }
 
   PageUsageStatsTest() : m_(mi_heap_get_backing()) {
+    InitTLJsonHeap(&m_);
   }
 
   void SetUp() override {
@@ -191,7 +192,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

@@ -6,7 +6,6 @@
 #include "base/logging.h"
 #include "facade/facade_test.h"
 #include "server/command_registry.h"
-#include "server/json_family.h"
 #include "server/test_utils.h"
 
 using namespace testing;
@@ -19,25 +18,33 @@ namespace dfly {
 
 class JsonFamilyMemoryTest : public BaseFamilyTest {
  public:
-  static dfly::MiMemoryResource* GetMemoryResource() {
-    static thread_local mi_heap_t* heap = mi_heap_new();
-    static thread_local dfly::MiMemoryResource memory_resource{heap};
+  static MiMemoryResource* GetMemoryResource() {
+    thread_local mi_heap_t* heap = mi_heap_new();
+    thread_local MiMemoryResource memory_resource{heap};
     return &memory_resource;
   }
 
  protected:
+  void SetUp() override {
+    BaseFamilyTest::SetUp();
+    // Make the core running the thread use the same resource as the rest of the test. Although
+    // BaseFamilyTest initializes the heap on shards serving transactions, the core running the test
+    // needs this initialized explicitly.
+    InitTLJsonHeap(GetMemoryResource());
+  }
+
   auto GetJsonMemoryUsageFromDb(std::string_view key) {
     return Run({"MEMORY", "USAGE", key, "WITHOUTKEY"});
   }
 };
 
 size_t GetMemoryUsage() {
-  return static_cast<MiMemoryResource*>(JsonFamilyMemoryTest::GetMemoryResource())->used();
+  return JsonFamilyMemoryTest::GetMemoryResource()->used();
 }
 
 size_t GetJsonMemoryUsageFromString(std::string_view json_str) {
   size_t start = GetMemoryUsage();
-  auto json = dfly::JsonFromString(json_str, JsonFamilyMemoryTest::GetMemoryResource());
+  auto json = ParseJsonUsingShardHeap(json_str);
   if (!json) {
     return 0;
   }
@@ -111,9 +118,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;