From e19bb18fad80ed95fb68318fc6ec7f2f892942dd Mon Sep 17 00:00:00 2001
From: li_yunqing <liyunqing@kylinos.cn>
Date: Tue, 19 Mar 2024 21:06:07 +0800
Subject: [PATCH] recvfrom support timeout

---
 src/lstack/api/lstack_epoll.c           | 57 ++++++++++++++++---------
 src/lstack/api/lstack_wrap.c            |  1 -
 src/lstack/core/lstack_lwip.c           | 28 +++++-------
 src/lstack/core/lstack_protocol_stack.c |  2 +-
 src/lstack/include/posix/lstack_epoll.h | 10 +----
 5 files changed, 50 insertions(+), 48 deletions(-)

diff --git a/src/lstack/api/lstack_epoll.c b/src/lstack/api/lstack_epoll.c
index 7dbef9d..1b68c36 100644
--- a/src/lstack/api/lstack_epoll.c
+++ b/src/lstack/api/lstack_epoll.c
@@ -586,6 +586,41 @@ static void ms_to_timespec(struct timespec *timespec, int32_t timeout)
     timespec->tv_nsec = timespec->tv_nsec % SEC_TO_NSEC;
 }
 
+/**
+ * Block lstack thread
+ *
+ * @param wakeup
+ *  The pointer to the wakeup_poll.
+ * @param timeout
+ *  The time to wait, if 'timeout <= 0' will block until unlock
+ *
+ * @return
+ * - return '0' on unlock
+ * - return 'ETIMEDOUT' on timeout
+ */
+int32_t lstack_block_wait(struct wakeup_poll *wakeup, int32_t timeout)
+{
+    int ret = 0;
+    if (wakeup == NULL) {
+        return ret;
+    }
+
+    __atomic_store_n(&wakeup->in_wait, true, __ATOMIC_RELEASE);
+    if (timeout > 0) {
+        struct timespec timespec;
+        ms_to_timespec(&timespec, timeout);
+        ret = pthread_mutex_timedlock(&wakeup->wait, &timespec);
+    } else {
+        ret = pthread_mutex_lock(&wakeup->wait);
+    }
+
+    if (__atomic_load_n(&wakeup->in_wait, __ATOMIC_ACQUIRE)) {
+        __atomic_store_n(&wakeup->in_wait, false, __ATOMIC_RELEASE);
+    }
+
+    return ret;
+}
+
 int32_t lstack_rtc_epoll_wait(int32_t epfd, struct epoll_event* events, int32_t maxevents, int32_t timeout)
 {
     struct lwip_sock *sock = get_socket_by_fd(epfd);
@@ -645,7 +680,6 @@ int32_t lstack_rtw_epoll_wait(int32_t epfd, struct epoll_event* events, int32_t
     struct wakeup_poll *wakeup = sock->wakeup;
     int32_t kernel_num = 0;
     int32_t lwip_num = 0;
-    int32_t ret = 0;
 
     if (get_global_cfg_params()->app_bind_numa) {
         epoll_bind_statck(sock->wakeup);
@@ -669,15 +703,7 @@ int32_t lstack_rtw_epoll_wait(int32_t epfd, struct epoll_event* events, int32_t
         if (timeout == 0) {
             break;
         }
-
-        if (timeout < 0) {
-            ret = pthread_mutex_lock(&wakeup->wait);
-        } else {
-            struct timespec epoll_time;
-            ms_to_timespec(&epoll_time, timeout);
-            ret = pthread_mutex_timedlock(&wakeup->wait, &epoll_time);
-        }
-    } while (ret == 0);
+    } while (lstack_block_wait(wakeup, timeout) == 0);
 
     __atomic_store_n(&wakeup->in_wait, false, __ATOMIC_RELEASE);
     wakeup->stat.app_events += lwip_num;
@@ -857,7 +883,6 @@ int32_t lstack_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
 
     int32_t kernel_num = 0;
     int32_t lwip_num = 0;
-    int32_t ret;
 
     do {
         __atomic_store_n(&wakeup->in_wait, true, __ATOMIC_RELEASE);
@@ -881,15 +906,7 @@ int32_t lstack_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout)
         if (timeout == 0) {
             break;
         }
-
-        if (timeout < 0) {
-            ret = pthread_mutex_lock(&wakeup->wait);
-        } else {
-            struct timespec epoll_time;
-            ms_to_timespec(&epoll_time, timeout);
-            ret = pthread_mutex_timedlock(&wakeup->wait, &epoll_time);
-        }
-    } while (ret == 0);
+    } while (lstack_block_wait(wakeup, timeout) == 0);
 
     __atomic_store_n(&wakeup->in_wait, false, __ATOMIC_RELEASE);
     wakeup->stat.app_events += lwip_num;
diff --git a/src/lstack/api/lstack_wrap.c b/src/lstack/api/lstack_wrap.c
index 49bbf91..0dac82e 100644
--- a/src/lstack/api/lstack_wrap.c
+++ b/src/lstack/api/lstack_wrap.c
@@ -380,7 +380,6 @@ static bool unsupport_socket_optname(int32_t optname)
     if ((optname == SO_BROADCAST) ||
         (optname == SO_PROTOCOL) ||
         (optname == SO_SNDTIMEO) ||
-        (optname == SO_RCVTIMEO) ||
         (optname == SO_SNDBUF) ||
         (optname == SO_RCVBUF) ||
         (optname == SO_DONTROUTE)) {
diff --git a/src/lstack/core/lstack_lwip.c b/src/lstack/core/lstack_lwip.c
index 8aae433..d0dea80 100644
--- a/src/lstack/core/lstack_lwip.c
+++ b/src/lstack/core/lstack_lwip.c
@@ -859,39 +859,33 @@ static bool recv_break_for_err(struct lwip_sock *sock)
     return break_wait;
 }
 
-static void recv_block_wait(struct lwip_sock *sock)
-{
-    lstack_block_wait(sock->wakeup);
-}
-
 /*
  * return 0 on success, -1 on error
  * pbuf maybe NULL(tcp fin packet)
  */
 static int recv_ring_get_one(struct lwip_sock *sock, bool noblock, struct pbuf **pbuf)
 {
+    int32_t expect = 1; // only get one pbuf
+
     if (sock->recv_lastdata != NULL) {
         *pbuf = sock->recv_lastdata;
         sock->recv_lastdata = NULL;
         return 0;
     }
 
-    if (noblock) {
-        if (gazelle_ring_read(sock->recv_ring, (void **)pbuf, 1) != 1) {
-            errno = EAGAIN;
+    while (gazelle_ring_read(sock->recv_ring, (void **)pbuf, expect) != expect) {
+        if (noblock) {
+            GAZELLE_RETURN(EAGAIN);
+        }
+        if (recv_break_for_err(sock)) {
             return -1;
-        } else {
-            return 0;
         }
-    } else  {
-        while (gazelle_ring_read(sock->recv_ring, (void **)pbuf, 1) != 1) {
-            if (recv_break_for_err(sock)) {
-                return -1;
-            }
-            recv_block_wait(sock);
+        if (lstack_block_wait(sock->wakeup, sock->conn->recv_timeout) == ETIMEDOUT) {
+            noblock = true;
         }
-        return 0;
     }
+    
+    return 0;
 }
 
 /* return true: fin is read to user, false: pend fin */
diff --git a/src/lstack/core/lstack_protocol_stack.c b/src/lstack/core/lstack_protocol_stack.c
index a545b73..7c4af64 100644
--- a/src/lstack/core/lstack_protocol_stack.c
+++ b/src/lstack/core/lstack_protocol_stack.c
@@ -1270,7 +1270,7 @@ int32_t stack_broadcast_accept4(int32_t fd, struct sockaddr *addr, socklen_t *ad
         min_sock = get_min_accept_sock(fd);
     } else {
         while ((min_sock = get_min_accept_sock(fd)) == NULL) {
-            lstack_block_wait(sock->wakeup);
+            lstack_block_wait(sock->wakeup, 0);
 	}
     }
 
diff --git a/src/lstack/include/posix/lstack_epoll.h b/src/lstack/include/posix/lstack_epoll.h
index 7591f0f..59b5ef7 100644
--- a/src/lstack/include/posix/lstack_epoll.h
+++ b/src/lstack/include/posix/lstack_epoll.h
@@ -80,15 +80,7 @@ int32_t lstack_rtc_epoll_wait(int32_t epfd, struct epoll_event* events, int32_t
 int32_t lstack_poll(struct pollfd *fds, nfds_t nfds, int32_t timeout);
 int lstack_select(int maxfd, fd_set *readfds, fd_set *writefds, fd_set *exceptfds, struct timeval *timeval);
 
-static inline void lstack_block_wait(struct wakeup_poll *wakeup)
-{
-    if (wakeup == NULL) {
-        return;
-    }
-
-    __atomic_store_n(&wakeup->in_wait, true, __ATOMIC_RELEASE);
-    pthread_mutex_lock(&wakeup->wait);
-}
+int32_t lstack_block_wait(struct wakeup_poll *wakeup, int32_t timeout);
 
 static inline void lstack_block_wakeup(struct wakeup_poll *wakeup)
 {
-- 
2.33.0