tools/scylla-nodetool: status: fix token count for tablets

Currently, the token count column is always based on the vnodes, which
makes no sense for tablet keyspaces. If a tablet keyspace is provided as
the keyspace argument, don't print the vnode token count. If the user
provided a table argument as well, print the tablet count, otherwise
print "?".

test/nodetool/test_status.py

@@ -45,7 +45,7 @@ null_ownership_error = ("Non-system keyspaces don't have the same replication se
                         "effective ownership information is meaningless")
 
 
-def validate_status_output(res, keyspace, nodes, ownership, resolve, effective_ownership_unknown):
+def validate_status_output(res, keyspace, nodes, ownership, resolve, effective_ownership_unknown, token_count_unknown):
     datacenters = sorted(list(set([node.datacenter for node in nodes.values()])))
     load_multiplier = {"bytes": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4}
 
@@ -106,7 +106,10 @@ def validate_status_output(res, keyspace, nodes, ownership, resolve, effective_o
             else:
                 assert load_unit is not None
                 assert load == "{:.2f}".format(int(node.load) / load_multiplier[load_unit])
-            assert int(tokens) == len(node.tokens)
+            if token_count_unknown:
+                tokens == "?"
+            else:
+                assert int(tokens) == len(node.tokens)
             if effective_ownership_unknown:
                 assert owns == "?"
             else:
@@ -198,6 +201,11 @@ def _do_test_status(request, nodetool, status_query_target, node_list, resolve=N
 
     host_id_map = [{"key": ep, "value": node.host_id} for ep, node in nodes.items() if node.host_id is not None]
 
+    tokens_endpoint_params = {}
+    if keyspace_uses_tablets and table:
+        tokens_endpoint_params["keyspace"] = keyspace
+        tokens_endpoint_params["cf"] = table
+
     tokens_endpoint = []
     for ep, node in nodes.items():
         for token in node.tokens:
@@ -212,7 +220,8 @@ def _do_test_status(request, nodetool, status_query_target, node_list, resolve=N
         expected_request("GET", "/storage_service/nodes/leaving", response=leaving),
         expected_request("GET", "/storage_service/nodes/moving", response=moving),
         expected_request("GET", "/storage_service/load_map", response=load_map),
-        expected_request("GET", "/storage_service/tokens_endpoint", response=tokens_endpoint),
+        expected_request("GET", "/storage_service/tokens_endpoint", params=tokens_endpoint_params,
+                         response=tokens_endpoint),
         expected_request("GET", "/gossiper/endpoint/live", response=live),
         expected_request("GET", "/gossiper/endpoint/down", response=down),
         expected_request("GET", "/storage_service/host_id", response=host_id_map),
@@ -227,18 +236,21 @@ def _do_test_status(request, nodetool, status_query_target, node_list, resolve=N
                                  response={"message": f"std::runtime_error({null_ownership_error})", "code": 500}),
                 expected_request("GET", "/storage_service/ownership", multiple=expected_request.ANY,
                                  response=ownership_response)]
-    elif table is None:
+    else:
         if not uses_cassandra_nodetool:
             keyspaces_using_tablets = [keyspace] if keyspace_uses_tablets else []
             expected_requests.append(
-                expected_request("GET", "/storage_service/keyspaces", params={"replication": "tablets"}, multiple=expected_request.ONE, response=keyspaces_using_tablets))
-
-        if not keyspace_uses_tablets:
+                expected_request("GET", "/storage_service/keyspaces", params={"replication": "tablets"},
+                                 multiple=expected_request.ONE, response=keyspaces_using_tablets))
+        if table is None:
+            if not keyspace_uses_tablets:
+                expected_requests.append(
+                    expected_request("GET", f"/storage_service/ownership/{keyspace}",
+                                     multiple=expected_request.ONE, response=ownership_response))
+        else:
             expected_requests.append(
-                expected_request("GET", f"/storage_service/ownership/{keyspace}", multiple=expected_request.ONE, response=ownership_response))
-    else:
-        expected_requests.append(
-                expected_request("GET", f"/storage_service/ownership/{keyspace}", params={"cf": table}, response=ownership_response))
+                    expected_request("GET", f"/storage_service/ownership/{keyspace}", params={"cf": table},
+                                     response=ownership_response))
 
     for ep, node in nodes.items():
         expected_requests += [
@@ -262,7 +274,9 @@ def _do_test_status(request, nodetool, status_query_target, node_list, resolve=N
     res = nodetool(*args, expected_requests=expected_requests)
 
     effective_ownership_unknown = keyspace is None or (table is None and keyspace_uses_tablets)
-    validate_status_output(res.stdout, keyspace, nodes, ownership, bool(resolve), effective_ownership_unknown)
+    token_count_unknown = keyspace_uses_tablets and not table
+    validate_status_output(res.stdout, keyspace, nodes, ownership, bool(resolve), effective_ownership_unknown,
+                           token_count_unknown)
 
 
 def test_status_no_keyspace_single_dc(request, nodetool):

tools/scylla-nodetool.cc

@@ -1923,7 +1923,10 @@ void status_operation(scylla_rest_client& client, const bpo::variables_map& vm)
     const auto endpoint_load = rjson_to_map<size_t>(client.get("/storage_service/load_map"));
     const auto endpoint_host_id = rjson_to_map<sstring>(client.get("/storage_service/host_id"));
 
-    const auto is_effective_ownership_unknown = (!keyspace || (!table && keyspace_uses_tablets(client, *keyspace)));
+    const auto tablets_keyspace = keyspace && keyspace_uses_tablets(client, *keyspace);
+
+    const auto is_effective_ownership_unknown = (!keyspace || (!table && tablets_keyspace));
+
     const auto endpoint_ownership = is_effective_ownership_unknown
         ? std::map<sstring, float>{}
         : get_effective_ownership(client, *keyspace, table);
@@ -1931,7 +1934,12 @@ void status_operation(scylla_rest_client& client, const bpo::variables_map& vm)
     std::unordered_map<sstring, sstring> endpoint_rack;
     std::map<sstring, std::set<sstring>> dc_endpoints;
     std::unordered_map<sstring, size_t> endpoint_tokens;
-    const auto tokens_endpoint_res = client.get("/storage_service/tokens_endpoint");
+    std::unordered_map<sstring, sstring> tokens_endpoint_params;
+    if (tablets_keyspace && table) {
+        tokens_endpoint_params["keyspace"] = *keyspace;
+        tokens_endpoint_params["cf"] = *table;
+    }
+    const auto tokens_endpoint_res = client.get("/storage_service/tokens_endpoint", std::move(tokens_endpoint_params));
     for (const auto& te : tokens_endpoint_res.GetArray()) {
         const auto ep = sstring(rjson::to_string_view(te["value"]));
          // We are not printing the actual tokens, so it is enough just to count them.
@@ -1945,6 +1953,8 @@ void status_operation(scylla_rest_client& client, const bpo::variables_map& vm)
         dc_endpoints[dc].insert(ep);
     }
 
+    const bool token_count_unknown = tablets_keyspace && !table;
+
     for (const auto& [dc, endpoints] : dc_endpoints) {
         const auto dc_header = fmt::format("Datacenter: {}", dc);
         fmt::print("{}\n", dc_header);
@@ -1981,7 +1991,7 @@ void status_operation(scylla_rest_client& client, const bpo::variables_map& vm)
                     fmt::format("{}{}", status, state),
                     address,
                     load,
-                    endpoint_tokens.at(ep),
+                    token_count_unknown ? "?" : fmt::to_string(endpoint_tokens.at(ep)),
                     !is_effective_ownership_unknown ? format("{:.1f}%", endpoint_ownership.at(ep) * 100) : "?",
                     endpoint_host_id.contains(ep) ? endpoint_host_id.at(ep) : "?",
                     endpoint_rack.at(ep));