From 17f26075142ae912f67937099c53d93003118e1e Mon Sep 17 00:00:00 2001
From: Trond Norbye <tn202803@tor01>
Date: Tue, 6 Oct 2009 12:48:57 +0200
Subject: [PATCH] Bug #434843: Large multigets with binary protocol may hang
 client

---
 libmemcached/memcached_get.c      |  7 ++-
 libmemcached/memcached_response.c |  8 +++-
 tests/function.c                  | 77 +++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/libmemcached/memcached_get.c b/libmemcached/memcached_get.c
index 176be6ad..7ee714fe 100644
--- a/libmemcached/memcached_get.c
+++ b/libmemcached/memcached_get.c
@@ -338,6 +338,9 @@ static memcached_return simple_binary_mget(memcached_st *ptr,
       rc= MEMCACHED_SOME_ERRORS;
       continue;
     }
+    
+    /* We just want one pending response per server */
+    memcached_server_response_reset(&ptr->hosts[server_key]); 
     memcached_server_response_increment(&ptr->hosts[server_key]);    
     if ((x > 0 && x == ptr->io_key_prefetch) &&
         memcached_flush_buffers(ptr) != MEMCACHED_SUCCESS)
@@ -371,7 +374,6 @@ static memcached_return simple_binary_mget(memcached_st *ptr,
           memcached_io_reset(&ptr->hosts[x]);
           rc= MEMCACHED_SOME_ERRORS;
         }
-        memcached_server_response_increment(&ptr->hosts[x]);    
       }
     }
 
@@ -438,6 +440,8 @@ static memcached_return replication_binary_mget(memcached_st *ptr,
         success= false;
         continue;
       }
+      /* we just want one pending response per server */
+      memcached_server_response_reset(&ptr->hosts[server]); 
       memcached_server_response_increment(&ptr->hosts[server]);
     }
 
@@ -461,7 +465,6 @@ static memcached_return replication_binary_mget(memcached_st *ptr,
             dead_servers[x]= true;
             success= false;
           }
-          memcached_server_response_increment(&ptr->hosts[x]);
 
           /* mark all of the messages bound for this server as sent! */
           for (x= 0; x < number_of_keys; ++x)
diff --git a/libmemcached/memcached_response.c b/libmemcached/memcached_response.c
index ea13bb15..fe13ddd5 100644
--- a/libmemcached/memcached_response.c
+++ b/libmemcached/memcached_response.c
@@ -356,8 +356,14 @@ static memcached_return binary_read_one_response(memcached_server_st *ptr,
   {
     switch (header.response.opcode)
     {
-    case PROTOCOL_BINARY_CMD_GETK:
     case PROTOCOL_BINARY_CMD_GETKQ:
+      /*
+       * We didn't increment the response counter for the GETKQ packet
+       * (only the final NOOP), so we need to increment the counter again.
+       */ 
+      memcached_server_response_increment(ptr); 
+      /* FALLTHROUGH */
+    case PROTOCOL_BINARY_CMD_GETK:
       {
         uint16_t keylen= header.response.keylen;
         memcached_result_reset(result);
diff --git a/tests/function.c b/tests/function.c
index 844666d1..66f819de 100644
--- a/tests/function.c
+++ b/tests/function.c
@@ -4505,6 +4505,81 @@ static test_return regression_bug_434484(memcached_st *memc)
   return TEST_SUCCESS;
 }
 
+static test_return regression_bug_434843(memcached_st *memc)
+{
+  if (pre_binary(memc) != TEST_SUCCESS)
+    return TEST_SUCCESS;
+
+  memcached_return rc;
+  unsigned int counter= 0;
+  memcached_execute_function callbacks[1]= { [0]= &callback_counter };
+
+  /*
+   * I only want to hit only _one_ server so I know the number of requests I'm
+   * sending in the pipleine to the server. Let's try to do a multiget of
+   * 10240 (that should satisfy most users don't you tink?)
+   */
+  uint32_t number_of_hosts= memc->number_of_hosts;
+  memc->number_of_hosts= 1;
+  const int max_keys= 10240;
+  char **keys= calloc(max_keys, sizeof(char*));
+  size_t *key_length=calloc(max_keys, sizeof(size_t));
+
+  for (int x= 0; x < max_keys; ++x)
+  {
+     char k[251];
+     key_length[x]= snprintf(k, sizeof(k), "0200%u", x);
+     keys[x]= strdup(k);
+     assert(keys[x] != NULL);
+  }
+
+  /*
+   * Run two times.. the first time we should have 100% cache miss,
+   * and the second time we should have 100% cache hits
+   */
+  for (int y= 0; y < 2; ++y)
+  {
+    rc= memcached_mget(memc, (const char**)keys, key_length, max_keys);
+    assert(rc == MEMCACHED_SUCCESS);
+    rc= memcached_fetch_execute(memc, callbacks, (void *)&counter, 1);
+    if (y == 0)
+    {
+      /* The first iteration should give me a 100% cache miss. verify that*/
+      assert(counter == 0);
+      char blob[1024];
+      for (int x= 0; x < max_keys; ++x)
+      {
+        rc= memcached_add(memc, keys[x], key_length[x],
+                          blob, sizeof(blob), 0, 0);
+        assert(rc == MEMCACHED_SUCCESS || rc == MEMCACHED_BUFFERED);
+      }
+    }
+    else
+    {
+      /* Verify that we received all of the key/value pairs */
+      assert(counter == max_keys);
+    }
+  }
+
+  /* Release allocated resources */
+  for (int x= 0; x < max_keys; ++x)
+    free(keys[x]);
+  free(keys);
+  free(key_length);
+
+  memc->number_of_hosts= number_of_hosts;
+  return TEST_SUCCESS;
+}
+
+static test_return regression_bug_434843_buffered(memcached_st *memc)
+{
+  memcached_return rc;
+  rc= memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS, 1);
+  assert(rc == MEMCACHED_SUCCESS);
+
+  return regression_bug_434843(memc);
+}
+
 test_st udp_setup_server_tests[] ={
   {"set_udp_behavior_test", 0, set_udp_behavior_test},
   {"add_tcp_server_udp_client_test", 0, add_tcp_server_udp_client_test},
@@ -4666,6 +4741,8 @@ test_st replication_tests[]= {
  */
 test_st regression_tests[]= {
   {"lp:434484", 1, regression_bug_434484 },
+  {"lp:434843", 1, regression_bug_434843 },
+  {"lp:434843 buffered", 1, regression_bug_434843_buffered },
   {0, 0, 0}
 };
 
-- 
2.30.2