From e5c478286bc242a55241f3709393bbb66ed4ca8f Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Fri, 20 Dec 2024 15:22:30 -0800 Subject: [PATCH] support for client count command Signed-off-by: Sarthak Aggarwal --- src/commands.def | 59 +++++++++++- src/commands/client-count.json | 163 +++++++++++++++++++++++++++++++++ src/networking.c | 53 +++++++++++ tests/unit/introspection.tcl | 158 ++++++++++++++++++++++++++++++++ 4 files changed, 431 insertions(+), 2 deletions(-) create mode 100644 src/commands/client-count.json diff --git a/src/commands.def b/src/commands.def index 9089b994c3..5de1c7c696 100644 --- a/src/commands.def +++ b/src/commands.def @@ -1162,6 +1162,58 @@ struct COMMAND_ARG CLIENT_CAPA_Args[] = { {MAKE_ARG("capability",ARG_TYPE_STRING,-1,NULL,NULL,NULL,CMD_ARG_MULTIPLE,0,NULL)}, }; +/********** CLIENT COUNT ********************/ + +#ifndef SKIP_CMD_HISTORY_TABLE +/* CLIENT COUNT history */ +commandHistory CLIENT_COUNT_History[] = { +{"8.1.0","Introducing client count command with filters"}, +}; +#endif + +#ifndef SKIP_CMD_TIPS_TABLE +/* CLIENT COUNT tips */ +const char *CLIENT_COUNT_Tips[] = { +"nondeterministic_output", +}; +#endif + +#ifndef SKIP_CMD_KEY_SPECS_TABLE +/* CLIENT COUNT key specs */ +#define CLIENT_COUNT_Keyspecs NULL +#endif + +/* CLIENT COUNT client_type argument table */ +struct COMMAND_ARG CLIENT_COUNT_client_type_Subargs[] = { +{MAKE_ARG("normal",ARG_TYPE_PURE_TOKEN,-1,"NORMAL",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +{MAKE_ARG("master",ARG_TYPE_PURE_TOKEN,-1,"MASTER",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +{MAKE_ARG("replica",ARG_TYPE_PURE_TOKEN,-1,"REPLICA",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +{MAKE_ARG("pubsub",ARG_TYPE_PURE_TOKEN,-1,"PUBSUB",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +}; + +/* CLIENT COUNT skipme argument table */ +struct COMMAND_ARG CLIENT_COUNT_skipme_Subargs[] = { +{MAKE_ARG("yes",ARG_TYPE_PURE_TOKEN,-1,"YES",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +{MAKE_ARG("no",ARG_TYPE_PURE_TOKEN,-1,"NO",NULL,NULL,CMD_ARG_NONE,0,NULL)}, +}; + +/* CLIENT COUNT argument table */ +struct COMMAND_ARG CLIENT_COUNT_Args[] = { +{MAKE_ARG("client-id",ARG_TYPE_INTEGER,-1,"ID",NULL,"8.1.0",CMD_ARG_OPTIONAL|CMD_ARG_MULTIPLE,0,NULL)}, +{MAKE_ARG("client-type",ARG_TYPE_ONEOF,-1,"TYPE",NULL,NULL,CMD_ARG_OPTIONAL,4,NULL),.subargs=CLIENT_COUNT_client_type_Subargs}, +{MAKE_ARG("username",ARG_TYPE_STRING,-1,"USER",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("addr",ARG_TYPE_STRING,-1,"ADDR",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL),.display_text="ip:port"}, +{MAKE_ARG("laddr",ARG_TYPE_STRING,-1,"LADDR",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL),.display_text="ip:port"}, +{MAKE_ARG("skipme",ARG_TYPE_ONEOF,-1,"SKIPME",NULL,"8.1.0",CMD_ARG_OPTIONAL,2,NULL),.subargs=CLIENT_COUNT_skipme_Subargs}, +{MAKE_ARG("maxage",ARG_TYPE_INTEGER,-1,"MAXAGE",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("name",ARG_TYPE_STRING,-1,"NAME",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("minidle",ARG_TYPE_INTEGER,-1,"MINIDLE",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("flags",ARG_TYPE_STRING,-1,"FLAGS",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("pattern",ARG_TYPE_STRING,-1,"PATTERN",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("channel",ARG_TYPE_STRING,-1,"CHANNEL",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +{MAKE_ARG("shardchannel",ARG_TYPE_STRING,-1,"SHARDCHANNEL",NULL,"8.1.0",CMD_ARG_OPTIONAL,0,NULL)}, +}; + /********** CLIENT GETNAME ********************/ #ifndef SKIP_CMD_HISTORY_TABLE @@ -1289,6 +1341,7 @@ commandHistory CLIENT_KILL_History[] = { {"6.2.0","`LADDR` option."}, {"8.0.0","`MAXAGE` option."}, {"8.0.0","Replaced `master` `TYPE` with `primary`. `master` still supported for backward compatibility."}, +{"8.1.0","Added filters NAME, MINIDLE, FLAGS, PATTERN, CHANNEL and SHARDCHANNEL "}, }; #endif @@ -1358,6 +1411,7 @@ commandHistory CLIENT_LIST_History[] = { {"7.0.0","Added `resp`, `multi-mem`, `rbs` and `rbp` fields."}, {"7.0.3","Added `ssub` field."}, {"8.0.0","Replaced `master` `TYPE` with `primary`. `master` still supported for backward compatibility."}, +{"8.1.0","Added filters like USER, ADDR, LADDR, SKIPME, MAXAGE, NAME, MINIDLE, FLAGS, PATTERN, CHANNEL and SHARDCHANNEL "}, }; #endif @@ -1677,14 +1731,15 @@ struct COMMAND_ARG CLIENT_UNBLOCK_Args[] = { struct COMMAND_STRUCT CLIENT_Subcommands[] = { {MAKE_CMD("caching","Instructs the server whether to track the keys in the next request.","O(1)","6.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_CACHING_History,0,CLIENT_CACHING_Tips,0,clientCommand,3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_CACHING_Keyspecs,0,NULL,1),.args=CLIENT_CACHING_Args}, {MAKE_CMD("capa","A client claims its capability.","O(1)","8.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_CAPA_History,0,CLIENT_CAPA_Tips,0,clientCommand,-3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE,ACL_CATEGORY_CONNECTION,CLIENT_CAPA_Keyspecs,0,NULL,1),.args=CLIENT_CAPA_Args}, +{MAKE_CMD("count","Counts open connections.","O(N) where N is the number of client connections","8.1.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_COUNT_History,1,CLIENT_COUNT_Tips,1,clientCommand,-2,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_COUNT_Keyspecs,0,NULL,13),.args=CLIENT_COUNT_Args}, {MAKE_CMD("getname","Returns the name of the connection.","O(1)","2.6.9",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_GETNAME_History,0,CLIENT_GETNAME_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_GETNAME_Keyspecs,0,NULL,0)}, {MAKE_CMD("getredir","Returns the client ID to which the connection's tracking notifications are redirected.","O(1)","6.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_GETREDIR_History,0,CLIENT_GETREDIR_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_GETREDIR_Keyspecs,0,NULL,0)}, {MAKE_CMD("help","Returns helpful text about the different subcommands.","O(1)","5.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_HELP_History,0,CLIENT_HELP_Tips,0,clientCommand,2,CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_HELP_Keyspecs,0,NULL,0)}, {MAKE_CMD("id","Returns the unique client ID of the connection.","O(1)","5.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_ID_History,0,CLIENT_ID_Tips,0,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_ID_Keyspecs,0,NULL,0)}, {MAKE_CMD("import-source","Mark this client as an import source when server is in import mode.","O(1)","8.1.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_IMPORT_SOURCE_History,0,CLIENT_IMPORT_SOURCE_Tips,0,clientCommand,3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE,ACL_CATEGORY_CONNECTION,CLIENT_IMPORT_SOURCE_Keyspecs,0,NULL,1),.args=CLIENT_IMPORT_SOURCE_Args}, {MAKE_CMD("info","Returns information about the connection.","O(1)","6.2.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_INFO_History,0,CLIENT_INFO_Tips,1,clientCommand,2,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_INFO_Keyspecs,0,NULL,0)}, -{MAKE_CMD("kill","Terminates open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_KILL_History,7,CLIENT_KILL_Tips,0,clientCommand,-3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_KILL_Keyspecs,0,NULL,1),.args=CLIENT_KILL_Args}, -{MAKE_CMD("list","Lists open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_LIST_History,7,CLIENT_LIST_Tips,1,clientCommand,-2,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_LIST_Keyspecs,0,NULL,13),.args=CLIENT_LIST_Args}, +{MAKE_CMD("kill","Terminates open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_KILL_History,8,CLIENT_KILL_Tips,0,clientCommand,-3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_KILL_Keyspecs,0,NULL,1),.args=CLIENT_KILL_Args}, +{MAKE_CMD("list","Lists open connections.","O(N) where N is the number of client connections","2.4.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_LIST_History,8,CLIENT_LIST_Tips,1,clientCommand,-2,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_LIST_Keyspecs,0,NULL,13),.args=CLIENT_LIST_Args}, {MAKE_CMD("no-evict","Sets the client eviction mode of the connection.","O(1)","7.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_NO_EVICT_History,0,CLIENT_NO_EVICT_Tips,0,clientCommand,3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_NO_EVICT_Keyspecs,0,NULL,1),.args=CLIENT_NO_EVICT_Args}, {MAKE_CMD("no-touch","Controls whether commands sent by the client affect the LRU/LFU of accessed keys.","O(1)","7.2.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_NO_TOUCH_History,0,CLIENT_NO_TOUCH_Tips,0,clientCommand,3,CMD_NOSCRIPT|CMD_LOADING|CMD_STALE,ACL_CATEGORY_CONNECTION,CLIENT_NO_TOUCH_Keyspecs,0,NULL,1),.args=CLIENT_NO_TOUCH_Args}, {MAKE_CMD("pause","Suspends commands processing.","O(1)","3.0.0",CMD_DOC_NONE,NULL,NULL,"connection",COMMAND_GROUP_CONNECTION,CLIENT_PAUSE_History,1,CLIENT_PAUSE_Tips,0,clientCommand,-3,CMD_ADMIN|CMD_NOSCRIPT|CMD_LOADING|CMD_STALE|CMD_SENTINEL,ACL_CATEGORY_CONNECTION,CLIENT_PAUSE_Keyspecs,0,NULL,2),.args=CLIENT_PAUSE_Args}, diff --git a/src/commands/client-count.json b/src/commands/client-count.json new file mode 100644 index 0000000000..e62cd054ed --- /dev/null +++ b/src/commands/client-count.json @@ -0,0 +1,163 @@ +{ + "COUNT": { + "summary": "Counts open connections.", + "complexity": "O(N) where N is the number of client connections", + "group": "connection", + "since": "8.1.0", + "arity": -2, + "container": "CLIENT", + "function": "clientCommand", + "history": [ + [ + "8.1.0", + "Introducing client count command with filters" + ] + ], + "command_flags": [ + "ADMIN", + "NOSCRIPT", + "LOADING", + "STALE", + "SENTINEL" + ], + "acl_categories": [ + "CONNECTION" + ], + "command_tips": [ + "NONDETERMINISTIC_OUTPUT" + ], + "reply_schema": { + "type": "string", + "description": "Information and statistics about client connections" + }, + "arguments": [ + { + "name": "client-id", + "token": "ID", + "type": "integer", + "optional": true, + "multiple": true, + "since": "8.1.0" + }, + { + "token": "TYPE", + "name": "client-type", + "type": "oneof", + "optional": true, + "arguments": [ + { + "name": "normal", + "type": "pure-token", + "token": "normal" + }, + { + "name": "master", + "type": "pure-token", + "token": "master" + }, + { + "name": "replica", + "type": "pure-token", + "token": "replica" + }, + { + "name": "pubsub", + "type": "pure-token", + "token": "pubsub" + } + ] + }, + { + "token": "USER", + "name": "username", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "ADDR", + "name": "addr", + "display": "ip:port", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "LADDR", + "name": "laddr", + "display": "ip:port", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "SKIPME", + "name": "skipme", + "type": "oneof", + "optional": true, + "since": "8.1.0", + "arguments": [ + { + "name": "yes", + "type": "pure-token", + "token": "YES" + }, + { + "name": "no", + "type": "pure-token", + "token": "NO" + } + ] + }, + { + "token": "MAXAGE", + "name": "maxage", + "type": "integer", + "optional": true, + "since": "8.1.0" + }, + { + "token": "NAME", + "name": "name", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "MINIDLE", + "name": "minidle", + "type": "integer", + "optional": true, + "since": "8.1.0" + }, + { + "token": "FLAGS", + "name": "flags", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "PATTERN", + "name": "pattern", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "CHANNEL", + "name": "channel", + "type": "string", + "optional": true, + "since": "8.1.0" + }, + { + "token": "SHARDCHANNEL", + "name": "shardchannel", + "type": "string", + "optional": true, + "since": "8.1.0" + } + ] + } +} diff --git a/src/networking.c b/src/networking.c index 9cfda21edd..f534d9ebfd 100644 --- a/src/networking.c +++ b/src/networking.c @@ -4459,6 +4459,57 @@ static void clientCommandImportSource(client *c) { } } +void clientCountCommand(client *c) { + /* If there are additional arguments */ + if (c->argc > 2) { + clientFilter filter = {.ids = NULL, .max_age = 0, .addr = NULL, .laddr = NULL, .user = NULL, .type = -1, .skipme = 0}; + const int i = 2; // Start parsing filters from the third argument + + if (parseClientFiltersOrReply(c, i, &filter) != C_OK) { + /* Free filter resources on failure */ + freeClientFilter(&filter); + return; + } + + long long count = 0; + listIter li; + listNode *ln; + listRewind(server.clients, &li); + + /* Count clients that match the filter */ + while ((ln = listNext(&li)) != NULL) { + client *cl = listNodeValue(ln); + if (clientMatchesFilter(cl, filter)) { + count++; + } + } + + /* Free filter resources after use */ + freeClientFilter(&filter); + /* Return the count to the client */ + addReplyLongLong(c, count); + + } else if (c->argc == 2) { + /* No filters, just "CLIENT COUNT" */ + long long count = 0; + listIter li; + listNode *ln; + listRewind(server.clients, &li); + + /* Count all clients */ + while ((ln = listNext(&li)) != NULL) { + count++; + } + + /* Return the total count to the client */ + addReplyLongLong(c, count); + } else { + /* Invalid syntax */ + addReplyErrorObject(c, shared.syntaxerr); + return; + } +} + void clientCommand(client *c) { if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr, "help")) { clientCommandHelp(c); @@ -4498,6 +4549,8 @@ void clientCommand(client *c) { clientCommandCapa(c); } else if (!strcasecmp(c->argv[1]->ptr, "import-source")) { clientCommandImportSource(c); + } else if (!strcasecmp(c->argv[1]->ptr, "count")) { + clientCountCommand(c); } else { addReplySubcommandSyntaxError(c); } diff --git a/tests/unit/introspection.tcl b/tests/unit/introspection.tcl index e03770f600..92fcbe9172 100644 --- a/tests/unit/introspection.tcl +++ b/tests/unit/introspection.tcl @@ -417,6 +417,164 @@ start_server {tags {"introspection"}} { assert_error "ERR *greater than 0*" {r client list maxage -1} } + test {CLIENT COUNT: Count all clients} { + # Connect multiple clients + set c1 [valkey_client] + set c2 [valkey_client] + set c3 [valkey_client] + + # Assert the count matches the number of connected clients + assert {[r client count] == 4} + + # Close all clients + $c1 close + $c2 close + $c3 close + } + + test {CLIENT COUNT: Filter by specific ID} { + # Create two clients + set c1 [valkey_client] + set c2 [valkey_client] + + # Get the ID of the first client + set client_list [r client list] + regexp {id=(\d+)} $client_list -> id1 + + # Assert only the client with the matching ID is counted + assert {[r client count id $id1] == 1} + + # Close the clients + $c1 close + $c2 close + } + + test {CLIENT COUNT: Filter by maximum age} { + # Create two clients + set c1 [valkey_client] + set c2 [valkey_client] + + # Wait 2 seconds + after 2000 + + # Assert no clients younger than 1 second are counted + assert {[r client count maxage 1] == 3} + + # Close the clients + $c1 close + $c2 close + } + + test {CLIENT COUNT: Filter by client address} { + # Create a client + set c1 [valkey_client] + + # Get the client's address + set client_list [r client list] + regexp {addr=([^ ]+)} $client_list -> addr + + # Assert only the client with the matching address is counted + assert {[r client count addr $addr] == 1} + + # Close the client + $c1 close + } + + test {CLIENT COUNT: Filter by local address} { + # Create a client + set c1 [valkey_client] + + # Get the client's local address + set client_list [r client list] + regexp {laddr=([^ ]+)} $client_list -> laddr + + # Assert only the client with the matching local address is counted + assert {[r client count laddr $laddr] == 2} + + # Close the client + $c1 close + } + + test {CLIENT COUNT: Exclude current client} { + # Create two clients + set c1 [valkey_client] + set c2 [valkey_client] + + # Assert the correct number of clients is returned for both cases + assert {[r client count skipme no] == 3} + assert {[r client count skipme yes] == 2} + + # Close the clients + $c1 close + $c2 close + } + + test {CLIENT COUNT: Filter by user} { + # Create a user and assign it to a client + r acl setuser user1 on +@all >password + set c1 [valkey_client] + $c1 auth user1 password + set c2 [valkey_client] + + # Assert only the client associated with the user is counted + assert {[r client count user user1] == 1} + + # Close the clients + $c1 close + $c2 close + + # Delete the user + r acl deluser user1 + } + + test {CLIENT COUNT: Filter by type} { + # Create clients of different types + set c1 [valkey_client] + set c2 [valkey_client] + $c2 subscribe test_channel + + # Assert only normal clients are counted + assert {[r client count type normal] == 2} + assert {[r client count type pubsub] == 1} + + # Close the clients + $c1 close + $c2 close + } + + + test {CLIENT COUNT: Filter by client name} { + # Create a client and set its name + set c1 [valkey_client] + $c1 client setname test_client + + # Assert only the client with the matching name is counted + assert {[r client count name test_client] == 1} + + # Close the client + $c1 close + } + + test {CLIENT COUNT: Combine filters for name and flags} { + # Create multiple clients with different attributes + set c1 [valkey_client] + set c2 [valkey_client] + $c1 client setname client1 + $c2 client setname client2 + $c2 multi + + # Assert only clients matching both filters are counted + assert {[r client count name client1 flags x] == 0} + + + # Assert only clients matching both filters are counted + assert {[r client count name client2 flags x] == 1} + + # Close the clients + $c1 close + $c2 close + } + proc get_field_in_client_info {info field} { set info [string trim $info] foreach item [split $info " "] {