Bug fix: JSON document size reported by 'json.debug memory'

Signed-off-by: Joe Hu <jowhuw@amazon.com>

src/json/dom.cc

@@ -82,6 +82,10 @@ size_t dom_get_doc_size(const JDocument *doc) {
     return doc->size;
 }
 
+size_t dom_get_doc_size_no_opt(const JDocument *doc) {
+    return sizeof(JDocument) + doc->GetJValue().ComputeMallocedSize();
+}
+
 void dom_set_doc_size(JDocument *doc, const size_t size) {
     doc->size = size;
 }
@@ -96,6 +100,10 @@ void dom_set_bucket_id(JDocument *doc, const uint32_t bucket_id) {
 
 JsonUtilCode dom_parse(ValkeyModuleCtx *ctx, const char *json_buf, const size_t buf_len, JDocument **doc) {
     *doc = nullptr;
+    // begin tracking memory
+    int64_t begin_val = jsonstats_begin_track_mem();
+
+    {
         JParser parser;
         if (parser.Parse(json_buf, buf_len).HasParseError()) {
             return parser.GetParseErrorCode();
@@ -104,6 +112,12 @@ JsonUtilCode dom_parse(ValkeyModuleCtx *ctx, const char *json_buf, const size_t
         *doc = create_doc();
         (*doc)->SetJValue(parser.GetJValue());
         jsonstats_update_max_depth_ever_seen(parser.GetMaxDepth());
+    }
+
+    // end tracking memory
+    int64_t delta = jsonstats_end_track_mem(begin_val);
+    dom_set_doc_size(*doc, dom_get_doc_size(*doc) + delta);
+
     return JSONUTIL_SUCCESS;
 }
 
@@ -1188,25 +1202,6 @@ JsonUtilCode dom_reply_with_resp(ValkeyModuleCtx *ctx, JDocument *doc, const cha
     return JSONUTIL_SUCCESS;
 }
 
-STATIC size_t mem_size_internal(const JValue& v) {
-    size_t size = sizeof(v);  // data structure size
-    if (v.IsString()) {
-        size += v.IsShortString() ? 0 : v.GetStringLength();  // add scalar string value's length
-    } else if (v.IsDouble()) {
-        size += v.IsShortDouble() ? 0 : v.GetDoubleStringLength();
-    } else if (v.IsObject()) {
-        for (auto m = v.MemberBegin(); m != v.MemberEnd(); ++m) {
-            size += m.NodeSize() - sizeof(m->value);  // Overhead (not including the value, which gets added below)
-            size += m->name.GetStringLength();    // add key's length
-            size += mem_size_internal(m->value);  // add value's size
-        }
-    } else if (v.IsArray()) {
-        for (auto &m : v.GetArray())
-            size += mem_size_internal(m);  // add member's size
-    }
-    return size;
-}
-
 JsonUtilCode dom_mem_size(JDocument *doc, const char *path, jsn::vector<size_t> &vec, bool &is_v2_path,
                           bool default_path) {
     vec.clear();
@@ -1234,7 +1229,11 @@ JsonUtilCode dom_mem_size(JDocument *doc, const char *path, jsn::vector<size_t>
     }
 
     for (auto &v : selector.getResultSet()) {
-        vec.push_back(mem_size_internal(*v.first));
+        if (jsonutil_is_root_path(path)) {
+            vec.push_back(dom_get_doc_size_no_opt(static_cast<const JDocument*>(v.first)));
+        } else {
+            vec.push_back(sizeof(*v.first) + v.first->ComputeMallocedSize());
+        }
     }
     return JSONUTIL_SUCCESS;
 }

src/json/dom.h

@@ -211,9 +211,17 @@ JsonUtilCode dom_parse(ValkeyModuleCtx *ctx, const char *json_buf, const size_t
 /* Free a document object */
 void dom_free_doc(JDocument *doc);
 
-/* Get document size */
+/**
+ * Get document size with the optimization of returning the meta data.
+ */
 size_t dom_get_doc_size(const JDocument *doc);
 
+/**
+ * Get document size without optimization.
+ * Calculate the size by walking through the JSON tree.
+ */
+size_t dom_get_doc_size_no_opt(const JDocument *doc);
+
 /* Set document size */
 void dom_set_doc_size(JDocument *doc, const size_t size);
 

src/json/json.cc

@@ -102,11 +102,15 @@ size_t json_get_max_query_string_size() {
 }
 
 #define CHECK_DOCUMENT_SIZE_LIMIT(ctx, new_doc_size) \
-if (!(ValkeyModule_GetContextFlags(ctx) & VALKEYMODULE_CTX_FLAGS_REPLICATED) && \
+if (ctx != nullptr && !(ValkeyModule_GetContextFlags(ctx) & VALKEYMODULE_CTX_FLAGS_REPLICATED) && \
     json_get_max_document_size() > 0 && (new_doc_size > json_get_max_document_size())) { \
     return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(JSONUTIL_DOCUMENT_SIZE_LIMIT_EXCEEDED)); \
 }
 
+#define END_TRACKING_MEMORY(ctx, cmd, doc, orig_doc_size, begin_val) \
+int64_t delta = jsonstats_end_track_mem(begin_val); \
+dom_set_doc_size(doc, orig_doc_size + delta);
+
 // module config params
 // NOTE: We save a copy of the value for each config param instead of pointer address, because the compiler does
 // not allow casting const pointer to pointer.
@@ -328,9 +332,11 @@ STATIC JsonUtilCode parseSimpleCmdArgs(ValkeyModuleString **argv, const int argc
 }
 
 STATIC JsonUtilCode parseNumIncrOrMultByCmdArgs(ValkeyModuleString **argv, const int argc,
-                                                ValkeyModuleString **key, const char **path, JValue *jvalue) {
+                                                ValkeyModuleString **key, const char **path,
+                                                JValue *jvalue, size_t *jvalue_size) {
     *key = nullptr;
     *path = nullptr;
+    *jvalue_size = 0;
 
     // we need exactly 4 arguments
     if (argc != 4) return JSONUTIL_WRONG_NUM_ARGS;
@@ -345,6 +351,7 @@ STATIC JsonUtilCode parseNumIncrOrMultByCmdArgs(ValkeyModuleString **argv, const
         return JSONUTIL_VALUE_NOT_NUMBER;
     }
     *jvalue = parser.GetJValue();
+    *jvalue_size = parser.GetJValueSize();
     return JSONUTIL_SUCCESS;
 }
 
@@ -577,9 +584,7 @@ int Command_JsonSet(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
         if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
         // end tracking memory
-        int64_t delta = jsonstats_end_track_mem(begin_val);
+        END_TRACKING_MEMORY(ctx, "JSON.SET", doc, 0, begin_val)
-        size_t doc_size = dom_get_doc_size(doc) + delta;
-        dom_set_doc_size(doc, doc_size);
 
         if (json_is_instrument_enabled_insert() || json_is_instrument_enabled_update()) {
             size_t len;
@@ -595,13 +600,14 @@ int Command_JsonSet(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
         ValkeyModule_ModuleTypeSetValue(key, DocumentType, doc);
 
         // update stats
-        jsonstats_update_stats_on_insert(doc, true, 0, doc_size, doc_size);
+        size_t new_doc_size = dom_get_doc_size(doc);
+        jsonstats_update_stats_on_insert(doc, true, 0, new_doc_size, new_doc_size);
     } else {
         // fetch doc object from Valkey dict
         JDocument *doc = static_cast<JDocument*>(ValkeyModule_ModuleTypeGetValue(key));
         if (doc == nullptr) return ValkeyModule_ReplyWithError(ctx, ERRMSG_JSON_DOCUMENT_NOT_FOUND);
-
         size_t orig_doc_size = dom_get_doc_size(doc);
+
         rc = dom_set_value(ctx, doc, args.path, args.json, args.json_len, args.is_create_only, args.is_update_only);
         if (rc != JSONUTIL_SUCCESS) {
             if (rc == JSONUTIL_NX_XX_CONDITION_NOT_SATISFIED)
@@ -610,11 +616,10 @@ int Command_JsonSet(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
         }
 
         // end tracking memory
-        int64_t delta = jsonstats_end_track_mem(begin_val);
+        END_TRACKING_MEMORY(ctx, "JSON.SET", doc, orig_doc_size, begin_val)
-        size_t new_doc_size = dom_get_doc_size(doc) + delta;
-        dom_set_doc_size(doc, new_doc_size);
 
         // update stats
+        size_t new_doc_size = dom_get_doc_size(doc);
         jsonstats_update_stats_on_update(doc, orig_doc_size, new_doc_size, args.json_len);
     }
 
@@ -782,12 +787,10 @@ int Command_JsonDel(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc) {
     }
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.DEL", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, new_doc_size, abs(delta));
+    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, dom_get_doc_size(doc), abs(delta));
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -837,7 +840,8 @@ int Command_JsonNumIncrBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     ValkeyModuleString *key_str;
     const char *path;
     JValue jvalue;
-    JsonUtilCode rc = parseNumIncrOrMultByCmdArgs(argv, argc, &key_str, &path, &jvalue);
+    size_t jvalue_size;
+    JsonUtilCode rc = parseNumIncrOrMultByCmdArgs(argv, argc, &key_str, &path, &jvalue, &jvalue_size);
     if (rc != JSONUTIL_SUCCESS) {
         if (rc == JSONUTIL_WRONG_NUM_ARGS) return ValkeyModule_WrongArity(ctx);
         return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
@@ -849,6 +853,10 @@ int Command_JsonNumIncrBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
 
     // fetch doc object from Valkey dict
     JDocument *doc = static_cast<JDocument*>(ValkeyModule_ModuleTypeGetValue(key));
+    size_t orig_doc_size = dom_get_doc_size(doc);
+
+    // begin tracking memory
+    int64_t begin_val = jsonstats_begin_track_mem();
 
     // increment the value at path
     jsn::vector<double> vec;
@@ -856,6 +864,12 @@ int Command_JsonNumIncrBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     rc = dom_increment_by(doc, path, &jvalue, vec, is_v2_path);
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
+    // end tracking memory
+    END_TRACKING_MEMORY(ctx, "JSON.NUMINCRBY", doc, orig_doc_size, begin_val)
+
+    // update stats
+    jsonstats_update_stats_on_update(doc, orig_doc_size, dom_get_doc_size(doc), jvalue_size);
+
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
 
@@ -872,7 +886,8 @@ int Command_JsonNumMultBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     ValkeyModuleString *key_str;
     const char *path;
     JValue jvalue;
-    JsonUtilCode rc = parseNumIncrOrMultByCmdArgs(argv, argc, &key_str, &path, &jvalue);
+    size_t jvalue_size;
+    JsonUtilCode rc = parseNumIncrOrMultByCmdArgs(argv, argc, &key_str, &path, &jvalue, &jvalue_size);
     if (rc != JSONUTIL_SUCCESS) {
         if (rc == JSONUTIL_WRONG_NUM_ARGS) return ValkeyModule_WrongArity(ctx);
         return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
@@ -884,6 +899,10 @@ int Command_JsonNumMultBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
 
     // fetch doc object from Valkey dict
     JDocument *doc = static_cast<JDocument*>(ValkeyModule_ModuleTypeGetValue(key));
+    size_t orig_doc_size = dom_get_doc_size(doc);
+
+    // begin tracking memory
+    int64_t begin_val = jsonstats_begin_track_mem();
 
     // multiply the value at path
     jsn::vector<double> vec;
@@ -891,6 +910,12 @@ int Command_JsonNumMultBy(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     rc = dom_multiply_by(doc, path, &jvalue, vec, is_v2_path);
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
+    // end tracking memory
+    END_TRACKING_MEMORY(ctx, "JSON.NUMMULTBY", doc, orig_doc_size, begin_val)
+
+    // update stats
+    jsonstats_update_stats_on_update(doc, orig_doc_size, dom_get_doc_size(doc), jvalue_size);
+
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
 
@@ -1022,12 +1047,10 @@ int Command_JsonStrAppend(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.STRAPPEND", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_update(doc, orig_doc_size, new_doc_size, json_len);
+    jsonstats_update_stats_on_update(doc, orig_doc_size, dom_get_doc_size(doc), json_len);
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -1091,6 +1114,10 @@ int Command_JsonToggle(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc
 
     // fetch doc object from Valkey dict
     JDocument *doc = static_cast<JDocument*>(ValkeyModule_ModuleTypeGetValue(key));
+    size_t orig_doc_size = dom_get_doc_size(doc);
+
+    // begin tracking memory
+    int64_t begin_val = jsonstats_begin_track_mem();
 
     // toggle the boolean value at this path
     jsn::vector<int> vec;
@@ -1098,6 +1125,9 @@ int Command_JsonToggle(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc
     rc = dom_toggle(doc, path, vec, is_v2_path);
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
+    // end tracking memory
+    END_TRACKING_MEMORY(ctx, "JSON.TOGGLE", doc, orig_doc_size, begin_val)
+
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
 
@@ -1302,12 +1332,10 @@ int Command_JsonArrAppend(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.STRAPPEND", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_update(doc, orig_doc_size, new_doc_size, args.total_json_len);
+    jsonstats_update_stats_on_update(doc, orig_doc_size, dom_get_doc_size(doc), args.total_json_len);
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -1379,12 +1407,10 @@ int Command_JsonArrPop(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc
     }
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.ARRPOP", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, new_doc_size, abs(delta));
+    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, dom_get_doc_size(doc), abs(delta));
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -1426,12 +1452,10 @@ int Command_JsonArrInsert(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int a
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.ARRINSERT", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_insert(doc, false, orig_doc_size, new_doc_size, args.total_json_len);
+    jsonstats_update_stats_on_insert(doc, false, orig_doc_size, dom_get_doc_size(doc), args.total_json_len);
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -1475,12 +1499,10 @@ int Command_JsonArrTrim(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int arg
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.ARRTRIM", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, new_doc_size, abs(delta));
+    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, dom_get_doc_size(doc), abs(delta));
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -1522,12 +1544,10 @@ int Command_JsonClear(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, int argc)
     if (rc != JSONUTIL_SUCCESS) return ValkeyModule_ReplyWithError(ctx, jsonutil_code_to_message(rc));
 
     // end tracking memory
-    int64_t delta = jsonstats_end_track_mem(begin_val);
+    END_TRACKING_MEMORY(ctx, "JSON.CLEAR", doc, orig_doc_size, begin_val)
-    size_t new_doc_size = dom_get_doc_size(doc) + delta;
-    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, new_doc_size, abs(delta));
+    jsonstats_update_stats_on_delete(doc, false, orig_doc_size, dom_get_doc_size(doc), abs(delta));
 
     // replicate the command
     ValkeyModule_ReplicateVerbatim(ctx);
@@ -2082,21 +2102,22 @@ void *DocumentType_RdbLoad(ValkeyModuleIO *rdb, int encver) {
         return nullptr;
     }
 
-    // begin tracking memory
     JDocument *doc;
+    // begin tracking memory
     int64_t begin_val = jsonstats_begin_track_mem();
     JsonUtilCode rc = dom_load(&doc, rdb, encver);
+    // end tracking memory
     int64_t delta = jsonstats_end_track_mem(begin_val);
     if (rc != JSONUTIL_SUCCESS) {
         ValkeyModule_Assert(delta == 0);
         return nullptr;
     }
-    // end tracking memory
+    size_t orig_doc_size = 0;
-    size_t doc_size = dom_get_doc_size(doc) + delta;
+    size_t new_doc_size = orig_doc_size + delta;
-    dom_set_doc_size(doc, doc_size);
+    dom_set_doc_size(doc, new_doc_size);
 
     // update stats
-    jsonstats_update_stats_on_insert(doc, true, 0, doc_size, doc_size);
+    jsonstats_update_stats_on_insert(doc, true, orig_doc_size, new_doc_size, new_doc_size);
     return doc;
 }
 

src/rapidjson/document.h

@@ -984,6 +984,51 @@ public:
         }
     }
 
+    /**
+     * Compute malloc'ed size of the subtree
+     */
+    size_t ComputeMallocedSize() const {
+        size_t mem_size = 0;
+        switch(data_.f.flags) {
+            case kArrayFlag:
+            {
+                GenericValue *e = GetElementsPointer();
+                mem_size += ValkeyModule_MallocSize(e);
+                for (GenericValue *v = e; v != e + data_.a.size; ++v) {
+                    mem_size += v->ComputeMallocedSize();
+                }
+            }
+                break;
+
+            case kObjectHTFlag:
+                mem_size += ValkeyModule_MallocSize(GetMembersPointerHT());
+                for (auto &m : GetObject()) {
+                    mem_size += m.value.ComputeMallocedSize();
+                }
+                break;
+
+            case kObjectVecFlag:
+                mem_size += ValkeyModule_MallocSize(GetMembersPointerVec());
+                for (auto &m : GetObject()) {
+                    mem_size += m.value.ComputeMallocedSize();
+                }
+                break;
+
+            case kCopyStringFlag:
+            case kNumberDoubleFlag:
+                mem_size += ValkeyModule_MallocSize(const_cast<Ch*>(GetStringPointer()));
+                break;
+
+            case kHandleFlag:
+                // KeyTable handles are not accounted to JSON memory consumption
+                break;
+
+            default:
+                break;
+        }
+        return mem_size;
+    }
+
     //@}
 
     //!@name Assignment operators

tst/integration/data/webxml.json

@@ -0,0 +1,88 @@
+{"web-app": {
+  "servlet": [   
+    {
+      "servlet-name": "cofaxCDS",
+      "servlet-class": "org.cofax.cds.CDSServlet",
+      "init-param": {
+        "configGlossary:installationAt": "Philadelphia, PA",
+        "configGlossary:adminEmail": "ksm@pobox.com",
+        "configGlossary:poweredBy": "Cofax",
+        "configGlossary:poweredByIcon": "/images/cofax.gif",
+        "configGlossary:staticPath": "/content/static",
+        "templateProcessorClass": "org.cofax.WysiwygTemplate",
+        "templateLoaderClass": "org.cofax.FilesTemplateLoader",
+        "templatePath": "templates",
+        "templateOverridePath": "",
+        "defaultListTemplate": "listTemplate.htm",
+        "defaultFileTemplate": "articleTemplate.htm",
+        "useJSP": false,
+        "jspListTemplate": "listTemplate.jsp",
+        "jspFileTemplate": "articleTemplate.jsp",
+        "cachePackageTagsTrack": 200,
+        "cachePackageTagsStore": 200,
+        "cachePackageTagsRefresh": 60,
+        "cacheTemplatesTrack": 100,
+        "cacheTemplatesStore": 50,
+        "cacheTemplatesRefresh": 15,
+        "cachePagesTrack": 200,
+        "cachePagesStore": 100,
+        "cachePagesRefresh": 10,
+        "cachePagesDirtyRead": 10,
+        "searchEngineListTemplate": "forSearchEnginesList.htm",
+        "searchEngineFileTemplate": "forSearchEngines.htm",
+        "searchEngineRobotsDb": "WEB-INF/robots.db",
+        "useDataStore": true,
+        "dataStoreClass": "org.cofax.SqlDataStore",
+        "redirectionClass": "org.cofax.SqlRedirection",
+        "dataStoreName": "cofax",
+        "dataStoreDriver": "com.microsoft.jdbc.sqlserver.SQLServerDriver",
+        "dataStoreUrl": "jdbc:microsoft:sqlserver://LOCALHOST:1433;DatabaseName=goon",
+        "dataStoreUser": "sa",
+        "dataStorePassword": "dataStoreTestQuery",
+        "dataStoreTestQuery": "SET NOCOUNT ON;select test='test';",
+        "dataStoreLogFile": "/usr/local/tomcat/logs/datastore.log",
+        "dataStoreInitConns": 10,
+        "dataStoreMaxConns": 100,
+        "dataStoreConnUsageLimit": 100,
+        "dataStoreLogLevel": "debug",
+        "maxUrlLength": 500}},
+    {
+      "servlet-name": "cofaxEmail",
+      "servlet-class": "org.cofax.cds.EmailServlet",
+      "init-param": {
+      "mailHost": "mail1",
+      "mailHostOverride": "mail2"}},
+    {
+      "servlet-name": "cofaxAdmin",
+      "servlet-class": "org.cofax.cds.AdminServlet"},
+ 
+    {
+      "servlet-name": "fileServlet",
+      "servlet-class": "org.cofax.cds.FileServlet"},
+    {
+      "servlet-name": "cofaxTools",
+      "servlet-class": "org.cofax.cms.CofaxToolsServlet",
+      "init-param": {
+        "templatePath": "toolstemplates/",
+        "log": 1,
+        "logLocation": "/usr/local/tomcat/logs/CofaxTools.log",
+        "logMaxSize": "",
+        "dataLog": 1,
+        "dataLogLocation": "/usr/local/tomcat/logs/dataLog.log",
+        "dataLogMaxSize": "",
+        "removePageCache": "/content/admin/remove?cache=pages&id=",
+        "removeTemplateCache": "/content/admin/remove?cache=templates&id=",
+        "fileTransferFolder": "/usr/local/tomcat/webapps/content/fileTransferFolder",
+        "lookInContext": 1,
+        "adminGroupID": 4,
+        "betaServer": true}}],
+  "servlet-mapping": {
+    "cofaxCDS": "/",
+    "cofaxEmail": "/cofaxutil/aemail/*",
+    "cofaxAdmin": "/admin/*",
+    "fileServlet": "/static/*",
+    "cofaxTools": "/tools/*"},
+ 
+  "taglib": {
+    "taglib-uri": "cofax.tld",
+    "taglib-location": "/WEB-INF/tlds/cofax.tld"}}}

tst/integration/test_json_basic.py

@@ -1,4 +1,4 @@
-from utils_json import DEFAULT_MAX_PATH_LIMIT, DEFAULT_MAX_DOCUMENT_SIZE, \
+from utils_json import DEFAULT_MAX_PATH_LIMIT, SIZE_64MB, \
     DEFAULT_WIKIPEDIA_COMPACT_PATH, DEFAULT_WIKIPEDIA_PATH, \
     JSON_INFO_METRICS_SECTION, JSON_INFO_NAMES
 from valkey.exceptions import ResponseError, NoPermissionError
@@ -144,10 +144,8 @@ class TestJsonBasic(JsonTestCase):
 
     def setup_data(self):
         client = self.server.get_new_client()
-        client.config_set(
+        client.config_set('json.max-path-limit', DEFAULT_MAX_PATH_LIMIT)
-            'json.max-path-limit', DEFAULT_MAX_PATH_LIMIT)
+        client.config_set('json.max-document-size', SIZE_64MB)
-        client.config_set(
-            'json.max-document-size', DEFAULT_MAX_DOCUMENT_SIZE)
         # Need the following line when executing the test against a running Valkey.
         # Otherwise, data from previous test cases will interfere current test case.
         client.execute_command("FLUSHDB")
@@ -2716,12 +2714,11 @@ class TestJsonBasic(JsonTestCase):
                 'JSON.DEBUG MEMORY', wikipedia, '.', 'extra')
         assert str(e.value).find('wrong number of arguments') >= 0
 
-        # Test shared path
+        # Verify the document size calculated by the "per document memory tracking" machinery matches the size
-        no_shared_mem = client.execute_command(
+        # calculated by the method of walking the JSON tree.
-            'JSON.DEBUG', 'MEMORY', wikipedia)
+        metadate_val = client.execute_command('JSON.DEBUG','MEMORY',wikipedia)
-        with_shared_mem = client.execute_command(
+        exp_val = client.execute_command('JSON.DEBUG','MEMORY',wikipedia,'.')
-            'JSON.DEBUG', 'MEMORY', wikipedia, '.')
+        assert exp_val == metadate_val
-        assert with_shared_mem > no_shared_mem
 
     def test_json_duplicate_keys(self):
         client = self.server.get_new_client()
@@ -4058,3 +4055,71 @@ class TestJsonBasic(JsonTestCase):
 
         for i in range(len(output)):
             assert subcmd_dict[output[i][0].decode('ascii')] == output[i][1]
+
+    def test_doc_mem_size_with_numincrby(self):
+        """
+        Verify the correctness of json document size.
+        """
+        client = self.server.get_new_client()
+
+        # Set JSON key
+        key = k1
+        json = "{\"a\":1955439684,\"b\":5398,\"c\":[],\"d\":{\"e\":0,\"f\":2.7233281135559082,\"g\":1732953600,\"h\":false,\"i\":true}}"
+        cmd = f'JSON.SET {key} . {json}'
+        assert b'OK' == client.execute_command(cmd)
+
+        for _ in range(0, 10000):
+            # Increment the value by a large double and then set it to 0, so that the value flips between 0 and double.
+            # Since double is stored as string, each flip changes the memory size of the value.
+            cmd = f'json.numincrby {key} .d.e 7.055081799084578998899889890890'
+            client.execute_command(cmd)
+
+            cmd = f'json.set {key} .d.e 0'
+            assert b'OK' == client.execute_command(cmd)
+
+            # This value is the meta data updated by the "per key memory tracking" machinery
+            cmd = f'json.debug memory {key}'
+            metadata_val = client.execute_command(cmd)
+            assert metadata_val < 1000
+
+            # This value the malloc'ed size calculated by walking through the tree
+            cmd = f'json.debug memory {key} .'
+            exp_val = client.execute_command(cmd)
+            assert exp_val == metadata_val
+
+    def test_verify_doc_mem_size(self):
+        """
+        Verify per key doc memory size is correct.
+        """
+        client = self.server.get_new_client()
+        src_dir = os.getenv('SOURCE_DIR')
+        data_dir = f"{src_dir}/tst/integration/data/"
+        file_names = os.listdir(data_dir)
+        for file_name in file_names:
+            with open(f'{data_dir}/{file_name}', 'r') as f:
+                logging.info(f'Loading {file_name}')
+                data = f.read()
+            base_name = os.path.splitext(file_name)[0]
+            key_name = base_name
+            logging.info(f"Setting key {key_name}")
+            client.execute_command('json.set', key_name, '.', data)
+
+        # Scan keyspace
+        count = 0
+        cursor = 0
+        while True:
+            result = client.execute_command("SCAN", cursor, "TYPE", "ReJSON-RL")
+            for key in result[1]:
+                logging.info(f"Verifying memory size of key {key}")
+                # This value is the meta data updated by the "per key memory tracking" machinery
+                metadata_val = client.execute_command("JSON.DEBUG", "MEMORY", key)
+                # This value the malloc'ed size calculated by walking through the tree
+                exp_val = client.execute_command("JSON.DEBUG", "MEMORY", key, ".")
+                assert exp_val == metadata_val
+                assert metadata_val < SIZE_64MB
+                count += 1
+
+            if result[0] == 0:
+                break
+            cursor = result[0]
+        logging.info(f"Verified {count} json keys")

tst/integration/utils_json.py

@@ -18,7 +18,7 @@ JSON_INFO_NAMES = {
     'total_malloc_bytes_used':      JSON_MODULE_NAME + "_total_malloc_bytes_used",
     'memory_traps_enabled':         JSON_MODULE_NAME + "_memory_traps_enabled",
 }
-DEFAULT_MAX_DOCUMENT_SIZE = 64*1024*1024
+SIZE_64MB = 64 * 1024 * 1024
 DEFAULT_MAX_PATH_LIMIT = 128
 DEFAULT_WIKIPEDIA_PATH = 'data/wikipedia.json'
 DEFAULT_WIKIPEDIA_COMPACT_PATH = 'data/wikipedia_compact.json'