chore: switch json object to pmr allocator (#2483)

* chore: switch json object to pmr allocator

1. Move json object creation code to the shard-thread inside rdb_load
2. Now json_family never references "json" type, always dfly::JsonType
3. Switch JsonType to pmr::json.
4. Make sure we pass the correct memory_resource when creating json object from string.

Signed-off-by: Roman Gershman <[email protected]>
---------

Signed-off-by: Roman Gershman <[email protected]>

src/core/compact_object_test.cc

@@ -577,14 +577,15 @@ TEST_F(CompactObjectTest, JsonTypeWithPathTest) {
   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
-  auto f = [](const std::string& /*path*/, JsonType& book) {
+  auto f = [](const auto& /*path*/, JsonType& book) {
     if (book.at("category") == "memoir" && !book.contains("price")) {
       book.try_emplace("price", 140.0);
     }
   };
   JsonType* json = cobj_.GetJson();
   ASSERT_TRUE(json != nullptr);
-  jsonpath::json_replace(*json, "$.books[*]", f);
+  auto allocator_set = jsoncons::combine_allocators(json->get_allocator());
+  jsonpath::json_replace(allocator_set, *json, "$.books[*]"sv, f);
 
   // Check whether we've changed the entry for json in place
   // we should have prices only for memoir books

src/core/json_object.cc

@@ -1,28 +1,26 @@
-// Copyright 2022, DragonflyDB authors.  All rights reserved.
+// Copyright 2023, DragonflyDB authors.  All rights reserved.
 // See LICENSE for licensing terms.
 //
 
 #include "core/json_object.h"
-extern "C" {
-#include "redis/object.h"
-#include "redis/zmalloc.h"
-}
-#include <jsoncons/json.hpp>
 
 #include "base/logging.h"
+#include "core/compact_object.h"
 
-namespace dfly {
+using namespace jsoncons;
+using namespace std;
 
-std::optional<JsonType> JsonFromString(std::string_view input) {
-  using namespace jsoncons;
+namespace dfly {
 
-  std::error_code ec;
+optional<JsonType> JsonFromString(string_view input) {
+  error_code ec;
   auto JsonErrorHandler = [](json_errc ec, const ser_context&) {
     VLOG(1) << "Error while decode JSON: " << make_error_code(ec).message();
     return false;
   };
 
-  json_decoder<JsonType> decoder;
+  json_decoder<JsonType> decoder(
+      std::pmr::polymorphic_allocator<char>{CompactObj::memory_resource()});
   json_parser parser(basic_json_decode_options<char>{}, JsonErrorHandler);
 
   parser.update(input);
@@ -31,7 +29,7 @@ std::optional<JsonType> JsonFromString(std::string_view input) {
   if (decoder.is_valid()) {
     return decoder.get_result();
   }
-  return {};
+  return nullopt;
 }
 
 }  // namespace dfly

src/core/json_object.h

@@ -4,32 +4,26 @@
 
 #pragma once
 
+#include <jsoncons/json.hpp>
+#include <jsoncons_ext/jsonpath/jsonpath.hpp>
 #include <memory>
 #include <optional>
 #include <string_view>
 
-extern "C" {
-#include "redis/redis_aux.h"
-}
-
-// Note about this file - once we have the issue with jsonpath in jsoncons resolved
-// we would add the implementation for the allocator here as well. Right now this
-// file is a little bit empty, but for external "users" such as json_family they
-// should include this when creating JSON object from string that we're getting
-// from the commands.
-namespace jsoncons {
-struct sorted_policy;
-template <typename CharT, typename Policy, typename Allocator> class basic_json;
-}  // namespace jsoncons
-
 namespace dfly {
 
 // This is temporary, there is an issue right now with jsoncons about using jsonpath
 // with custom allocator. once it would resolved, we would change this to use custom allocator
 // that allocate memory from mimalloc
-using JsonType = jsoncons::basic_json<char, jsoncons::sorted_policy, std::allocator<char>>;
+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);
 
+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>>(
+      jsoncons::allocator_set<JsonType::allocator_type, std::allocator<char>>(), path, ec);
+}
+
 }  // namespace dfly

src/core/json_test.cc

@@ -35,12 +35,17 @@ TEST_F(JsonTest, Basic) {
     }
 )";
 
-  json j = json::parse(data);
+  pmr::json j = pmr::json::parse(data);
   EXPECT_TRUE(j.contains("reputons"));
   jsonpath::json_replace(j, "$.reputons[*].rating", 1.1);
   EXPECT_EQ(1.1, j["reputons"][0]["rating"].as_double());
 }
 
+TEST_F(JsonTest, SetEmpty) {
+  pmr::json dest{json_object_arg};  // crashes on UB without the tag.
+  dest["bar"] = "foo";
+}
+
 TEST_F(JsonTest, Query) {
   json j = R"(
 {"a":{}, "b":{"a":1}, "c":{"a":1, "b":2}}

src/core/sorted_map.cc

@@ -1054,7 +1054,7 @@ uint64_t SortedMap::DfImpl::Scan(uint64_t cursor,
 /***************************************************************************/
 /* SortedMap */
 /***************************************************************************/
-SortedMap::SortedMap(PMR_NS::memory_resource* mr) : impl_(RdImpl()), mr_res_(mr) {
+SortedMap::SortedMap(PMR_NS::memory_resource* mr) : impl_(RdImpl()) {
   if (absl::GetFlag(FLAGS_use_zset_tree)) {
     impl_ = DfImpl();
   }

src/core/sorted_map.h

@@ -268,7 +268,6 @@ class SortedMap {
   };
 
   std::variant<RdImpl, DfImpl> impl_;
-  PMR_NS::memory_resource* mr_res_;
 };
 
 // Used by CompactObject.