nginx_upstream_hash/ngx_http_upstream_hash_module.c

   1 /*
   2  * Hash a variable to choose an upstream server.
   3  *
   4  * Copyright (C) Evan Miller
   5  *
   6  * This module can be distributed under the same terms as Nginx itself.
   7  */
   8
   9
  10 #include <ngx_config.h>
  11 #include <ngx_core.h>
  12 #include <ngx_http.h>
  13
  14 #define ngx_bitvector_index(index) index / (8 * sizeof(uintptr_t))
  15 #define ngx_bitvector_bit(index) (uintptr_t) 1 << index % (8 * sizeof(uintptr_t))
  16
  17 typedef struct {
  18     struct sockaddr                *sockaddr;
  19     socklen_t                       socklen;
  20     ngx_str_t                       name;
  21 } ngx_http_upstream_hash_peer_t;
  22
  23 typedef struct {
  24     ngx_uint_t                        number;
  25     ngx_http_upstream_hash_peer_t     peer[0];
  26 } ngx_http_upstream_hash_peers_t;
  27
  28 typedef struct {
  29     ngx_http_upstream_hash_peers_t   *peers;
  30     ngx_uint_t                        hash;
  31     ngx_str_t                         current_key;
  32     ngx_str_t                         original_key;
  33     ngx_uint_t                        try_i;
  34     uintptr_t                         tried[1];
  35 } ngx_http_upstream_hash_peer_data_t;
  36
  37
  38 static ngx_int_t ngx_http_upstream_init_hash_peer(ngx_http_request_t *r,
  39     ngx_http_upstream_srv_conf_t *us);
  40 static ngx_int_t ngx_http_upstream_get_hash_peer(ngx_peer_connection_t *pc,
  41     void *data);
  42 static void ngx_http_upstream_free_hash_peer(ngx_peer_connection_t *pc,
  43     void *data, ngx_uint_t state);
  44 static char *ngx_http_upstream_hash(ngx_conf_t *cf, ngx_command_t *cmd,
  45     void *conf);
  46 static char *ngx_http_upstream_hash_again(ngx_conf_t *cf, ngx_command_t *cmd,
  47     void *conf);
  48 static ngx_int_t ngx_http_upstream_init_hash(ngx_conf_t *cf,
  49     ngx_http_upstream_srv_conf_t *us);
  50 static ngx_uint_t ngx_http_upstream_hash_crc32(u_char *keydata, size_t keylen);
  51
  52
  53 static ngx_command_t  ngx_http_upstream_hash_commands[] = {
  54     { ngx_string("hash"),
  55       NGX_HTTP_UPS_CONF|NGX_CONF_TAKE1,
  56       ngx_http_upstream_hash,
  57       0,
  58       0,
  59       NULL },
  60
  61     { ngx_string("hash_again"),
  62       NGX_HTTP_UPS_CONF|NGX_CONF_TAKE1,
  63       ngx_http_upstream_hash_again,
  64       0,
  65       0,
  66       NULL },
  67
  68       ngx_null_command
  69 };
  70
  71
  72 static ngx_http_module_t  ngx_http_upstream_hash_module_ctx = {
  73     NULL,                                  /* preconfiguration */
  74     NULL,                                  /* postconfiguration */
  75
  76     NULL,                                  /* create main configuration */
  77     NULL,                                  /* init main configuration */
  78
  79     NULL,                                  /* create server configuration */
  80     NULL,                                  /* merge server configuration */
  81
  82     NULL,                                  /* create location configuration */
  83     NULL                                   /* merge location configuration */
  84 };
  85
  86
  87 ngx_module_t  ngx_http_upstream_hash_module = {
  88     NGX_MODULE_V1,
  89     &ngx_http_upstream_hash_module_ctx,    /* module context */
  90     ngx_http_upstream_hash_commands,       /* module directives */
  91     NGX_HTTP_MODULE,                       /* module type */
  92     NULL,                                  /* init master */
  93     NULL,                                  /* init module */
  94     NULL,                                  /* init process */
  95     NULL,                                  /* init thread */
  96     NULL,                                  /* exit thread */
  97     NULL,                                  /* exit process */
  98     NULL,                                  /* exit master */
  99     NGX_MODULE_V1_PADDING
 100 };
 101
 102
 103 static ngx_int_t
 104 ngx_http_upstream_init_hash(ngx_conf_t *cf, ngx_http_upstream_srv_conf_t *us)
 105 {
 106     ngx_uint_t                       i, j, n;
 107     ngx_http_upstream_server_t      *server;
 108     ngx_http_upstream_hash_peers_t  *peers;
 109
 110     us->peer.init = ngx_http_upstream_init_hash_peer;
 111
 112     if (!us->servers) {
 113
 114         return NGX_ERROR;
 115     }
 116
 117     server = us->servers->elts;
 118
 119     for (n = 0, i = 0; i < us->servers->nelts; i++) {
 120         n += server[i].naddrs;
 121     }
 122
 123     peers = ngx_pcalloc(cf->pool, sizeof(ngx_http_upstream_hash_peers_t)
 124             + sizeof(ngx_http_upstream_hash_peer_t) * n);
 125
 126     if (peers == NULL) {
 127         return NGX_ERROR;
 128     }
 129
 130     peers->number = n;
 131
 132     /* one hostname can have multiple IP addresses in DNS */
 133     for (n = 0, i = 0; i < us->servers->nelts; i++) {
 134         for (j = 0; j < server[i].naddrs; j++, n++) {
 135             peers->peer[n].sockaddr = server[i].addrs[j].sockaddr;
 136             peers->peer[n].socklen = server[i].addrs[j].socklen;
 137             peers->peer[n].name = server[i].addrs[j].name;
 138         }
 139     }
 140
 141     us->peer.data = peers;
 142
 143     return NGX_OK;
 144 }
 145
 146
 147 static ngx_int_t
 148 ngx_http_upstream_init_hash_peer(ngx_http_request_t *r,
 149     ngx_http_upstream_srv_conf_t *us)
 150 {
 151     ngx_http_upstream_hash_peer_data_t     *uhpd;
 152
 153     ngx_str_t val;
 154
 155     if (ngx_http_script_run(r, &val, us->lengths, 0, us->values) == NULL) {
 156         return NGX_ERROR;
 157     }
 158
 159     uhpd = ngx_pcalloc(r->pool, sizeof(ngx_http_upstream_hash_peer_data_t)
 160             + sizeof(uintptr_t) *
 161                 ((ngx_http_upstream_hash_peers_t *)us->peer.data)->number /
 162                     (8 * sizeof(uintptr_t)));
 163     if (uhpd == NULL) {
 164         return NGX_ERROR;
 165     }
 166
 167     r->upstream->peer.data = uhpd;
 168
 169     uhpd->peers = us->peer.data;
 170
 171     r->upstream->peer.free = ngx_http_upstream_free_hash_peer;
 172     r->upstream->peer.get = ngx_http_upstream_get_hash_peer;
 173     r->upstream->peer.tries = us->retries + 1;
 174
 175     /* must be big enough for the retry keys */
 176     if ((uhpd->current_key.data = ngx_pcalloc(r->pool, NGX_ATOMIC_T_LEN + val.len)) == NULL) {
 177         return NGX_ERROR;
 178     }
 179
 180     ngx_memcpy(uhpd->current_key.data, val.data, val.len);
 181     uhpd->current_key.len = val.len;
 182     uhpd->original_key = val;
 183
 184     ngx_log_debug1(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
 185                    "upstream_hash: hashing \"%V\"", &val);
 186
 187     uhpd->hash = ngx_http_upstream_hash_crc32(val.data, val.len);
 188
 189     return NGX_OK;
 190 }
 191
 192
 193 static ngx_int_t
 194 ngx_http_upstream_get_hash_peer(ngx_peer_connection_t *pc, void *data)
 195 {
 196     ngx_http_upstream_hash_peer_data_t  *uhpd = data;
 197     ngx_http_upstream_hash_peer_t       *peer;
 198     ngx_uint_t                           peer_index;
 199
 200     ngx_log_debug1(NGX_LOG_DEBUG_HTTP, pc->log, 0,
 201                    "upstream_hash: get upstream request hash peer try %ui", pc->tries);
 202
 203     pc->cached = 0;
 204     pc->connection = NULL;
 205
 206     peer_index = uhpd->hash % uhpd->peers->number;
 207
 208     peer = &uhpd->peers->peer[peer_index];
 209
 210     ngx_log_debug2(NGX_LOG_DEBUG_HTTP, pc->log, 0,
 211                    "upstream_hash: chose peer %ui w/ hash %ui", peer_index, uhpd->hash);
 212
 213     pc->sockaddr = peer->sockaddr;
 214     pc->socklen = peer->socklen;
 215     pc->name = &peer->name;
 216
 217     return NGX_OK;
 218 }
 219
 220 /* retry implementation is PECL memcache compatible */
 221 static void
 222 ngx_http_upstream_free_hash_peer(ngx_peer_connection_t *pc, void *data,
 223     ngx_uint_t state)
 224 {
 225     ngx_http_upstream_hash_peer_data_t  *uhpd = data;
 226     ngx_uint_t                           current;
 227
 228     ngx_log_debug1(NGX_LOG_DEBUG_HTTP, pc->log, 0,
 229             "upstream_hash: free upstream hash peer try %ui", pc->tries);
 230
 231     if (state & (NGX_PEER_FAILED|NGX_PEER_NEXT)
 232             && --pc->tries)
 233     {
 234         current = uhpd->hash % uhpd->peers->number;
 235
 236         uhpd->tried[ngx_bitvector_index(current)] |= ngx_bitvector_bit(current);
 237
 238         do {
 239             uhpd->current_key.len = ngx_sprintf(uhpd->current_key.data, "%d%V",
 240                     ++uhpd->try_i, &uhpd->original_key) - uhpd->current_key.data;
 241             ngx_log_debug1(NGX_LOG_DEBUG_HTTP, pc->log, 0,
 242                     "upstream_hash: hashing \"%V\"", &uhpd->current_key);
 243             uhpd->hash += ngx_http_upstream_hash_crc32(uhpd->current_key.data,
 244                     uhpd->current_key.len);
 245             current = uhpd->hash % uhpd->peers->number;
 246         } while ((uhpd->tried[ngx_bitvector_index(current)] & ngx_bitvector_bit(current))
 247                 && --pc->tries);
 248     }
 249 }
 250
 251 /* bit-shift, bit-mask, and non-zero requirement are for libmemcache compatibility */
 252 static ngx_uint_t
 253 ngx_http_upstream_hash_crc32(u_char *keydata, size_t keylen)
 254 {
 255     ngx_uint_t crc32 = (ngx_crc32_short(keydata, keylen) >> 16) & 0x7fff;
 256     return crc32 ? crc32 : 1;
 257 }
 258
 259 static char *
 260 ngx_http_upstream_hash(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
 261 {
 262     ngx_http_upstream_srv_conf_t  *uscf;
 263     ngx_http_script_compile_t      sc;
 264     ngx_str_t                     *value;
 265     ngx_array_t                   *vars_lengths, *vars_values;
 266
 267     value = cf->args->elts;
 268
 269     ngx_memzero(&sc, sizeof(ngx_http_script_compile_t));
 270
 271     vars_lengths = NULL;
 272     vars_values = NULL;
 273
 274     sc.cf = cf;
 275     sc.source = &value[1];
 276     sc.lengths = &vars_lengths;
 277     sc.values = &vars_values;
 278     sc.complete_lengths = 1;
 279     sc.complete_values = 1;
 280
 281     if (ngx_http_script_compile(&sc) != NGX_OK) {
 282         return NGX_CONF_ERROR;
 283     }
 284
 285     uscf = ngx_http_conf_get_module_srv_conf(cf, ngx_http_upstream_module);
 286
 287     uscf->peer.init_upstream = ngx_http_upstream_init_hash;
 288
 289     uscf->flags = NGX_HTTP_UPSTREAM_CREATE;
 290
 291     uscf->values = vars_values->elts;
 292     uscf->lengths = vars_lengths->elts;
 293
 294     return NGX_CONF_OK;
 295 }
 296
 297 static char *
 298 ngx_http_upstream_hash_again(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
 299 {
 300     ngx_http_upstream_srv_conf_t  *uscf;
 301     ngx_int_t n;
 302
 303     ngx_str_t *value;
 304
 305     uscf = ngx_http_conf_get_module_srv_conf(cf, ngx_http_upstream_module);
 306
 307     value = cf->args->elts;
 308
 309     n = ngx_atoi(value[1].data, value[1].len);
 310
 311     if (n == NGX_ERROR || n < 0) {
 312         return "invalid number";
 313     }
 314
 315     uscf->retries = n;
 316
 317     return NGX_CONF_OK;
 318 }