refactor: ensure exit code reference is passed to all workers (#14814)

4 files changed, 30 insertions, 29 deletions

diff --git a/runtime/ops/os.rs b/runtime/ops/os.rs
index 37da410ca..f138bae31 100644
--- a/runtime/ops/os.rs
+++ b/runtime/ops/os.rs
@@ -2,6 +2,7 @@
 
 use super::utils::into_string;
 use crate::permissions::Permissions;
+use crate::worker::ExitCode;
 use deno_core::error::{type_error, AnyError};
 use deno_core::op;
 use deno_core::url::Url;
@@ -10,11 +11,8 @@ use deno_core::OpState;
 use serde::Serialize;
 use std::collections::HashMap;
 use std::env;
-use std::sync::atomic::AtomicI32;
-use std::sync::atomic::Ordering::Relaxed;
-use std::sync::Arc;
 
-pub fn init(maybe_exit_code: Option<Arc<AtomicI32>>) -> Extension {
+pub fn init(exit_code: ExitCode) -> Extension {
   Extension::builder()
     .ops(vec![
       op_env::decl(),
@@ -33,8 +31,7 @@ pub fn init(maybe_exit_code: Option<Arc<AtomicI32>>) -> Extension {
       op_system_memory_info::decl(),
     ])
     .state(move |state| {
-      let exit_code = maybe_exit_code.clone().unwrap_or_default();
-      state.put::<Arc<AtomicI32>>(exit_code);
+      state.put::<ExitCode>(exit_code.clone());
       Ok(())
     })
     .build()
@@ -105,12 +102,12 @@ fn op_delete_env(state: &mut OpState, key: String) -> Result<(), AnyError> {
 
 #[op]
 fn op_set_exit_code(state: &mut OpState, code: i32) {
-  state.borrow_mut::<Arc<AtomicI32>>().store(code, Relaxed);
+  state.borrow_mut::<ExitCode>().set(code);
 }
 
 #[op]
 fn op_exit(state: &mut OpState) {
-  let code = state.borrow::<Arc<AtomicI32>>().load(Relaxed);
+  let code = state.borrow::<ExitCode>().get();
   std::process::exit(code)
 }
 
diff --git a/runtime/ops/worker_host.rs b/runtime/ops/worker_host.rs
index d869e5871..9860f7cb5 100644
--- a/runtime/ops/worker_host.rs
+++ b/runtime/ops/worker_host.rs
@@ -11,6 +11,7 @@ use crate::web_worker::WebWorkerHandle;
 use crate::web_worker::WebWorkerType;
 use crate::web_worker::WorkerControlEvent;
 use crate::web_worker::WorkerId;
+use crate::worker::ExitCode;
 use crate::worker::FormatJsErrorFn;
 use deno_core::error::AnyError;
 use deno_core::futures::future::LocalFutureObj;
@@ -27,7 +28,6 @@ use log::debug;
 use std::cell::RefCell;
 use std::collections::HashMap;
 use std::rc::Rc;
-use std::sync::atomic::AtomicI32;
 use std::sync::Arc;
 
 pub struct CreateWebWorkerArgs {
@@ -37,7 +37,7 @@ pub struct CreateWebWorkerArgs {
   pub permissions: Permissions,
   pub main_module: ModuleSpecifier,
   pub worker_type: WebWorkerType,
-  pub maybe_exit_code: Option<Arc<AtomicI32>>,
+  pub exit_code: ExitCode,
 }
 
 pub type CreateWebWorkerCb = dyn Fn(CreateWebWorkerArgs) -> (WebWorker, SendableWebWorkerHandle)
@@ -171,11 +171,7 @@ fn op_create_worker(
     parent_permissions.clone()
   };
   let parent_permissions = parent_permissions.clone();
-  // `try_borrow` here, because worker might have been started without
-  // access to `Deno` namespace.
-  // TODO(bartlomieju): can a situation happen when parent doesn't
-  // have access to `exit_code` but the child does?
-  let maybe_exit_code = state.try_borrow::<Arc<AtomicI32>>().cloned();
+  let exit_code = state.borrow::<ExitCode>().clone();
   let worker_id = state.take::<WorkerId>();
   let create_web_worker_cb = state.take::<CreateWebWorkerCbHolder>();
   state.put::<CreateWebWorkerCbHolder>(create_web_worker_cb.clone());
@@ -211,7 +207,7 @@ fn op_create_worker(
         permissions: worker_permissions,
         main_module: module_specifier.clone(),
         worker_type,
-        maybe_exit_code,
+        exit_code,
       });
 
     // Send thread safe handle from newly created worker to host thread
diff --git a/runtime/web_worker.rs b/runtime/web_worker.rs
index a2f2ffc5d..eb0aaf916 100644
--- a/runtime/web_worker.rs
+++ b/runtime/web_worker.rs
@@ -6,6 +6,7 @@ use crate::ops;
 use crate::ops::io::Stdio;
 use crate::permissions::Permissions;
 use crate::tokio_util::run_basic;
+use crate::worker::ExitCode;
 use crate::worker::FormatJsErrorFn;
 use crate::BootstrapOptions;
 use deno_broadcast_channel::InMemoryBroadcastChannel;
@@ -40,7 +41,6 @@ use std::cell::RefCell;
 use std::fmt;
 use std::rc::Rc;
 use std::sync::atomic::AtomicBool;
-use std::sync::atomic::AtomicI32;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::task::Context;
@@ -335,7 +335,7 @@ pub struct WebWorkerOptions {
   pub broadcast_channel: InMemoryBroadcastChannel,
   pub shared_array_buffer_store: Option<SharedArrayBufferStore>,
   pub compiled_wasm_module_store: Option<CompiledWasmModuleStore>,
-  pub maybe_exit_code: Option<Arc<AtomicI32>>,
+  pub exit_code: ExitCode,
   pub stdio: Stdio,
 }
 
@@ -421,11 +421,7 @@ impl WebWorker {
         unstable,
         options.unsafely_ignore_certificate_errors.clone(),
       ),
-      ops::os::init(Some(
-        options
-          .maybe_exit_code
-          .expect("Worker has access to OS ops but exit code was not passed."),
-      )),
+      ops::os::init(options.exit_code),
       ops::permissions::init(),
       ops::process::init(),
       ops::spawn::init(),
diff --git a/runtime/worker.rs b/runtime/worker.rs
index 4c38d232f..94c05cd82 100644
--- a/runtime/worker.rs
+++ b/runtime/worker.rs
@@ -35,6 +35,18 @@ use std::task::Poll;
 
 pub type FormatJsErrorFn = dyn Fn(&JsError) -> String + Sync + Send;
 
+#[derive(Clone)]
+pub struct ExitCode(Arc<AtomicI32>);
+
+impl ExitCode {
+  pub fn get(&self) -> i32 {
+    self.0.load(Relaxed)
+  }
+
+  pub fn set(&mut self, code: i32) {
+    self.0.store(code, Relaxed);
+  }
+}
 /// This worker is created and used by almost all
 /// subcommands in Deno executable.
 ///
@@ -45,6 +57,7 @@ pub type FormatJsErrorFn = dyn Fn(&JsError) -> String + Sync + Send;
 pub struct MainWorker {
   pub js_runtime: JsRuntime,
   should_break_on_first_statement: bool,
+  exit_code: ExitCode,
 }
 
 pub struct WorkerOptions {
@@ -98,6 +111,7 @@ impl MainWorker {
         Ok(())
       })
       .build();
+    let exit_code = ExitCode(Arc::new(AtomicI32::new(0)));
 
     // Internal modules
     let mut extensions: Vec<Extension> = vec![
@@ -147,7 +161,7 @@ impl MainWorker {
         unstable,
         options.unsafely_ignore_certificate_errors.clone(),
       ),
-      ops::os::init(None),
+      ops::os::init(exit_code.clone()),
       ops::permissions::init(),
       ops::process::init(),
       ops::signal::init(),
@@ -181,6 +195,7 @@ impl MainWorker {
     Self {
       js_runtime,
       should_break_on_first_statement: options.should_break_on_first_statement,
+      exit_code,
     }
   }
 
@@ -314,11 +329,8 @@ impl MainWorker {
 
   /// Return exit code set by the executed code (either in main worker
   /// or one of child web workers).
-  pub fn get_exit_code(&mut self) -> i32 {
-    let op_state_rc = self.js_runtime.op_state();
-    let op_state = op_state_rc.borrow();
-    let exit_code = op_state.borrow::<Arc<AtomicI32>>().load(Relaxed);
-    exit_code
+  pub fn get_exit_code(&self) -> i32 {
+    self.exit_code.get()
   }
 
   /// Dispatches "load" event to the JavaScript runtime.