From ef990009558e4b5b66081ad94ee09caebfa15314 Mon Sep 17 00:00:00 2001
From: LongYinan <lynweklm@gmail.com>
Date: Thu, 12 Mar 2020 22:17:39 +0800
Subject: [PATCH] fix(executor): use Arc to ensure thread safe

---
 .github/workflows/linux.yaml   |  2 +-
 .github/workflows/macos.yaml   |  2 +-
 .github/workflows/windows.yaml |  2 +-
 napi/src/executor.rs           | 50 ++++++++++++++++------------------
 test_module/package.json       |  2 +-
 test_module/tests.js           | 14 ----------
 6 files changed, 28 insertions(+), 44 deletions(-)
 delete mode 100644 test_module/tests.js

diff --git a/.github/workflows/linux.yaml b/.github/workflows/linux.yaml
index 04ba0838..516c7854 100644
--- a/.github/workflows/linux.yaml
+++ b/.github/workflows/linux.yaml
@@ -68,7 +68,7 @@ jobs:
           yarn
           cd test_module
           yarn build
-          node tests.js
+          yarn test
 
       - name: Clear the cargo caches
         run: |
diff --git a/.github/workflows/macos.yaml b/.github/workflows/macos.yaml
index 21ae38aa..c020d38c 100644
--- a/.github/workflows/macos.yaml
+++ b/.github/workflows/macos.yaml
@@ -68,7 +68,7 @@ jobs:
           yarn
           cd test_module
           yarn build
-          node tests.js
+          yarn test
 
       - name: Clear the cargo caches
         run: |
diff --git a/.github/workflows/windows.yaml b/.github/workflows/windows.yaml
index c664c133..8fa93ca4 100644
--- a/.github/workflows/windows.yaml
+++ b/.github/workflows/windows.yaml
@@ -74,7 +74,7 @@ jobs:
           yarn
           cd test_module
           yarn build
-          node tests.js
+          yarn test
         env:
           RUST_BACKTRACE: 1
 
diff --git a/napi/src/executor.rs b/napi/src/executor.rs
index f7650a5f..728cbaa1 100644
--- a/napi/src/executor.rs
+++ b/napi/src/executor.rs
@@ -4,7 +4,7 @@ use std::future::Future;
 use std::mem;
 use std::os::raw::c_void;
 use std::pin::Pin;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::task::{Context, RawWaker, RawWakerVTable, Waker};
 
 pub struct LibuvExecutor {
@@ -46,7 +46,6 @@ unsafe fn drop_uv_async(uv_async_t_ptr: *const ()) {
 struct Task<'a> {
   future: Pin<Box<dyn Future<Output = ()>>>,
   context: Context<'a>,
-  resolved: AtomicBool,
 }
 
 impl LibuvExecutor {
@@ -68,46 +67,45 @@ impl LibuvExecutor {
         &UV_ASYNC_V_TABLE,
       ));
       let context = Context::from_waker(&waker);
-      let task = Box::leak(Box::new(Task {
+      let mut task = Box::new(Task {
         future: Box::pin(future),
         context,
-        resolved: AtomicBool::new(false),
-      }));
-      sys::uv_handle_set_data(
-        uv_async_t_ref as *mut _ as *mut sys::uv_handle_t,
-        task as *mut _ as *mut c_void,
-      );
-      task.poll_future();
+      });
+      if !task.as_mut().poll_future() {
+        let arc_task = Arc::new(task);
+        sys::uv_handle_set_data(
+          uv_async_t_ref as *mut _ as *mut sys::uv_handle_t,
+          Arc::into_raw(arc_task) as *mut c_void,
+        );
+      }
     }
   }
 }
 
 impl<'a> Task<'a> {
   fn poll_future(&mut self) -> bool {
-    if self.resolved.load(Ordering::Relaxed) {
-      return true;
-    }
     match self.future.as_mut().poll(&mut self.context) {
-      Poll::Ready(_) => {
-        while !self.resolved.swap(true, Ordering::Relaxed) {}
-        true
-      }
+      Poll::Ready(_) => true,
       Poll::Pending => false,
     }
   }
 }
 
 unsafe extern "C" fn poll_future(handle: *mut sys::uv_async_t) {
-  let data_ptr = sys::uv_handle_get_data(handle as *mut sys::uv_handle_t) as *mut Task;
-  let mut task = Box::from_raw(data_ptr);
-  if !task.as_mut().poll_future() {
-    Box::leak(task);
+  let data_ptr = sys::uv_handle_get_data(handle as *mut sys::uv_handle_t) as *mut Box<Task>;
+  let mut task = Arc::from_raw(data_ptr);
+  if let Some(mut_task) = Arc::get_mut(&mut task) {
+    if mut_task.poll_future() {
+      Arc::into_raw(task);
+    } else {
+      sys::uv_close(
+        handle as *mut sys::uv_handle_t,
+        Some(drop_handle_after_close),
+      );
+    };
   } else {
-    sys::uv_close(
-      handle as *mut sys::uv_handle_t,
-      Some(drop_handle_after_close),
-    );
-  };
+    Arc::into_raw(task);
+  }
 }
 
 unsafe extern "C" fn drop_handle_after_close(handle: *mut sys::uv_handle_t) {
diff --git a/test_module/package.json b/test_module/package.json
index e18741d9..96097214 100644
--- a/test_module/package.json
+++ b/test_module/package.json
@@ -4,6 +4,6 @@
   "scripts": {
     "build": "cargo build && node ../scripts/napi.js",
     "build-release": "cargo build --release && node ../scripts/napi.js --release",
-    "test": "node ./tests.js"
+    "test": "node ./index.js"
   }
 }
diff --git a/test_module/tests.js b/test_module/tests.js
deleted file mode 100644
index 21ca1f75..00000000
--- a/test_module/tests.js
+++ /dev/null
@@ -1,14 +0,0 @@
-const { platform } = require('os')
-const { fork } = require('child_process')
-
-fork('./index.js', {
-  stdio: 'inherit',
-}).on('exit', (code) => {
-  if (code !== 0) {
-    if (code === 3221225477 && platform() === 'win32') {
-      console.error(code)
-      process.exit(0)
-    }
-    process.exit(code)
-  }
-})