token_metadata: update_topology: make endpoint_dc_rack arg optional

It's better to pass a disengaged optional when
the caller doesn't have the information rather than
passing the default dc_rack location so the latter
will never implicitly override a known endpoint dc/rack location.

Signed-off-by: Benny Halevy <bhalevy@scylladb.com>

Closes #15300

locator/token_metadata.cc

@@ -118,8 +118,8 @@ public:
         return _bootstrap_tokens;
     }
 
-    void update_topology(inet_address ep, endpoint_dc_rack dr, std::optional<node::state> opt_st, std::optional<shard_id> shard_count = std::nullopt) {
-        _topology.add_or_update_endpoint(ep, std::nullopt, std::move(dr), std::move(opt_st), std::move(shard_count));
+    void update_topology(inet_address ep, std::optional<endpoint_dc_rack> opt_dr, std::optional<node::state> opt_st, std::optional<shard_id> shard_count = std::nullopt) {
+        _topology.add_or_update_endpoint(ep, std::nullopt, std::move(opt_dr), std::move(opt_st), std::move(shard_count));
     }
 
     /**
@@ -933,8 +933,8 @@ token_metadata::get_bootstrap_tokens() const {
 }
 
 void
-token_metadata::update_topology(inet_address ep, endpoint_dc_rack dr, std::optional<node::state> opt_st, std::optional<shard_id> shard_count) {
-    _impl->update_topology(ep, std::move(dr), std::move(opt_st), std::move(shard_count));
+token_metadata::update_topology(inet_address ep, std::optional<endpoint_dc_rack> opt_dr, std::optional<node::state> opt_st, std::optional<shard_id> shard_count) {
+    _impl->update_topology(ep, std::move(opt_dr), std::move(opt_st), std::move(shard_count));
 }
 
 boost::iterator_range<token_metadata::tokens_iterator>

locator/token_metadata.hh

@@ -133,7 +133,7 @@ public:
     /**
      * Update or add endpoint given its inet_address and endpoint_dc_rack.
      */
-    void update_topology(inet_address ep, endpoint_dc_rack dr, std::optional<node::state> opt_st = std::nullopt,
+    void update_topology(inet_address ep, std::optional<endpoint_dc_rack> opt_dr, std::optional<node::state> opt_st = std::nullopt,
                          std::optional<shard_id> shard_count = std::nullopt);
     /**
      * Creates an iterable range of the sorted tokens starting at the token t

locator/types.hh

@@ -31,6 +31,6 @@ struct endpoint_dc_rack {
     bool operator==(const endpoint_dc_rack&) const = default;
 };
 
-using dc_rack_fn = seastar::noncopyable_function<endpoint_dc_rack(inet_address)>;
+using dc_rack_fn = seastar::noncopyable_function<std::optional<endpoint_dc_rack>(inet_address)>;
 
 } // namespace locator

service/storage_service.cc

@@ -3508,15 +3508,26 @@ std::unordered_set<locator::token> storage_service::get_tokens_for(inet_address
     return ret;
 }
 
-locator::endpoint_dc_rack storage_service::get_dc_rack_for(inet_address endpoint) {
-    auto* dc = _gossiper.get_application_state_ptr(endpoint, gms::application_state::DC);
-    auto* rack = _gossiper.get_application_state_ptr(endpoint, gms::application_state::RACK);
+std::optional<locator::endpoint_dc_rack> storage_service::get_dc_rack_for(const gms::endpoint_state& ep_state) {
+    auto* dc = ep_state.get_application_state_ptr(gms::application_state::DC);
+    auto* rack = ep_state.get_application_state_ptr(gms::application_state::RACK);
+    if (!dc || !rack) {
+        return std::nullopt;
+    }
     return locator::endpoint_dc_rack{
-        .dc = dc ? dc->value() : locator::endpoint_dc_rack::default_location.dc,
-        .rack = rack ? rack->value() : locator::endpoint_dc_rack::default_location.rack,
+        .dc = dc->value(),
+        .rack = rack->value(),
     };
 }
 
+std::optional<locator::endpoint_dc_rack> storage_service::get_dc_rack_for(inet_address endpoint) {
+    auto eps = _gossiper.get_endpoint_state_ptr(endpoint);
+    if (!eps) {
+        return std::nullopt;
+    }
+    return get_dc_rack_for(*eps);
+}
+
 void endpoint_lifecycle_notifier::register_subscriber(endpoint_lifecycle_subscriber* subscriber)
 {
     _subscribers.add(subscriber);
@@ -3892,7 +3903,7 @@ storage_service::prepare_replacement_info(std::unordered_set<gms::inet_address>
         }
     }
 
-    auto dc_rack = get_dc_rack_for(replace_address);
+    auto dc_rack = get_dc_rack_for(replace_address).value_or(locator::endpoint_dc_rack::default_location);
 
     if (!replace_host_id) {
         replace_host_id = _gossiper.get_host_id(replace_address);

service/storage_service.hh

@@ -480,7 +480,8 @@ private:
     future<> do_update_system_peers_table(gms::inet_address endpoint, const application_state& state, const versioned_value& value);
 
     std::unordered_set<token> get_tokens_for(inet_address endpoint);
-    locator::endpoint_dc_rack get_dc_rack_for(inet_address endpoint);
+    std::optional<locator::endpoint_dc_rack> get_dc_rack_for(const gms::endpoint_state& ep_state);
+    std::optional<locator::endpoint_dc_rack> get_dc_rack_for(inet_address endpoint);
 private:
     // Should be serialized under token_metadata_lock.
     future<> replicate_to_all_cores(mutable_token_metadata_ptr tmptr) noexcept;

test/boost/storage_proxy_test.cc

@@ -55,7 +55,7 @@ SEASTAR_TEST_CASE(test_get_restricted_ranges) {
             {
                 // Ring with minimum token
                 auto tmptr = locator::make_token_metadata_ptr(locator::token_metadata::config{});
-                tmptr->update_topology(gms::inet_address("10.0.0.1"), {"dc1", "rack1"});
+                tmptr->update_topology(gms::inet_address("10.0.0.1"), locator::endpoint_dc_rack{"dc1", "rack1"});
                 tmptr->update_normal_tokens(std::unordered_set<dht::token>({dht::minimum_token()}), gms::inet_address("10.0.0.1")).get();
 
                 check(tmptr, dht::partition_range::make_singular(ring[0]), {
@@ -69,9 +69,9 @@ SEASTAR_TEST_CASE(test_get_restricted_ranges) {
 
             {
                 auto tmptr = locator::make_token_metadata_ptr(locator::token_metadata::config{});
-                tmptr->update_topology(gms::inet_address("10.0.0.1"), {"dc1", "rack1"});
+                tmptr->update_topology(gms::inet_address("10.0.0.1"), locator::endpoint_dc_rack{"dc1", "rack1"});
                 tmptr->update_normal_tokens(std::unordered_set<dht::token>({ring[2].token()}), gms::inet_address("10.0.0.1")).get();
-                tmptr->update_topology(gms::inet_address("10.0.0.2"), {"dc1", "rack1"});
+                tmptr->update_topology(gms::inet_address("10.0.0.2"), locator::endpoint_dc_rack{"dc1", "rack1"});
                 tmptr->update_normal_tokens(std::unordered_set<dht::token>({ring[5].token()}), gms::inet_address("10.0.0.2")).get();
 
                 check(tmptr, dht::partition_range::make_singular(ring[0]), {