Allow MEMORY MALLOC-STATS and MEMORY PURGE during loading phase (#1317)

- Enable investigation of memory issues during loading
- Previously, all memory commands were rejected with LOADING error
(except memory help)
- `MEMORY MALLOC-STATS` and `MEMORTY PURGE` are now allowed
as they don't depend on the dataset
- `MEMORY STATS` and `MEMORY USAGE KEY` remain disallowed

Fixes #1299

Signed-off-by: Guillaume Koenig <[email protected]>
Signed-off-by: Binbin <[email protected]>
Co-authored-by: Binbin <[email protected]>

src/commands.def

@@ -7320,8 +7320,8 @@ struct COMMAND_ARG MEMORY_USAGE_Args[] = {
 struct COMMAND_STRUCT MEMORY_Subcommands[] = {
 {MAKE_CMD("doctor","Outputs a memory problems report.","O(1)","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_DOCTOR_History,0,MEMORY_DOCTOR_Tips,3,memoryCommand,2,0,0,MEMORY_DOCTOR_Keyspecs,0,NULL,0)},
 {MAKE_CMD("help","Returns helpful text about the different subcommands.","O(1)","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_HELP_History,0,MEMORY_HELP_Tips,0,memoryCommand,2,CMD_LOADING|CMD_STALE,0,MEMORY_HELP_Keyspecs,0,NULL,0)},
-{MAKE_CMD("malloc-stats","Returns the allocator statistics.","Depends on how much memory is allocated, could be slow","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_MALLOC_STATS_History,0,MEMORY_MALLOC_STATS_Tips,3,memoryCommand,2,0,0,MEMORY_MALLOC_STATS_Keyspecs,0,NULL,0)},
-{MAKE_CMD("purge","Asks the allocator to release memory.","Depends on how much memory is allocated, could be slow","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_PURGE_History,0,MEMORY_PURGE_Tips,2,memoryCommand,2,0,0,MEMORY_PURGE_Keyspecs,0,NULL,0)},
+{MAKE_CMD("malloc-stats","Returns the allocator statistics.","Depends on how much memory is allocated, could be slow","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_MALLOC_STATS_History,0,MEMORY_MALLOC_STATS_Tips,3,memoryCommand,2,CMD_LOADING,0,MEMORY_MALLOC_STATS_Keyspecs,0,NULL,0)},
+{MAKE_CMD("purge","Asks the allocator to release memory.","Depends on how much memory is allocated, could be slow","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_PURGE_History,0,MEMORY_PURGE_Tips,2,memoryCommand,2,CMD_LOADING,0,MEMORY_PURGE_Keyspecs,0,NULL,0)},
 {MAKE_CMD("stats","Returns details about memory usage.","O(1)","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_STATS_History,0,MEMORY_STATS_Tips,3,memoryCommand,2,0,0,MEMORY_STATS_Keyspecs,0,NULL,0)},
 {MAKE_CMD("usage","Estimates the memory usage of a key.","O(N) where N is the number of samples.","4.0.0",CMD_DOC_NONE,NULL,NULL,"server",COMMAND_GROUP_SERVER,MEMORY_USAGE_History,0,MEMORY_USAGE_Tips,0,memoryCommand,-3,CMD_READONLY,0,MEMORY_USAGE_Keyspecs,1,NULL,2),.args=MEMORY_USAGE_Args},
 {0}

src/commands/memory-malloc-stats.json

@@ -12,6 +12,9 @@
             "REQUEST_POLICY:ALL_SHARDS",
             "RESPONSE_POLICY:SPECIAL"
         ],
+        "command_flags": [
+            "LOADING"
+        ],
         "reply_schema": {
             "type": "string",
             "description": "The memory allocator's internal statistics report."

src/commands/memory-purge.json

@@ -11,6 +11,9 @@
             "REQUEST_POLICY:ALL_SHARDS",
             "RESPONSE_POLICY:ALL_SUCCEEDED"
         ],
+        "command_flags": [
+            "LOADING"
+        ],
         "reply_schema": {
             "const": "OK"
         }

tests/unit/introspection.tcl

@@ -1042,6 +1042,49 @@ test {config during loading} {
     }
 } {} {external:skip}
 
+test {MEMORY commands during loading} {
+    start_server [list overrides [list key-load-delay 50 loading-process-events-interval-bytes 1024]] {
+        # Set up some initial data
+        r debug populate 100000 key 1000
+
+        # Save and restart
+        r save
+        restart_server 0 false false
+
+        # At this point, keys are loaded one at time, busy looping 50usec
+        # between each. Further, other events are processed every 1024 bytes
+        # of RDB. We're sending all our commands deferred, so they have a
+        # chance to be processed all at once between loading two keys.
+
+        set rd [valkey_deferring_client]
+
+        # Allowed during loading
+        $rd memory help
+        $rd memory malloc-stats
+        $rd memory purge
+
+        # Disallowed during loading (because directly dependent on the dataset)
+        $rd memory doctor
+        $rd memory stats
+        $rd memory usage key:1
+
+        # memory help
+        assert_match {{MEMORY <subcommand> *}} [$rd read]
+        # memory malloc-stats
+        assert_match {*alloc*} [$rd read]
+        # memory purge
+        assert_match OK [$rd read]
+        # memory doctor
+        assert_error {*LOADING*} {$rd read}
+        # memory stats
+        assert_error {*LOADING*} {$rd read}
+        # memory usage key:1
+        assert_error {*LOADING*} {$rd read}
+
+        $rd close
+    }
+} {} {external:skip}
+
 test {CONFIG REWRITE handles rename-command properly} {
     start_server {tags {"introspection"} overrides {rename-command {flushdb badger}}} {
         assert_error {ERR unknown command*} {r flushdb}