# HG changeset patch
# User Maxim Dounin <mdounin@mdounin.ru>
# Date 1224799419 -14400
# Node ID bef88ba0b37850c114dfa121027b7900caab0aaf
# Parent  e7d1b49e0611754d6dc6628af01b439523ca7d93
Keepalive: distinguish between upstream servers by default.

Option 'single' may be used if single pool of connections desired for some
reason.

diff --git a/ngx_http_upstream_keepalive_module.c b/ngx_http_upstream_keepalive_module.c
--- a/ngx_http_upstream_keepalive_module.c
+++ b/ngx_http_upstream_keepalive_module.c
@@ -12,11 +12,16 @@
 typedef struct {
     ngx_queue_t                        queue;
     ngx_connection_t                  *connection;
+
+    socklen_t                          socklen;
+    struct sockaddr_storage            sockaddr;
+
 } ngx_http_upstream_keepalive_cache_t;
 
 
 typedef struct {
     ngx_uint_t                         max_cached;
+    ngx_uint_t                         single;       /* unsigned:1 */
 
     ngx_queue_t                        cache;
     ngx_queue_t                        free;
@@ -56,7 +61,7 @@ static char *ngx_http_upstream_keepalive
 static ngx_command_t  ngx_http_upstream_keepalive_commands[] = {
 
     { ngx_string("keepalive"),
-      NGX_HTTP_UPS_CONF|NGX_CONF_NOARGS|NGX_CONF_TAKE1,
+      NGX_HTTP_UPS_CONF|NGX_CONF_TAKE12,
       ngx_http_upstream_keepalive,
       0,
       0,
@@ -179,15 +184,16 @@ ngx_http_upstream_get_keepalive_peer(ngx
     ngx_http_upstream_keepalive_peer_data_t  *kp = data;
     ngx_http_upstream_keepalive_cache_t      *item;
 
-    ngx_queue_t       *q;
+    ngx_int_t          rc;
+    ngx_queue_t       *q, *cache;
     ngx_connection_t  *c;
 
     ngx_log_debug0(NGX_LOG_DEBUG_HTTP, pc->log, 0,
                    "get keepalive peer");
 
-    /* XXX single pool of cached connections */
+    /* single pool of cached connections */
 
-    if (!ngx_queue_empty(&kp->conf->cache)) {
+    if (kp->conf->single && !ngx_queue_empty(&kp->conf->cache)) {
 
         q = ngx_queue_head(&kp->conf->cache);
         ngx_queue_remove(q);
@@ -208,7 +214,43 @@ ngx_http_upstream_get_keepalive_peer(ngx
         return NGX_DONE;
     }
 
-    return kp->original_get_peer(pc, kp->data);
+    rc = kp->original_get_peer(pc, kp->data);
+
+    if (kp->conf->single || rc != NGX_OK) {
+        return rc;
+    }
+
+    /* search cache for suitable connection */
+
+    cache = &kp->conf->cache;
+
+    for (q = ngx_queue_head(cache);
+         q != ngx_queue_sentinel(cache);
+         q = ngx_queue_next(q))
+    {
+        item = ngx_queue_data(q, ngx_http_upstream_keepalive_cache_t, queue);
+        c = item->connection;
+
+        if (ngx_memn2cmp((u_char *) &item->sockaddr, (u_char *) pc->sockaddr,
+                         item->socklen, pc->socklen)
+            == 0)
+        {
+            ngx_queue_remove(q);
+            ngx_queue_insert_head(&kp->conf->free, q);
+
+            c->idle = 0;
+            c->log = pc->log;
+            c->read->log = pc->log;
+            c->write->log = pc->log;
+
+            pc->connection = c;
+            pc->cached = 1;
+
+            return NGX_DONE;
+        }
+    }
+
+    return NGX_OK;
 }
 
 
@@ -268,6 +310,9 @@ ngx_http_upstream_free_keepalive_peer(ng
         c->log = ngx_cycle->log;
         c->read->log = ngx_cycle->log;
         c->write->log = ngx_cycle->log;
+
+        item->socklen = pc->socklen;
+        ngx_memcpy(&item->sockaddr, pc->sockaddr, pc->socklen);
     }
 
     return kp->original_free_peer(pc, kp->data, state);
@@ -367,16 +412,22 @@ ngx_http_upstream_keepalive(ngx_conf_t *
 
     value = cf->args->elts;
 
-    for (i = 1; i < cf->args->nelts; i++) {
-
-        if (ngx_strncmp(value[i].data, "cached=", 7) == 0) {
-            n = ngx_atoi(&value[i].data[7], value[i].len - 7);
+    n = ngx_atoi(value[1].data, value[1].len);
 
-            if (n == NGX_ERROR || n == 0) {
-                goto invalid;
-            }
+    if (n == NGX_ERROR || n == 0) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "invalid value \"%V\" in \"%V\" directive",
+                           &value[1], &cmd->name);
+        return NGX_CONF_ERROR;
+    }
 
-            kcf->max_cached = n;
+    kcf->max_cached = n;
+
+    for (i = 2; i < cf->args->nelts; i++) {
+
+        if (ngx_strcmp(value[i].data, "single") == 0) {
+
+            kcf->single = 1;
 
             continue;
         }
diff --git a/t/memcached-keepalive.t b/t/memcached-keepalive.t
--- a/t/memcached-keepalive.t
+++ b/t/memcached-keepalive.t
@@ -20,7 +20,7 @@ select STDOUT; $| = 1;
 eval { require Cache::Memcached; };
 plain(skip_all => 'Cache::Memcached not installed') if $@;
 
-my $t = Test::Nginx->new()->has('rewrite')->has_daemon('memcached')->plan(7)
+my $t = Test::Nginx->new()->has('rewrite')->has_daemon('memcached')->plan(10)
 	->write_file_expand('nginx.conf', <<'EOF');
 
 master_process off;
@@ -39,7 +39,25 @@ http {
 
     upstream memd {
         server 127.0.0.1:8081;
-        keepalive;
+        keepalive 1;
+    }
+
+    upstream memd2 {
+        server 127.0.0.1:8081;
+        server 127.0.0.1:8082;
+        keepalive 1 single;
+    }
+
+    upstream memd3 {
+        server 127.0.0.1:8081;
+        server 127.0.0.1:8082;
+        keepalive 1;
+    }
+
+    upstream memd4 {
+        server 127.0.0.1:8081;
+        server 127.0.0.1:8082;
+        keepalive 10;
     }
 
     server {
@@ -56,20 +74,39 @@ http {
             memcached_next_upstream  not_found;
             memcached_pass memd;
         }
+
+        location /memd2 {
+            set $memcached_key "/";
+            memcached_pass memd2;
+        }
+
+        location /memd3 {
+            set $memcached_key "/";
+            memcached_pass memd3;
+        }
+
+        location /memd4 {
+            set $memcached_key "/";
+            memcached_pass memd4;
+        }
     }
 }
 
 EOF
 
 $t->run_daemon('memcached', '-l', '127.0.0.1', '-p', '8081');
+$t->run_daemon('memcached', '-l', '127.0.0.1', '-p', '8082');
 $t->run();
 
 ###############################################################################
 
-my $memd = Cache::Memcached->new(servers => [ '127.0.0.1:8081' ]);
-$memd->set('/', 'SEE-THIS');
+my $memd1 = Cache::Memcached->new(servers => [ '127.0.0.1:8081' ]);
+my $memd2 = Cache::Memcached->new(servers => [ '127.0.0.1:8082' ]);
 
-my $total = $memd->stats()->{total}->{total_connections};
+$memd1->set('/', 'SEE-THIS');
+$memd2->set('/', 'SEE-THIS');
+
+my $total = $memd1->stats()->{total}->{total_connections};
 
 like(http_get('/'), qr/SEE-THIS/, 'keepalive memcached request');
 like(http_get('/notfound'), qr/404/, 'keepalive memcached not found');
@@ -79,6 +116,48 @@ like(http_get('/'), qr/SEE-THIS/, 'keepa
 like(http_get('/'), qr/SEE-THIS/, 'keepalive memcached request again');
 like(http_get('/'), qr/SEE-THIS/, 'keepalive memcached request again');
 
-is($memd->stats()->{total}->{total_connections}, $total + 1, 'keepalive used');
+is($memd1->stats()->{total}->{total_connections}, $total + 1,
+	'only one connection used');
+
+# two backends with 'single' option - should establish only one connection
+
+$total = $memd1->stats()->{total}->{total_connections} +
+	$memd2->stats()->{total}->{total_connections};
+
+http_get('/memd2');
+http_get('/memd2');
+http_get('/memd2');
+
+is($memd1->stats()->{total}->{total_connections} +
+	$memd2->stats()->{total}->{total_connections}, $total + 1,
+	'only one connection with two backends and single');
+
+$total = $memd1->stats()->{total}->{total_connections} +
+	$memd2->stats()->{total}->{total_connections};
+
+# two backends without 'single' option and maximum number of cached
+# connections set to 1 - should establish new connection on each request
+
+http_get('/memd3');
+http_get('/memd3');
+http_get('/memd3');
+
+is($memd1->stats()->{total}->{total_connections} +
+	$memd2->stats()->{total}->{total_connections}, $total + 3,
+	'3 connections should be established');
+
+# two backends without 'single' option and maximum number of cached
+# connections set to 10 - should establish only two connections (1 per backend)
+
+$total = $memd1->stats()->{total}->{total_connections} +
+        $memd2->stats()->{total}->{total_connections};
+
+http_get('/memd4');
+http_get('/memd4');
+http_get('/memd4');
+
+is($memd1->stats()->{total}->{total_connections} +
+	$memd2->stats()->{total}->{total_connections}, $total + 2,
+	'connection per backend');
 
 ###############################################################################