v0.2.0 (#50)

Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#50

envd-rs/src/rpc/process_service.rs

@@ -14,14 +14,14 @@ use crate::state::AppState;
 
 pub struct ProcessServiceImpl {
     state: Arc<AppState>,
-    processes: DashMap<u32, Arc<ProcessHandle>>,
+    processes: Arc<DashMap<u32, Arc<ProcessHandle>>>,
 }
 
 impl ProcessServiceImpl {
     pub fn new(state: Arc<AppState>) -> Self {
         Self {
             state,
-            processes: DashMap::new(),
+            processes: Arc::new(DashMap::new()),
         }
     }
 
@@ -131,11 +131,21 @@ impl ProcessServiceImpl {
 
         self.processes.insert(spawned.handle.pid, Arc::clone(&spawned.handle));
 
-        let processes = self.processes.clone();
+        let processes = Arc::clone(&self.processes);
         let pid = spawned.handle.pid;
+        // Subscribe before checking cached_end so the prune cannot be lost to a
+        // race: a short-lived process can exit and broadcast its end event
+        // before this task runs. A broadcast receiver only sees messages sent
+        // after subscribe(), so a late subscribe would miss the event forever
+        // (recv() never returns Closed either — the handle keeps end_tx alive
+        // until it leaves the map, which only this task does). The waiter sets
+        // ended before sending end_tx, so cached_end() is a reliable fallback.
         let mut cleanup_end_rx = spawned.handle.subscribe_end();
+        let already_ended = spawned.handle.cached_end().is_some();
         tokio::spawn(async move {
-            let _ = cleanup_end_rx.recv().await;
+            if !already_ended {
+                let _ = cleanup_end_rx.recv().await;
+            }
             processes.remove(&pid);
         });
 
@@ -199,12 +209,28 @@ impl Process for ProcessServiceImpl {
                         match data {
                             Ok(ev) => yield Ok(make_data_start_response(ev)),
                             Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
-                            Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                            Err(tokio::sync::broadcast::error::RecvError::Closed) => {
+                                // Data channel closed: the process ended and its
+                                // handle was dropped. The end event is published
+                                // before the handle drop, so it is still buffered
+                                // — emit it rather than losing the exit code.
+                                if let Ok(end) = end_rx.try_recv() {
+                                    yield Ok(make_end_start_response(end));
+                                }
+                                break;
+                            }
                         }
                     }
                     end = end_rx.recv() => {
-                        while let Ok(ev) = data_rx.try_recv() {
-                            yield Ok(make_data_start_response(ev));
+                        // Process ended. The waiter joins the output readers
+                        // before sending this event, so every byte is already
+                        // in the data channel — drain it fully before the end.
+                        loop {
+                            match data_rx.try_recv() {
+                                Ok(ev) => yield Ok(make_data_start_response(ev)),
+                                Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
+                                Err(_) => break,
+                            }
                         }
                         if let Ok(end) = end {
                             yield Ok(make_end_start_response(end));
@@ -268,15 +294,35 @@ impl Process for ProcessServiceImpl {
                                     });
                                 }
                                 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
-                                Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                                Err(tokio::sync::broadcast::error::RecvError::Closed) => {
+                                    // Data channel closed: the process ended and
+                                    // its handle was dropped. The end event is
+                                    // published before the handle drop, so it is
+                                    // still buffered — emit it rather than losing
+                                    // the exit code.
+                                    if let Ok(end) = end_rx.try_recv() {
+                                        yield Ok(ConnectResponse {
+                                            event: buffa::MessageField::some(make_end_event(end)),
+                                            ..Default::default()
+                                        });
+                                    }
+                                    break;
+                                }
                             }
                         }
                         end = end_rx.recv() => {
-                            while let Ok(ev) = data_rx.try_recv() {
-                                yield Ok(ConnectResponse {
-                                    event: buffa::MessageField::some(make_data_event(ev)),
-                                    ..Default::default()
-                                });
+                            // Process ended. The waiter joins the output readers
+                            // before sending this event, so every byte is already
+                            // in the data channel — drain it fully before the end.
+                            loop {
+                                match data_rx.try_recv() {
+                                    Ok(ev) => yield Ok(ConnectResponse {
+                                        event: buffa::MessageField::some(make_data_event(ev)),
+                                        ..Default::default()
+                                    }),
+                                    Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
+                                    Err(_) => break,
+                                }
                             }
                             if let Ok(end) = end {
                                 yield Ok(ConnectResponse {