feat: add 'drain' method (#46)

Author: spacewander

lib/resty/apisix/stream/xrpc/socket.lua

@@ -65,28 +65,24 @@ local function get_tcp_socket(cosocket)
 end
 
 
--- read the given length of data to a buffer in C land and return the buffer address
--- return error if the read data is less than given length
-local function read_buf(cosocket, len)
-    if len <= 0 then
-        error("bad length: length of data should be positive, got " .. len, 2)
-    end
-
-    if len > 4 * 1024 * 1024 then
-        error("bad length: length of data too big, got " .. len, 2)
-    end
-
+local function _read(cosocket, len, single_buf)
     local r = get_request()
     if not r then
         error("no request found", 2)
     end
 
     local u = get_tcp_socket(cosocket)
+
+    local buf
+    if single_buf then
+        buf = resbuf
+    end
+
     local errbuf = get_string_buf(ERR_BUF_SIZE)
     local errbuf_size = get_size_ptr()
     errbuf_size[0] = ERR_BUF_SIZE
 
-    local rc = socket_tcp_read(r, u, resbuf, len, errbuf, errbuf_size)
+    local rc = socket_tcp_read(r, u, buf, len, errbuf, errbuf_size)
     if rc == FFI_DONE then
         error(ffi_str(errbuf, errbuf_size[0]), 2)
     end
@@ -97,7 +93,11 @@ local function read_buf(cosocket, len)
         end
 
         if rc >= 0 then
-            return resbuf[0]
+            if single_buf then
+                return resbuf[0]
+            end
+
+            return true
         end
 
         assert(rc == FFI_AGAIN)
@@ -107,8 +107,35 @@ local function read_buf(cosocket, len)
         errbuf = get_string_buf(ERR_BUF_SIZE)
         errbuf_size = get_size_ptr()
         errbuf_size[0] = ERR_BUF_SIZE
-        rc = socket_tcp_get_read_result(r, u, resbuf, len, errbuf, errbuf_size)
+        rc = socket_tcp_get_read_result(r, u, buf, len, errbuf, errbuf_size)
+    end
+end
+
+
+-- read the given length of data to a buffer in C land and return the buffer address
+-- return error if the read data is less than given length
+local function read(cosocket, len)
+    if len <= 0 then
+        error("bad length: length of data should be positive, got " .. len, 2)
+    end
+
+    if len > 4 * 1024 * 1024 then
+        error("bad length: length of data too big, got " .. len, 2)
     end
+
+    return _read(cosocket, len, true)
+end
+
+
+-- work like `read` but don't return the buffer address and don't guarantee all the data is
+-- in the same buffer.
+-- return error if the read data is less than given length
+local function drain(cosocket, len)
+    if len <= 0 then
+        error("bad length: length of data should be positive, got " .. len, 2)
+    end
+
+    return _read(cosocket, len, false)
 end
 
 
@@ -172,7 +199,8 @@ local function patch_methods(sk)
     copy.receiveany = nil
     copy.receiveuntil = nil
 
-    copy.read = read_buf
+    copy.read = read
+    copy.drain = drain
     copy.move = move
 
     return {__index = copy}

patch/1.19.9/ngx_stream_lua-xrpc.patch

@@ -1,5 +1,5 @@
 diff --git src/ngx_stream_lua_socket_tcp.c src/ngx_stream_lua_socket_tcp.c
-index 7fcfb45..3178588 100644
+index 7fcfb45..13e942d 100644
 --- src/ngx_stream_lua_socket_tcp.c
 +++ src/ngx_stream_lua_socket_tcp.c
 @@ -234,6 +234,41 @@ enum {
@@ -44,7 +44,7 @@ index 7fcfb45..3178588 100644
 
  static char ngx_stream_lua_raw_req_socket_metatable_key;
  static char ngx_stream_lua_tcp_socket_metatable_key;
-@@ -6005,6 +6040,582 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
+@@ -6005,6 +6040,605 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
  }
 
 
@@ -213,7 +213,10 @@ index 7fcfb45..3178588 100644
 +        return NGX_ERROR;
 +    }
 +
-+    ngx_stream_lua_ffi_socket_push_res(r, ctx, u, buf, len);
++    if (buf != NULL) {
++        ngx_stream_lua_ffi_socket_push_res(r, ctx, u, buf, len);
++    }
++
 +    return NGX_OK;
 +}
 +
@@ -263,63 +266,83 @@ index 7fcfb45..3178588 100644
 +
 +    lctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
 +
-+    if (u->bufs_in == NULL) {
-+        size_t buf_len = len > u->conf->buffer_size ? len : u->conf->buffer_size;
++    if (buf != NULL) {
++        if (u->bufs_in == NULL) {
++            size_t buf_len = len > u->conf->buffer_size ? len : u->conf->buffer_size;
++
++            ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
++                           "stream lua tcp socket allocate new new buf of size %uz",
++                           buf_len);
++
++            u->bufs_in =
++                ngx_stream_lua_chain_get_free_buf(r->connection->log,
++                                                  r->pool,
++                                                  &lctx->free_recv_bufs,
++                                                  buf_len);
++
++            if (u->bufs_in == NULL) {
++                *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "no memory") - errbuf;
++                return NGX_DONE;
++            }
++
++            u->buf_in = u->bufs_in;
++            u->buffer = *u->buf_in->buf;
++
++        } else {
++            size_t           remain_space = u->buffer.end - u->buffer.pos;
++            size_t           remain_data = u->buffer.last - u->buffer.pos;
++            size_t           buf_len;
++            u_char          *pos;
++            ngx_chain_t     *cl, *tmp_cl;
++
++            if (remain_space < len) {
++                buf_len = len > u->conf->buffer_size ? len : u->conf->buffer_size;
++
++                ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
++                               "stream lua tcp socket allocate new new buf of size %uz",
++                               buf_len);
++
++                cl = ngx_stream_lua_chain_get_free_buf(r->connection->log,
++                                                       r->pool,
++                                                       &lctx->free_recv_bufs,
++                                                       buf_len);
++                if (cl == NULL) {
++                    *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "no memory") - errbuf;
++                    return NGX_DONE;
++                }
++
++                for (tmp_cl = u->bufs_in; tmp_cl->next; tmp_cl = tmp_cl->next) {}
++                tmp_cl->next = cl;
++                u->buf_in = cl;
++
++                u->buffer.last = u->buffer.pos;
++                pos = u->buffer.pos;
++                u->buffer = *cl->buf;
 +
++                if (remain_data > 0) {
++                    u->buffer.last = ngx_copy(u->buffer.last, pos, remain_data);
++                }
++            }
++        }
++
++    } else if (u->bufs_in == NULL) {
 +        ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
 +                       "stream lua tcp socket allocate new new buf of size %uz",
-+                       buf_len);
++                       u->conf->buffer_size);
 +
 +        u->bufs_in =
 +            ngx_stream_lua_chain_get_free_buf(r->connection->log,
 +                                              r->pool,
 +                                              &lctx->free_recv_bufs,
-+                                              buf_len);
++                                              u->conf->buffer_size);
 +
 +        if (u->bufs_in == NULL) {
-+            *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "no memory")
-+                           - errbuf;
++            *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "no memory") - errbuf;
 +            return NGX_DONE;
 +        }
 +
 +        u->buf_in = u->bufs_in;
 +        u->buffer = *u->buf_in->buf;
-+
-+    } else {
-+        size_t           remain_space = u->buffer.end - u->buffer.pos;
-+        size_t           remain_data = u->buffer.last - u->buffer.pos;
-+        size_t           buf_len;
-+        u_char          *pos;
-+        ngx_chain_t     *cl, *tmp_cl;
-+
-+        if (remain_space < len) {
-+            buf_len = len > u->conf->buffer_size ? len : u->conf->buffer_size;
-+
-+            ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
-+                           "stream lua tcp socket allocate new new buf of size %uz",
-+                           buf_len);
-+
-+            cl = ngx_stream_lua_chain_get_free_buf(r->connection->log,
-+                                                   r->pool,
-+                                                   &lctx->free_recv_bufs,
-+                                                   buf_len);
-+            if (cl == NULL) {
-+                *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "no memory") - errbuf;
-+                return NGX_DONE;
-+            }
-+
-+            for (tmp_cl = u->bufs_in; tmp_cl->next; tmp_cl = tmp_cl->next) {}
-+            tmp_cl->next = cl;
-+            u->buf_in = cl;
-+
-+            u->buffer.last = u->buffer.pos;
-+            pos = u->buffer.pos;
-+            u->buffer = *cl->buf;
-+
-+            if (remain_data > 0) {
-+                u->buffer.last = ngx_copy(u->buffer.last, pos, remain_data);
-+            }
-+        }
 +    }
 +
 +    dd("tcp receive: buf_in: %p, bufs_in: %p", u->buf_in, u->bufs_in);

t/stream/xrpc/downstream.t

@@ -256,3 +256,46 @@ qr/stream lua tcp socket allocate new new buf of size \d+/
 --- grep_error_log_out
 stream lua tcp socket allocate new new buf of size 128
 stream lua tcp socket allocate new new buf of size 144
+
+
+
+=== TEST 10: drain & move
+--- stream_config
+    server {
+        listen 1995;
+        content_by_lua_block {
+            local sk = ngx.req.socket(true)
+            local exp = {
+                "hell",
+                "o wo",
+                "r",
+            }
+            for i = 1, 3 do
+                local data = sk:receiveany(128)
+                if data ~= exp[i] then
+                    ngx.log(ngx.ERR, "actual: ", data, ", expected: ", exp[i])
+                end
+            end
+        }
+    }
+--- stream_server_config
+    content_by_lua_block {
+        local ffi = require("ffi")
+        local ds = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        ds:settimeout(5)
+
+        local us = require("resty.apisix.stream.xrpc.socket").upstream.socket()
+        us:settimeout(50)
+        assert(us:connect("127.0.0.1", 1995))
+
+        assert(ds:drain(4))
+        assert(us:move(ds))
+        ngx.sleep(0.01)
+        assert(ds:drain(4))
+        assert(us:move(ds))
+        ngx.sleep(0.01)
+        assert(ds:drain(1))
+        assert(us:move(ds))
+    }
+--- stream_request
+hello world

t/stream/xrpc/fuzzing.t

@@ -173,3 +173,47 @@ __DATA__
 --- stream_request
 --- stream_response eval
 "123456789" x 128
+
+
+
+=== TEST 5: move (read & drain)
+--- stream_config
+    lua_socket_buffer_size 128;
+    server {
+        listen 1995;
+        content_by_lua_block {
+            local s = ("123456789"):rep(1280)
+            ngx.say(s)
+        }
+    }
+--- stream_server_config
+    lua_socket_buffer_size 128;
+    content_by_lua_block {
+        math.randomseed(ngx.time())
+
+        local ffi = require("ffi")
+        local sk = require("resty.apisix.stream.xrpc.socket").upstream.socket()
+        local dsk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        assert(sk:connect("127.0.0.1", 1995))
+        sk:settimeout(500)
+        local total = 9 * 1280
+        while total > 0 do
+            local len = math.random(1, 512)
+            if len > total then
+                len = total
+            end
+            total = total - len
+            if math.random(1, 2) == 1 then
+                assert(sk:read(len))
+            else
+                assert(sk:drain(len))
+            end
+            if total % 2 == 0 then
+                dsk:move(sk)
+            end
+        end
+    }
+--- stream_request
+--- stream_response eval
+"123456789" x 1280
+--- timeout: 5