fix(serverless): drop isolates in the same thread as they were created (

#483)

.changeset/fast-geese-help.md

@@ -0,0 +1,5 @@
+---
+'@lagon/serverless': patch
+---
+
+Drop and exit isolates in the same thread as they were created

crates/serverless/package.json

@@ -5,6 +5,7 @@
   "private": true,
   "scripts": {
     "build": "cargo build",
-    "lint": "cargo clippy -- -Dwarnings --no-deps"
+    "lint": "cargo clippy -- -Dwarnings --no-deps",
+    "debug": "cargo build && rust-lldb ../../target/debug/lagon-serverless"
   }
 }

crates/serverless/src/deployments/cache.rs

@@ -5,14 +5,19 @@ use std::{
     time::{Duration, Instant},
 };
 
-use log::info;
+use log::{error, info};
 use tokio::sync::RwLock;
+use tokio_util::task::LocalPoolHandle;
 
 use crate::ISOLATES;
 
 const CACHE_TASK_INTERVAL: Duration = Duration::from_secs(60);
 
-pub fn run_cache_clear_task(last_requests: Arc<RwLock<HashMap<String, Instant>>>) {
+pub fn run_cache_clear_task(
+    last_requests: Arc<RwLock<HashMap<String, Instant>>>,
+    thread_ids: Arc<RwLock<HashMap<String, usize>>>,
+    pool: LocalPoolHandle,
+) {
     let isolates_cache_seconds = Duration::from_secs(
         env::var("LAGON_ISOLATES_CACHE_SECONDS")
             .expect("LAGON_ISOLATES_CACHE_SECONDS is not set")
@@ -28,8 +33,6 @@ pub fn run_cache_clear_task(last_requests: Arc<RwLock<HashMap<String, Instant>>>
             let mut isolates_to_clear = Vec::new();
             let last_requests_reader = last_requests.read().await;
 
-            info!("Running cache clear task");
-
             for (hostname, last_request) in last_requests_reader.iter() {
                 if now.duration_since(*last_request) > isolates_cache_seconds {
                     isolates_to_clear.push(hostname.clone());
@@ -43,20 +46,37 @@ pub fn run_cache_clear_task(last_requests: Arc<RwLock<HashMap<String, Instant>>>
             // Drop the read lock because we now acquire a write lock
             drop(last_requests_reader);
 
-            let mut thread_isolates = ISOLATES.write().await;
             let mut last_requests = last_requests.write().await;
+            let thread_ids = thread_ids.read().await;
 
             for hostname in isolates_to_clear {
-                for isolates in thread_isolates.values_mut() {
+                if let Some(thread_id) = thread_ids.get(&hostname) {
                     last_requests.remove(&hostname);
 
-                    if let Some(isolate) = isolates.remove(&hostname) {
-                        let metadata = isolate.get_metadata();
+                    let thread_id = *thread_id;
+
+                    // The isolate is implicitely dropped when the block after `remove()` ends
+                    //
+                    // An isolate must be dropped (which will call `exit()` and terminate the
+                    // execution) in the same thread as it was created in
+                    match pool.spawn_pinned_by_idx(move || async move {
+                        let mut thread_isolates = ISOLATES.write().await;
+                        let thread_isolates = thread_isolates.get_mut(&thread_id).unwrap();
+
+                        if let Some(isolate) = thread_isolates.remove(&hostname) {
+                            let metadata = isolate.get_metadata();
 
-                        if let Some((deployment, ..)) = metadata.as_ref() {
-                            info!(deployment = deployment; "Clearing deployment from cache due to expiration");
+                            if let Some((deployment, ..)) = metadata.as_ref() {
+                                info!(deployment = deployment; "Clearing deployment from cache due to expiration");
+                            }
                         }
-                    }
+                    }, thread_id)
+                        .await {
+                            Ok(_) => {},
+                            Err(err) => {
+                                error!("Failed to clear deployment from cache: {}", err);
+                            }
+                        };
                 }
             }
         }

crates/serverless/src/main.rs

@@ -396,11 +396,15 @@ async fn main() -> Result<()> {
     let deployments = get_deployments(conn, bucket.clone()).await?;
     let redis = listen_pub_sub(bucket.clone(), Arc::clone(&deployments));
     let last_requests = Arc::new(RwLock::new(HashMap::new()));
-    run_cache_clear_task(Arc::clone(&last_requests));
-
     let pool = LocalPoolHandle::new(POOL_SIZE);
     let thread_ids = Arc::new(RwLock::new(HashMap::new()));
 
+    run_cache_clear_task(
+        Arc::clone(&last_requests),
+        Arc::clone(&thread_ids),
+        pool.clone(),
+    );
+
     let server = Server::bind(&addr).serve(make_service_fn(move |conn: &AddrStream| {
         let deployments = Arc::clone(&deployments);
         let pool = pool.clone();