From f4abd3b3fd5004405cb186981b93f5d40e4648db Mon Sep 17 00:00:00 2001
From: jiangheng <jiangheng14@huawei.com>
Date: Sun, 10 Dec 2023 17:47:56 +0800
Subject: [PATCH] fix close can't exit

---
 src/lstack/core/lstack_lwip.c           | 12 ++++++++----
 src/lstack/core/lstack_protocol_stack.c | 25 +++++++++++++++----------
 src/lstack/core/lstack_thread_rpc.c     |  6 +++---
 src/lstack/include/lstack_lwip.h        |  4 ++--
 4 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/lstack/core/lstack_lwip.c b/src/lstack/core/lstack_lwip.c
index c4b1ebc..73a6f12 100644
--- a/src/lstack/core/lstack_lwip.c
+++ b/src/lstack/core/lstack_lwip.c
@@ -646,14 +646,18 @@ bool do_lwip_replenish_sendring(struct protocol_stack *stack, struct lwip_sock *
     return replenish_again;
 }
 
-bool do_lwip_send(struct protocol_stack *stack, int32_t fd, struct lwip_sock *sock,
-                  size_t len, int32_t flags)
+int do_lwip_send(struct protocol_stack *stack, int32_t fd, struct lwip_sock *sock,
+                 size_t len, int32_t flags)
 {
+    ssize_t ret;
     /* send all send_ring, so len set lwip send max. */
     if (NETCONN_IS_UDP(sock)) {
-        (void)lwip_send(fd, sock, len, flags);
+        ret = lwip_send(fd, sock, len, flags);
     } else {
-        (void)lwip_send(fd, sock, UINT16_MAX, flags);
+        ret = lwip_send(fd, sock, UINT16_MAX, flags);
+    }
+    if (ret < 0 && (errno == ENOTCONN || errno == ECONNRESET || errno == ECONNABORTED)) {
+        return -1;
     }
 
     return do_lwip_replenish_sendring(stack, sock);
diff --git a/src/lstack/core/lstack_protocol_stack.c b/src/lstack/core/lstack_protocol_stack.c
index f61e7a8..8dbd9ad 100644
--- a/src/lstack/core/lstack_protocol_stack.c
+++ b/src/lstack/core/lstack_protocol_stack.c
@@ -699,10 +699,10 @@ void stack_close(struct rpc_msg *msg)
     struct protocol_stack *stack = get_protocol_stack_by_fd(fd);
     struct lwip_sock *sock = get_socket(fd);
     
-    if (sock && NETCONN_IS_DATAOUT(sock)) {
+    if (sock && __atomic_load_n(&sock->call_num, __ATOMIC_ACQUIRE) > 0) {
         msg->recall_flag = 1;
         rpc_call(&stack->rpc_queue, msg); /* until stack_send recall finish */
-	return;
+        return;
     }
     
     msg->result = lwip_close(fd);
@@ -860,27 +860,32 @@ void stack_send(struct rpc_msg *msg)
     int32_t fd = msg->args[MSG_ARG_0].i;
     size_t len = msg->args[MSG_ARG_1].size;
     struct protocol_stack *stack = (struct protocol_stack *)msg->args[MSG_ARG_3].p;
-    bool replenish_again;
+    int replenish_again;
 
     struct lwip_sock *sock = get_socket(fd);
     if (sock == NULL) {
         msg->result = -1;
         LSTACK_LOG(ERR, LSTACK, "get sock error! fd=%d, len=%ld\n", fd, len);
-        rpc_msg_free(msg);
+        __sync_fetch_and_sub(&sock->call_num, 1);
         return;
     }
 
     replenish_again = do_lwip_send(stack, sock->conn->socket, sock, len, 0);
-    __sync_fetch_and_sub(&sock->call_num, 1);
-    if (!NETCONN_IS_DATAOUT(sock) && !replenish_again) {
+    if (replenish_again < 0) {
+        __sync_fetch_and_sub(&sock->call_num, 1);
         return;
-    } else {
-        if (__atomic_load_n(&sock->call_num, __ATOMIC_ACQUIRE) == 0) {
-	    msg->recall_flag = 1;
+    }
+
+    if (NETCONN_IS_DATAOUT(sock) || replenish_again > 0) {
+        if (__atomic_load_n(&sock->call_num, __ATOMIC_ACQUIRE) == 1) {
+            msg->recall_flag = 1;
             rpc_call(&stack->rpc_queue, msg);
-             __sync_fetch_and_add(&sock->call_num, 1);
+            return;
         }
     }
+
+    __sync_fetch_and_sub(&sock->call_num, 1);
+    return;
 }
 
 /* any protocol stack thread receives arp packet and sync it to other threads so that it can have the arp table */
diff --git a/src/lstack/core/lstack_thread_rpc.c b/src/lstack/core/lstack_thread_rpc.c
index 4aceee6..0b2a62a 100644
--- a/src/lstack/core/lstack_thread_rpc.c
+++ b/src/lstack/core/lstack_thread_rpc.c
@@ -109,14 +109,14 @@ void poll_rpc_msg(struct protocol_stack *stack, uint32_t max_num)
         }
 
         if (!msg->recall_flag) {
-	    if (msg->sync_flag) {
+            if (msg->sync_flag) {
                 pthread_spin_unlock(&msg->lock);
             } else {
                 rpc_msg_free(msg);
             }
         } else {
-	  msg->recall_flag = 0;
-	}
+            msg->recall_flag = 0;
+        }
     }
 }
 
diff --git a/src/lstack/include/lstack_lwip.h b/src/lstack/include/lstack_lwip.h
index 4a13204..a11489c 100644
--- a/src/lstack/include/lstack_lwip.h
+++ b/src/lstack/include/lstack_lwip.h
@@ -50,8 +50,8 @@ ssize_t do_lwip_read_from_stack(int32_t fd, void *buf, size_t len, int32_t flags
 
 void do_lwip_read_recvlist(struct protocol_stack *stack, uint32_t max_num);
 void do_lwip_add_recvlist(int32_t fd);
-bool do_lwip_send(struct protocol_stack *stack, int32_t fd, struct lwip_sock *sock,
-                  size_t len, int32_t flags);
+int do_lwip_send(struct protocol_stack *stack, int32_t fd, struct lwip_sock *sock,
+                 size_t len, int32_t flags);
 
 uint32_t do_lwip_get_conntable(struct gazelle_stat_lstack_conn_info *conn, uint32_t max_num);
 uint32_t do_lwip_get_connnum(void);
-- 
2.27.0

