tools/scylla-sstable: add scylla sstable shard-of command

when migrating to the uuid-based identifiers, the mapping from the
integer-based generation to the shard-id is preserved. we used to have
"gen % smp_count" for calculating the shard which is responsible to host
a given sstable. despite that this is not a documented behavior, this is
handy when we try to correlate an sstable to a shard, typically when
looking at a performance issue.

in this change, a new subcommand is added to expose the connection
between the sstable and its "owner" shards.

Fixes #16343
Signed-off-by: Kefu Chai <kefu.chai@scylladb.com>

Closes scylladb/scylladb#16345

docs/operating-scylla/admin-tools/scylla-sstable.rst

@@ -648,6 +648,22 @@ Note that levels are cumulative - each contains all the checks of the previous l
 By default, the strictest level is used.
 This can be relaxed, for example, if you want to produce intentionally corrupt SStables for tests.
 
+shard-of
+^^^^^^^^
+
+Pint out the shards which own the specified SSTables.
+
+The content is dumped in JSON, using the following schema:
+
+.. code-block:: none
+    :class: hide-copy-button
+
+    $ROOT := { "$sstable_path": $SHARD_IDS, ... }
+
+    $SHARD_IDS := [$SHARD_ID, ...]
+
+    $SHARD_ID := Uint
+
 script
 ^^^^^^
 

test/cql-pytest/test_tools.py

@@ -8,6 +8,8 @@
 
 import contextlib
 import glob
+import itertools
+import functools
 import json
 import nodetool
 import os
@@ -18,6 +20,7 @@ import random
 import re
 import shutil
 import util
+from typing import Iterable, Type, Union
 
 
 def simple_no_clustering_table(cql, keyspace):
@@ -865,3 +868,83 @@ def test_scrub_validate_mode(scylla_path, scrub_workdir, scrub_schema_file, scru
 
         # Check that validate did not move the bad sstable into qurantine
         assert os.path.exists(scrub_bad_sstable)
+
+
+def _to_cql3_type(t: Type) -> str:
+    # map from Python type to Cassandra type, only a small subset is supported
+    py_to_cql3_type = {int: "Int32Type",
+                       str: "UTF8Type",
+                       bool: "BooleanType"}
+    return py_to_cql3_type[t]
+
+
+KeyType = Union[int, str, bool]
+
+
+def _serialize_value(scylla_path: str, value: KeyType) -> str:
+    return subprocess.check_output([scylla_path,
+                                    "types", "serialize", "--full-compound",
+                                    "-t", _to_cql3_type(type(value)),
+                                    "--",
+                                    str(value)]).strip().decode()
+
+
+def _shard_of_values(scylla_path: str, shards: int, *values: list[KeyType]) -> int:
+    args = [scylla_path, "types", "shardof",
+            "--full-compound",
+            "--shards", str(shards)]
+    for value in values:
+        args.extend(['-t', _to_cql3_type(type(value))])
+    serialized = ''.join(_serialize_value(scylla_path, v) for v in values)
+    args.extend(['--', serialized])
+    output = subprocess.check_output(args).strip().decode()
+    # the output looks like:
+    # (file_instance, 2021-03-27, c61a3321-0459-41c3-8e56-75255feb0196): token: -5043005771368701888, shard: 1
+    shard = output.rsplit(':', 1)[-1]
+    return int(shard)
+
+
+def _generate_key_for_shard(scylla_path: str, shards: int, shard_id: int) -> Iterable[int]:
+    # this only works with the table with a single integer pk. if we want to
+    # be more general, we could use a randomized generator to enumerate all
+    # possible pk combinations.
+    for pk in itertools.count(start=0, step=1):
+        if _shard_of_values(scylla_path, shards, pk) == shard_id:
+            yield pk
+
+
+def _simple_table_with_keys(cql, keyspace: str, keys: Iterable[int]) -> tuple[str, str]:
+    table = util.unique_name()
+    schema = (f"CREATE TABLE {keyspace}.{table} (pk int PRIMARY KEY, v int) "
+              "WITH compaction = {'class': 'NullCompactionStrategy'}")
+    cql.execute(schema)
+
+    for pk in keys:
+        cql.execute(f"INSERT INTO {keyspace}.{table} (pk, v) VALUES ({pk}, 0)")
+    nodetool.flush(cql, f"{keyspace}.{table}")
+
+    return table, schema
+
+
+def test_scylla_sstable_shard_of(cql, test_keyspace, scylla_path, scylla_data_dir) -> None:
+    # cql-pytest/run.py::run_scylla_cmd() passes "--smp 2" to scylla, so we
+    # need to be consistent with it to get the correct sstable-shard mapping
+    scylla_option_smp = 2
+    shards = scylla_option_smp
+    num_keys = 42
+    for shard_id in range(shards):
+        all_keys_for_shard = _generate_key_for_shard(scylla_path, shards, shard_id)
+        keys = itertools.islice(all_keys_for_shard, num_keys)
+        table_factory = functools.partial(_simple_table_with_keys, keys=keys)
+        with scylla_sstable(table_factory, cql, test_keyspace, scylla_data_dir) as (schema_file, sstables):
+            out = subprocess.check_output([scylla_path,
+                                           "sstable", "shard-of",
+                                           "--schema-file", schema_file,
+                                           "--shards", str(shards)] +
+                                          sstables)
+            # all sstables contains the rows with the keys deliberately
+            # created for specified shard
+            sstables_json = json.loads(out)['sstables']
+            expected_json = [shard_id]
+            for actual_json in sstables_json.values():
+                assert actual_json == expected_json

tools/scylla-sstable.cc

@@ -2553,6 +2553,43 @@ void sstable_consumer_operation(schema_ptr schema, reader_permit permit, const s
     consumer->consume_stream_end().get();
 }
 
+void shard_of_operation(schema_ptr, reader_permit,
+                        const std::vector<sstables::shared_sstable>& sstables,
+                        sstables::sstables_manager& sstable_manager,
+                        const bpo::variables_map& vm) {
+    if (!vm.count("shards")) {
+        throw std::invalid_argument("missing required option '--shards'");
+    }
+    unsigned shard_count = vm["shards"].as<unsigned>();
+    unsigned ignore_msb_bits = vm["ignore-msb-bits"].as<unsigned>();
+
+    json_writer writer;
+    writer.StartStream();
+    for (auto& sst : sstables) {
+        // sst was loaded with the smp::count as its shard_count but that's not
+        // necessarily identical to the "shards" specified in the command line.
+        // reload the sst with the specified shard_count and ignore_msb_bits
+        auto schema = schema_builder(sst->get_schema()).with_sharder(
+            shard_count, ignore_msb_bits).build();
+        auto new_sst = sstable_manager.make_sstable(
+            schema,
+            sst->get_storage().prefix(),
+            data_dictionary::storage_options{},
+            sst->generation(),
+            sstable_state::normal,
+            sst->get_version());
+        new_sst->load(schema->get_sharder(), sstables::sstable_open_config{}).get();
+
+        writer.Key(sst->get_filename());
+        writer.StartArray();
+        for (unsigned shard_id : new_sst->get_shards_for_this_sstable()) {
+            writer.Uint(shard_id);
+        }
+        writer.EndArray();
+    }
+    writer.EndStream();
+}
+
 const std::vector<operation_option> global_options {
     typed_option<sstring>("schema-file", "schema.cql", "file containing the schema description"),
     typed_option<sstring>("keyspace", "keyspace name"),
@@ -2816,6 +2853,15 @@ for more information on this operation, including the API documentation.
                 typed_option<program_options::string_map>("script-arg", {}, "parameter(s) for the script"),
             }},
             script_operation},
+/* shard-of */
+    {{"shard-of",
+            "Print out the shard which 'owns' the sstable",
+            "Print out the intersection(s) of the shard-ranges and the partition ranges",
+            {
+                typed_option<unsigned>("shards", "the number of shards the source scylla instance has"),
+                typed_option<unsigned>("ignore-msb-bits", 12u, "'murmur3_partitioner_ignore_msb_bits' set by scylla.yaml"),
+            }},
+            shard_of_operation},
 };
 
 } // anonymous namespace