Nginx模块开发（十二）（续二）：upstream负载均衡模块的例子

upstream

nginx模块一贯被分成三大类：handler、filter和upstream。前面的章节中，读者已经了解了handler、filter。通过开发这两类模块，可以使nginx轻松搞定任何单机的工作。现在，我们终于迎来了upstream。有了这一利器，nginx将可以跨越单机的限制，在网络进行数据的接收、处理和转发。

听起来让人热血沸腾吧？啊哈？没感觉？

假如上班时只能按老大的号令行事，相互之间不能有沟通，所有的人都会觉得憋屈吧。nginx模块也是如此，不能横向联系的nginx模块只能为某个终端节点提供单一功能，而一旦加上upstream，nginx模块就具备了战略价值，有能力成为一个网络应用的关键组件。当然，一个网络应用的关键组件往往一开始都会考虑通过这样那样的高级开发语言编写，因为开发比较方便，但当考虑到性能以后，以及这些高级语言为了实现高性能所要付出的代价以后，nginx的upstream模块就会呈现出极大的吸引力，因为他天生就快。作为一个附带的好处，nginx模块的运维相当简单。

upstream基本接口

upstream从本质上说还是一个handler，只是他不自己产生内容，而是通过请求一个后端服务器来得到内容，就好像逆流而上一般，所以才称为“upstream”。这个请求的过程已经被封装到nginx内部，所以开发一个upstream模块只需要提供若干回调函数，完成请求的构造和响应的解析等等具体的工作。

这些回调函数如下表所示：

create_request

生成发送到后端服务器的请求缓冲（缓冲链）。

reinit_request

在某台后端服务器出错的情况，nginx会尝试另一台后端服务器。nginx选定新的服务器以后，会先调用此函数，以重新初始化upstream模块的工作状态，然后重连其他服务器。

process_header

处理后端服务器返回的信息头部。所谓头部是与upstream server通信的协议规定的，比如HTTP协议的header部分，或者memcache协议的响应状态部分。

abort_request

在客户端放弃请求时被调用。不需要在函数中实现关闭后端服务器连接的功能，系统会自动完成关闭连接的步骤，所以一般此函数不会进行任何具体工作。

finalize_request

正常完成与后端服务器的请求后调用该函数，与abort_request相同，一般也不会进行任何具体工作。

input_filter

处理后端服务器返回的响应正文。nginx默认的input_filter会将收到的内容封装成为缓冲区链ngx_chain。该链由upstream的out_bufs指针域定位，所以我们可以在模块以外通过该指针得到后端服务器返回的正文数据。memcache模块实现了自己的input_filter，我们在后面会具体分析这个模块。

input_filter_init

初始化input filter的上下文。nginx默认的input_filter_init直接返回。

upstream分析：memcached模块

通过上一节，我们准备好了作料，从这一节开始我们可以来学习做upstream这盘菜了。我们先来分析nginx中最简单的memcache模块。

memcache是一款高性能的分布式cache系统，得到了非常广泛的应用。memcache定义了一套私有通信协议，使得不能通过HTTP请求来访问memcache。但协议本身简单高效，而且memcache使用广泛，所以大部分现代开发语言和平台都提供了memcache支持，方便开发者使用memcache。

nginx提供了ngx_http_memcached模块，提供了直接从memcache读取数据的功能，而将更新数据的操作交给其他应用。作为web服务器，这是可以接受的。

现在我们来分解ngx_http_memcached模块，一窥upstream的奥秘。

Handler模块？

初看memcached模块，大家可能觉得并无特别之处。如果稍微细看，甚至觉得有点像handler模块，当大家看到这段代码以后，必定疑惑为什么会跟handler模块一模一样。

clcf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);

clcf->handler = ngx_http_memcached_handler;

因为upstream模块使用的就是handler模块的接入方式。同时，upstream模块的指令系统的设计也是遵循handler模块的基本规则：配置该模块才会执行该模块。

{ ngx_string("memcached_pass"),

  NGX_HTTP_LOC_CONF|NGX_HTTP_LIF_CONF|NGX_CONF_TAKE1,

  ngx_http_memcached_pass,

  NGX_HTTP_LOC_CONF_OFFSET,

  0,

  NULL }

所以大家觉得眼熟是好事，说明大家对Handler的写法已经很熟悉了。

Upstream模块！

那么，upstream模块的特别之处究竟在哪里呢？答案是就在模块的处理函数实现中。upstream模块的处理函数进行的操作都含有一个固定的流程，就像例行公事一样。在memcached的例子中，我们观察ngx_http_memcached_handler的代码，可以发现，这个固定的操作是：

1. 创建upstream数据结构。

if (ngx_http_upstream_create(r) != NGX_OK) {

    return NGX_HTTP_INTERNAL_SERVER_ERROR;

}

2. 设置模块的tag和schema。schema现在只会用于日志，tag会用于buf_chain管理。

u = r->upstream;

ngx_str_set(&u->schema, "memcached://");

u->output.tag = (ngx_buf_tag_t) &ngx_http_memcached_module;

3. 设置upstream的后端服务器列表数据结构。

mlcf = ngx_http_get_module_loc_conf(r, ngx_http_memcached_module);

u->conf = &mlcf->upstream;

4. 设置upstream回调函数。在这里列出的代码稍稍调整了代码顺序。

u->create_request = ngx_http_memcached_create_request;

u->reinit_request = ngx_http_memcached_reinit_request;

u->process_header = ngx_http_memcached_process_header;

u->abort_request = ngx_http_memcached_abort_request;

u->finalize_request = ngx_http_memcached_finalize_request;

u->input_filter_init = ngx_http_memcached_filter_init;

u->input_filter = ngx_http_memcached_filter;

5. 创建并设置upstream环境数据结构。

ctx = ngx_palloc(r->pool, sizeof(ngx_http_memcached_ctx_t));

if (ctx == NULL) {

    return NGX_HTTP_INTERNAL_SERVER_ERROR;

}

ctx->rest = NGX_HTTP_MEMCACHED_END;

ctx->request = r;

ngx_http_set_ctx(r, ctx, ngx_http_memcached_module);

u->input_filter_ctx = ctx;

6. 完成upstream初始化并进行收尾工作。

r->main->count++;

ngx_http_upstream_init(r);

return NGX_DONE;

任何upstream模块，简单如memcached，复杂如proxy、fastcgi都是如此。不同的upstream模块在这6步中的最大差别会出现在第2、3、4、5上。其中第2、4两步很容易理解，不同的模块设置的标志和使用的回调函数肯定不同。第5步也比较直白，只有第3步是最为晦涩的，不同的模块在取得后端服务器列表时，策略的差距相当大，有如memcached这样简单明了的，也有如proxy那样逻辑复杂的。但是，我们先把这一块放一放，等我们把memcached剖析清楚了，再单独讨论下。

第6步是一个常态。将count加1，然后返回NGX_DONE，使得nginx结束当前请求的处理，但是不关闭与客户段的连接。这样说可能有些晦涩，但是nginx是将upstream请求作为一个另外的请求来处理的，目前的客户端请求处理到这里已经结束了，处理的结果就是需要建立一个upstream请求来完成后续的工作。而upstream请求和原本客户端请求之间是一对一的关系，nginx会使用内建的ngx_event_pipe来完成将upstream请求得到的响应写回到客户端。

这个设计有优势也有缺陷。优势就是简化模块开发，可以将精力集中在模块逻辑上，缺陷同样明显，一对一的设计很多时候都不能满足复杂逻辑的需要。对于这一点，将会在后面的原理篇来阐述。

回调函数

前面剖析了memcached模块的骨架，现在开始逐个解决每个回调函数。

1. ngx_http_memcached_create_request：很简单的按照设置的内容生成一个key，接着生成一个“get $key”的请求，放在r->upstream->request_bufs里面。

2. ngx_http_memcached_reinit_request：无事可做。

3. ngx_http_memcached_abort_request：无事可做。

4. ngx_http_memcached_finalize_request：无事可做。

5. ngx_http_memcached_process_header：终于可以找点事情做了。对于memcache协议来说，header被定义为第一行文本的内容，可以找到一段代码说明问题。

for (p = u->buffer.pos; p < u->buffer.last; p++) {

    if (*p == LF) {

        goto found;

    }

}

如果在已读入的缓冲数据中没有发现LF字符，函数返回NGX_AGAIN，表示需要继续读取数据。nginx在收到新的数据以后会再次调用该函数。nginx处理后端服务器的header时只会使用一块buffer，所有数据都在这块buffer中，所以解析是不需要考虑头部信息跨越多个buffer的情况。

在process_header中进行的一件比较重要的事情是将后端服务器返回的状态翻译成返回给客户端的状态。例如，在ngx_http_memcached_process_header中，有这样几段代码：

r->headers_out.content_length_n = ngx_atoof(len, p - len - 1);

u->headers_in.status_n = 200;

u->state->status = 200;

u->headers_in.status_n = 404;

u->state->status = 404;

u->state是用于upstream相关的变量。比如u->status->status将用于计算变量“upstream_status”。u->headers_in将被作为返回给客户端的响应返回状态码。而第一行则是设置返回给客户端的响应的长度。

还需要注意的一件事情就是不能忘记了在处理完头部信息以后将读指针pos后移，否则这段数据也会被复制到返回给客户端的响应的正文中。

u->buffer.pos = p + 1;

process_header函数完成头部的正确处理，应该返回NGX_OK。如果返回NGX_AGAIN，表示数据未读完整，需要继续从后端服务器读取数据。返回NGX_DECLINED无意义，其他任何返回值都被认为是出错状态，nginx将结束upstream请求并返回错误信息。

6. ngx_http_memcached_filter_init：基本上无事可做，只是修正一下从后端服务器收到的内容长度。因为在处理header时没有加上这部分长度。

7. ngx_http_memcached_filter：memcached模块是少有的带有处理正文的回调函数的模块。因为memcached模块需要过滤正文末尾CRLF "END" CRLF，所以实现了自己的filter回调函数。处理正文的实际意义是将从后端服务器收到的正文有效内容封装成ngx_chain_t，并加在u->out_bufs末尾。nginx并不进行数据拷贝，而是建立ngx_buf_t数据结构指向这些数据内存区，然后由ngx_chain_t组织这些buf。这种实现避免了大量内存搬迁，也是nginx高效的奥秘之一。

本节小结

在这一节里，大家对upstream模块的基本组成有了一些认识。upstream模块是从handler模块发展而来，指令系统和模块生效方式与handler模块无异。不同之处在于，upstream模块在handler函数中设置众多回调函数。实际工作都是由这些回调函数完成的。每个回调函数都是在upstream的某个固定阶段执行，各司其职，大部分回调函数一般不会真正用到。upstream最重要的回调函数是create_request、process_header和input_filter，他们共同实现了与后端服务器的协议的解析部分。

近又没时间去摆弄那本nginx的书的，所以这个例子虽然是准备放在书里面的，还是先拿出来分享下算了。

这是一个counter的例子，能够统计upstream中每个server目前保持的连接数量。当然，是每个worker自己计数的。模块正常工作的基础是原本的负载均衡算法data数据结构与round robin兼容，比如ip hash等等，为什么呢，因为如果想更通用的话，只能给nginx打patch了。另外，对程序的分析以后在书里写吧。

#include <ngx_config.h>

#include <ngx_core.h>
#include <ngx_http.h>


typedef struct {
ngx_http_upstream_rr_peer_t *peer;
ngx_uint_t count;
} ngx_http_upstream_count_peer;

typedef struct {
ngx_http_upstream_peer_t peer;
ngx_uint_t nelts;
ngx_http_upstream_count_peer *peers;
} ngx_http_upstream_count_conf_t;

typedef struct {
ngx_peer_connection_t pc;
ngx_http_upstream_count_conf_t *conf;

unsigned index:16;
unsigned counted:1;
} ngx_http_upstream_count_ctx_t;


static ngx_int_t ngx_http_upstream_init_count_peer(ngx_http_request_t *r,
ngx_http_upstream_srv_conf_t *us);
static char *ngx_http_upstream_count(ngx_conf_t *cf, ngx_command_t *cmd,
void *conf);
static void *ngx_http_upstream_count_create_conf(ngx_conf_t *cf);
static void ngx_http_upstream_free_count_peer(ngx_peer_connection_t *pc,
void *data, ngx_uint_t state);
static ngx_int_t ngx_http_upstream_get_count_peer(ngx_peer_connection_t *pc,
void *data);
static ngx_int_t ngx_http_upstream_init_count(ngx_conf_t *cf,
ngx_http_upstream_srv_conf_t *us);

static ngx_int_t ngx_http_upstream_count_info_handler(ngx_http_request_t *r);
static char *ngx_http_upstream_count_info(ngx_conf_t *cf, ngx_command_t *cmd,
void *conf);


static ngx_command_t ngx_http_upstream_count_commands[] = {

{ ngx_string("count"),
NGX_HTTP_UPS_CONF|NGX_CONF_NOARGS,
ngx_http_upstream_count,
0,
0,
NULL },

{ ngx_string("count_info"),
NGX_HTTP_LOC_CONF|NGX_CONF_NOARGS,
ngx_http_upstream_count_info,
0,
0,
NULL },

ngx_null_command
};


static ngx_http_module_t ngx_http_upstream_count_module_ctx = {
NULL,
NULL,

NULL,
NULL,

ngx_http_upstream_count_create_conf,
NULL,

NULL,
NULL
};


ngx_module_t ngx_http_upstream_count_module = {
NGX_MODULE_V1,
&ngx_http_upstream_count_module_ctx,
ngx_http_upstream_count_commands,
NGX_HTTP_MODULE,
NULL,
NULL,
NULL,
NULL,
NULL,
NULL,
NULL,
NGX_MODULE_V1_PADDING
};


static char *
ngx_http_upstream_count(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
ngx_http_upstream_srv_conf_t *uscf;
ngx_http_upstream_count_conf_t *cscf;

uscf = ngx_http_conf_get_module_srv_conf(cf, ngx_http_upstream_module);
cscf = ngx_http_conf_get_module_srv_conf(cf, ngx_http_upstream_count_module);

cscf->peer.init_upstream = uscf->peer.init_upstream;
uscf->peer.init_upstream = ngx_http_upstream_init_count;

return NGX_CONF_OK;
}


static ngx_int_t
ngx_http_upstream_init_count(ngx_conf_t *cf, ngx_http_upstream_srv_conf_t *us)
{
ngx_uint_t i;
ngx_http_upstream_rr_peers_t *peers;
ngx_http_upstream_count_conf_t *cscf;

cscf = us->srv_conf[ngx_http_upstream_count_module.ctx_index];

if (cscf->peer.init_upstream == NULL) {
cscf->peer.init_upstream = ngx_http_upstream_init_round_robin;
}

if (cscf->peer.init_upstream(cf, us) != NGX_OK) {
return NGX_ERROR;
}

cscf->peer.init = us->peer.init;
us->peer.init = ngx_http_upstream_init_count_peer;

peers = us->peer.data;

cscf->nelts = peers->number;

cscf->peers = ngx_pcalloc(cf->pool,
sizeof(ngx_http_upstream_count_peer) * cscf->nelts);
if (cscf->peers == NULL) {
return NGX_ERROR;
}

for (i = 0; i < cscf->nelts; i++) {
cscf->peers[i].peer = &peers->peer[i];
}

return NGX_OK;
}


static ngx_int_t
ngx_http_upstream_init_count_peer(ngx_http_request_t *r,
ngx_http_upstream_srv_conf_t *us)
{
ngx_int_t rc;
ngx_http_upstream_count_ctx_t *ctx;
ngx_http_upstream_count_conf_t *cscf;

cscf = us->srv_conf[ngx_http_upstream_count_module.ctx_index];

ctx = ngx_palloc(r->pool, sizeof(ngx_http_upstream_count_ctx_t));
if (ctx == NULL) {
return NGX_ERROR;
}

rc = cscf->peer.init(r, us);
if (rc != NGX_OK) {
return rc;
}

ctx->pc.data = r->upstream->peer.data;
r->upstream->peer.data = ctx;

ctx->pc.get = r->upstream->peer.get;
r->upstream->peer.get = ngx_http_upstream_get_count_peer;

ctx->pc.free = r->upstream->peer.free;
r->upstream->peer.free = ngx_http_upstream_free_count_peer;

ctx->conf = cscf;

return NGX_OK;
}


static void
ngx_http_upstream_free_count_peer(ngx_peer_connection_t *pc, void *data,
ngx_uint_t state)
{
ngx_http_upstream_count_ctx_t *ctx;

ctx = data;

if (ctx->pc.free) {
ctx->pc.free(pc, ctx->pc.data, state);
}

if (ctx->counted) {
--ctx->conf->peers[ctx->index].count;
ctx->counted = 0;
}
}


static ngx_int_t
ngx_http_upstream_get_count_peer(ngx_peer_connection_t *pc, void *data)
{
ngx_int_t rc;
ngx_uint_t i;
ngx_http_upstream_count_ctx_t *ctx;

ctx = data;

rc = ctx->pc.get(pc, ctx->pc.data);
if (rc != NGX_OK && rc != NGX_DONE) {
return rc;
}

for (i = 0; i < ctx->conf->nelts; i++) {
if (pc->name == &ctx->conf->peers[i].peer->name) {
ctx->index = i;
ctx->counted = 1;
++ctx->conf->peers[i].count;
break;
}
}

return rc;
}


static void *
ngx_http_upstream_count_create_conf(ngx_conf_t *cf)
{
ngx_http_upstream_count_conf_t *conf;

conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_upstream_count_conf_t));
if (conf == NULL) {
return NULL;
}

return conf;
}


static char *
ngx_http_upstream_count_info(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
ngx_http_core_loc_conf_t *clcf;

clcf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
clcf->handler = ngx_http_upstream_count_info_handler;

return NGX_CONF_OK;
}


static ngx_int_t
ngx_http_upstream_count_info_handler(ngx_http_request_t *r)
{
u_char *p, *b;
ngx_str_t value;
ngx_uint_t i, k, count, len;
ngx_chain_t *cl;
ngx_http_upstream_srv_conf_t **uscfp, *us;
ngx_http_upstream_main_conf_t *umcf;
ngx_http_upstream_count_conf_t *cscf;

umcf = ngx_http_get_module_main_conf(r, ngx_http_upstream_module);
uscfp = umcf->upstreams.elts;
us = NULL;

value.len = 0;

if (r->args.len != 0) {
ngx_http_arg(r, (u_char *) "us", 2, &value);
}

if (value.len == 0) {
return NGX_HTTP_BAD_REQUEST;
}

if (umcf->upstreams.nelts != 0) {

len = 0;

for (i = 0; i < umcf->upstreams.nelts; i++) {
if (ngx_strncasecmp(value.data,
uscfp[i]->host.data, value.len) == 0
&& value.len == uscfp[i]->host.len)
{
us = uscfp[i];
break;
}
}

cscf = us->srv_conf[ngx_http_upstream_count_module.ctx_index];

if (cscf == NULL) {
return NGX_HTTP_NOT_FOUND;
}

len += sizeof("worker id: ");
for (k = 0, count = ngx_pid; count; k++, count /= 10) ;
len += k;

for (i = 0; i < cscf->nelts; i++) {
len += cscf->peers[i].peer->name.len + 2;

for (k = 0, count = cscf->peers[i].count;
count; k++, count /= 10) ;

len += k == 0 ? 2 : (k + 1);
}

if (len == 0) {
return NGX_HTTP_NO_CONTENT;
}

p = b = ngx_palloc(r->pool, len);
if (b == NULL) {
return NGX_ERROR;
}

p = ngx_slprintf(p, b + len, "worker id: %P\n", ngx_pid);

for (i = 0; i < cscf->nelts; i++) {
p = ngx_slprintf(p, b + len, "%V: ",
&cscf->peers[i].peer->name);
p = ngx_slprintf(p, b + len, "%ui\n",
cscf->peers[i].count);
}

cl = ngx_alloc_chain_link(r->pool);
if (cl == NULL) {
return NGX_ERROR;
}

cl->next = NULL;
cl->buf = ngx_calloc_buf(r->pool);
if (cl->buf == NULL) {
return NGX_ERROR;
}

cl->buf->pos = b;
cl->buf->last = b + len;
cl->buf->last_buf = 1;
cl->buf->memory = 1;

r->headers_out.content_length_n = len;
r->headers_out.status = NGX_HTTP_OK;

if (ngx_http_send_header(r) != NGX_OK) {
return NGX_ERROR;
}

return ngx_http_output_filter(r, cl);
}

return NGX_HTTP_NO_CONTENT;

}

配置示例：

upstream  a {

     ip_hash;

     server 10.0.0.1:7001;

     server 10.0.0.2:7001;

     count;

}

server {

     location / {

           proxy http://a;

     }

     location =/us {

           count_info;

     }

}

使用curl 'http://127.0.0.1/us?us=a'可以看到a这个upstream中的各server的连接数。