diff --git a/locator/token_metadata.cc b/locator/token_metadata.cc index abc7507859..7b009a278a 100644 --- a/locator/token_metadata.cc +++ b/locator/token_metadata.cc @@ -118,8 +118,8 @@ public: return _bootstrap_tokens; } - void update_topology(inet_address ep, endpoint_dc_rack dr, std::optional opt_st, std::optional 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 opt_dr, std::optional opt_st, std::optional 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 opt_st, std::optional 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 opt_dr, std::optional opt_st, std::optional shard_count) { + _impl->update_topology(ep, std::move(opt_dr), std::move(opt_st), std::move(shard_count)); } boost::iterator_range diff --git a/locator/token_metadata.hh b/locator/token_metadata.hh index 97f76db34b..e8abaa8ace 100644 --- a/locator/token_metadata.hh +++ b/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 opt_st = std::nullopt, + void update_topology(inet_address ep, std::optional opt_dr, std::optional opt_st = std::nullopt, std::optional shard_count = std::nullopt); /** * Creates an iterable range of the sorted tokens starting at the token t diff --git a/locator/types.hh b/locator/types.hh index 2de382130b..4625062ead 100644 --- a/locator/types.hh +++ b/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; +using dc_rack_fn = seastar::noncopyable_function(inet_address)>; } // namespace locator diff --git a/service/storage_service.cc b/service/storage_service.cc index 4f4f83d10f..7fda75265b 100644 --- a/service/storage_service.cc +++ b/service/storage_service.cc @@ -3508,15 +3508,26 @@ std::unordered_set 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 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 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 } } - 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); diff --git a/service/storage_service.hh b/service/storage_service.hh index 869387b80f..9bc321840a 100644 --- a/service/storage_service.hh +++ b/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 get_tokens_for(inet_address endpoint); - locator::endpoint_dc_rack get_dc_rack_for(inet_address endpoint); + std::optional get_dc_rack_for(const gms::endpoint_state& ep_state); + std::optional 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; diff --git a/test/boost/storage_proxy_test.cc b/test/boost/storage_proxy_test.cc index 49c68faa67..908ba2dd05 100644 --- a/test/boost/storage_proxy_test.cc +++ b/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::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({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({ring[5].token()}), gms::inet_address("10.0.0.2")).get(); check(tmptr, dht::partition_range::make_singular(ring[0]), {