grpc 1.71.0 → 1.73.0

data/src/core/client_channel/global_subchannel_pool.cc

@@ -26,12 +26,13 @@
 
 namespace grpc_core {
 
-RefCountedPtr<GlobalSubchannelPool> GlobalSubchannelPool::instance() {
-  static GlobalSubchannelPool* p = new GlobalSubchannelPool();
-  return p->RefAsSubclass<GlobalSubchannelPool>();
+RefCountedPtr<LegacyGlobalSubchannelPool>
+LegacyGlobalSubchannelPool::instance() {
+  static LegacyGlobalSubchannelPool* p = new LegacyGlobalSubchannelPool();
+  return p->RefAsSubclass<LegacyGlobalSubchannelPool>();
 }
 
-RefCountedPtr<Subchannel> GlobalSubchannelPool::RegisterSubchannel(
+RefCountedPtr<Subchannel> LegacyGlobalSubchannelPool::RegisterSubchannel(
     const SubchannelKey& key, RefCountedPtr<Subchannel> constructed) {
   MutexLock lock(&mu_);
   auto it = subchannel_map_.find(key);
@@ -43,8 +44,8 @@ RefCountedPtr<Subchannel> GlobalSubchannelPool::RegisterSubchannel(
   return constructed;
 }
 
-void GlobalSubchannelPool::UnregisterSubchannel(const SubchannelKey& key,
-                                                Subchannel* subchannel) {
+void LegacyGlobalSubchannelPool::UnregisterSubchannel(const SubchannelKey& key,
+                                                      Subchannel* subchannel) {
   MutexLock lock(&mu_);
   auto it = subchannel_map_.find(key);
   // delete only if key hasn't been re-registered to a different subchannel
@@ -54,7 +55,7 @@ void GlobalSubchannelPool::UnregisterSubchannel(const SubchannelKey& key,
   }
 }
 
-RefCountedPtr<Subchannel> GlobalSubchannelPool::FindSubchannel(
+RefCountedPtr<Subchannel> LegacyGlobalSubchannelPool::FindSubchannel(
     const SubchannelKey& key) {
   MutexLock lock(&mu_);
   auto it = subchannel_map_.find(key);
@@ -62,4 +63,64 @@ RefCountedPtr<Subchannel> GlobalSubchannelPool::FindSubchannel(
   return it->second->RefIfNonZero();
 }
 
+RefCountedPtr<GlobalSubchannelPool> GlobalSubchannelPool::instance() {
+  static GlobalSubchannelPool* p = new GlobalSubchannelPool();
+  return p->RefAsSubclass<GlobalSubchannelPool>();
+}
+
+RefCountedPtr<Subchannel> GlobalSubchannelPool::RegisterSubchannel(
+    const SubchannelKey& key, RefCountedPtr<Subchannel> constructed) {
+  auto shard_index = ShardIndex(key);
+  auto& write_shard = write_shards_[shard_index];
+  auto& read_shard = read_shards_[shard_index];
+  SubchannelMap old_map1;
+  SubchannelMap old_map2;
+  MutexLock lock(&write_shard.mu);
+  auto* existing = write_shard.map.Lookup(key);
+  if (existing != nullptr) return (*existing)->RefIfNonZero();
+  old_map1 = std::exchange(write_shard.map,
+                           write_shard.map.Add(key, constructed->WeakRef()));
+  MutexLock lock_read(&read_shard.mu);
+  old_map2 = std::exchange(read_shard.map, write_shard.map);
+  return constructed;
+}
+
+void GlobalSubchannelPool::UnregisterSubchannel(const SubchannelKey& key,
+                                                Subchannel* subchannel) {
+  auto shard_index = ShardIndex(key);
+  auto& write_shard = write_shards_[shard_index];
+  auto& read_shard = read_shards_[shard_index];
+  SubchannelMap old_map1;
+  SubchannelMap old_map2;
+  MutexLock lock(&write_shard.mu);
+  auto* existing = write_shard.map.Lookup(key);
+  // delete only if key hasn't been re-registered to a different subchannel
+  // between strong-unreffing and unregistration of subchannel.
+  if (existing == nullptr || existing->get() != subchannel) return;
+  old_map1 = std::exchange(write_shard.map, write_shard.map.Remove(key));
+  MutexLock lock_read(&read_shard.mu);
+  old_map2 = std::exchange(read_shard.map, write_shard.map);
+}
+
+RefCountedPtr<Subchannel> GlobalSubchannelPool::FindSubchannel(
+    const SubchannelKey& key) {
+  auto shard_index = ShardIndex(key);
+  auto& read_shard = read_shards_[shard_index];
+  read_shard.mu.Lock();
+  auto map = read_shard.map;
+  read_shard.mu.Unlock();
+  auto* subchannel = map.Lookup(key);
+  if (subchannel == nullptr) return nullptr;
+  return (*subchannel)->RefIfNonZero();
+}
+
+size_t GlobalSubchannelPool::ShardIndex(const SubchannelKey& key) {
+  absl::string_view addr(key.address().addr, key.address().len);
+  return absl::HashOf(addr) % kShards;
+}
+
+GlobalSubchannelPool::GlobalSubchannelPool() = default;
+
+GlobalSubchannelPool::~GlobalSubchannelPool() = default;
+
 }  // namespace grpc_core

data/src/core/client_channel/global_subchannel_pool.h

@@ -32,10 +32,10 @@ namespace grpc_core {
 
 // The global subchannel pool. It shares subchannels among channels. There
 // should be only one instance of this class.
-class GlobalSubchannelPool final : public SubchannelPoolInterface {
+class LegacyGlobalSubchannelPool final : public SubchannelPoolInterface {
  public:
   // Gets the singleton instance.
-  static RefCountedPtr<GlobalSubchannelPool> instance();
+  static RefCountedPtr<LegacyGlobalSubchannelPool> instance();
 
   // Implements interface methods.
   RefCountedPtr<Subchannel> RegisterSubchannel(
@@ -48,8 +48,8 @@ class GlobalSubchannelPool final : public SubchannelPoolInterface {
       ABSL_LOCKS_EXCLUDED(mu_);
 
  private:
-  GlobalSubchannelPool() {}
-  ~GlobalSubchannelPool() override {}
+  LegacyGlobalSubchannelPool() {}
+  ~LegacyGlobalSubchannelPool() override {}
 
   // A map from subchannel key to subchannel.
   std::map<SubchannelKey, Subchannel*> subchannel_map_ ABSL_GUARDED_BY(mu_);
@@ -57,6 +57,39 @@ class GlobalSubchannelPool final : public SubchannelPoolInterface {
   Mutex mu_;
 };
 
+// The global subchannel pool. It shares subchannels among channels. There
+// should be only one instance of this class.
+class GlobalSubchannelPool final : public SubchannelPoolInterface {
+ public:
+  // Gets the singleton instance.
+  static RefCountedPtr<GlobalSubchannelPool> instance();
+
+  // Implements interface methods.
+  RefCountedPtr<Subchannel> RegisterSubchannel(
+      const SubchannelKey& key, RefCountedPtr<Subchannel> constructed) override;
+  void UnregisterSubchannel(const SubchannelKey& key,
+                            Subchannel* subchannel) override;
+  RefCountedPtr<Subchannel> FindSubchannel(const SubchannelKey& key) override;
+
+ private:
+  GlobalSubchannelPool();
+  ~GlobalSubchannelPool() override;
+
+  static const size_t kShards = 127;
+
+  using SubchannelMap = AVL<SubchannelKey, WeakRefCountedPtr<Subchannel>>;
+  struct LockedMap {
+    Mutex mu;
+    SubchannelMap map ABSL_GUARDED_BY(mu);
+  };
+  using ShardedMap = std::array<LockedMap, kShards>;
+
+  static size_t ShardIndex(const SubchannelKey& key);
+
+  ShardedMap write_shards_;
+  ShardedMap read_shards_;
+};
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_CLIENT_CHANNEL_GLOBAL_SUBCHANNEL_POOL_H

data/src/core/client_channel/lb_metadata.h

@@ -22,7 +22,7 @@
 #include <vector>
 
 #include "absl/strings/string_view.h"
-#include "src/core/lib/transport/metadata_batch.h"
+#include "src/core/call/metadata_batch.h"
 #include "src/core/load_balancing/lb_policy.h"
 
 namespace grpc_core {

data/src/core/client_channel/load_balanced_call_destination.cc

@@ -15,12 +15,12 @@
 #include "src/core/client_channel/load_balanced_call_destination.h"
 
 #include "absl/log/log.h"
+#include "src/core/call/status_util.h"
 #include "src/core/client_channel/client_channel.h"
 #include "src/core/client_channel/client_channel_internal.h"
 #include "src/core/client_channel/lb_metadata.h"
 #include "src/core/client_channel/subchannel.h"
 #include "src/core/config/core_configuration.h"
-#include "src/core/lib/channel/status_util.h"
 #include "src/core/lib/promise/loop.h"
 #include "src/core/telemetry/call_tracer.h"
 

data/src/core/client_channel/load_balanced_call_destination.h

@@ -16,9 +16,9 @@
 #define GRPC_SRC_CORE_CLIENT_CHANNEL_LOAD_BALANCED_CALL_DESTINATION_H
 
 #include "absl/functional/any_invocable.h"
+#include "src/core/call/call_destination.h"
 #include "src/core/client_channel/client_channel.h"
 #include "src/core/lib/promise/context.h"
-#include "src/core/lib/transport/call_destination.h"
 #include "src/core/load_balancing/lb_policy.h"
 
 namespace grpc_core {

data/src/core/client_channel/retry_filter_legacy_call_data.cc

@@ -24,11 +24,12 @@
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
+#include "src/core/call/metadata_batch.h"
+#include "src/core/call/status_util.h"
 #include "src/core/client_channel/client_channel_internal.h"
 #include "src/core/client_channel/retry_service_config.h"
 #include "src/core/client_channel/retry_throttle.h"
 #include "src/core/lib/channel/channel_stack.h"
-#include "src/core/lib/channel/status_util.h"
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/iomgr/call_combiner.h"
 #include "src/core/lib/iomgr/closure.h"
@@ -39,7 +40,6 @@
 #include "src/core/lib/slice/slice.h"
 #include "src/core/lib/slice/slice_buffer.h"
 #include "src/core/lib/transport/error_utils.h"
-#include "src/core/lib/transport/metadata_batch.h"
 #include "src/core/lib/transport/transport.h"
 #include "src/core/util/backoff.h"
 #include "src/core/util/construct_destruct.h"
@@ -53,6 +53,8 @@
 
 namespace grpc_core {
 
+using TaskHandle = grpc_event_engine::experimental::EventEngine::TaskHandle;
+
 //
 // RetryFilter::LegacyCallData::CallStackDestructionBarrier
 //
@@ -220,7 +222,7 @@ void RetryFilter::LegacyCallData::CallAttempt::MaybeSwitchToFastPath() {
   // If we've already switched to fast path, there's nothing to do here.
   if (calld_->committed_call_ != nullptr) return;
   // If the perAttemptRecvTimeout timer is pending, we can't switch yet.
-  if (per_attempt_recv_timer_handle_.has_value()) return;
+  if (per_attempt_recv_timer_handle_ != TaskHandle::kInvalid) return;
   // If there are still send ops to replay, we can't switch yet.
   if (HaveSendOpsToReplay()) return;
   // If we started an internal batch for recv_trailing_metadata but have not
@@ -643,31 +645,33 @@ void RetryFilter::LegacyCallData::CallAttempt::OnPerAttemptRecvTimerLocked(
       << "chand=" << calld->chand_ << " calld=" << calld
       << " attempt=" << call_attempt
       << ": perAttemptRecvTimeout timer fired: error=" << StatusToString(error)
-      << ", per_attempt_recv_timer_handle_.has_value()="
-      << call_attempt->per_attempt_recv_timer_handle_.has_value();
+      << ", per_attempt_recv_timer_handle_ is valid ="
+      << (call_attempt->per_attempt_recv_timer_handle_ != TaskHandle::kInvalid);
   CallCombinerClosureList closures;
-  call_attempt->per_attempt_recv_timer_handle_.reset();
-  // Cancel this attempt.
-  // TODO(roth): When implementing hedging, we should not cancel the
-  // current attempt.
-  call_attempt->MaybeAddBatchForCancelOp(
-      grpc_error_set_int(
-          GRPC_ERROR_CREATE("retry perAttemptRecvTimeout exceeded"),
-          StatusIntProperty::kRpcStatus, GRPC_STATUS_CANCELLED),
-      &closures);
-  // Check whether we should retry.
-  if (call_attempt->ShouldRetry(/*status=*/std::nullopt,
-                                /*server_pushback_ms=*/std::nullopt)) {
-    // Mark current attempt as abandoned.
-    call_attempt->Abandon();
-    // We are retrying.  Start backoff timer.
-    calld->StartRetryTimer(/*server_pushback=*/std::nullopt);
-  } else {
-    // Not retrying, so commit the call.
-    calld->RetryCommit(call_attempt);
-    // If retry state is no longer needed, switch to fast path for
-    // subsequent batches.
-    call_attempt->MaybeSwitchToFastPath();
+  if (call_attempt->per_attempt_recv_timer_handle_ != TaskHandle::kInvalid) {
+    call_attempt->per_attempt_recv_timer_handle_ = TaskHandle::kInvalid;
+    // Cancel this attempt.
+    // TODO(roth): When implementing hedging, we should not cancel the
+    // current attempt.
+    call_attempt->MaybeAddBatchForCancelOp(
+        grpc_error_set_int(
+            GRPC_ERROR_CREATE("retry perAttemptRecvTimeout exceeded"),
+            StatusIntProperty::kRpcStatus, GRPC_STATUS_CANCELLED),
+        &closures);
+    // Check whether we should retry.
+    if (call_attempt->ShouldRetry(/*status=*/std::nullopt,
+                                  /*server_pushback_ms=*/std::nullopt)) {
+      // Mark current attempt as abandoned.
+      call_attempt->Abandon();
+      // We are retrying.  Start backoff timer.
+      calld->StartRetryTimer(/*server_pushback=*/std::nullopt);
+    } else {
+      // Not retrying, so commit the call.
+      calld->RetryCommit(call_attempt);
+      // If retry state is no longer needed, switch to fast path for
+      // subsequent batches.
+      call_attempt->MaybeSwitchToFastPath();
+    }
   }
   closures.RunClosures(calld->call_combiner_);
   call_attempt->Unref(DEBUG_LOCATION, "OnPerAttemptRecvTimer");
@@ -676,16 +680,16 @@ void RetryFilter::LegacyCallData::CallAttempt::OnPerAttemptRecvTimerLocked(
 
 void RetryFilter::LegacyCallData::CallAttempt::
     MaybeCancelPerAttemptRecvTimer() {
-  if (per_attempt_recv_timer_handle_.has_value()) {
+  if (per_attempt_recv_timer_handle_ != TaskHandle::kInvalid) {
     GRPC_TRACE_LOG(retry, INFO)
         << "chand=" << calld_->chand_ << " calld=" << calld_
         << " attempt=" << this << ": cancelling perAttemptRecvTimeout timer";
     if (calld_->chand_->event_engine()->Cancel(
-            *per_attempt_recv_timer_handle_)) {
+            per_attempt_recv_timer_handle_)) {
       Unref(DEBUG_LOCATION, "OnPerAttemptRecvTimer");
       GRPC_CALL_STACK_UNREF(calld_->owning_call_, "OnPerAttemptRecvTimer");
     }
-    per_attempt_recv_timer_handle_.reset();
+    per_attempt_recv_timer_handle_ = TaskHandle::kInvalid;
   }
 }
 
@@ -1481,7 +1485,6 @@ RetryFilter::LegacyCallData::LegacyCallData(RetryFilter* chand,
               .set_max_backoff(retry_policy_ == nullptr
                                    ? Duration::Zero()
                                    : retry_policy_->max_backoff())),
-      path_(CSliceRef(args.path)),
       deadline_(args.deadline),
       arena_(args.arena),
       owning_call_(args.call_stack),
@@ -1497,7 +1500,6 @@ RetryFilter::LegacyCallData::LegacyCallData(RetryFilter* chand,
 
 RetryFilter::LegacyCallData::~LegacyCallData() {
   FreeAllCachedSendOpData();
-  CSliceUnref(path_);
   // Make sure there are no remaining pending batches.
   for (size_t i = 0; i < GPR_ARRAY_SIZE(pending_batches_); ++i) {
     CHECK_EQ(pending_batches_[i].batch, nullptr);
@@ -1550,13 +1552,13 @@ void RetryFilter::LegacyCallData::StartTransportStreamOpBatch(
       return;
     }
     // Cancel retry timer if needed.
-    if (retry_timer_handle_.has_value()) {
+    if (retry_timer_handle_ != TaskHandle::kInvalid) {
       GRPC_TRACE_LOG(retry, INFO) << "chand=" << chand_ << " calld=" << this
                                   << ": cancelling retry timer";
-      if (chand_->event_engine()->Cancel(*retry_timer_handle_)) {
+      if (chand_->event_engine()->Cancel(retry_timer_handle_)) {
         GRPC_CALL_STACK_UNREF(owning_call_, "OnRetryTimer");
       }
-      retry_timer_handle_.reset();
+      retry_timer_handle_ = TaskHandle::kInvalid;
       FreeAllCachedSendOpData();
     }
     // We have no call attempt, so there's nowhere to send the cancellation
@@ -1570,7 +1572,7 @@ void RetryFilter::LegacyCallData::StartTransportStreamOpBatch(
   PendingBatch* pending = PendingBatchesAdd(batch);
   // If the timer is pending, yield the call combiner and wait for it to
   // run, since we don't want to start another call attempt until it does.
-  if (retry_timer_handle_.has_value()) {
+  if (retry_timer_handle_ != TaskHandle::kInvalid) {
     GRPC_CALL_COMBINER_STOP(call_combiner_,
                             "added pending batch while retry timer pending");
     return;
@@ -1627,9 +1629,9 @@ void RetryFilter::LegacyCallData::StartTransportStreamOpBatch(
 OrphanablePtr<ClientChannelFilter::FilterBasedLoadBalancedCall>
 RetryFilter::LegacyCallData::CreateLoadBalancedCall(
     absl::AnyInvocable<void()> on_commit, bool is_transparent_retry) {
-  grpc_call_element_args args = {owning_call_,     nullptr,   path_,
-                                 /*start_time=*/0, deadline_, arena_,
-                                 call_combiner_};
+  grpc_call_element_args args = {owning_call_,     nullptr,
+                                 /*start_time=*/0, deadline_,
+                                 arena_,           call_combiner_};
   return chand_->client_channel()->CreateLoadBalancedCall(
       args, pollent_,
       // This callback holds a ref to the CallStackDestructionBarrier
@@ -1912,8 +1914,8 @@ void RetryFilter::LegacyCallData::OnRetryTimer() {
 void RetryFilter::LegacyCallData::OnRetryTimerLocked(
     void* arg, grpc_error_handle /*error*/) {
   auto* calld = static_cast<RetryFilter::LegacyCallData*>(arg);
-  if (calld->retry_timer_handle_.has_value()) {
-    calld->retry_timer_handle_.reset();
+  if (calld->retry_timer_handle_ != TaskHandle::kInvalid) {
+    calld->retry_timer_handle_ = TaskHandle::kInvalid;
     calld->CreateCallAttempt(/*is_transparent_retry=*/false);
   }
   GRPC_CALL_STACK_UNREF(calld->owning_call_, "OnRetryTimer");

data/src/core/client_channel/retry_filter_legacy_call_data.h

@@ -27,6 +27,7 @@
 
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/any_invocable.h"
+#include "src/core/call/metadata_batch.h"
 #include "src/core/client_channel/client_channel_filter.h"
 #include "src/core/client_channel/retry_filter.h"
 #include "src/core/client_channel/retry_service_config.h"
@@ -39,7 +40,6 @@
 #include "src/core/lib/iomgr/polling_entity.h"
 #include "src/core/lib/resource_quota/arena.h"
 #include "src/core/lib/slice/slice_buffer.h"
-#include "src/core/lib/transport/metadata_batch.h"
 #include "src/core/lib/transport/transport.h"
 #include "src/core/util/backoff.h"
 #include "src/core/util/debug_location.h"
@@ -256,8 +256,9 @@ class RetryFilter::LegacyCallData final {
     bool lb_call_committed_ = false;
 
     grpc_closure on_per_attempt_recv_timer_;
-    std::optional<grpc_event_engine::experimental::EventEngine::TaskHandle>
-        per_attempt_recv_timer_handle_;
+    grpc_event_engine::experimental::EventEngine::TaskHandle
+        per_attempt_recv_timer_handle_ =
+            grpc_event_engine::experimental::EventEngine::TaskHandle::kInvalid;
 
     // BatchData.batch.payload points to this.
     grpc_transport_stream_op_batch_payload batch_payload_;
@@ -373,7 +374,6 @@ class RetryFilter::LegacyCallData final {
   const internal::RetryMethodConfig* retry_policy_ = nullptr;
   BackOff retry_backoff_;
 
-  grpc_slice path_;  // Request path.
   Timestamp deadline_;
   Arena* arena_;
   grpc_call_stack* owning_call_;
@@ -410,8 +410,8 @@ class RetryFilter::LegacyCallData final {
   bool retry_codepath_started_ : 1;
   bool sent_transparent_retry_not_seen_by_server_ : 1;
   int num_attempts_completed_ = 0;
-  std::optional<grpc_event_engine::experimental::EventEngine::TaskHandle>
-      retry_timer_handle_;
+  grpc_event_engine::experimental::EventEngine::TaskHandle retry_timer_handle_ =
+      grpc_event_engine::experimental::EventEngine::TaskHandle::kInvalid;
   grpc_closure retry_closure_;
 
   // Cached data for retrying send ops.

data/src/core/client_channel/retry_interceptor.h

@@ -15,12 +15,12 @@
 #ifndef GRPC_SRC_CORE_CLIENT_CHANNEL_RETRY_INTERCEPTOR_H
 #define GRPC_SRC_CORE_CLIENT_CHANNEL_RETRY_INTERCEPTOR_H
 
+#include "src/core/call/interception_chain.h"
 #include "src/core/call/request_buffer.h"
 #include "src/core/client_channel/client_channel_args.h"
 #include "src/core/client_channel/retry_service_config.h"
 #include "src/core/client_channel/retry_throttle.h"
 #include "src/core/filter/filter_args.h"
-#include "src/core/lib/transport/interception_chain.h"
 #include "src/core/util/backoff.h"
 
 namespace grpc_core {

data/src/core/client_channel/retry_service_config.cc

@@ -30,9 +30,9 @@
 #include "absl/log/log.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
+#include "src/core/call/status_util.h"
 #include "src/core/config/core_configuration.h"
 #include "src/core/lib/channel/channel_args.h"
-#include "src/core/lib/channel/status_util.h"
 #include "src/core/util/json/json_channel_args.h"
 
 // As per the retry design, we do not allow more than 5 retry attempts.

data/src/core/client_channel/retry_service_config.h

@@ -25,9 +25,9 @@
 #include <optional>
 
 #include "absl/strings/string_view.h"
+#include "src/core/call/status_util.h"
 #include "src/core/config/core_configuration.h"
 #include "src/core/lib/channel/channel_args.h"
-#include "src/core/lib/channel/status_util.h"
 #include "src/core/service_config/service_config_parser.h"
 #include "src/core/util/json/json.h"
 #include "src/core/util/json/json_args.h"

data/src/core/client_channel/retry_throttle.cc

@@ -49,32 +49,12 @@ T ClampedAdd(std::atomic<T>& value, T delta, T min, T max) {
 // ServerRetryThrottleData
 //
 
-ServerRetryThrottleData::ServerRetryThrottleData(
-    uintptr_t max_milli_tokens, uintptr_t milli_token_ratio,
-    ServerRetryThrottleData* old_throttle_data)
+ServerRetryThrottleData::ServerRetryThrottleData(uintptr_t max_milli_tokens,
+                                                 uintptr_t milli_token_ratio,
+                                                 uintptr_t milli_tokens)
     : max_milli_tokens_(max_milli_tokens),
-      milli_token_ratio_(milli_token_ratio) {
-  uintptr_t initial_milli_tokens = max_milli_tokens;
-  // If there was a pre-existing entry for this server name, initialize
-  // the token count by scaling proportionately to the old data.  This
-  // ensures that if we're already throttling retries on the old scale,
-  // we will start out doing the same thing on the new one.
-  if (old_throttle_data != nullptr) {
-    double token_fraction =
-        static_cast<double>(
-            old_throttle_data->milli_tokens_.load(std::memory_order_relaxed)) /
-        static_cast<double>(old_throttle_data->max_milli_tokens_);
-    initial_milli_tokens =
-        static_cast<uintptr_t>(token_fraction * max_milli_tokens);
-  }
-  milli_tokens_.store(initial_milli_tokens, std::memory_order_relaxed);
-  // If there was a pre-existing entry, mark it as stale and give it a
-  // pointer to the new entry, which is its replacement.
-  if (old_throttle_data != nullptr) {
-    Ref().release();  // Ref held by pre-existing entry.
-    old_throttle_data->replacement_.store(this, std::memory_order_release);
-  }
-}
+      milli_token_ratio_(milli_token_ratio),
+      milli_tokens_(milli_tokens) {}
 
 ServerRetryThrottleData::~ServerRetryThrottleData() {
   ServerRetryThrottleData* replacement =
@@ -84,6 +64,11 @@ ServerRetryThrottleData::~ServerRetryThrottleData() {
   }
 }
 
+void ServerRetryThrottleData::SetReplacement(
+    RefCountedPtr<ServerRetryThrottleData> replacement) {
+  replacement_.store(replacement.release(), std::memory_order_release);
+}
+
 void ServerRetryThrottleData::GetReplacementThrottleDataIfNeeded(
     ServerRetryThrottleData** throttle_data) {
   while (true) {
@@ -133,20 +118,31 @@ RefCountedPtr<ServerRetryThrottleData> ServerRetryThrottleMap::GetDataForServer(
     const std::string& server_name, uintptr_t max_milli_tokens,
     uintptr_t milli_token_ratio) {
   MutexLock lock(&mu_);
-  auto it = map_.find(server_name);
-  ServerRetryThrottleData* throttle_data =
-      it == map_.end() ? nullptr : it->second.get();
+  auto& throttle_data = map_[server_name];
   if (throttle_data == nullptr ||
       throttle_data->max_milli_tokens() != max_milli_tokens ||
       throttle_data->milli_token_ratio() != milli_token_ratio) {
     // Entry not found, or found with old parameters.  Create a new one.
-    it = map_.emplace(server_name,
-                      MakeRefCounted<ServerRetryThrottleData>(
-                          max_milli_tokens, milli_token_ratio, throttle_data))
-             .first;
-    throttle_data = it->second.get();
+    auto old_throttle_data = std::move(throttle_data);
+    uintptr_t initial_milli_tokens = max_milli_tokens;
+    // If there was a pre-existing entry for this server name, initialize
+    // the token count by scaling proportionately to the old data.  This
+    // ensures that if we're already throttling retries on the old scale,
+    // we will start out doing the same thing on the new one.
+    if (old_throttle_data != nullptr) {
+      double token_fraction =
+          static_cast<double>(old_throttle_data->milli_tokens()) /
+          static_cast<double>(old_throttle_data->max_milli_tokens());
+      initial_milli_tokens =
+          static_cast<uintptr_t>(token_fraction * max_milli_tokens);
+    }
+    throttle_data = MakeRefCounted<ServerRetryThrottleData>(
+        max_milli_tokens, milli_token_ratio, initial_milli_tokens);
+    if (old_throttle_data != nullptr) {
+      old_throttle_data->SetReplacement(throttle_data);
+    }
   }
-  return throttle_data->Ref();
+  return throttle_data;
 }
 
 }  // namespace internal

data/src/core/client_channel/retry_throttle.h

@@ -34,13 +34,14 @@
 namespace grpc_core {
 namespace internal {
 
+class ServerRetryThrottleMap;
+
 /// Tracks retry throttling data for an individual server name.
 class ServerRetryThrottleData final
     : public RefCounted<ServerRetryThrottleData> {
  public:
   ServerRetryThrottleData(uintptr_t max_milli_tokens,
-                          uintptr_t milli_token_ratio,
-                          ServerRetryThrottleData* old_throttle_data);
+                          uintptr_t milli_token_ratio, uintptr_t milli_tokens);
   ~ServerRetryThrottleData() override;
 
   /// Records a failure.  Returns true if it's okay to send a retry.
@@ -51,8 +52,15 @@ class ServerRetryThrottleData final
 
   uintptr_t max_milli_tokens() const { return max_milli_tokens_; }
   uintptr_t milli_token_ratio() const { return milli_token_ratio_; }
+  intptr_t milli_tokens() const {
+    return milli_tokens_.load(std::memory_order_relaxed);
+  }
 
  private:
+  friend ServerRetryThrottleMap;
+
+  void SetReplacement(RefCountedPtr<ServerRetryThrottleData> replacement);
+
   void GetReplacementThrottleDataIfNeeded(
       ServerRetryThrottleData** throttle_data);