From e0bba18a02139a2d0deaff8811743e3f37c69480 Mon Sep 17 00:00:00 2001
From: YuQing <384681@qq.com>
Date: Wed, 21 Oct 2020 09:32:49 +0800
Subject: [PATCH] use atomic operation to avoid reentrance

---
 src/sf_define.h | 15 ++++++++-------
 src/sf_nio.c    | 15 ++++++++++++---
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/sf_define.h b/src/sf_define.h
index a3b1ce0..65ba983 100644
--- a/src/sf_define.h
+++ b/src/sf_define.h
@@ -26,13 +26,14 @@
 #define SF_DEF_MIN_BUFF_SIZE  (64 * 1024)
 #define SF_DEF_MAX_BUFF_SIZE  (64 * 1024)
 
-#define SF_NIO_STAGE_INIT        0  //set ioevent
-#define SF_NIO_STAGE_CONNECT     1  //do connect  (client only)
-#define SF_NIO_STAGE_HANDSHAKE   2  //notify the thread to handshake (client only)
-#define SF_NIO_STAGE_RECV        3  //do recv
-#define SF_NIO_STAGE_SEND        4  //do send
-#define SF_NIO_STAGE_FORWARDED   5  //deal the forwarded request
-#define SF_NIO_STAGE_CONTINUE    6  //notify the thread continue deal
+#define SF_NIO_STAGE_NONE        0
+#define SF_NIO_STAGE_INIT        1  //set ioevent
+#define SF_NIO_STAGE_CONNECT     2  //do connect  (client only)
+#define SF_NIO_STAGE_HANDSHAKE   3  //notify the thread to handshake (client only)
+#define SF_NIO_STAGE_RECV        4  //do recv
+#define SF_NIO_STAGE_SEND        5  //do send
+#define SF_NIO_STAGE_FORWARDED   6  //deal the forwarded request
+#define SF_NIO_STAGE_CONTINUE    7  //notify the thread continue deal
 #define SF_NIO_STAGE_CLOSE     127  //cleanup the task
 
 #define SF_NIO_TASK_STAGE_FETCH(task)  task->nio_stages.current
diff --git a/src/sf_nio.c b/src/sf_nio.c
index e7659e3..4c965a0 100644
--- a/src/sf_nio.c
+++ b/src/sf_nio.c
@@ -247,8 +247,10 @@ static int sf_connect_server(struct fast_task_info *task)
 static int sf_nio_deal_task(struct fast_task_info *task)
 {
     int result;
+    int stage;
 
-    switch (task->nio_stages.notify) {
+    stage = __sync_add_and_fetch(&task->nio_stages.notify, 0);
+    switch (stage) {
         case SF_NIO_STAGE_INIT:
             task->nio_stages.current = SF_NIO_STAGE_RECV;
             result = sf_nio_init(task);
@@ -283,7 +285,7 @@ static int sf_nio_deal_task(struct fast_task_info *task)
         default:
             logError("file: "__FILE__", line: %d, "
                     "client ip: %s, invalid notify stage: %d",
-                    __LINE__, task->client_ip, task->nio_stages.notify);
+                    __LINE__, task->client_ip, stage);
             result = -EINVAL;
             break;
     }
@@ -292,6 +294,8 @@ static int sf_nio_deal_task(struct fast_task_info *task)
         ioevent_add_to_deleted_list(task);
     }
 
+    __sync_bool_compare_and_swap(&task->nio_stages.notify,
+            stage, SF_NIO_STAGE_NONE);
     return result;
 }
 
@@ -301,8 +305,13 @@ int sf_nio_notify(struct fast_task_info *task, const int stage)
     int result;
     bool notify;
 
+    if (!__sync_bool_compare_and_swap(&task->nio_stages.notify,
+                SF_NIO_STAGE_NONE, stage))
+    {
+        return EAGAIN;
+    }
+
     PTHREAD_MUTEX_LOCK(&task->thread_data->waiting_queue.lock);
-    task->nio_stages.notify = stage;
     task->next = NULL;
 
     if (task->thread_data->waiting_queue.tail == NULL) {