server: Parse JSON on destination shard (#6061)

abhijat · web-flow · commit 978961ff793b · 2025-11-17T09:56:58.000Z
* server: Parse JSON on shard

JSON is parsed on shard to make sure that the shard's thread local
memory resource is used to allocate. This is done in preparation to use
a stateless allocator.

With a stateless allocator, it will not be possible to safely parse the
search value on coordinator shard, because currently the coordinator
shard uses the default memory resource and the destination shard uses
mimalloc.

It can be possible that the coordinator shard does not have mimalloc
heap, in which case it will not be able to parse using such a heap.

To avoid this the parsing is pushed to destination shard, and the error
handling is delayed until the transaction hop finishes.

* core: Introduce short lived json type

This is the default heap allocated json type from jsoncons, using
std::allocator. It is renamed to clearly indicate its use in a non db
context.

* server: Use std allocated JSON where applicable

Signed-off-by: Abhijat Malviya &lt;abhijat@dragonflydb.io&gt;
diff --git a/src/core/json/json_object.cc b/src/core/json/json_object.cc
@@ -7,12 +7,12 @@
 #include "base/logging.h"
 
 using namespace jsoncons;
-using namespace std;
 
-namespace dfly {
+namespace {
 
-optional<JsonType> JsonFromString(string_view input, PMR_NS::memory_resource* mr) {
-  error_code ec;
+template <typename T>
+std::optional<T> ParseWithDecoder(std::string_view input, json_decoder<T>&& decoder) {
+  std::error_code ec;
   auto JsonErrorHandler = [](json_errc ec, const ser_context&) {
     VLOG(1) << "Error while decode JSON: " << make_error_code(ec).message();
     return false;
@@ -28,10 +28,9 @@ optional<JsonType> JsonFromString(string_view input, PMR_NS::memory_resource* mr
 
   /* The maximum possible JSON nesting depth is either the specified json_nesting_depth_limit or
      half of the input size. Since nesting a JSON object requires at least 2 characters. */
-  auto parser_options = jsoncons::json_options{}.max_nesting_depth(
+  auto parser_options = json_options{}.max_nesting_depth(
       std::min(json_nesting_depth_limit, uint32_t(input.size() / 2)));
 
-  json_decoder<JsonType> decoder(std::pmr::polymorphic_allocator<char>{mr});
   json_parser parser(parser_options, JsonErrorHandler);
 
   parser.update(input);
@@ -40,7 +39,19 @@ optional<JsonType> JsonFromString(string_view input, PMR_NS::memory_resource* mr
   if (!ec && decoder.is_valid()) {
     return decoder.get_result();
   }
-  return nullopt;
+  return std::nullopt;
+}
+
+}  // namespace
+
+namespace dfly {
+
+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}});
 }
 
 JsonType DeepCopyJSON(const JsonType* j, PMR_NS::memory_resource* mr) {
diff --git a/src/core/json/json_object.h b/src/core/json/json_object.h
@@ -21,9 +21,21 @@
 
 namespace dfly {
 
+using ShortLivedJSON = jsoncons::json;
 using JsonType = jsoncons::pmr::json;
 
-// Build a json object from string. If the string is not legal json, will return nullopt
+// A helper type to use in template functions which are expected to work with both ShortLivedJSON
+// and JsonType
+template <typename Allocator>
+using JsonWithAllocator = jsoncons::basic_json<char, jsoncons::sorted_policy, Allocator>;
+
+// Parses string into JSON. Any allocatons are done using the std allocator. This method should be
+// used for generic JSON parsing, in particular, it should not be used to parse objects which will
+// be stored in the db, as the backing storage is not managed by mimalloc.
+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);
 
 // Deep copy a JSON object, by first serializing it to a string and then deserializing the string.
diff --git a/src/server/cluster/cluster_config.cc b/src/server/cluster/cluster_config.cc
@@ -127,7 +127,7 @@ shared_ptr<ClusterConfig> ClusterConfig::CreateFromConfig(string_view my_id,
 namespace {
 constexpr string_view kInvalidConfigPrefix = "Invalid JSON cluster config: "sv;
 
-template <typename T> optional<T> ReadNumeric(const JsonType& obj) {
+template <typename T> optional<T> ReadNumeric(const ShortLivedJSON& obj) {
   if (!obj.is_number()) {
     LOG(ERROR) << kInvalidConfigPrefix << "object is not a number " << obj;
     return nullopt;
@@ -136,7 +136,7 @@ template <typename T> optional<T> ReadNumeric(const JsonType& obj) {
   return obj.as<T>();
 }
 
-optional<SlotRanges> GetClusterSlotRanges(const JsonType& slots) {
+optional<SlotRanges> GetClusterSlotRanges(const ShortLivedJSON& slots) {
   if (!slots.is_array()) {
     LOG(ERROR) << kInvalidConfigPrefix << "slot_ranges is not an array " << slots;
     return nullopt;
@@ -162,7 +162,7 @@ optional<SlotRanges> GetClusterSlotRanges(const JsonType& slots) {
   return SlotRanges(ranges);
 }
 
-optional<ClusterExtendedNodeInfo> ParseClusterNode(const JsonType& json) {
+optional<ClusterExtendedNodeInfo> ParseClusterNode(const ShortLivedJSON& json) {
   if (!json.is_object()) {
     LOG(ERROR) << kInvalidConfigPrefix << "node config is not an object " << json;
     return nullopt;
@@ -221,7 +221,7 @@ optional<ClusterExtendedNodeInfo> ParseClusterNode(const JsonType& json) {
   return node;
 }
 
-optional<std::vector<MigrationInfo>> ParseMigrations(const JsonType& json) {
+optional<std::vector<MigrationInfo>> ParseMigrations(const ShortLivedJSON& json) {
   std::vector<MigrationInfo> res;
   if (json.is_null()) {
     return res;
@@ -251,7 +251,7 @@ optional<std::vector<MigrationInfo>> ParseMigrations(const JsonType& json) {
   return res;
 }
 
-optional<ClusterShardInfos> BuildClusterConfigFromJson(const JsonType& json) {
+optional<ClusterShardInfos> BuildClusterConfigFromJson(const ShortLivedJSON& json) {
   std::vector<ClusterShardInfo> config;
 
   if (!json.is_array()) {
@@ -309,7 +309,7 @@ optional<ClusterShardInfos> BuildClusterConfigFromJson(const JsonType& json) {
 /* static */
 shared_ptr<ClusterConfig> ClusterConfig::CreateFromConfig(string_view my_id,
                                                           std::string_view json_str) {
-  optional<JsonType> json_config = JsonFromString(json_str, PMR_NS::get_default_resource());
+  optional<ShortLivedJSON> json_config = JsonFromString(json_str);
   if (!json_config.has_value()) {
     LOG(ERROR) << "Can't parse JSON for ClusterConfig " << json_str;
     return nullptr;
diff --git a/src/server/json_family.cc b/src/server/json_family.cc
@@ -235,11 +235,12 @@ void Send(const std::vector<std::string>& vec, RedisReplyBuilder* rb) {
   Send(vec.begin(), vec.end(), rb);
 }
 
-void Send(const JsonType& value, RedisReplyBuilder* rb) {
+template <typename Allocator>
+void Send(const JsonWithAllocator<Allocator>& value, RedisReplyBuilder* rb) {
   if (value.is_double()) {
     Send(value.as_double(), rb);
   } else if (value.is_number()) {
-    Send(value.as_integer<long>(), rb);
+    Send(value.template as_integer<long>(), rb);
   } else if (value.is_bool()) {
     rb->SendSimpleString(value.as_bool() ? "true" : "false");
   } else if (value.is_null()) {
@@ -341,9 +342,7 @@ void SendJsonString(const OpResult<string>& result, RedisReplyBuilder* rb) {
   if (result) {
     const string& json_str = result.value();
     if (rb->IsResp3()) {
-      std::optional<JsonType> parsed_json =
-          JsonFromString(json_str, PMR_NS::get_default_resource());
-      if (parsed_json) {
+      if (const std::optional<ShortLivedJSON> parsed_json = JsonFromString(json_str)) {
         Send(parsed_json.value(), rb);
         return;
       }
@@ -470,18 +469,13 @@ std::optional<std::string> ConvertJsonPathToJsonPointer(string_view json_path) {
   return pointer;
 }
 
-// Use this method on the coordinator thread
-std::optional<JsonType> JsonFromString(std::string_view input) {
-  return dfly::JsonFromString(input, PMR_NS::get_default_resource());
-}
-
 /* Use this method on the shard thread
 
    If you do memory tracking, make sure to initialize it before calling this method, and reset the
    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 dfly::JsonFromString(input, CompactObj::memory_resource());
+  return JsonFromString(input, CompactObj::memory_resource());
 }
 
 OpStatus SetFullJson(const OpArgs& op_args, string_view key, string_view json_str) {
@@ -1307,7 +1301,19 @@ auto OpArrTrim(const OpArgs& op_args, string_view key, const WrappedJsonPath& pa
 // Returns numeric vector that represents the new length of the array at each path.
 OpResult<JsonCallbackResult<OptSize>> OpArrInsert(const OpArgs& op_args, string_view key,
                                                   const WrappedJsonPath& json_path, int index,
-                                                  const vector<JsonType>& new_values) {
+                                                  const vector<string_view>& new_values) {
+  vector<JsonType> parsed_values;
+  parsed_values.reserve(new_values.size());
+
+  for (const auto& nv : new_values) {
+    const optional<JsonType> v = ShardJsonFromString(nv);
+    if (!v) {
+      return OpStatus::SYNTAX_ERR;
+    }
+
+    parsed_values.emplace_back(std::move(*v));
+  }
+
   bool out_of_boundaries_encountered = false;
 
   // Insert user-supplied value into the supplied index that should be valid.
@@ -1336,7 +1342,7 @@ OpResult<JsonCallbackResult<OptSize>> OpArrInsert(const OpArgs& op_args, string_
     }
 
     auto it = GetJsonArrayIterator(val, insert_before_index);
-    for (auto& new_val : new_values) {
+    for (auto& new_val : parsed_values) {
       it = val->insert(it, new_val);
       it++;
     }
@@ -1350,14 +1356,26 @@ OpResult<JsonCallbackResult<OptSize>> OpArrInsert(const OpArgs& op_args, string_
   return res;
 }
 
-auto OpArrAppend(const OpArgs& op_args, string_view key, const WrappedJsonPath& path,
-                 const vector<JsonType>& append_values) {
+OpResult<JsonCallbackResult<optional<optional<unsigned long>>>> OpArrAppend(
+    const OpArgs& op_args, string_view key, const WrappedJsonPath& path,
+    const vector<string_view>& append_values) {
+  vector<JsonType> parsed_values;
+  parsed_values.reserve(append_values.size());
+
+  for (const auto& v : append_values) {
+    const optional<JsonType> parsed = ShardJsonFromString(v);
+    if (!parsed) {
+      return OpStatus::SYNTAX_ERR;
+    }
+    parsed_values.emplace_back(std::move(*parsed));
+  }
+
   auto cb = [&](std::optional<std::string_view>,
                 JsonType* val) -> MutateCallbackResult<std::optional<std::size_t>> {
     if (!val->is_array()) {
       return {};
     }
-    for (auto& new_val : append_values) {
+    for (auto& new_val : parsed_values) {
       val->emplace_back(new_val);
     }
     return {false, val->size()};
@@ -1368,8 +1386,15 @@ auto OpArrAppend(const OpArgs& op_args, string_view key, const WrappedJsonPath&
 // Returns a numeric vector representing each JSON value first index of the JSON scalar.
 // An index value of -1 represents unfound in the array.
 // JSON scalar has types of string, boolean, null, and number.
-auto OpArrIndex(const OpArgs& op_args, string_view key, const WrappedJsonPath& json_path,
-                const JsonType& search_val, int start_index, int end_index) {
+OpResult<JsonCallbackResult<optional<long>>> OpArrIndex(const OpArgs& op_args, string_view key,
+                                                        const WrappedJsonPath& json_path,
+                                                        string_view search_val, int start_index,
+                                                        int end_index) {
+  const optional<JsonType> search_value_json = ShardJsonFromString(search_val);
+  if (!search_value_json) {
+    return OpStatus::SYNTAX_ERR;
+  }
+
   auto cb = [&](const string_view&, const JsonType& val) -> std::optional<long> {
     if (!val.is_array()) {
       return std::nullopt;
@@ -1403,7 +1428,7 @@ auto OpArrIndex(const OpArgs& op_args, string_view key, const WrappedJsonPath& j
     size_t pos = -1;
     auto it = GetJsonArrayIterator(val, pos_start_index);
     while (it != val.array_range().end()) {
-      if (JsonAreEquals(search_val, *it)) {
+      if (JsonAreEquals(search_value_json, *it)) {
         pos = pos_start_index;
         break;
       }
@@ -1710,8 +1735,7 @@ void JsonFamily::Resp(CmdArgList args, const CommandContext& cmd_cntx) {
   };
 
   auto result = cmd_cntx.tx->ScheduleSingleHopT(std::move(cb));
-  auto* rb = static_cast<RedisReplyBuilder*>(builder);
-  reply_generic::Send(result, rb);
+  reply_generic::Send(result, builder);
 }
 
 void JsonFamily::Debug(CmdArgList args, const CommandContext& cmd_cntx) {
@@ -1826,11 +1850,7 @@ void JsonFamily::ArrIndex(CmdArgList args, const CommandContext& cmd_cntx) {
   auto* builder = static_cast<RedisReplyBuilder*>(cmd_cntx.rb);
   WrappedJsonPath json_path = GET_OR_SEND_UNEXPECTED(ParseJsonPath(path));
 
-  optional<JsonType> search_value = JsonFromString(parser.Next());
-  if (!search_value) {
-    builder->SendError(kSyntaxErr);
-    return;
-  }
+  string_view search_value = parser.Next();
 
   int start_index = 0;
   if (parser.HasNext()) {
@@ -1851,12 +1871,11 @@ void JsonFamily::ArrIndex(CmdArgList args, const CommandContext& cmd_cntx) {
   }
 
   auto cb = [&](Transaction* t, EngineShard* shard) {
-    return OpArrIndex(t->GetOpArgs(shard), key, json_path, *search_value, start_index, end_index);
+    return OpArrIndex(t->GetOpArgs(shard), key, json_path, search_value, start_index, end_index);
   };
 
   auto result = cmd_cntx.tx->ScheduleSingleHopT(std::move(cb));
-  auto* rb = static_cast<RedisReplyBuilder*>(builder);
-  reply_generic::Send(result, rb);
+  reply_generic::Send(result, builder);
 }
 
 void JsonFamily::ArrInsert(CmdArgList args, const CommandContext& cmd_cntx) {
@@ -1873,15 +1892,9 @@ void JsonFamily::ArrInsert(CmdArgList args, const CommandContext& cmd_cntx) {
 
   WrappedJsonPath json_path = GET_OR_SEND_UNEXPECTED(ParseJsonPath(path));
 
-  vector<JsonType> new_values;
+  vector<string_view> new_values;
   for (size_t i = 3; i < args.size(); i++) {
-    optional<JsonType> val = JsonFromString(ArgS(args, i));
-    if (!val) {
-      builder->SendError(kSyntaxErr);
-      return;
-    }
-
-    new_values.emplace_back(std::move(*val));
+    new_values.emplace_back(ArgS(args, i));
   }
 
   auto cb = [&](Transaction* t, EngineShard* shard) {
@@ -1900,26 +1913,17 @@ void JsonFamily::ArrAppend(CmdArgList args, const CommandContext& cmd_cntx) {
   auto* builder = static_cast<RedisReplyBuilder*>(cmd_cntx.rb);
   WrappedJsonPath json_path = GET_OR_SEND_UNEXPECTED(ParseJsonPath(path));
 
-  vector<JsonType> append_values;
-
-  // TODO: there is a bug here, because we parse json using the allocator from
-  // the coordinator thread, and we pass it to the shard thread, which is not safe.
+  vector<string_view> append_values;
   for (size_t i = 2; i < args.size(); ++i) {
-    optional<JsonType> converted_val = JsonFromString(ArgS(args, i));
-    if (!converted_val) {
-      builder->SendError(kSyntaxErr);
-      return;
-    }
-    append_values.emplace_back(converted_val);
+    append_values.emplace_back(ArgS(args, i));
   }
 
   auto cb = [&](Transaction* t, EngineShard* shard) {
     return OpArrAppend(t->GetOpArgs(shard), key, json_path, append_values);
   };
 
   auto result = cmd_cntx.tx->ScheduleSingleHopT(std::move(cb));
-  auto* rb = static_cast<RedisReplyBuilder*>(builder);
-  reply_generic::Send(result, rb);
+  reply_generic::Send(result, builder);
 }
 
 void JsonFamily::ArrTrim(CmdArgList args, const CommandContext& cmd_cntx) {
@@ -2000,7 +2004,7 @@ void JsonFamily::StrAppend(CmdArgList args, const CommandContext& cmd_cntx) {
   WrappedJsonPath json_path = GET_OR_SEND_UNEXPECTED(ParseJsonPath(path));
 
   // We try parsing the value into json string object first.
-  optional<JsonType> parsed_json = JsonFromString(value);
+  optional<ShortLivedJSON> parsed_json = JsonFromString(value);
   if (!parsed_json || !parsed_json->is_string()) {
     return builder->SendError("expected string value", kSyntaxErrType);
   };
@@ -2011,8 +2015,7 @@ void JsonFamily::StrAppend(CmdArgList args, const CommandContext& cmd_cntx) {
   };
 
   auto result = cmd_cntx.tx->ScheduleSingleHopT(std::move(cb));
-  auto* rb = static_cast<RedisReplyBuilder*>(builder);
-  reply_generic::Send(result, rb);
+  reply_generic::Send(result, builder);
 }
 
 void JsonFamily::ObjKeys(CmdArgList args, const CommandContext& cmd_cntx) {
