Fix token kind comparison in decorated_key::tri_compare

When comparing decorated_key with ring_position, we need to account for
the token kind. decorated_key tokens are always token_kind::key, but
ring_position tokens can be before_all_keys or after_all_keys.

The previous version incorrectly compared only _data fields, which would
produce wrong results when the ring_position token had a different kind.

Co-authored-by: tgrabiec <283695+tgrabiec@users.noreply.github.com>

compaction/leveled_manifest.hh

@@ -329,13 +329,13 @@ public:
         auto it = candidates.begin();
         auto& first_sstable = *it;
         it++;
-        dht::token first = first_sstable->get_first_decorated_key()._token;
-        dht::token last = first_sstable->get_last_decorated_key()._token;
+        dht::token first = first_sstable->get_first_decorated_key().token();
+        dht::token last = first_sstable->get_last_decorated_key().token();
         while (it != candidates.end()) {
             auto& candidate_sstable = *it;
             it++;
-            dht::token first_candidate = candidate_sstable->get_first_decorated_key()._token;
-            dht::token last_candidate = candidate_sstable->get_last_decorated_key()._token;
+            dht::token first_candidate = candidate_sstable->get_first_decorated_key().token();
+            dht::token last_candidate = candidate_sstable->get_last_decorated_key().token();
 
             first = first <= first_candidate? first : first_candidate;
             last = last >= last_candidate ? last : last_candidate;
@@ -345,7 +345,7 @@ public:
 
     template <typename T>
     static std::vector<sstables::shared_sstable> overlapping(const schema& s, const sstables::shared_sstable& sstable, const T& others) {
-        return overlapping(s, sstable->get_first_decorated_key()._token, sstable->get_last_decorated_key()._token, others);
+        return overlapping(s, sstable->get_first_decorated_key().token(), sstable->get_last_decorated_key().token(), others);
     }
 
     /**
@@ -359,7 +359,7 @@ public:
         auto range = ::wrapping_interval<dht::token>::make(start, end);
 
         for (auto& candidate : sstables) {
-            auto candidate_range = ::wrapping_interval<dht::token>::make(candidate->get_first_decorated_key()._token, candidate->get_last_decorated_key()._token);
+            auto candidate_range = ::wrapping_interval<dht::token>::make(candidate->get_first_decorated_key().token(), candidate->get_last_decorated_key().token());
 
             if (range.overlaps(candidate_range, dht::token_comparator())) {
                 overlapped.push_back(candidate);

dht/decorated_key.hh

@@ -30,11 +30,13 @@ namespace dht {
 // Total ordering defined by comparators is compatible with Origin's ordering.
 class decorated_key {
 public:
-    dht::token _token;
+    // Store only the token data as int64_t to avoid the bloat of storing
+    // token_kind, which is always token_kind::key for decorated_key.
+    int64_t _token_data;
     partition_key _key;
 
     decorated_key(dht::token t, partition_key k)
-        : _token(std::move(t))
+        : _token_data(t._data)
         , _key(std::move(k)) {
     }
 
@@ -56,8 +58,8 @@ public:
     std::strong_ordering tri_compare(const schema& s, const decorated_key& other) const;
     std::strong_ordering tri_compare(const schema& s, const ring_position& other) const;
 
-    const dht::token& token() const noexcept {
-        return _token;
+    dht::token token() const noexcept {
+        return dht::token(_token_data);
     }
 
     const partition_key& key() const {
@@ -65,7 +67,7 @@ public:
     }
 
     size_t external_memory_usage() const {
-        return _key.external_memory_usage() + _token.external_memory_usage();
+        return _key.external_memory_usage();
     }
 
     size_t memory_usage() const {
@@ -102,6 +104,6 @@ template <> struct fmt::formatter<dht::decorated_key> {
     constexpr auto parse(format_parse_context& ctx) { return ctx.begin(); }
     template <typename FormatContext>
     auto format(const dht::decorated_key& dk, FormatContext& ctx) const {
-        return fmt::format_to(ctx.out(), "{{key: {}, token: {}}}", dk._key, dk._token);
+        return fmt::format_to(ctx.out(), "{{key: {}, token: {}}}", dk._key, dk.token());
     }
 };

dht/i_partitioner.cc

@@ -95,7 +95,7 @@ std::unique_ptr<dht::i_partitioner> make_partitioner(sstring partitioner_name) {
 
 bool
 decorated_key::equal(const schema& s, const decorated_key& other) const {
-    if (_token == other._token) {
+    if (_token_data == other._token_data) {
         return _key.legacy_equal(s, other._key);
     }
     return false;
@@ -103,7 +103,7 @@ decorated_key::equal(const schema& s, const decorated_key& other) const {
 
 std::strong_ordering
 decorated_key::tri_compare(const schema& s, const decorated_key& other) const {
-    auto r = _token <=> other._token;
+    auto r = _token_data <=> other._token_data;
     if (r != 0) {
         return r;
     } else {
@@ -113,13 +113,24 @@ decorated_key::tri_compare(const schema& s, const decorated_key& other) const {
 
 std::strong_ordering
 decorated_key::tri_compare(const schema& s, const ring_position& other) const {
-    auto r = _token <=> other.token();
-    if (r != 0) {
-        return r;
-    } else if (other.has_key()) {
-        return _key.legacy_tri_compare(s, *other.key());
+    // decorated_key tokens are always of token_kind::key, so we need to 
+    // account for ring_position tokens that might be before_all_keys or after_all_keys
+    const auto& other_token = other.token();
+    if (other_token._kind == token_kind::key) [[likely]] {
+        auto r = _token_data <=> other_token._data;
+        if (r != 0) {
+            return r;
+        } else if (other.has_key()) {
+            return _key.legacy_tri_compare(s, *other.key());
+        }
+        return 0 <=> other.relation_to_keys();
+    } else if (other_token._kind == token_kind::before_all_keys) {
+        // decorated_key (token_kind::key) > before_all_keys
+        return std::strong_ordering::greater;
+    } else {
+        // decorated_key (token_kind::key) < after_all_keys
+        return std::strong_ordering::less;
     }
-    return 0 <=> other.relation_to_keys();
 }
 
 bool

dht/ring_position.hh

@@ -93,12 +93,12 @@ public:
     { }
 
     ring_position(const dht::decorated_key& dk)
-        : _token(dk._token)
+        : _token(dk.token())
         , _key(std::make_optional(dk._key))
     { }
 
     ring_position(dht::decorated_key&& dk)
-        : _token(std::move(dk._token))
+        : _token(dk.token())
         , _key(std::make_optional(std::move(dk._key)))
     { }
 

mutation/mutation.cc

@@ -316,7 +316,7 @@ auto fmt::formatter<mutation>::format(const mutation& m, fmt::format_context& ct
         ++column_iterator;
     }
 
-    return fmt::format_to(out, "token: {}}}, {}\n}}", dk._token, mutation_partition::printer(s, m.partition()));
+    return fmt::format_to(out, "token: {}}}, {}\n}}", dk.token(), mutation_partition::printer(s, m.partition()));
 }
 
 namespace mutation_json {

mutation/mutation.hh

@@ -126,7 +126,7 @@ public:
     const partition_key& key() const { return _ptr->_dk._key; };
     const dht::decorated_key& decorated_key() const { return _ptr->_dk; };
     dht::ring_position ring_position() const { return { decorated_key() }; }
-    const dht::token& token() const { return _ptr->_dk._token; }
+    dht::token token() const { return _ptr->_dk.token(); }
     const schema_ptr& schema() const { return _ptr->_schema; }
     const mutation_partition& partition() const { return _ptr->_p; }
     mutation_partition& partition() { return _ptr->_p; }

sstables/trie/bti_key_translation.cc

@@ -43,7 +43,7 @@ lazy_comparable_bytes_from_ring_position::lazy_comparable_bytes_from_ring_positi
     , _weight(bound_weight::equal)
     , _pk(std::move(dk._key))
 {
-    init_first_fragment(dk._token);
+    init_first_fragment(dk.token());
 }
 
 void lazy_comparable_bytes_from_ring_position::init_first_fragment(dht::token dht_token) {

test/boost/compound_test.cc

@@ -201,7 +201,7 @@ SEASTAR_THREAD_TEST_CASE(test_conversion_to_legacy_form_same_token_singular) {
     auto key1 = partition_key::from_single_value(*s1, b);
     auto dk1 = partitioner.decorate_key(*s1, key1);
 
-    BOOST_REQUIRE_EQUAL(dk._token, dk1._token);
+    BOOST_REQUIRE_EQUAL(dk.token(), dk1.token());
 }
 
 SEASTAR_THREAD_TEST_CASE(test_conversion_to_legacy_form_same_token_two_components) {
@@ -223,7 +223,7 @@ SEASTAR_THREAD_TEST_CASE(test_conversion_to_legacy_form_same_token_two_component
     auto key1 = partition_key::from_single_value(*s1, b);
     auto dk1 = partitioner.decorate_key(*s1, key1);
 
-    BOOST_REQUIRE_EQUAL(dk._token, dk1._token);
+    BOOST_REQUIRE_EQUAL(dk.token(), dk1.token());
 }
 
 SEASTAR_THREAD_TEST_CASE(test_legacy_ordering_of_singular) {

test/boost/partitioner_test.cc

@@ -120,7 +120,7 @@ SEASTAR_THREAD_TEST_CASE(test_decorated_key_is_compatible_with_origin) {
     auto dk = partitioner.decorate_key(*s, key);
 
     // Expected value was taken from Origin
-    BOOST_REQUIRE_EQUAL(dk._token, token_from_long(4958784316840156970));
+    BOOST_REQUIRE_EQUAL(dk.token(), token_from_long(4958784316840156970));
     BOOST_REQUIRE(dk._key.equal(*s, key));
 }
 

test/boost/sstable_compaction_test.cc

@@ -653,8 +653,8 @@ static bool key_range_overlaps(table_for_tests& cf, const dht::decorated_key& a,
 }
 
 static bool sstable_overlaps(const lw_shared_ptr<replica::column_family>& cf, sstables::shared_sstable candidate1, sstables::shared_sstable candidate2) {
-    auto range1 = wrapping_interval<dht::token>::make(candidate1->get_first_decorated_key()._token, candidate1->get_last_decorated_key()._token);
-    auto range2 = wrapping_interval<dht::token>::make(candidate2->get_first_decorated_key()._token, candidate2->get_last_decorated_key()._token);
+    auto range1 = wrapping_interval<dht::token>::make(candidate1->get_first_decorated_key().token(), candidate1->get_last_decorated_key().token());
+    auto range2 = wrapping_interval<dht::token>::make(candidate2->get_first_decorated_key().token(), candidate2->get_last_decorated_key().token());
     return range1.overlaps(range2, dht::token_comparator());
 }
 

test/boost/sstable_datafile_test.cc

@@ -2986,8 +2986,8 @@ SEASTAR_TEST_CASE(test_index_fast_forwarding_after_eof) {
             sst->load(sst->get_schema()->get_sharder()).get();
         }
 
-        const auto t1 = muts.front().decorated_key()._token;
-        const auto t2 = muts.back().decorated_key()._token;
+        const auto t1 = muts.front().decorated_key().token();
+        const auto t2 = muts.back().decorated_key().token();
         dht::partition_range_vector prs;
 
         prs.emplace_back(dht::ring_position::starting_at(dht::token{t1.raw() - 200}), dht::ring_position::ending_at(dht::token{t1.raw() - 100}));

test/lib/cql_assertions.cc

@@ -322,7 +322,7 @@ future<> require_column_has_value(cql_test_env& e, const sstring& table_name,
     auto ckey = clustering_key::from_deeply_exploded(*schema, ck);
     auto exp = expected.type()->decompose(expected);
     auto dk = dht::decorate_key(*schema, pkey);
-    auto shard = cf.get_effective_replication_map()->shard_for_reads(*schema, dk._token);
+    auto shard = cf.get_effective_replication_map()->shard_for_reads(*schema, dk.token());
     return e.db().invoke_on(shard, [&e, dk = std::move(dk),
                                   ckey = std::move(ckey),
                                   column_name = std::move(column_name),