From d693a028e89c97a73e0d39c1fc829875453a05d7 Mon Sep 17 00:00:00 2001
From: maosiping <maosiping@huawei.com>
Date: Fri, 3 Dec 2021 10:52:47 +0800
Subject: [PATCH] =?UTF-8?q?=E8=93=9D=E7=89=99=E7=BD=91=E7=BB=9C=E4=BB=A3?=
 =?UTF-8?q?=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: maosiping <maosiping@huawei.com>
---
 .../distributed_net/distributed_net_core.c    | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/core/distributed_net/distributed_net_core.c b/src/core/distributed_net/distributed_net_core.c
index eb98e1c..21bce08 100644
--- a/src/core/distributed_net/distributed_net_core.c
+++ b/src/core/distributed_net/distributed_net_core.c
@@ -36,6 +36,29 @@
 #include "lwip/distributed_net/udp_transmit.h"
 #include "lwip/priv/sockets_priv.h"
 
+#define DISTRIBUTED_NET_SET_NONBLOCK_FLAG(flags, sock) \
+  do { \
+    if ((flags) & (O_NONBLOCK)) { \
+      make_nonblock((sock), 1); \
+    } \
+  } while (0)
+
+#define DISTRIBUTED_NET_RESET_NONBLOCK_FLAG(flags, sock) \
+  do { \
+    if ((flags) & (O_NONBLOCK)) { \
+      make_nonblock((sock), 0); \
+    } \
+  } while (0)
+
+static int make_nonblock(int sock, int nonblock)
+{
+  int flags = lwip_fcntl(sock, F_GETFL, 0);
+  if (nonblock) {
+    return lwip_fcntl(sock, F_SETFL, flags | O_NONBLOCK);
+  }
+  return lwip_fcntl(sock, F_SETFL, flags & (~O_NONBLOCK));
+}
+
 int distributed_net_connect(int sock, const struct sockaddr *addr, socklen_t addr_len)
 {
   CHECK_PARA(SOCKET_TO_INDEX(sock) >= 0 && SOCKET_TO_INDEX(sock) < NUM_SOCKETS, EBADF);
@@ -51,7 +74,10 @@ int distributed_net_connect(int sock, const struct sockaddr *addr, socklen_t add
 
   (void)memset_s(&addr_in, sizeof(addr_in), 0, sizeof(addr_in));
   INIT_SOCK_ADDR(&addr_in, LOCAL_SERVER_IP, get_local_tcp_server_port());
+  int flags = lwip_fcntl(sock, F_GETFL, 0);
+  DISTRIBUTED_NET_RESET_NONBLOCK_FLAG(flags, sock);
   if (lwip_connect_internal(sock, (struct sockaddr *)&addr_in, sizeof(addr_in)) < 0) {
+    DISTRIBUTED_NET_SET_NONBLOCK_FLAG(flags, sock);
     return -1;
   }
   set_distributed_net_socket(sock);
@@ -66,8 +92,10 @@ int distributed_net_connect(int sock, const struct sockaddr *addr, socklen_t add
 
   if (lwip_send(sock, &data, sizeof(data), 0) < 0) {
     reset_distributed_net_socket(sock);
+    DISTRIBUTED_NET_SET_NONBLOCK_FLAG(flags, sock);
     return -1;
   }
+  DISTRIBUTED_NET_SET_NONBLOCK_FLAG(flags, sock);
   return 0;
 }
 
-- 
GitLab