diff --git a/machine_interface/src/function_driver/thread_utils.rs b/machine_interface/src/function_driver/thread_utils.rs index f9174601..b5d49242 100644 --- a/machine_interface/src/function_driver/thread_utils.rs +++ b/machine_interface/src/function_driver/thread_utils.rs @@ -147,6 +147,11 @@ fn run_thread(core_id: u8, mut queue: impl EngineWorkQueue) { 'engine: loop { // TODO catch unwind so we can always return an error or shut down gracefully let (args, debt) = waker::manual_pull(&mut queue); + + if !debt.is_alive() { + continue 'engine; + } + match args { WorkToDo::FunctionArguments { function_id: _, @@ -230,6 +235,10 @@ fn run_thread(core_id: u8, mut queue: impl EngineWorkQueue) { } } + if !debt.is_alive() { + continue 'engine; + } + recorder.record(RecordPoint::EngineStart); let result = engine_state.run( diff --git a/machine_interface/src/promise.rs b/machine_interface/src/promise.rs index dee0264c..25e45677 100644 --- a/machine_interface/src/promise.rs +++ b/machine_interface/src/promise.rs @@ -9,7 +9,7 @@ use core::{ task::{Poll, Waker}, }; use dandelion_commons::{err_dandelion, DandelionError, DandelionResult, PromiseError}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; static WAKER_INDEX: u8 = 0b0000_0001; static DEBT_ALIVE: u8 = 0b0001_0000; @@ -17,11 +17,13 @@ static PROMISE_ALIVE: u8 = 0b0010_0000; static CONTENT_SET: u8 = 0b0100_0000; static ALIVE: u8 = DEBT_ALIVE | PROMISE_ALIVE; +type AbortHandle = Box; + // TODO replace 2 wakers with a futures::task::AtomicWaker struct PromiseData { /// Abort handle, only to be called once, as long as this value ///non null that means the function has not been aborted or terminated on it's own - abort_handle: AtomicPtr ()>, + abort_handle: Mutex>, /// Points to raw box of the results the engine has put in there results: Cell>, /// TODO replace Option with only waker when Waker::noop stabilizes, @@ -118,7 +120,7 @@ impl PromiseBuffer { let data_ptr = self.internal.get_promise_data()?; let data = unsafe { &mut (&mut *data_ptr).data }; let default = ManuallyDrop::new(PromiseData { - abort_handle: AtomicPtr::new(ptr::null_mut()), + abort_handle: Mutex::new(None), results: Cell::new(err_dandelion!(DandelionError::PromiseError( PromiseError::Default ))), @@ -151,9 +153,9 @@ impl Promise { } fn abort_internal(&mut self) { let data = unsafe { &(&*self.data).data }; - let abort_handle = data.abort_handle.swap(ptr::null_mut(), Ordering::SeqCst); - if !abort_handle.is_null() { - unsafe { (*abort_handle)() } + let abort_handle = data.abort_handle.lock().unwrap().take(); + if let Some(abort_handle) = abort_handle { + abort_handle() } } } @@ -217,8 +219,8 @@ impl Debt { pub fn fulfill(self, results: DandelionResult) { let data = unsafe { &(&*self.data).data }; - // make sure we are not aborted by this promise anymore - data.abort_handle.store(ptr::null_mut(), Ordering::SeqCst); + // prevent a later abort from invoking the callback + data.abort_handle.lock().unwrap().take(); // write a result data.results.set(results); let flags = data.flags.fetch_or(CONTENT_SET, Ordering::SeqCst); @@ -227,18 +229,20 @@ impl Debt { waker.wake(); } } - pub fn install_abort_handle(&self, handle: fn()) { + pub fn install_abort_handle(&self, handle: F) + where + F: FnOnce() + Send + 'static, + { let data = unsafe { &(&*self.data).data }; - data.abort_handle - .store(handle as *mut fn(), Ordering::SeqCst); + *data.abort_handle.lock().unwrap() = Some(Box::new(handle)); } } impl Drop for Debt { fn drop(&mut self) { let data = unsafe { &(&*self.data).data }; - // make sure we can't get aborted by this handle anymore - data.abort_handle.store(ptr::null_mut(), Ordering::SeqCst); + // prevent a later abort from invoking the callback + data.abort_handle.lock().unwrap().take(); // if promise is still alive, there is still a promise waiting for a result let flags = data.flags.load(Ordering::SeqCst); if flags & PROMISE_ALIVE == 1 { diff --git a/multinode/proto/multinode.proto b/multinode/proto/multinode.proto index d74affd7..e05e4a9b 100644 --- a/multinode/proto/multinode.proto +++ b/multinode/proto/multinode.proto @@ -126,6 +126,7 @@ message QueueMessage { RepeatedInvocations invocations = 2; // Message containing the work for one function, can be refused Invocation try_offload = 3; - // TODO add message to cancel invocation + // Message indicating remote invocation should not be scheduled + uint32 cancel_invocation = 4; } } diff --git a/multinode/src/client.rs b/multinode/src/client.rs index ce832115..d4567759 100644 --- a/multinode/src/client.rs +++ b/multinode/src/client.rs @@ -29,7 +29,7 @@ use machine_interface::{ }; use prost::bytes::{Bytes, BytesMut}; use std::{ - collections::{BTreeMap, BinaryHeap}, + collections::{BTreeMap, BTreeSet, BinaryHeap}, sync::Arc, time::{Duration, Instant, SystemTime}, }; @@ -154,6 +154,7 @@ enum QueueOption { Message(remote_message::RemoteMessage, Option), WorkAvailable, TryOffload(WorkToDo, machine_interface::promise::Debt, usize), + CancelRemote(u32), /// The connection to the remote node was lost, so the server logic should tear down. Disconnected, } @@ -244,6 +245,7 @@ async fn remote_queue_server_try_offload( /// The protocol logic handling for the remote queue server async fn remote_queue_server_logic( mut message_receiver: mpsc::Receiver, + local_sender: mpsc::Sender, message_sender: mpsc::Sender, queue: WorkQueue, export_registry: ExportRegistry, @@ -254,6 +256,7 @@ async fn remote_queue_server_logic( let mut waiting_for_work = false; let mut debt_map = BTreeMap::new(); + let mut cancelled_debt_ids = BTreeSet::new(); let mut free_debt_ids = BinaryHeap::new(); let mut max_debt_id = 0; @@ -287,8 +290,11 @@ async fn remote_queue_server_logic( // TODO: consider limiting the work we give to a node based on the know max capacity, // to limit potential stragglers if we know the node asked for more than it can handle (possibly because of race conditions) // do not give even more. - invocations.extend(work_found.into_iter().map(|(work, debt, composition_id)| - { + invocations.extend(work_found.into_iter().filter_map( + |(work, debt, composition_id)| { + if !debt.is_alive() { + return None; + } // there is some work so send it out // find the local function id to use let promise_id = if let Some(free_id) = free_debt_ids.pop() { @@ -333,6 +339,18 @@ async fn remote_queue_server_logic( ); let caching = caching; let function_id = function_id.to_string(); + let cancel_sender = local_sender.clone(); + debt.install_abort_handle(move || { + if cancel_sender + .try_send(QueueOption::CancelRemote(promise_id)) + .is_err() + { + warn!( + "Failed to enqueue cancellation for remote invocation {}", + promise_id + ); + } + }); debt_map.insert( promise_id, ( @@ -344,13 +362,14 @@ async fn remote_queue_server_logic( work, ), ); - Invocation { + Some(Invocation { metadata_sets, function_id, invocation_id: promise_id, caching, - } - })); + }) + }, + )); } if invocations.is_empty() { waiting_for_work = true; @@ -386,16 +405,31 @@ async fn remote_queue_server_logic( response, } = response; // TODO: handle failure - let ( + let Some(( composition_id, debt, mut recorder, start_epoch, remote_data_references, work, - ) = debt_map - .remove(&invocation_id) - .expect("Should always get back function response for a present debt"); + )) = debt_map.remove(&invocation_id) + else { + if cancelled_debt_ids.remove(&invocation_id) { + // The acknowledgement guarantees that the remote will not send a + // delayed response that could be mistaken for a later invocation. + free_debt_ids.push(invocation_id); + trace!( + "Received cancellation acknowledgement for remote invocation {}", + invocation_id + ); + } else { + warn!( + "Received response for unknown remote invocation {}", + invocation_id + ); + } + continue; + }; free_debt_ids.push(invocation_id); drop(remote_data_references); // remote did not do work, was a try offload request, reenqueu the work @@ -454,6 +488,9 @@ async fn remote_queue_server_logic( } } QueueOption::TryOffload(work, debt, composition_id) => { + if !debt.is_alive() { + continue; + } // if this node already sent enough work for the remote to be at capacity don't send more if debt_map.len() >= remote_num_cores as usize { queue.reenqueue(work, debt, composition_id); @@ -498,6 +535,18 @@ async fn remote_queue_server_logic( }); let caching = caching; let function_id = function_id.to_string(); + let cancel_sender = local_sender.clone(); + debt.install_abort_handle(move || { + if cancel_sender + .try_send(QueueOption::CancelRemote(promise_id)) + .is_err() + { + warn!( + "Failed to enqueue cancellation for remote invocation {}", + promise_id + ); + } + }); debt_map.insert( promise_id, ( @@ -533,6 +582,27 @@ async fn remote_queue_server_logic( } } } + QueueOption::CancelRemote(invocation_id) => { + if let Some(( + _composition_id, + _debt, + _recorder, + _start_epoch, + remote_data_references, + _work, + )) = debt_map.remove(&invocation_id) + { + drop(remote_data_references); + cancelled_debt_ids.insert(invocation_id); + if message_sender + .send(queue_message::QueueMessage::CancelInvocation(invocation_id)) + .await + .is_err() + { + break; + } + } + } QueueOption::Disconnected => { // The remote node disconnected, stop the loop and run the cleanup below. break; @@ -553,9 +623,10 @@ async fn remote_queue_server_logic( while let Some(message) = message_receiver.recv().await { match message { // Ignore messages that have no effect on the clean up - QueueOption::WorkAvailable | QueueOption::Disconnected | QueueOption::Message(_, _) => { - () - } + QueueOption::WorkAvailable + | QueueOption::Disconnected + | QueueOption::Message(_, _) + | QueueOption::CancelRemote(_) => (), // Reenqueue work that was tried to offload QueueOption::TryOffload(work, debt, composition_id) => { queue.reenqueue(work, debt, composition_id); @@ -642,12 +713,13 @@ pub async fn remote_queue_server( queue.add_remote_channel(node_id, offload_sender); spawn(remote_queue_server_try_offload( offload_receiver, - queue_option_sender, + queue_option_sender.clone(), queue.clone(), )); remote_queue_server_logic( queue_option_receiver, + queue_option_sender.clone(), queue_message_sender, queue, export_registry, @@ -669,7 +741,7 @@ pub enum PollingOption { IdleChanged, LocalCoreCountChanged(usize), // Results(RemoteMessage, Option<(Vec>, u64)>), - Results(remote_message::RemoteMessage), + Results(u32, remote_message::RemoteMessage), /// The connection to the remote node was lost, so the client logic should tear down. Disconnected, } @@ -789,6 +861,7 @@ async fn dispatcher_call( // fine, the master will reenqueue the work after detecting the disconnect. let _ = sender .send(PollingOption::Results( + invocation_id, remote_message::RemoteMessage::Response(Response { invocation_id, response: Some(response_message), @@ -819,7 +892,8 @@ async fn remote_queue_client_logic( // This makes sure we don't overfetch from this node. // TODO: think about the general issue of state synchronization between the dispatcher and multinode client adding things, // and the engines and server taking things. - let mut work_from_remote = 0; + let mut work_from_remote: usize = 0; + let mut cancelled_remote_invocations = BTreeSet::new(); while let Some(current_future) = receiver.recv().await { match current_future { @@ -839,12 +913,17 @@ async fn remote_queue_client_logic( debug_assert!(data_option.is_none()); remote_had_work = true; } + queue_message::QueueMessage::CancelInvocation(invocation_id) => { + trace!( + "Queue Client received cancellation for invocation {}", + invocation_id + ); + cancelled_remote_invocations.insert(invocation_id); + } // TODO for try offload decide when to refuse work queue_message::QueueMessage::TryOffload(invocation) => { trace!("Queue Client recieved try offload"); - work_from_remote += 1; - // mark remote as having work, so we ask for more as idle cores change let start_instance = Instant::now(); let start_time = @@ -856,6 +935,21 @@ async fn remote_queue_client_logic( metadata_sets, caching, } = invocation; + if cancelled_remote_invocations.remove(&invocation_id) { + trace!("Skipping canceled try-offload invocation {}", invocation_id); + if message_sender + .send(remote_message::RemoteMessage::Response(Response { + invocation_id, + response: None, + })) + .await + .is_err() + { + break; + } + continue; + } + work_from_remote += 1; let function_arc = Arc::new(function_id); let recorder = Recorder::new(function_arc.clone(), start_instance); let inputs = @@ -880,7 +974,6 @@ async fn remote_queue_client_logic( invocation_request_in_flight = false; // mark remote as having work, so we ask for more as idle cores change - work_from_remote += invocations.invocations.len(); let start_instance = Instant::now(); let start_time = std::time::SystemTime::elapsed(&std::time::SystemTime::UNIX_EPOCH) @@ -892,6 +985,21 @@ async fn remote_queue_client_logic( metadata_sets, caching, } = invocation; + if cancelled_remote_invocations.remove(&invocation_id) { + trace!("Skipping canceled invocation {}", invocation_id); + if message_sender + .send(remote_message::RemoteMessage::Response(Response { + invocation_id, + response: None, + })) + .await + .is_err() + { + break; + } + continue; + } + work_from_remote += 1; let function_arc = Arc::new(function_id); let recorder = Recorder::new(function_arc.clone(), start_instance); let inputs = proto_data_sets_to_composition_sets( @@ -922,13 +1030,31 @@ async fn remote_queue_client_logic( // The remote node disconnected, stop the loop so the caller can reconnect. break; } - PollingOption::Results(results) => { + PollingOption::Results(invocation_id, results) => { work_from_remote -= 1; - trace!("Queue Client sending out result, {} outstanding", { - work_from_remote - }); - if message_sender.send(results).await.is_err() { - break; + if cancelled_remote_invocations.remove(&invocation_id) { + trace!( + "Suppressing result for canceled remote invocation {}", + invocation_id + ); + if message_sender + .send(remote_message::RemoteMessage::Response(Response { + invocation_id, + response: None, + })) + .await + .is_err() + { + break; + } + } else { + trace!( + "Queue Client sending out result, {} outstanding", + work_from_remote + ); + if message_sender.send(results).await.is_err() { + break; + } } } // getting a notification so should poll the queue diff --git a/multinode/src/client/test.rs b/multinode/src/client/test.rs index c0dbc6e5..6da126fe 100644 --- a/multinode/src/client/test.rs +++ b/multinode/src/client/test.rs @@ -82,6 +82,7 @@ fn test_remote_queue_server() { let mut context = Context::from_waker(Waker::noop()); let mut server_future = Box::pin(remote_queue_server_logic( queue_option_receiver, + queue_option_sender.clone(), queue_message_sender, work_queue.clone(), ExportRegistry::new(1), @@ -470,6 +471,7 @@ fn test_remote_queue_client() { // send back a result, mark the queue as changed and poll to get it processed poll_option_sender .try_send(crate::client::PollingOption::Results( + invocation_id_3, remote_message::RemoteMessage::Response(Response { invocation_id: invocation_id_3, response: Some(response::Response::ErrorMsg(