Fix Segfault on Electron Exit (#501)

lib/native/binding.d.ts

@@ -41,6 +41,8 @@ export function native_memory_dump(): void;
 export function error_code_to_string(error_code: number): string;
 /** @internal */
 export function error_code_to_name(error_code: number): string;
+/** @internal */
+export function disable_threadsafe_function(): number;
 
 /* IO */
 /** @internal */

lib/native/binding.js

@@ -16,7 +16,7 @@ const CRuntimeType = Object.freeze({
 });
 
 function getCRuntime() {
-    if(platform() !== "linux") {
+    if (platform() !== "linux") {
         return CRuntimeType.NON_LINUX;
     }
 
@@ -74,5 +74,15 @@ if (binding == undefined) {
     throw new Error("AWS CRT binary not present in any of the following locations:\n\t" + search_paths.join('\n\t'));
 }
 
+import crt_native from "./binding"
+/** Electron will shutdown the node process on exit, which causes the threadsafe function to segfault. To prevent
+  * the segfault we disable the threadsafe function on node process exit. */
+if (process.versions.hasOwnProperty('electron')) {
+    process.on('exit', function () {
+        crt_native.disable_threadsafe_function();
+    });
+}
+
+
 export default binding;
-export { CRuntimeType , cRuntime };
+export { CRuntimeType, cRuntime };

source/module.c

@@ -25,6 +25,7 @@
 #include <aws/common/logging.h>
 #include <aws/common/mutex.h>
 #include <aws/common/ref_count.h>
+#include <aws/common/rw_lock.h>
 #include <aws/common/system_info.h>
 
 #include <aws/event-stream/event_stream.h>
@@ -47,6 +48,15 @@ AWS_STATIC_ASSERT(NAPI_VERSION >= 4);
 
 #define AWS_DEFINE_ERROR_INFO_CRT_NODEJS(CODE, STR) AWS_DEFINE_ERROR_INFO(CODE, STR, "aws-crt-nodejs")
 
+/* TODO:
+ * Hardcoded enum value for `napi_no_external_buffers_allowed`.
+ * The enum `napi_no_external_buffers_allowed` is introduced in node14.
+ * Use it for external buffer related changes after bump to node 14 */
+#define NAPI_NO_EXTERNAL_BUFFER_ENUM_VALUE 22
+
+static bool s_tsfn_enabled = false;
+static struct aws_rw_lock s_tsfn_lock;
+
 /* clang-format off */
 static struct aws_error_info s_errors[] = {
     AWS_DEFINE_ERROR_INFO_CRT_NODEJS(
@@ -743,6 +753,28 @@ int aws_napi_get_property_array_size(
     return AWS_OP_SUCCESS;
 }
 
+void s_aws_enable_threadsafe_function(void) {
+    aws_rw_lock_wlock(&s_tsfn_lock);
+    s_tsfn_enabled = true;
+    aws_rw_lock_wunlock(&s_tsfn_lock);
+}
+
+void s_aws_disable_threadsafe_function(void) {
+    aws_rw_lock_wlock(&s_tsfn_lock);
+    s_tsfn_enabled = false;
+    aws_rw_lock_wunlock(&s_tsfn_lock);
+}
+
+napi_value aws_napi_disable_threadsafe_function(napi_env env, napi_callback_info info) {
+    (void)info;
+    if (env == NULL) {
+        aws_raise_error(AWS_CRT_NODEJS_ERROR_THREADSAFE_FUNCTION_NULL_NAPI_ENV);
+        return NULL;
+    }
+    s_aws_disable_threadsafe_function();
+    return NULL;
+}
+
 void aws_napi_throw_last_error(napi_env env) {
     const int error_code = aws_last_error();
     napi_throw_error(env, aws_error_str(error_code), aws_error_debug_str(error_code));
@@ -975,19 +1007,25 @@ napi_status aws_napi_dispatch_threadsafe_function(
     size_t argc,
     napi_value *argv) {
 
-    napi_status call_status = napi_ok;
-    if (!this_ptr) {
-        AWS_NAPI_ENSURE(env, napi_get_undefined(env, &this_ptr));
+    aws_rw_lock_rlock(&s_tsfn_lock);
+    napi_status result = napi_ok;
+    if (s_tsfn_enabled) {
+        napi_status call_status = napi_ok;
+        if (!this_ptr) {
+            AWS_NAPI_ENSURE(env, napi_get_undefined(env, &this_ptr));
+        }
+        AWS_NAPI_CALL(env, napi_call_function(env, this_ptr, function, argc, argv, NULL), {
+            call_status = status;
+            s_handle_failed_callback(env, function, status);
+        });
+        /* main thread can exit now */
+        napi_unref_threadsafe_function(env, tsfn);
+        /* Must always decrement the ref count, or the function will be pinned */
+        napi_status release_status = napi_release_threadsafe_function(tsfn, napi_tsfn_release);
+        result = (call_status != napi_ok) ? call_status : release_status;
     }
-    AWS_NAPI_CALL(env, napi_call_function(env, this_ptr, function, argc, argv, NULL), {
-        call_status = status;
-        s_handle_failed_callback(env, function, status);
-    });
-    /* main thread can exit now */
-    napi_unref_threadsafe_function(env, tsfn);
-    /* Must always decrement the ref count, or the function will be pinned */
-    napi_status release_status = napi_release_threadsafe_function(tsfn, napi_tsfn_release);
-    return (call_status != napi_ok) ? call_status : release_status;
+    aws_rw_lock_runlock(&s_tsfn_lock);
+    return result;
 }
 
 napi_status aws_napi_create_threadsafe_function(
@@ -1018,30 +1056,45 @@ napi_status aws_napi_create_threadsafe_function(
 napi_status aws_napi_release_threadsafe_function(
     napi_threadsafe_function function,
     napi_threadsafe_function_release_mode mode) {
-    if (function) {
-        return napi_release_threadsafe_function(function, mode);
+    napi_status result = napi_ok;
+    aws_rw_lock_rlock(&s_tsfn_lock);
+    if (s_tsfn_enabled && function) {
+        result = napi_release_threadsafe_function(function, mode);
     }
-    return napi_ok;
+    aws_rw_lock_runlock(&s_tsfn_lock);
+    return result;
 }
 
 napi_status aws_napi_acquire_threadsafe_function(napi_threadsafe_function function) {
-    if (function) {
-        return napi_acquire_threadsafe_function(function);
+    napi_status result = napi_ok;
+    aws_rw_lock_rlock(&s_tsfn_lock);
+    if (s_tsfn_enabled && function) {
+        result = napi_acquire_threadsafe_function(function);
     }
-    return napi_ok;
+    aws_rw_lock_runlock(&s_tsfn_lock);
+    return result;
 }
 
 napi_status aws_napi_unref_threadsafe_function(napi_env env, napi_threadsafe_function function) {
-    if (function) {
-        return napi_unref_threadsafe_function(env, function);
+    napi_status result = napi_ok;
+    aws_rw_lock_rlock(&s_tsfn_lock);
+    if (s_tsfn_enabled && function) {
+        result = napi_unref_threadsafe_function(env, function);
     }
-    return napi_ok;
+    aws_rw_lock_runlock(&s_tsfn_lock);
+    return result;
 }
 
 napi_status aws_napi_queue_threadsafe_function(napi_threadsafe_function function, void *user_data) {
-    /* increase the ref count, gets decreased when the call completes */
-    AWS_NAPI_ENSURE(NULL, napi_acquire_threadsafe_function(function));
-    return napi_call_threadsafe_function(function, user_data, napi_tsfn_nonblocking);
+    napi_status result = napi_ok;
+    aws_rw_lock_rlock(&s_tsfn_lock);
+    if (s_tsfn_enabled && function) {
+        /* increase the ref count, gets decreased when the call completes */
+        AWS_NAPI_ENSURE(NULL, napi_acquire_threadsafe_function(function));
+        result = napi_call_threadsafe_function(function, user_data, napi_tsfn_nonblocking);
+    }
+    aws_rw_lock_runlock(&s_tsfn_lock);
+    return result;
 }
 
 AWS_STATIC_STRING_FROM_LITERAL(s_mem_tracing_env_var, "AWS_CRT_MEMORY_TRACING");
@@ -1145,6 +1198,7 @@ static void s_napi_context_finalize(napi_env env, void *user_data, void *finaliz
     --s_module_initialize_count;
 
     if (s_module_initialize_count == 0) {
+
         aws_client_bootstrap_release(s_default_client_bootstrap);
         s_default_client_bootstrap = NULL;
 
@@ -1163,6 +1217,8 @@ static void s_napi_context_finalize(napi_env env, void *user_data, void *finaliz
         aws_mqtt_library_clean_up();
 
         s_uninstall_crash_handler();
+        // clean up threadsafe function lock
+        aws_rw_lock_clean_up(&s_tsfn_lock);
     }
 
     struct aws_napi_context *ctx = user_data;
@@ -1178,7 +1234,6 @@ static void s_napi_context_finalize(napi_env env, void *user_data, void *finaliz
             aws_mem_tracer_destroy(ctx_allocator);
         }
     }
-
     aws_mutex_unlock(&s_module_lock);
 }
 
@@ -1230,6 +1285,9 @@ static bool s_create_and_register_function(
     struct aws_allocator *allocator = aws_napi_get_allocator();
 
     if (s_module_initialize_count == 0) {
+        aws_rw_lock_init(&s_tsfn_lock);
+        s_aws_enable_threadsafe_function();
+
         s_install_crash_handler();
 
         aws_mqtt_library_init(allocator);
@@ -1295,6 +1353,7 @@ static bool s_create_and_register_function(
     CREATE_AND_REGISTER_FN(native_memory_dump)
     CREATE_AND_REGISTER_FN(error_code_to_string)
     CREATE_AND_REGISTER_FN(error_code_to_name)
+    CREATE_AND_REGISTER_FN(disable_threadsafe_function)
 
     /* IO */
     CREATE_AND_REGISTER_FN(io_logging_enable)

source/module.h

@@ -309,6 +309,12 @@ napi_status aws_napi_unref_threadsafe_function(napi_env env, napi_threadsafe_fun
  */
 napi_status aws_napi_queue_threadsafe_function(napi_threadsafe_function function, void *user_data);
 
+/**
+ * Disable the thread safe function operations. The function will prevent any access to threadsafe function
+ * including acquire, release, function call and so on.
+ */
+napi_value aws_napi_disable_threadsafe_function(napi_env env, napi_callback_info info);
+
 /*
  * One of these will be allocated each time the module init function is called
  * Any global state that isn't thread safe or requires clean up should be stored