Stream: udp streams.

Previously, only one client packet could be processed in a udp stream session
even though multiple response packets were supported.  Now multiple packets
coming from the same client address and port are delivered to the same stream
session.

If it's required to maintain a single stream of data, nginx should be
configured in a way that all packets from a client are delivered to the same
worker.  On Linux and DragonFly BSD the "reuseport" parameter should be
specified for this.  Other systems do not currently provide appropriate
mechanisms.  For these systems a single stream of udp packets is only
guaranteed in single-worker configurations.

The proxy_response directive now specifies how many packets are expected in
response to a single client packet.
diff --git a/src/core/ngx_connection.c b/src/core/ngx_connection.c
index 878c6be..dc60679 100644
--- a/src/core/ngx_connection.c
+++ b/src/core/ngx_connection.c
@@ -72,6 +72,10 @@
 
     ngx_memcpy(ls->addr_text.data, text, len);
 
+#if !(NGX_WIN32)
+    ngx_rbtree_init(&ls->rbtree, &ls->sentinel, ngx_udp_rbtree_insert_value);
+#endif
+
     ls->fd = (ngx_socket_t) -1;
     ls->type = SOCK_STREAM;
 
diff --git a/src/core/ngx_connection.h b/src/core/ngx_connection.h
index e4dfe58..ef07558 100644
--- a/src/core/ngx_connection.h
+++ b/src/core/ngx_connection.h
@@ -51,6 +51,9 @@
     ngx_listening_t    *previous;
     ngx_connection_t   *connection;
 
+    ngx_rbtree_t        rbtree;
+    ngx_rbtree_node_t   sentinel;
+
     ngx_uint_t          worker;
 
     unsigned            open:1;
@@ -151,6 +154,8 @@
     ngx_ssl_connection_t  *ssl;
 #endif
 
+    ngx_udp_connection_t  *udp;
+
     struct sockaddr    *local_sockaddr;
     socklen_t           local_socklen;
 
diff --git a/src/core/ngx_core.h b/src/core/ngx_core.h
index 2069373..93ca917 100644
--- a/src/core/ngx_core.h
+++ b/src/core/ngx_core.h
@@ -27,6 +27,7 @@
 typedef struct ngx_thread_task_s     ngx_thread_task_t;
 typedef struct ngx_ssl_s             ngx_ssl_t;
 typedef struct ngx_ssl_connection_s  ngx_ssl_connection_t;
+typedef struct ngx_udp_connection_s  ngx_udp_connection_t;
 
 typedef void (*ngx_event_handler_pt)(ngx_event_t *ev);
 typedef void (*ngx_connection_handler_pt)(ngx_connection_t *c);
diff --git a/src/event/ngx_event.h b/src/event/ngx_event.h
index bcc2271..44f9e05 100644
--- a/src/event/ngx_event.h
+++ b/src/event/ngx_event.h
@@ -506,6 +506,8 @@
 void ngx_event_accept(ngx_event_t *ev);
 #if !(NGX_WIN32)
 void ngx_event_recvmsg(ngx_event_t *ev);
+void ngx_udp_rbtree_insert_value(ngx_rbtree_node_t *temp,
+    ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel);
 #endif
 ngx_int_t ngx_trylock_accept_mutex(ngx_cycle_t *cycle);
 ngx_int_t ngx_enable_accept_events(ngx_cycle_t *cycle);
diff --git a/src/event/ngx_event_udp.c b/src/event/ngx_event_udp.c
index 962b8e3..8b0ae55 100644
--- a/src/event/ngx_event_udp.c
+++ b/src/event/ngx_event_udp.c
@@ -12,13 +12,28 @@
 
 #if !(NGX_WIN32)
 
+struct ngx_udp_connection_s {
+    ngx_rbtree_node_t   node;
+    ngx_connection_t   *connection;
+    ngx_buf_t          *buffer;
+};
+
+
 static void ngx_close_accepted_udp_connection(ngx_connection_t *c);
+static ssize_t ngx_udp_shared_recv(ngx_connection_t *c, u_char *buf,
+    size_t size);
+static ngx_int_t ngx_insert_udp_connection(ngx_connection_t *c);
+static void ngx_delete_udp_connection(void *data);
+static ngx_connection_t *ngx_lookup_udp_connection(ngx_listening_t *ls,
+    struct sockaddr *sockaddr, socklen_t socklen,
+    struct sockaddr *local_sockaddr, socklen_t local_socklen);
 
 
 void
 ngx_event_recvmsg(ngx_event_t *ev)
 {
     ssize_t            n;
+    ngx_buf_t          buf;
     ngx_log_t         *log;
     ngx_err_t          err;
     socklen_t          socklen, local_socklen;
@@ -215,6 +230,43 @@
 
 #endif
 
+        c = ngx_lookup_udp_connection(ls, sockaddr, socklen, local_sockaddr,
+                                      local_socklen);
+
+        if (c) {
+
+#if (NGX_DEBUG)
+            if (c->log->log_level & NGX_LOG_DEBUG_EVENT) {
+                ngx_log_handler_pt  handler;
+
+                handler = c->log->handler;
+                c->log->handler = NULL;
+
+                ngx_log_debug2(NGX_LOG_DEBUG_EVENT, c->log, 0,
+                               "recvmsg: fd:%d n:%z", c->fd, n);
+
+                c->log->handler = handler;
+            }
+#endif
+
+            ngx_memzero(&buf, sizeof(ngx_buf_t));
+
+            buf.pos = buffer;
+            buf.last = buffer + n;
+
+            rev = c->read;
+
+            c->udp->buffer = &buf;
+            rev->ready = 1;
+
+            rev->handler(rev);
+
+            c->udp->buffer = NULL;
+            rev->ready = 0;
+
+            goto next;
+        }
+
 #if (NGX_STAT_STUB)
         (void) ngx_atomic_fetch_add(ngx_stat_accepted, 1);
 #endif
@@ -257,6 +309,7 @@
 
         *log = ls->log;
 
+        c->recv = ngx_udp_shared_recv;
         c->send = ngx_udp_send;
         c->send_chain = ngx_udp_send_chain;
 
@@ -344,11 +397,18 @@
         }
 #endif
 
+        if (ngx_insert_udp_connection(c) != NGX_OK) {
+            ngx_close_accepted_udp_connection(c);
+            return;
+        }
+
         log->data = NULL;
         log->handler = NULL;
 
         ls->handler(c);
 
+    next:
+
         if (ngx_event_flags & NGX_USE_KQUEUE_EVENT) {
             ev->available -= n;
         }
@@ -373,4 +433,209 @@
 #endif
 }
 
+
+static ssize_t
+ngx_udp_shared_recv(ngx_connection_t *c, u_char *buf, size_t size)
+{
+    ssize_t     n;
+    ngx_buf_t  *b;
+
+    if (c->udp == NULL || c->udp->buffer == NULL) {
+        return NGX_AGAIN;
+    }
+
+    b = c->udp->buffer;
+
+    n = ngx_min(b->last - b->pos, (ssize_t) size);
+
+    ngx_memcpy(buf, b->pos, n);
+
+    c->udp->buffer = NULL;
+    c->read->ready = 0;
+
+    return n;
+}
+
+
+void
+ngx_udp_rbtree_insert_value(ngx_rbtree_node_t *temp,
+    ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
+{
+    ngx_int_t               rc;
+    ngx_connection_t       *c, *ct;
+    ngx_rbtree_node_t     **p;
+    ngx_udp_connection_t   *udp, *udpt;
+
+    for ( ;; ) {
+
+        if (node->key < temp->key) {
+
+            p = &temp->left;
+
+        } else if (node->key > temp->key) {
+
+            p = &temp->right;
+
+        } else { /* node->key == temp->key */
+
+            udp = (ngx_udp_connection_t *) node;
+            c = udp->connection;
+
+            udpt = (ngx_udp_connection_t *) temp;
+            ct = udpt->connection;
+
+            rc = ngx_cmp_sockaddr(c->sockaddr, c->socklen,
+                                  ct->sockaddr, ct->socklen, 1);
+
+            if (rc == 0 && c->listening->wildcard) {
+                rc = ngx_cmp_sockaddr(c->local_sockaddr, c->local_socklen,
+                                      ct->local_sockaddr, ct->local_socklen, 1);
+            }
+
+            p = (rc < 0) ? &temp->left : &temp->right;
+        }
+
+        if (*p == sentinel) {
+            break;
+        }
+
+        temp = *p;
+    }
+
+    *p = node;
+    node->parent = temp;
+    node->left = sentinel;
+    node->right = sentinel;
+    ngx_rbt_red(node);
+}
+
+
+static ngx_int_t
+ngx_insert_udp_connection(ngx_connection_t *c)
+{
+    uint32_t               hash;
+    ngx_pool_cleanup_t    *cln;
+    ngx_udp_connection_t  *udp;
+
+    if (c->udp) {
+        return NGX_OK;
+    }
+
+    udp = ngx_pcalloc(c->pool, sizeof(ngx_udp_connection_t));
+    if (udp == NULL) {
+        return NGX_ERROR;
+    }
+
+    udp->connection = c;
+
+    ngx_crc32_init(hash);
+    ngx_crc32_update(&hash, (u_char *) c->sockaddr, c->socklen);
+
+    if (c->listening->wildcard) {
+        ngx_crc32_update(&hash, (u_char *) c->local_sockaddr, c->local_socklen);
+    }
+
+    ngx_crc32_final(hash);
+
+    udp->node.key = hash;
+
+    cln = ngx_pool_cleanup_add(c->pool, 0);
+    if (cln == NULL) {
+        return NGX_ERROR;
+    }
+
+    cln->data = c;
+    cln->handler = ngx_delete_udp_connection;
+
+    ngx_rbtree_insert(&c->listening->rbtree, &udp->node);
+
+    c->udp = udp;
+
+    return NGX_OK;
+}
+
+
+static void
+ngx_delete_udp_connection(void *data)
+{
+    ngx_connection_t  *c = data;
+
+    ngx_rbtree_delete(&c->listening->rbtree, &c->udp->node);
+}
+
+
+static ngx_connection_t *
+ngx_lookup_udp_connection(ngx_listening_t *ls, struct sockaddr *sockaddr,
+    socklen_t socklen, struct sockaddr *local_sockaddr, socklen_t local_socklen)
+{
+    uint32_t               hash;
+    ngx_int_t              rc;
+    ngx_connection_t      *c;
+    ngx_rbtree_node_t     *node, *sentinel;
+    ngx_udp_connection_t  *udp;
+
+#if (NGX_HAVE_UNIX_DOMAIN)
+
+    if (sockaddr->sa_family == AF_UNIX) {
+        struct sockaddr_un *saun = (struct sockaddr_un *) sockaddr;
+
+        if (socklen <= (socklen_t) offsetof(struct sockaddr_un, sun_path)
+            || saun->sun_path[0] == '\0')
+        {
+            ngx_log_debug0(NGX_LOG_DEBUG_EVENT, ngx_cycle->log, 0,
+                           "unbound unix socket");
+            return NULL;
+        }
+    }
+
+#endif
+
+    node = ls->rbtree.root;
+    sentinel = ls->rbtree.sentinel;
+
+    ngx_crc32_init(hash);
+    ngx_crc32_update(&hash, (u_char *) sockaddr, socklen);
+
+    if (ls->wildcard) {
+        ngx_crc32_update(&hash, (u_char *) local_sockaddr, local_socklen);
+    }
+
+    ngx_crc32_final(hash);
+
+    while (node != sentinel) {
+
+        if (hash < node->key) {
+            node = node->left;
+            continue;
+        }
+
+        if (hash > node->key) {
+            node = node->right;
+            continue;
+        }
+
+        /* hash == node->key */
+
+        udp = (ngx_udp_connection_t *) node;
+
+        c = udp->connection;
+
+        rc = ngx_cmp_sockaddr(sockaddr, socklen,
+                              c->sockaddr, c->socklen, 1);
+
+        if (rc == 0 && ls->wildcard) {
+            rc = ngx_cmp_sockaddr(local_sockaddr, local_socklen,
+                                  c->local_sockaddr, c->local_socklen, 1);
+        }
+
+        if (rc == 0) {
+            return c;
+        }
+
+        node = (rc < 0) ? node->left : node->right;
+    }
+
+    return NULL;
+}
+
 #endif
diff --git a/src/stream/ngx_stream_proxy_module.c b/src/stream/ngx_stream_proxy_module.c
index 30572cd..792bb77 100644
--- a/src/stream/ngx_stream_proxy_module.c
+++ b/src/stream/ngx_stream_proxy_module.c
@@ -377,6 +377,8 @@
 
     s->log_handler = ngx_stream_proxy_log_error;
 
+    u->requests = 1;
+
     u->peer.log = c->log;
     u->peer.log_error = NGX_ERROR_ERR;
 
@@ -398,21 +400,19 @@
         return;
     }
 
-    if (c->type == SOCK_STREAM) {
-        p = ngx_pnalloc(c->pool, pscf->buffer_size);
-        if (p == NULL) {
-            ngx_stream_proxy_finalize(s, NGX_STREAM_INTERNAL_SERVER_ERROR);
-            return;
-        }
+    p = ngx_pnalloc(c->pool, pscf->buffer_size);
+    if (p == NULL) {
+        ngx_stream_proxy_finalize(s, NGX_STREAM_INTERNAL_SERVER_ERROR);
+        return;
+    }
 
-        u->downstream_buf.start = p;
-        u->downstream_buf.end = p + pscf->buffer_size;
-        u->downstream_buf.pos = p;
-        u->downstream_buf.last = p;
+    u->downstream_buf.start = p;
+    u->downstream_buf.end = p + pscf->buffer_size;
+    u->downstream_buf.pos = p;
+    u->downstream_buf.last = p;
 
-        if (c->read->ready) {
-            ngx_post_event(c->read, &ngx_posted_events);
-        }
+    if (c->read->ready) {
+        ngx_post_event(c->read, &ngx_posted_events);
     }
 
     if (pscf->upstream_value) {
@@ -829,7 +829,6 @@
 
         cl->buf->tag = (ngx_buf_tag_t) &ngx_stream_proxy_module;
         cl->buf->flush = 1;
-        cl->buf->last_buf = (c->type == SOCK_DGRAM);
 
         cl->next = u->upstream_out;
         u->upstream_out = cl;
@@ -871,17 +870,12 @@
         u->proxy_protocol = 0;
     }
 
-    if (c->type == SOCK_DGRAM && pscf->responses == 0) {
-        pc->read->ready = 0;
-        pc->read->eof = 1;
-    }
-
     u->connected = 1;
 
     pc->read->handler = ngx_stream_proxy_upstream_handler;
     pc->write->handler = ngx_stream_proxy_upstream_handler;
 
-    if (pc->read->ready || pc->read->eof) {
+    if (pc->read->ready) {
         ngx_post_event(pc->read, &ngx_posted_events);
     }
 
@@ -1280,6 +1274,7 @@
 ngx_stream_proxy_process_connection(ngx_event_t *ev, ngx_uint_t from_upstream)
 {
     ngx_connection_t             *c, *pc;
+    ngx_log_handler_pt            handler;
     ngx_stream_session_t         *s;
     ngx_stream_upstream_t        *u;
     ngx_stream_proxy_srv_conf_t  *pscf;
@@ -1328,25 +1323,37 @@
                      * with unspecified number of responses
                      */
 
-                    pc->read->ready = 0;
-                    pc->read->eof = 1;
+                    handler = c->log->handler;
+                    c->log->handler = NULL;
 
-                    ngx_stream_proxy_process(s, 1, 0);
+                    ngx_log_error(NGX_LOG_INFO, c->log, 0,
+                                  "udp timed out"
+                                  ", packets from/to client:%ui/%ui"
+                                  ", bytes from/to client:%O/%O"
+                                  ", bytes from/to upstream:%O/%O",
+                                  u->requests, u->responses,
+                                  s->received, c->sent, u->received,
+                                  pc ? pc->sent : 0);
+
+                    c->log->handler = handler;
+
+                    ngx_stream_proxy_finalize(s, NGX_STREAM_OK);
                     return;
                 }
 
                 ngx_connection_error(pc, NGX_ETIMEDOUT, "upstream timed out");
 
-                if (u->received == 0) {
-                    ngx_stream_proxy_next_upstream(s);
-                    return;
-                }
+                pc->read->error = 1;
 
-            } else {
-                ngx_connection_error(c, NGX_ETIMEDOUT, "connection timed out");
+                ngx_stream_proxy_finalize(s, NGX_STREAM_BAD_GATEWAY);
+
+                return;
             }
 
+            ngx_connection_error(c, NGX_ETIMEDOUT, "connection timed out");
+
             ngx_stream_proxy_finalize(s, NGX_STREAM_OK);
+
             return;
         }
 
@@ -1453,7 +1460,7 @@
     ssize_t                       n;
     ngx_buf_t                    *b;
     ngx_int_t                     rc;
-    ngx_uint_t                    flags;
+    ngx_uint_t                    flags, *packets;
     ngx_msec_t                    delay;
     ngx_chain_t                  *cl, **ll, **out, **busy;
     ngx_connection_t             *c, *pc, *src, *dst;
@@ -1489,6 +1496,7 @@
         b = &u->upstream_buf;
         limit_rate = pscf->download_rate;
         received = &u->received;
+        packets = &u->responses;
         out = &u->downstream_out;
         busy = &u->downstream_busy;
         recv_action = "proxying and reading from upstream";
@@ -1500,6 +1508,7 @@
         b = &u->downstream_buf;
         limit_rate = pscf->upload_rate;
         received = &s->received;
+        packets = &u->requests;
         out = &u->upstream_out;
         busy = &u->upstream_busy;
         recv_action = "proxying and reading from client";
@@ -1516,11 +1525,6 @@
                 rc = ngx_stream_top_filter(s, *out, from_upstream);
 
                 if (rc == NGX_ERROR) {
-                    if (c->type == SOCK_DGRAM && !from_upstream) {
-                        ngx_stream_proxy_next_upstream(s);
-                        return;
-                    }
-
                     ngx_stream_proxy_finalize(s, NGX_STREAM_OK);
                     return;
                 }
@@ -1565,11 +1569,6 @@
             }
 
             if (n == NGX_ERROR) {
-                if (c->type == SOCK_DGRAM && u->received == 0) {
-                    ngx_stream_proxy_next_upstream(s);
-                    return;
-                }
-
                 src->read->eof = 1;
                 n = 0;
             }
@@ -1591,12 +1590,6 @@
                     }
                 }
 
-                if (c->type == SOCK_DGRAM && ++u->responses == pscf->responses)
-                {
-                    src->read->ready = 0;
-                    src->read->eof = 1;
-                }
-
                 for (ll = out; *ll; ll = &(*ll)->next) { /* void */ }
 
                 cl = ngx_chain_get_free_buf(c->pool, &u->free);
@@ -1616,6 +1609,7 @@
                 cl->buf->last_buf = src->read->eof;
                 cl->buf->flush = 1;
 
+                (*packets)++;
                 *received += n;
                 b->last += n;
                 do_write = 1;
@@ -1629,15 +1623,38 @@
 
     c->log->action = "proxying connection";
 
-    if (src->read->eof && dst && (dst->read->eof || !dst->buffered)) {
+    if (c->type == SOCK_DGRAM
+        && pscf->responses != NGX_MAX_INT32_VALUE
+        && u->responses >= pscf->responses * u->requests
+        && !src->buffered && dst && !dst->buffered)
+    {
         handler = c->log->handler;
         c->log->handler = NULL;
 
         ngx_log_error(NGX_LOG_INFO, c->log, 0,
-                      "%s%s disconnected"
+                      "udp done"
+                      ", packets from/to client:%ui/%ui"
                       ", bytes from/to client:%O/%O"
                       ", bytes from/to upstream:%O/%O",
-                      src->type == SOCK_DGRAM ? "udp " : "",
+                      u->requests, u->responses,
+                      s->received, c->sent, u->received, pc ? pc->sent : 0);
+
+        c->log->handler = handler;
+
+        ngx_stream_proxy_finalize(s, NGX_STREAM_OK);
+        return;
+    }
+
+    if (c->type == SOCK_STREAM
+        && src->read->eof && dst && (dst->read->eof || !dst->buffered))
+    {
+        handler = c->log->handler;
+        c->log->handler = NULL;
+
+        ngx_log_error(NGX_LOG_INFO, c->log, 0,
+                      "%s disconnected"
+                      ", bytes from/to client:%O/%O"
+                      ", bytes from/to upstream:%O/%O",
                       from_upstream ? "upstream" : "client",
                       s->received, c->sent, u->received, pc ? pc->sent : 0);
 
@@ -1739,6 +1756,7 @@
 static void
 ngx_stream_proxy_finalize(ngx_stream_session_t *s, ngx_uint_t rc)
 {
+    ngx_uint_t              state;
     ngx_connection_t       *pc;
     ngx_stream_upstream_t  *u;
 
@@ -1768,7 +1786,15 @@
     }
 
     if (u->peer.free && u->peer.sockaddr) {
-        u->peer.free(&u->peer, u->peer.data, 0);
+        state = 0;
+
+        if (pc && pc->type == SOCK_DGRAM
+            && (pc->read->error || pc->write->error))
+        {
+            state = NGX_PEER_FAILED;
+        }
+
+        u->peer.free(&u->peer, u->peer.data, state);
         u->peer.sockaddr = NULL;
     }
 
diff --git a/src/stream/ngx_stream_upstream.h b/src/stream/ngx_stream_upstream.h
index 73947f4..741ef4e 100644
--- a/src/stream/ngx_stream_upstream.h
+++ b/src/stream/ngx_stream_upstream.h
@@ -128,6 +128,7 @@
 
     off_t                              received;
     time_t                             start_sec;
+    ngx_uint_t                         requests;
     ngx_uint_t                         responses;
 
     ngx_str_t                          ssl_name;