router: plumb through max_host_selection_count

include/envoy/router/router.h

@@ -175,6 +175,12 @@ class RetryPolicy {
    * @return list of RetryHostPredicates to use
    */
   virtual std::vector<Upstream::RetryHostPredicateSharedPtr> retryHostPredicates() const PURE;
+
+  /**
+   * Number of times host selection should be reattempted when selecting a host
+   * for a retry attempt.
+   */
+  virtual uint32_t hostSelectionMaxAttempts() const PURE;
 };
 
 /**
@@ -225,6 +231,11 @@ class RetryState {
    * @return whether host selection should be reattempted
    */
   virtual bool shouldSelectAnotherHost(const Upstream::Host& host) PURE;
+
+  /**
+   * @return how many times host selection should be reattempted during host selection.
+   */
+  virtual uint32_t hostSelectionMaxAttempts() const PURE;
 };
 
 typedef std::unique_ptr<RetryState> RetryStatePtr;

source/common/http/async_client_impl.h

@@ -125,6 +125,7 @@ class AsyncStreamImpl : public AsyncClient::Stream,
     std::vector<Upstream::RetryHostPredicateSharedPtr> retryHostPredicates() const override {
       return {};
     }
+    uint32_t hostSelectionMaxAttempts() const override { return 1; }
     uint32_t numRetries() const override { return 0; }
     uint32_t retryOn() const override { return 0; }
   };

source/common/router/config_impl.cc

@@ -54,6 +54,11 @@ RetryPolicyImpl::RetryPolicyImpl(const envoy::api::v2::route::RouteAction& confi
         host_predicate.name())
         ->createHostPredicate(*this, host_predicate.config());
   }
+
+  auto host_selection_attempts = config.retry_policy().host_selection_retry_max_attempts();
+  if (host_selection_attempts) {
+    host_selection_attempts_ = host_selection_attempts;
+  }
 }
 
 CorsPolicyImpl::CorsPolicyImpl(const envoy::api::v2::route::CorsPolicy& config) {

source/common/router/config_impl.h

@@ -194,6 +194,7 @@ class RetryPolicyImpl : public RetryPolicy, Upstream::RetryHostPredicateFactoryC
   std::vector<Upstream::RetryHostPredicateSharedPtr> retryHostPredicates() const override {
     return retry_host_predicates_;
   }
+  uint32_t hostSelectionMaxAttempts() const override { return host_selection_attempts_; }
 
   // Upstream::RetryHostPredicateFactoryCallbacks
   void addHostPredicate(Upstream::RetryHostPredicateSharedPtr predicate) override {
@@ -205,6 +206,7 @@ class RetryPolicyImpl : public RetryPolicy, Upstream::RetryHostPredicateFactoryC
   uint32_t num_retries_{};
   uint32_t retry_on_{};
   std::vector<Upstream::RetryHostPredicateSharedPtr> retry_host_predicates_;
+  uint32_t host_selection_attempts_{1};
 };
 
 /**

source/common/router/retry_state_impl.cc

@@ -74,6 +74,7 @@ RetryStateImpl::RetryStateImpl(const RetryPolicy& route_policy, Http::HeaderMap&
   const uint32_t base = runtime_.snapshot().getInteger("upstream.base_retry_backoff_ms", 25);
   // Cap the max interval to 10 times the base interval to ensure reasonable backoff intervals.
   backoff_strategy_ = std::make_unique<JitteredBackOffStrategy>(base, base * 10, random_);
+  host_selection_max_attempts_ = route_policy.hostSelectionMaxAttempts();
 }
 
 RetryStateImpl::~RetryStateImpl() { resetRetry(); }

source/common/router/retry_state_impl.h

@@ -50,6 +50,8 @@ class RetryStateImpl : public RetryState {
                   [&host](auto predicate) { predicate->onHostAttempted(host); });
   }
 
+  uint32_t hostSelectionMaxAttempts() const override { return host_selection_max_attempts_; }
+
 private:
   RetryStateImpl(const RetryPolicy& route_policy, Http::HeaderMap& request_headers,
                  const Upstream::ClusterInfo& cluster, Runtime::Loader& runtime,
@@ -72,6 +74,7 @@ class RetryStateImpl : public RetryState {
   Upstream::ResourcePriority priority_;
   BackOffStrategyPtr backoff_strategy_;
   std::vector<Upstream::RetryHostPredicateSharedPtr> retry_host_predicates_;
+  uint32_t host_selection_max_attempts_;
 };
 
 } // namespace Router

source/common/router/router.h

@@ -215,6 +215,14 @@ class Filter : Logger::Loggable<Logger::Id::router>,
     return retry_state_->shouldSelectAnotherHost(host);
   }
 
+  uint32_t hostSelectionRetryCount() const override {
+    if (!is_retry_) {
+      return 1;
+    }
+
+    return retry_state_->hostSelectionMaxAttempts();
+  }
+
   /**
    * Set a computed cookie to be sent with the downstream headers.
    * @param key supplies the size of the cookie

test/common/router/retry_state_impl_test.cc

@@ -27,6 +27,11 @@ class RouterRetryStateImplTest : public testing::Test {
         .WillByDefault(Return(true));
   }
 
+  void setup() {
+    Http::TestHeaderMapImpl headers;
+    setup(headers);
+  }
+
   void setup(Http::HeaderMap& request_headers) {
     state_ = RetryStateImpl::create(policy_, request_headers, cluster_, runtime_, random_,
                                     dispatcher_, Upstream::ResourcePriority::Default);
@@ -393,6 +398,15 @@ TEST_F(RouterRetryStateImplTest, Backoff) {
   EXPECT_EQ(1UL, cluster_.stats().upstream_rq_retry_success_.value());
 }
 
+TEST_F(RouterRetryStateImplTest, HostSelectionAttempts) {
+  policy_.host_selection_max_attempts_ = 2;
+  policy_.retry_on_ = RetryPolicy::RETRY_ON_CONNECT_FAILURE;
+
+  setup();
+
+  EXPECT_EQ(2, state_->hostSelectionMaxAttempts());
+}
+
 TEST_F(RouterRetryStateImplTest, Cancel) {
   // Cover the case where we start a retry, and then we get destructed. This is how the router
   // uses the implementation in the cancel case.

test/common/router/router_test.cc

@@ -1697,6 +1697,72 @@ TEST_F(RouterTest, RetryUpstreamGrpcCancelled) {
   EXPECT_TRUE(verifyHostUpstreamStats(1, 1));
 }
 
+// Verifies that the initial host is select with max host count of one, but during retries
+// RetryPolicy will be consulted.
+TEST_F(RouterTest, RetryRespsectsMaxHostSelectionCount) {
+  router_.reject_all_hosts_ = true;
+
+  NiceMock<Http::MockStreamEncoder> encoder1;
+  Http::StreamDecoder* response_decoder = nullptr;
+  EXPECT_CALL(cm_.conn_pool_, newStream(_, _))
+      .WillOnce(Invoke([&](Http::StreamDecoder& decoder, Http::ConnectionPool::Callbacks& callbacks)
+                           -> Http::ConnectionPool::Cancellable* {
+        response_decoder = &decoder;
+        callbacks.onPoolReady(encoder1, cm_.conn_pool_.host_);
+        return nullptr;
+      }));
+  expectResponseTimerCreate();
+
+  Http::TestHeaderMapImpl headers{{"x-envoy-retry-on", "5xx"}, {"x-envoy-internal", "true"}};
+  HttpTestUtility::addDefaultHeaders(headers);
+  router_.decodeHeaders(headers, false);
+
+  ON_CALL(*router_.retry_state_, hostSelectionMaxAttempts()).WillByDefault(Return(3));
+  // The router should accept any host at this point, since we're not in a retry.
+  EXPECT_EQ(1, router_.hostSelectionRetryCount());
+
+  Buffer::InstancePtr body_data(new Buffer::OwnedImpl("hello"));
+  EXPECT_CALL(*router_.retry_state_, enabled()).WillOnce(Return(true));
+  EXPECT_EQ(Http::FilterDataStatus::StopIterationAndBuffer, router_.decodeData(*body_data, false));
+
+  Http::TestHeaderMapImpl trailers{{"some", "trailer"}};
+  router_.decodeTrailers(trailers);
+
+  // 5xx response.
+  router_.retry_state_->expectRetry();
+  Http::HeaderMapPtr response_headers1(new Http::TestHeaderMapImpl{{":status", "503"}});
+  EXPECT_CALL(encoder1.stream_, resetStream(Http::StreamResetReason::LocalReset));
+  EXPECT_CALL(cm_.conn_pool_.host_->outlier_detector_, putHttpResponseCode(503));
+  response_decoder->decodeHeaders(std::move(response_headers1), false);
+  EXPECT_TRUE(verifyHostUpstreamStats(0, 1));
+
+  // We expect the 5xx response to kick off a new request.
+  NiceMock<Http::MockStreamEncoder> encoder2;
+  EXPECT_CALL(cm_.conn_pool_, newStream(_, _))
+      .WillOnce(Invoke([&](Http::StreamDecoder& decoder, Http::ConnectionPool::Callbacks& callbacks)
+                           -> Http::ConnectionPool::Cancellable* {
+        response_decoder = &decoder;
+        callbacks.onPoolReady(encoder2, cm_.conn_pool_.host_);
+        return nullptr;
+      }));
+  ON_CALL(callbacks_, decodingBuffer()).WillByDefault(Return(body_data.get()));
+  EXPECT_CALL(encoder2, encodeHeaders(_, false));
+  EXPECT_CALL(encoder2, encodeData(_, false));
+  EXPECT_CALL(encoder2, encodeTrailers(_));
+  router_.retry_state_->callback_();
+
+  // Now that we're triggered a retry, we should see the configured number of host selections.
+  EXPECT_EQ(3, router_.hostSelectionRetryCount());
+
+  // Normal response.
+  EXPECT_CALL(*router_.retry_state_, shouldRetry(_, _, _)).WillOnce(Return(RetryStatus::No));
+  EXPECT_CALL(cm_.conn_pool_.host_->health_checker_, setUnhealthy()).Times(0);
+  Http::HeaderMapPtr response_headers2(new Http::TestHeaderMapImpl{{":status", "200"}});
+  EXPECT_CALL(cm_.conn_pool_.host_->outlier_detector_, putHttpResponseCode(200));
+  response_decoder->decodeHeaders(std::move(response_headers2), true);
+  EXPECT_TRUE(verifyHostUpstreamStats(1, 1));
+}
+
 // Verifies that the initial request accepts any host, but during retries
 // RetryPolicy will be consulted.
 TEST_F(RouterTest, RetryRespectsRetryHostPredicate) {

test/mocks/router/mocks.h

@@ -70,10 +70,12 @@ class TestRetryPolicy : public RetryPolicy {
   uint32_t numRetries() const override { return num_retries_; }
   uint32_t retryOn() const override { return retry_on_; }
   MOCK_CONST_METHOD0(retryHostPredicates, std::vector<Upstream::RetryHostPredicateSharedPtr>());
+  uint32_t hostSelectionMaxAttempts() const override { return host_selection_max_attempts_; }
 
   std::chrono::milliseconds per_try_timeout_{0};
   uint32_t num_retries_{};
   uint32_t retry_on_{};
+  uint32_t host_selection_max_attempts_;
 };
 
 class MockRetryState : public RetryState {
@@ -89,6 +91,7 @@ class MockRetryState : public RetryState {
                                         DoRetryCallback callback));
   MOCK_METHOD1(onHostAttempted, void(Upstream::HostDescriptionConstSharedPtr));
   MOCK_METHOD1(shouldSelectAnotherHost, bool(const Upstream::Host& host));
+  MOCK_CONST_METHOD0(hostSelectionMaxAttempts, uint32_t());
 
   DoRetryCallback callback_;
 };