Notify TCP writers when the final synchronous send drains

lib/tcp_io.c

@@ -282,6 +282,14 @@ static int flush_write_buf(struct tcp_conn* tc) {
     return 0;
 }
 
+static void notify_flushed(struct tcp_conn* tc) {
+    if (tc->destroyed || tc->write_buf || tc->write_queue->head || !tc->on_flushed) return;
+    void (*cb)(struct tcp_conn*, void*) = tc->on_flushed;
+    tc->on_flushed = NULL;
+    DEBUG_DEBUG(DEBUG_CATEGORY_SOCKET, "tcp_io: flushed fd=%d", (int)tc->sock);
+    cb(tc, tc->arg);
+}
+
 static void write_queue_fetch_cb(struct ll_queue* q, void* arg) {
     struct tcp_conn* tc = (struct tcp_conn*)arg;
 
@@ -300,11 +308,7 @@ static void write_queue_fetch_cb(struct ll_queue* q, void* arg) {
         tc->write_monitor = 0;
         DEBUG_TRACE(DEBUG_CATEGORY_SOCKET, "tcp_io: wq_fetch POLLOUT OFF fd=%d (queue empty)", (int)tc->sock);
         queue_resume_callback(q);
-        if (tc->on_flushed) {
+        notify_flushed(tc);
-            void (*cb)(struct tcp_conn*, void*) = tc->on_flushed;
-            tc->on_flushed = NULL;
-            cb(tc, tc->arg);
-        }
         return;
     }
 
@@ -334,6 +338,7 @@ static void write_queue_fetch_cb(struct ll_queue* q, void* arg) {
         memory_pool_free(tc->data_pool, e->dgram);
         queue_entry_free(e);
         queue_resume_callback(q);
+        notify_flushed(tc);
     } else if (n > 0) {
         size_t rem = e->len - (size_t)n;
         tc->write_buf = memory_pool_alloc(tc->data_pool);
@@ -404,6 +409,7 @@ static void write_cb(socket_t sock, void* arg) {
             tc->write_monitor = 0;
         }
         queue_resume_callback(tc->write_queue);
+        notify_flushed(tc);
     }
 }
 

tests/Makefile.am

@@ -636,3 +636,7 @@ test_services_LDADD = $(LIBUTUN) $(CRYPTO_LIBS) $(COMMON_LIBS)
 check_PROGRAMS += test_proxy_packets
 test_proxy_packets_SOURCES = test_proxy_packets.c
 test_proxy_packets_LDADD = $(LIBUTUN) $(CRYPTO_LIBS) $(COMMON_LIBS)
+
+check_PROGRAMS += test_tcp_io_flush
+test_tcp_io_flush_SOURCES = test_tcp_io_flush.c
+test_tcp_io_flush_LDADD = $(LIBUTUN) $(CRYPTO_LIBS) $(COMMON_LIBS)

tests/test_tcp_io_flush.c

@@ -0,0 +1,42 @@
+// on_flushed должен приходить после последнего send даже без активного POLLOUT.
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include "tcp_io.h"
+#include "debug_config.h"
+#include "mem.h"
+#define CHECK(x) do { if (!(x)) { fprintf(stderr, "FAIL line %d: %s\n", __LINE__, #x); exit(1); } } while (0)
+static int flushed;
+static void on_flushed(struct tcp_conn* tc, void* arg) { (void)tc; (void)arg; flushed++; }
+static void on_error(struct tcp_conn* tc, int err, void* arg) { (void)tc; (void)arg; fprintf(stderr, "socket error %d\n", err); exit(1); }
+int main(void) {
+    debug_config_init(); debug_set_level(DEBUG_LEVEL_WARN);
+    socket_platform_init();
+    struct UASYNC* ua = uasync_create(); CHECK(ua);
+    socket_t listener = socket(AF_INET, SOCK_STREAM, 0); CHECK(listener != SOCKET_INVALID);
+    struct sockaddr_in addr = {.sin_family = AF_INET, .sin_addr.s_addr = htonl(INADDR_LOOPBACK)};
+    CHECK(bind(listener, (struct sockaddr*)&addr, sizeof(addr)) == 0 && listen(listener, 1) == 0);
+    socklen_t len = sizeof(addr); CHECK(getsockname(listener, (struct sockaddr*)&addr, &len) == 0);
+    socket_t peer = socket(AF_INET, SOCK_STREAM, 0); CHECK(peer != SOCKET_INVALID);
+    CHECK(connect(peer, (struct sockaddr*)&addr, sizeof(addr)) == 0);
+    socket_t fd = accept(listener, NULL, NULL); CHECK(fd != SOCKET_INVALID);
+    socket_close_wrapper(listener); socket_set_nonblocking(fd); socket_set_nonblocking(peer);
+    struct tcp_conn* tc = tcp_conn_create(ua, fd, 4096, 4096, 8, 4, 0, NULL, on_error, NULL); CHECK(tc);
+    tc->connected = 1;
+    for (int n = 0; n < 3; n++) {
+        for (int i = 0; i < 12; i++) uasync_poll(ua, 1);
+        CHECK(tc->write_monitor == 0);
+        struct ll_entry* e = queue_entry_new_from_pool(tc->entry_pool); CHECK(e);
+        e->dgram = memory_pool_alloc(tc->data_pool); CHECK(e->dgram);
+        memcpy(e->dgram, "last", 4); e->len = 4;
+        tcp_conn_set_flushed(tc, on_flushed); queue_data_put(tc->write_queue, e);
+        for (int i = 0; i < 12; i++) uasync_poll(ua, 1);
+        CHECK(flushed == n + 1);
+        char bytes[4]; CHECK(recv(peer, bytes, sizeof(bytes), 0) == 4 && memcmp(bytes, "last", 4) == 0);
+    }
+    tcp_conn_destroy(tc); socket_close_wrapper(peer);
+    for (int i = 0; i < 12; i++) uasync_poll(ua, 1);
+    uasync_destroy(ua, 0); CHECK(u_get_allocated_count() == 0);
+    puts("[PASS] tcp_io: final synchronous send notifies flushed exactly once");
+    return 0;
+}