From a735fd09aa59df805865331b3cbda5e0458c9fac Mon Sep 17 00:00:00 2001
From: yinbin <yinbin8@huawei.com>
Date: Tue, 19 Nov 2024 19:36:26 +0800
Subject: [PATCH] socket: init wakeup in blocking socket

---
 src/lstack/api/lstack_epoll.c     | 16 ++++++++++++----
 src/lstack/api/lstack_rtw_api.c   |  9 +++++++++
 src/lstack/core/lstack_lwip.c     | 15 ++++++++++++++-
 src/lstack/include/lstack_epoll.h |  2 ++
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/lstack/api/lstack_epoll.c b/src/lstack/api/lstack_epoll.c
index acbf393..ff9cccf 100644
--- a/src/lstack/api/lstack_epoll.c
+++ b/src/lstack/api/lstack_epoll.c
@@ -872,21 +872,30 @@ static int poll_init(struct wakeup_poll *wakeup, struct pollfd *fds, nfds_t nfds
     return 0;
 }
 
-int32_t lstack_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
+struct wakeup_poll* poll_construct_wakeup(void)
 {
     static PER_THREAD struct wakeup_poll *wakeup = NULL;
     if (wakeup == NULL) {
         wakeup = calloc(1, sizeof(struct wakeup_poll));
         if (wakeup == NULL) {
             LSTACK_LOG(ERR, LSTACK, "calloc failed errno=%d\n", errno);
-            GAZELLE_RETURN(EINVAL);
+            return NULL;
         }
 
         if (init_poll_wakeup_data(wakeup) < 0) {
             free(wakeup);
-            GAZELLE_RETURN(EINVAL);
+            return NULL;
         }
     }
+    return wakeup;
+}
+
+int32_t lstack_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
+{
+    struct wakeup_poll *wakeup = poll_construct_wakeup();
+    if (wakeup == NULL) {
+        GAZELLE_RETURN(EINVAL);
+    }
 
     if (poll_init(wakeup, fds, nfds) < 0) {
         free(wakeup);
@@ -1015,4 +1024,3 @@ int lstack_select(int maxfd, fd_set *readfds, fd_set *writefds, fd_set *exceptfd
 
     return event_num;
 }
-
diff --git a/src/lstack/api/lstack_rtw_api.c b/src/lstack/api/lstack_rtw_api.c
index 1b02e2a..6d0bd05 100644
--- a/src/lstack/api/lstack_rtw_api.c
+++ b/src/lstack/api/lstack_rtw_api.c
@@ -136,6 +136,15 @@ static struct lwip_sock *get_min_accept_sock(int fd)
     struct lwip_sock *min_sock = NULL;
 
     while (sock) {
+        if (!netconn_is_nonblocking(sock->conn)) {
+            if (sock->wakeup == NULL) {
+                sock->wakeup = poll_construct_wakeup();
+                if (sock->wakeup == NULL) {
+                    return NULL;
+                }
+                sock->epoll_events = POLLIN | POLLERR;
+            }
+        }
         if (!NETCONN_IS_ACCEPTIN(sock)) {
             sock = sock->listen_next;
             continue;
diff --git a/src/lstack/core/lstack_lwip.c b/src/lstack/core/lstack_lwip.c
index bb261d2..65acae2 100644
--- a/src/lstack/core/lstack_lwip.c
+++ b/src/lstack/core/lstack_lwip.c
@@ -993,6 +993,7 @@ static bool recv_break_for_err(struct lwip_sock *sock)
 static int recv_ring_get_one(struct lwip_sock *sock, bool noblock, struct pbuf **pbuf)
 {
     int32_t expect = 1; // only get one pbuf
+    int ret = 0;
     uint64_t time_stamp = sys_now_us();
 
     if (sock->recv_lastdata != NULL) {
@@ -1008,8 +1009,20 @@ static int recv_ring_get_one(struct lwip_sock *sock, bool noblock, struct pbuf *
         if (recv_break_for_err(sock)) {
             return -1;
         }
-        if (lstack_block_wait(sock->wakeup, sock->conn->recv_timeout) == ETIMEDOUT) {
+        if (unlikely(sock->wakeup == NULL)) {
+            sock->wakeup = poll_construct_wakeup();
+            if (sock->wakeup == NULL) {
+                return -1;
+            }
+            sock->epoll_events = POLLIN | POLLERR;
+        }
+
+        ret = lstack_block_wait(sock->wakeup, sock->conn->recv_timeout);
+        if (ret == ETIMEDOUT) {
             noblock = true;
+        } else if (ret != 0 && errno == EINTR) {
+            /* SIGALRM signal may interrupt blocking */
+            return ret;
         }
     }
 
diff --git a/src/lstack/include/lstack_epoll.h b/src/lstack/include/lstack_epoll.h
index e7ae26b..cad9aed 100644
--- a/src/lstack/include/lstack_epoll.h
+++ b/src/lstack/include/lstack_epoll.h
@@ -76,6 +76,8 @@ int lstack_select(int maxfd, fd_set *readfds, fd_set *writefds, fd_set *exceptfd
 
 int32_t lstack_block_wait(struct wakeup_poll *wakeup, int32_t timeout);
 
+struct wakeup_poll* poll_construct_wakeup(void);
+
 static inline void lstack_block_wakeup(struct wakeup_poll *wakeup)
 {
     if (wakeup && __atomic_load_n(&wakeup->in_wait, __ATOMIC_ACQUIRE)) {
-- 
2.33.0