feat: add has_pending_data (#51)

Author: spacewander

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

@@ -43,12 +43,18 @@ ngx_stream_lua_ffi_socket_tcp_get_send_result(ngx_stream_lua_request_t *r,
 void
 ngx_stream_lua_ffi_socket_tcp_reset_read_buf(ngx_stream_lua_request_t *r,
     ngx_stream_lua_socket_tcp_upstream_t *u);
+
+int
+ngx_stream_lua_ffi_socket_tcp_has_pending_data(ngx_stream_lua_request_t *r,
+    ngx_stream_lua_socket_tcp_upstream_t *u,
+    u_char *errbuf, size_t *errbuf_size);
 ]]
 local socket_tcp_read = C.ngx_stream_lua_ffi_socket_tcp_read_buf
 local socket_tcp_get_read_result = C.ngx_stream_lua_ffi_socket_tcp_get_read_buf_result
 local socket_tcp_move = C.ngx_stream_lua_ffi_socket_tcp_send_from_socket
 local socket_tcp_get_move_result = C.ngx_stream_lua_ffi_socket_tcp_get_send_result
 local socket_tcp_reset_read_buf = C.ngx_stream_lua_ffi_socket_tcp_reset_read_buf
+local socket_tcp_has_pending_data = C.ngx_stream_lua_ffi_socket_tcp_has_pending_data
 
 
 local ERR_BUF_SIZE = 256
@@ -178,6 +184,29 @@ local function drain(cosocket, len)
 end
 
 
+-- has_pending_data check if there is unread data in the given socket.
+-- return false if there is no pending data, and return true if there may be any pending data.
+-- we require it to be called after any read methods called successfully.
+local function has_pending_data(cosocket)
+    local r = get_request()
+    if not r then
+        error("no request found", 2)
+    end
+
+    local u = get_tcp_socket(cosocket)
+
+    local errbuf = get_string_buf(ERR_BUF_SIZE)
+    local errbuf_size = get_size_ptr()
+    errbuf_size[0] = ERR_BUF_SIZE
+
+    local rc = socket_tcp_has_pending_data(r, u, errbuf, errbuf_size)
+    if rc == FFI_ERROR then
+        return nil, ffi_str(errbuf, errbuf_size[0])
+    end
+    return rc == FFI_AGAIN
+end
+
+
 -- move the buffers from src cosocket to dst cosocket. The buffers are from previous one or multiple
 -- read calls. It is equal to send multiple read buffer in the src cosocket to the dst cosocket.
 local function move(dst, src)
@@ -256,6 +285,7 @@ local function patch_methods(sk)
     copy.read_line = read_line
     copy.move = move
     copy.reset_read_buf = reset_read_buf
+    copy.has_pending_data = has_pending_data
 
     return {__index = copy}
 end

patch/1.19.9/ngx_stream_lua-xrpc.patch

@@ -81,7 +81,7 @@ index e65d572..25175a5 100644
  #endif /* _NGX_STREAM_LUA_INPUT_FILTERS_H_INCLUDED_ */
 
 diff --git src/ngx_stream_lua_socket_tcp.c src/ngx_stream_lua_socket_tcp.c
-index 7fcfb45..a21561b 100644
+index 7fcfb45..9981b4b 100644
 --- src/ngx_stream_lua_socket_tcp.c
 +++ src/ngx_stream_lua_socket_tcp.c
 @@ -234,6 +234,41 @@ enum {
@@ -158,7 +158,7 @@ index 7fcfb45..a21561b 100644
  static ngx_int_t
  ngx_stream_lua_socket_tcp_read(ngx_stream_lua_request_t *r,
      ngx_stream_lua_socket_tcp_upstream_t *u)
-@@ -6005,6 +6065,647 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
+@@ -6005,6 +6065,672 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
  }
 
 
@@ -278,6 +278,7 @@ index 7fcfb45..a21561b 100644
 +            p = ngx_strerror(u->socket_errno, errbuf, *errbuf_size);
 +            /* for compatibility with LuaSocket */
 +            ngx_strlow(errbuf, errbuf, p - errbuf);
++            *errbuf_size = p - errbuf;
 +
 +        } else {
 +            *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "error")
@@ -585,6 +586,30 @@ index 7fcfb45..a21561b 100644
 +}
 +
 +
++int
++ngx_stream_lua_ffi_socket_tcp_has_pending_data(ngx_stream_lua_request_t *r,
++    ngx_stream_lua_socket_tcp_upstream_t *u,
++    u_char *errbuf, size_t *errbuf_size)
++{
++    /* skip many input checks as we require glance to be called
++     * after any read methods called successfully */
++
++    if (u == NULL || u->bufs_in == NULL) {
++        *errbuf_size = ngx_snprintf(errbuf, *errbuf_size, "closed") - errbuf;
++        return NGX_ERROR;
++    }
++
++    /* no remain data */
++    if ((u->buffer.last == u->buffer.pos && u->buffer.last != u->buffer.end)
++        || u->eof /* EOF reached */)
++    {
++        return NGX_OK;
++    }
++
++    return NGX_AGAIN;
++}
++
++
 +static void
 +ngx_stream_lua_ffi_socket_write_error_retval_handler(
 +    ngx_stream_lua_request_t *r, ngx_stream_lua_socket_tcp_upstream_t *u,

t/stream/xrpc/misc.t

@@ -0,0 +1,71 @@
+use t::APISIX_NGINX 'no_plan';
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: has_pending_data
+--- stream_server_config
+    content_by_lua_block {
+        local sk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        local p, err, len = sk:read_line(128)
+        if err then
+            ngx.say(err)
+            return
+        end
+        ngx.say(sk:has_pending_data())
+    }
+--- stream_request eval
+"hello world\r\n" x 2
+--- stream_response
+true
+
+
+
+=== TEST 2: has_pending_data, all are read
+--- stream_server_config
+    content_by_lua_block {
+        local sk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        local p, err, len = sk:read_line(128)
+        if err then
+            ngx.say(err)
+            return
+        end
+        ngx.say(sk:has_pending_data())
+    }
+--- stream_request eval
+"hello world\r\n"
+--- stream_response
+false
+
+
+
+=== TEST 3: has_pending_data, multiple read
+--- stream_server_config
+    content_by_lua_block {
+        local sk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        assert(sk:read(4))
+        assert(sk:drain(7))
+        ngx.say(sk:has_pending_data())
+    }
+--- stream_request eval
+"hello world"
+--- stream_response
+false
+
+
+
+=== TEST 4: has_pending_data, buffer is greater than lua_socket_buffer_size
+--- stream_server_config
+    lua_socket_buffer_size 128;
+    content_by_lua_block {
+        local sk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        assert(sk:read(132))
+        -- In this case, the has_pending_data has to return true as
+        -- there is no way to know if there is pending data without a read
+        ngx.say(sk:has_pending_data())
+    }
+--- stream_request eval
+"1234" x 33
+--- stream_response
+true