From d47dfceb34293ee7172bc540f69f97771a79a58e Mon Sep 17 00:00:00 2001
From: Tomasz Grabiec <tgrabiec@scylladb.com>
Date: Wed, 20 Mar 2024 13:53:25 +0100
Subject: [PATCH] tablets: tablet_sharder: Allow working with non-local host

Will be used in tests.
---
 locator/tablet_sharder.hh | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/locator/tablet_sharder.hh b/locator/tablet_sharder.hh
index 657bedf923..07d500c1a9 100644
--- a/locator/tablet_sharder.hh
+++ b/locator/tablet_sharder.hh
@@ -21,6 +21,7 @@ class tablet_sharder : public dht::sharder {
     const token_metadata& _tm;
     table_id _table;
     mutable const tablet_map* _tmap = nullptr;
+    host_id _host;
 private:
     // Tablet map is lazily initialized to avoid exceptions during effective_replication_map construction
     // in case tablet mapping is not yet available in token metadata at the time the table is constructed.
@@ -98,9 +99,10 @@ private:
         return get_shard(tinfo.replicas, host);
     }
 public:
-    tablet_sharder(const token_metadata& tm, table_id table)
+    tablet_sharder(const token_metadata& tm, table_id table, std::optional<host_id> host = std::nullopt)
             : _tm(tm)
             , _table(table)
+            , _host(host.value_or(tm.get_my_id()))
     { }
 
     virtual ~tablet_sharder() = default;
@@ -111,7 +113,7 @@ public:
         // FIXME: Consider throwing when there is no owning shard on the current host rather than returning 0.
         // It's a coordination mistake to route requests to non-owners. Topology coordinator should synchronize
         // with request coordinators before moving the shard away.
-        auto shard = shard_for_reads(tid, _tm.get_my_id()).value_or(0);
+        auto shard = shard_for_reads(tid, _host).value_or(0);
         tablet_logger.trace("[{}] shard_of({}) = {}, tablet={}", _table, t, shard, tid);
         return shard;
     }
@@ -119,17 +121,16 @@ public:
     virtual dht::shard_replica_set shard_for_writes(const token& t) const override {
         ensure_tablet_map();
         auto tid = _tmap->get_tablet_id(t);
-        auto shards = shard_for_writes(tid, _tm.get_my_id());
+        auto shards = shard_for_writes(tid, _host);
         tablet_logger.trace("[{}] shard_for_writes({}) = {}, tablet={}", _table, t, shards, tid);
         return shards;
     }
 
     virtual std::optional<dht::shard_and_token> next_shard(const token& t) const override {
         ensure_tablet_map();
-        auto me = _tm.get_my_id();
         std::optional<tablet_id> tb = _tmap->get_tablet_id(t);
         while ((tb = _tmap->next_tablet(*tb))) {
-            auto r = shard_for_reads(*tb, me);
+            auto r = shard_for_reads(*tb, _host);
             auto next = _tmap->get_first_token(*tb);
             tablet_logger.trace("[{}] token_for_next_shard({}) = {{{}, {}}}, tablet={}", _table, t, next, r, *tb);
             return dht::shard_and_token{r.value_or(0), next};