From 16e9036568d9c53aae0e9580571cfe6f693b10e1 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Sat, 4 Nov 2017 15:46:47 +0900
Subject: [PATCH] nghttpx: Fix affinity retry

---
 src/shrpx_client_handler.cc             | 30 ++++++++++++++++++++-----
 src/shrpx_http_downstream_connection.cc | 14 ++++++++----
 src/shrpx_http_downstream_connection.h  |  3 +++
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/src/shrpx_client_handler.cc b/src/shrpx_client_handler.cc
index a49d3b72..17d600c0 100644
--- a/src/shrpx_client_handler.cc
+++ b/src/shrpx_client_handler.cc
@@ -1057,11 +1057,31 @@ ClientHandler::get_downstream_connection(int &err, Downstream *downstream) {
       it = std::begin(affinity_hash);
     }
 
+    auto aff_idx = std::distance(std::begin(affinity_hash), it);
     auto idx = (*it).idx;
+    auto addr = &shared_addr->addrs[idx];
 
-    auto &addr = shared_addr->addrs[idx];
-    if (addr.proto == PROTO_HTTP2) {
-      auto http2session = select_http2_session_with_affinity(group, &addr);
+    if (addr->connect_blocker->blocked()) {
+      size_t i;
+      for (i = aff_idx + 1; i != aff_idx; ++i) {
+        if (i == shared_addr->affinity_hash.size()) {
+          i = 0;
+        }
+        addr = &shared_addr->addrs[shared_addr->affinity_hash[i].idx];
+        if (addr->connect_blocker->blocked()) {
+          continue;
+        }
+        break;
+      }
+      if (i == aff_idx) {
+        err = -1;
+        return nullptr;
+      }
+      aff_idx = i;
+    }
+
+    if (addr->proto == PROTO_HTTP2) {
+      auto http2session = select_http2_session_with_affinity(group, addr);
 
       auto dconn = make_unique<Http2DownstreamConnection>(http2session);
 
@@ -1070,11 +1090,11 @@ ClientHandler::get_downstream_connection(int &err, Downstream *downstream) {
       return std::move(dconn);
     }
 
-    auto &dconn_pool = addr.dconn_pool;
+    auto &dconn_pool = addr->dconn_pool;
     auto dconn = dconn_pool->pop_downstream_connection();
 
     if (!dconn) {
-      dconn = make_unique<HttpDownstreamConnection>(group, idx, conn_.loop,
+      dconn = make_unique<HttpDownstreamConnection>(group, aff_idx, conn_.loop,
                                                     worker_);
     }
 
diff --git a/src/shrpx_http_downstream_connection.cc b/src/shrpx_http_downstream_connection.cc
index 0e063705..3b26ef72 100644
--- a/src/shrpx_http_downstream_connection.cc
+++ b/src/shrpx_http_downstream_connection.cc
@@ -270,10 +270,16 @@ int HttpDownstreamConnection::initiate_connection() {
         assert(addr->dns);
       } else {
         assert(addr_ == nullptr);
-        addr = &addrs[next_downstream];
-
-        if (++next_downstream >= addrs.size()) {
-          next_downstream = 0;
+        if (shared_addr->affinity.type == AFFINITY_NONE) {
+          addr = &addrs[next_downstream];
+          if (++next_downstream >= addrs.size()) {
+            next_downstream = 0;
+          }
+        } else {
+          addr = &addrs[shared_addr->affinity_hash[next_downstream].idx];
+          if (++next_downstream >= shared_addr->affinity_hash.size()) {
+            next_downstream = 0;
+          }
         }
 
         if (addr->proto != PROTO_HTTP1) {
diff --git a/src/shrpx_http_downstream_connection.h b/src/shrpx_http_downstream_connection.h
index 94e1b964..94178858 100644
--- a/src/shrpx_http_downstream_connection.h
+++ b/src/shrpx_http_downstream_connection.h
@@ -108,6 +108,9 @@ private:
   std::unique_ptr<DNSQuery> dns_query_;
   IOControl ioctrl_;
   http_parser response_htp_;
+  // Index to backend address.  If client affinity is enabled, it is
+  // the index to affinity_hash.  Otherwise, it is the index to the
+  // backend addresses.
   ssize_t initial_addr_idx_;
   // true if first write of reused connection succeeded.  For
   // convenience, this is initialized as true.