diff --git a/dispatcher/src/queue.rs b/dispatcher/src/queue.rs index 5d78377f..fb1de73b 100644 --- a/dispatcher/src/queue.rs +++ b/dispatcher/src/queue.rs @@ -500,8 +500,13 @@ impl WorkQueue { self.remote_nodes.lock().unwrap().insert(node_id, channel); } + /// Removes the offload channel for a remote node, e.g. when it has disconnected. + pub fn remove_remote_channel(&self, node_id: u64) { + self.remote_nodes.lock().unwrap().remove(&node_id); + } + /// Put work back into queue after trying to offload without success. - pub async fn reenqueue(&self, work: WorkToDo, debt: Debt) { + pub fn reenqueue(&self, work: WorkToDo, debt: Debt) { self.push(work, debt, false); } diff --git a/dispatcher/src/queue/policy/data_locality.rs b/dispatcher/src/queue/policy/data_locality.rs index 03c080fe..5df61bc2 100644 --- a/dispatcher/src/queue/policy/data_locality.rs +++ b/dispatcher/src/queue/policy/data_locality.rs @@ -54,8 +54,23 @@ pub fn prepare_io_element( if max_size * 2 > total_input_size { let maybe_sender = remote_nodes.lock().unwrap().get(node_id).cloned(); if let Some(node_sender) = maybe_sender { - node_sender.send((work, debt)).unwrap(); - return None; + // If the remote node has disconnected its receiver is gone, in which case + // we recover the work and fall back to executing it locally. + match node_sender.send((work, debt)) { + Ok(()) => return None, + Err(mpsc::error::SendError((work, debt))) => { + // since the sender does not work, remove it. + remote_nodes.lock().unwrap().remove(&node_id); + return Some(( + work, + debt, + IOElementData { + remote_data, + total_input_size, + }, + )); + } + } } } } diff --git a/multinode/src/client.rs b/multinode/src/client.rs index 483783e2..96b1ce79 100644 --- a/multinode/src/client.rs +++ b/multinode/src/client.rs @@ -20,7 +20,7 @@ use dispatcher::{ dispatcher::Dispatcher, queue::{get_engine_flag, WorkQueue}, }; -use log::{error, trace, warn}; +use log::{error, info, trace, warn}; use machine_interface::{ composition::{CompositionSet, RemoteData}, function_driver::{WorkDone, WorkToDo}, @@ -40,7 +40,10 @@ use tokio::{ TcpStream, }, spawn, - sync::{mpsc, watch, Notify}, + sync::{ + mpsc::{self, error::SendError}, + watch, Notify, + }, }; #[cfg(test)] @@ -51,14 +54,15 @@ const _: () = assert!(size_of::() == size_of::()); // TODO ADDITIONAL_DATA_BUFFER and data_buffer are currently used only to carry IoData // We should consider removing this when recursive resolution of IoData is implemented, // as then all sets will be exchanged via the remote data server. -// TODO handle connection failure /// To send a message between nodes, always first send the length of the message, /// then the message, so the other side knows when one message ends. +/// Returns an error if the underlying connection failed, so the caller can tear down +/// the connection instead of panicking. async fn send_message( metadata_buffer: &Bytes, mut sender: impl AsyncWriteExt + Unpin, data_buffer: Option<(Vec>, u64)>, -) { +) -> std::io::Result<()> { let metadata_size: u32 = metadata_buffer.len().try_into().unwrap(); let flags = match data_buffer { Some((_, total_size)) => { @@ -70,8 +74,8 @@ async fn send_message( let packed_metadata = pack_metadata_size_and_flags(metadata_size, flags); - sender.write_u64(packed_metadata).await.unwrap(); - sender.write_all(metadata_buffer).await.unwrap(); + sender.write_u64(packed_metadata).await?; + sender.write_all(metadata_buffer).await?; // Code for sending data along with the request if needed // Keeping for later when we want to send small items along with requests / responses. @@ -108,21 +112,28 @@ async fn send_message( // } // } // } - sender.flush().await.unwrap(); + sender.flush().await } -// TODO handle connection failure -// For small messages we are expecting repeteatly, could have spezial read function with permanent preallocated buffers -// Issue: serialization does not give constant sizes, so would need to find an upper bound first -async fn receive_message(mut receiver: impl AsyncReadExt + Unpin) -> (Bytes, Option) { - let packed_metadata = receiver.read_u64().await.unwrap(); +/// For small messages we are expecting repeteatly, could have spezial read function with permanent preallocated buffers +/// Issue: serialization does not give constant sizes, so would need to find an upper bound first +/// Returns an error if the underlying connection failed (or was closed mid-message), +/// so the caller can tear down the connection instead of panicking. +async fn receive_message( + mut receiver: impl AsyncReadExt + Unpin, +) -> std::io::Result<(Bytes, Option)> { + let packed_metadata = receiver.read_u64().await?; let (metadata_size, _) = unpack_metadata_size_and_flags(packed_metadata); trace!("strart receiving: {}", metadata_size); // new buffer with size of message let mut metadata_buffer = BytesMut::with_capacity(metadata_size as usize); while metadata_buffer.len() < metadata_size as usize { - receiver.read_buf(&mut metadata_buffer).await.unwrap(); + // A read of 0 bytes means the peer closed the connection before sending the + // full message, treat it as an unexpected end of file. + if receiver.read_buf(&mut metadata_buffer).await? == 0 { + return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof)); + } } trace!("finish receiving"); @@ -138,19 +149,24 @@ async fn receive_message(mut receiver: impl AsyncReadExt + Unpin) -> (Bytes, Opt // } else { // return (metadata_buffer.freeze(), None); // } - (metadata_buffer.freeze(), None) + Ok((metadata_buffer.freeze(), None)) } enum QueueOption { Message(remote_message::RemoteMessage, Option), WorkAvailable, TryOffload(WorkToDo, machine_interface::promise::Debt), + /// The connection to the remote node was lost, so the server logic should tear down. + Disconnected, } async fn remote_queue_sever_notification(receiver: Arc, sender: mpsc::Sender) { loop { receiver.notified().await; - sender.send(QueueOption::WorkAvailable).await.unwrap(); + if sender.send(QueueOption::WorkAvailable).await.is_err() { + // logic loop has shut down, nothing left to notify + break; + } } } @@ -160,20 +176,30 @@ async fn remote_queue_server_receiver( sender: mpsc::Sender, ) { loop { - let (message_buffer, _) = receive_message(&mut socket).await; + let message_buffer = match receive_message(&mut socket).await { + Ok((message_buffer, _)) => message_buffer, + Err(_) => { + // The connection was lost, inform the logic loop so it can clean up. + let _ = sender.send(QueueOption::Disconnected).await; + break; + } + }; let message = deserialize_remote_message(message_buffer) .unwrap() .remote_message .unwrap(); - sender + if sender .send(QueueOption::Message(message, None)) .await - .unwrap(); + .is_err() + { + break; + } } } -/// The sender docket handling for the remote queue server -/// check if we can unite this and the reciever with the other one, by using traits +/// The sender socket handling for the remote queue server. +/// TODO: check if we can unite this and the receiver with the other client, by using traits. async fn remote_queue_server_sender( mut socket: OwnedWriteHalf, mut receiver: mpsc::Receiver, @@ -182,7 +208,13 @@ async fn remote_queue_server_sender( let message_buffer = serialize_queue_message(QueueMessage { queue_message: Some(queue_message), }); - send_message(&message_buffer, &mut socket, None).await; + if send_message(&message_buffer, &mut socket, None) + .await + .is_err() + { + // connection lost, the receiver side will trigger the teardown + break; + } } } @@ -190,12 +222,21 @@ async fn remote_queue_server_sender( async fn remote_queue_server_try_offload( mut queue_receiver: mpsc::UnboundedReceiver<(WorkToDo, Debt)>, sender: mpsc::Sender, + queue: WorkQueue, ) { while let Some((work, debt)) = queue_receiver.recv().await { - sender - .send(QueueOption::TryOffload(work, debt)) - .await - .unwrap(); + if let Err(send_err) = sender.send(QueueOption::TryOffload(work, debt)).await { + if let SendError(QueueOption::TryOffload(w, d)) = send_err { + queue.reenqueue(w, d); + } + break; + } + } + // can't send anymore so make sure the channel does not have things added to it. + queue_receiver.close(); + // drain the channel, reqenque all the work + while let Some((work, debt)) = queue_receiver.recv().await { + queue.reenqueue(work, debt); } } @@ -314,18 +355,24 @@ async fn remote_queue_server_logic( waiting_for_work = true; trace!("No work available"); // there is no work, so send message accordingly - message_sender + if message_sender .send(queue_message::QueueMessage::NoWork(true)) .await - .unwrap(); + .is_err() + { + break; + } } else { invocations_running += invocations.len(); - message_sender + if message_sender .send(queue_message::QueueMessage::Invocations( RepeatedInvocations { invocations }, )) .await - .unwrap(); + .is_err() + { + break; + } } } remote_message::RemoteMessage::Response(response) => { @@ -370,7 +417,7 @@ async fn remote_queue_server_logic( debt.fulfill(result) } else { // did not get response so need to reenqueue the work - queue.reenqueue(work, debt).await; + queue.reenqueue(work, debt); } } remote_message::RemoteMessage::NodeUpdate(node_update) => { @@ -401,7 +448,7 @@ async fn remote_queue_server_logic( QueueOption::TryOffload(work, debt) => { // if this node already sent enough work for the remote to be at capacity don't send more if invocations_running >= remote_num_cores as usize { - queue.reenqueue(work, debt).await; + queue.reenqueue(work, debt); continue; } invocations_running += 1; @@ -461,21 +508,57 @@ async fn remote_queue_server_logic( metadata_sets, caching, }); - message_sender.send(try_offload_message).await.unwrap(); + if message_sender.send(try_offload_message).await.is_err() { + break; + } } QueueOption::WorkAvailable => { trace!("Queue Server received work available notification"); if waiting_for_work { waiting_for_work = false; - message_sender + if message_sender .send(queue_message::QueueMessage::NoWork(false)) .await - .unwrap(); + .is_err() + { + break; + } } } + QueueOption::Disconnected => { + // The remote node disconnected, stop the loop and run the cleanup below. + break; + } } } - warn!("Arrived at end of remtote_queue_server_logic, which should stay in the loop forever"); + + // The connection to the remote node was lost (or the channel was closed). + // Close the message receiver so all senders will also shut down and no more messages can be enqueued. + info!("Lost connection to worker node {}", node_id); + message_receiver.close(); + + // Undo the bookkeeping for this node + let _ = queue.remove_remote_cores(remote_num_cores as usize); + queue.remove_remote_channel(node_id); + + // Drain any remaining messages. + 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(_, _) => { + () + } + // Reenqueue work that was tried to offload + QueueOption::TryOffload(work, debt) => { + queue.reenqueue(work, debt); + } + } + } + + // Recover any work that was offloaded but never completed, so it can be re-scheduled locally or on another node. + for (_promise_id, (debt, _recorder, _start_epoch, _remote_data_references, work)) in debt_map { + queue.reenqueue(work, debt); + } } /// Handler for one remote node, polling the local queue for them. @@ -497,33 +580,48 @@ pub async fn remote_queue_server( // First ask for the information about the other node // Currently not using engine information trace!("Queue Server wait for initial message"); - let (node_info_buffer, node_info_data) = receive_message(&mut read_socket).await; - debug_assert!(node_info_data.is_none()); + let node_info_buffer = match receive_message(&mut read_socket).await { + Ok((node_info_buffer, node_info_data)) => { + debug_assert!(node_info_data.is_none()); + node_info_buffer + } + Err(err) => { + // The connection dropped before we could identify the node, nothing to clean up. + warn!( + "Failed to receive initial message from worker node: {}", + err + ); + return; + } + }; let NodeInfo { version, id: node_id, num_local_cores, } = deserialize_node_info(node_info_buffer).unwrap(); assert_eq!(version, 1); - trace!("Queue Server received initial message"); + info!( + "Established connection to worker node {} ({} cores)", + node_id, num_local_cores + ); // tell the queue about the remote cores queue.add_remote_cores(num_local_cores as usize); // start sender loop let (queue_message_sender, queue_message_reciever) = mpsc::channel(64); - spawn(remote_queue_server_sender( + let sender_handle = spawn(remote_queue_server_sender( write_socket, queue_message_reciever, )); // start receiver loop let (queue_option_sender, queue_option_receiver) = mpsc::channel(64); - spawn(remote_queue_server_receiver( + let receiver_handle = spawn(remote_queue_server_receiver( read_socket, queue_option_sender.clone(), )); // spawn notificaiton loop - spawn(remote_queue_sever_notification( + let notification_handle = spawn(remote_queue_sever_notification( queue.queueing_notifier(), queue_option_sender.clone(), )); @@ -533,6 +631,7 @@ pub async fn remote_queue_server( spawn(remote_queue_server_try_offload( offload_receiver, queue_option_sender, + queue.clone(), )); remote_queue_server_logic( @@ -545,6 +644,12 @@ pub async fn remote_queue_server( num_local_cores, ) .await; + + // The logic loop returned because the connection was lost, stop the helper tasks so they do not + // linger waiting on a dead socket or closed channels. + sender_handle.abort(); + receiver_handle.abort(); + notification_handle.abort(); } pub enum PollingOption { @@ -553,6 +658,8 @@ pub enum PollingOption { LocalCoreCountChanged(usize), // Results(RemoteMessage, Option<(Vec>, u64)>), Results(remote_message::RemoteMessage), + /// The connection to the remote node was lost, so the client logic should tear down. + Disconnected, } async fn remote_queue_client_receiver( @@ -560,13 +667,23 @@ async fn remote_queue_client_receiver( sender: mpsc::Sender, ) { loop { - let (message_buffer, _) = receive_message(&mut socket).await; + let message_buffer = match receive_message(&mut socket).await { + Ok((message_buffer, _)) => message_buffer, + Err(_) => { + // The connection was lost, inform the logic loop so it can clean up. + let _ = sender.send(PollingOption::Disconnected).await; + break; + } + }; let message = deserialize_queue_message(message_buffer) .and_then(|message| Ok(message.queue_message.unwrap())); - sender + if sender .send(PollingOption::Message(message, None)) .await - .unwrap(); + .is_err() + { + break; + } } } @@ -578,7 +695,13 @@ async fn remote_queue_client_sender( let message_buffer = serialize_remote_message(RemoteMessage { remote_message: Some(remote_message), }); - send_message(&message_buffer, &mut socket, None).await; + if send_message(&message_buffer, &mut socket, None) + .await + .is_err() + { + // connection lost, the receiver side will trigger the teardown + break; + } } } @@ -592,10 +715,13 @@ async fn remote_queue_client_queue_state( receiver.changed().await.unwrap(); trace!("received local queue state"); let queue_state = *receiver.borrow_and_update(); - sender + if sender .send(PollingOption::QueueStateChanged(queue_state)) .await - .unwrap(); + .is_err() + { + break; + } } } @@ -606,10 +732,13 @@ async fn remote_queue_client_core_count( loop { receiver.changed().await.unwrap(); let num_local_cores = *receiver.borrow_and_update(); - sender + if sender .send(PollingOption::LocalCoreCountChanged(num_local_cores)) .await - .unwrap(); + .is_err() + { + break; + } } } @@ -639,15 +768,16 @@ async fn dispatcher_call( } Err(err) => proto::response::Response::ErrorMsg(err.error.to_string()), }; - sender + // If the connection was lost in the meantime the logic loop is gone; dropping the result is + // fine, the master will reenqueue the work after detecting the disconnect. + let _ = sender .send(PollingOption::Results( remote_message::RemoteMessage::Response(Response { invocation_id, response: Some(response_message), }), )) - .await - .unwrap(); + .await; } async fn remote_queue_client_logic( @@ -697,12 +827,15 @@ async fn remote_queue_client_logic( engine_capacity: (num_local_cores - occupancy) as u32, }) .collect(); - message_sender + if message_sender .send(remote_message::RemoteMessage::WorkRequest( RepeatedEngines { engines }, )) .await - .unwrap(); + .is_err() + { + break; + } remote_had_work = false; } else { // Only set to true, if we did not send out a message already, @@ -779,13 +912,21 @@ async fn remote_queue_client_logic( } } PollingOption::Message(Err(error), _) => { - // TODO: recover from message reception failure - panic!("Receiving remote queue message faied with: {}", error); + // A malformed message means the connection is unusable, tear it down so + // the caller can re-establish it instead of crashing the whole node. + error!("Receiving remote queue message failed with: {}", error); + break; + } + PollingOption::Disconnected => { + // The remote node disconnected, stop the loop so the caller can reconnect. + break; } PollingOption::Results(results) => { trace!("Queue Client sending out result"); work_from_remote -= 1; - message_sender.send(results).await.unwrap(); + if message_sender.send(results).await.is_err() { + break; + } let occupancy = std::cmp::max(queue_state, work_from_remote); if remote_had_work && occupancy < num_local_cores { let engines: Vec<_> = EngineType::iter() @@ -796,12 +937,15 @@ async fn remote_queue_client_logic( .collect(); trace!("Asking for more work"); - message_sender + if message_sender .send(remote_message::RemoteMessage::WorkRequest( RepeatedEngines { engines }, )) .await - .unwrap(); + .is_err() + { + break; + } trace!("Finished sending the message asking for more work"); // set false, to avoid double sending if multiple cores become idle, but did not have a response in between remote_had_work = false; @@ -826,12 +970,15 @@ async fn remote_queue_client_logic( .collect(); trace!("Asking for more work"); - message_sender + if message_sender .send(remote_message::RemoteMessage::WorkRequest( RepeatedEngines { engines }, )) .await - .unwrap(); + .is_err() + { + break; + } trace!("Finished sending the message asking for more work"); // set false, to avoid double sending if multiple cores become idle, but did not have a response in between remote_had_work = false; @@ -840,12 +987,15 @@ async fn remote_queue_client_logic( PollingOption::LocalCoreCountChanged(new_core_number) => { trace!("Sending new local core count: {}", new_core_number); num_local_cores = new_core_number; - message_sender + if message_sender .send(remote_message::RemoteMessage::NodeUpdate(NodeUpdate { num_local_cores: new_core_number as u64, })) .await - .unwrap(); + .is_err() + { + break; + } // check if we now want to get more work let occupancy = std::cmp::max(queue_state, work_from_remote); if remote_had_work && occupancy < num_local_cores { @@ -857,12 +1007,15 @@ async fn remote_queue_client_logic( .collect(); trace!("Asking for more work"); - message_sender + if message_sender .send(remote_message::RemoteMessage::WorkRequest( RepeatedEngines { engines }, )) .await - .unwrap(); + .is_err() + { + break; + } trace!("Finished sending the message asking for more work"); // set false, to avoid double sending if multiple cores become idle, but did not have a response in between remote_had_work = false; @@ -870,7 +1023,8 @@ async fn remote_queue_client_logic( } } } - warn!("Arrived at end of remote_qeueu_client_logic, which should stay in the loop forever"); + // Reaching here means the connection was lost; the caller will attempt to reconnect. + trace!("remote_queue_client_logic exited, connection to remote was lost"); } /// Client to ask for work from a remote queue. @@ -898,7 +1052,14 @@ pub async fn remote_queue_client( let (read_socket, mut write_socket) = socket.into_split(); - send_message(&node_info_buffer, &mut write_socket, None).await; + if send_message(&node_info_buffer, &mut write_socket, None) + .await + .is_err() + { + // Could not even send the initial message, let the caller retry the connection. + warn!("Failed to send initial message to remote queue, connection lost"); + return; + } trace!("Queue Client sent out initial message"); // Create a second copy of the watcher to wait on asynchronously, marke changed to check once in the beginning, @@ -909,23 +1070,23 @@ pub async fn remote_queue_client( // start sender loop let (remote_message_sender, remote_message_reciever) = mpsc::channel(64); - spawn(remote_queue_client_sender( + let sender_handle = spawn(remote_queue_client_sender( write_socket, remote_message_reciever, )); // start receiver loop let (poll_option_sender, poll_option_receiver) = mpsc::channel(64); - spawn(remote_queue_client_receiver( + let receiver_handle = spawn(remote_queue_client_receiver( read_socket, poll_option_sender.clone(), )); // start core count loop - spawn(remote_queue_client_core_count( + let core_count_handle = spawn(remote_queue_client_core_count( local_core_watcher, poll_option_sender.clone(), )); // spawn queue state loop - spawn(remote_queue_client_queue_state( + let queue_state_handle = spawn(remote_queue_client_queue_state( local_queue_state, poll_option_sender.clone(), )); @@ -951,4 +1112,11 @@ pub async fn remote_queue_client( local_core_count, ) .await; + + // The logic loop returned because the connection was lost, stop the helper tasks so they do not + // linger waiting on a dead socket or closed channels. + sender_handle.abort(); + receiver_handle.abort(); + core_count_handle.abort(); + queue_state_handle.abort(); } diff --git a/multinode/src/data.rs b/multinode/src/data.rs index 5ba968e4..6582b724 100644 --- a/multinode/src/data.rs +++ b/multinode/src/data.rs @@ -214,6 +214,20 @@ impl ExportRegistry { Ok(()) } + /// Drops all exported data. Used when the connection to the node that manages these + /// contexts is lost, so the worker does not hold on to contexts that will never be + /// fetched or explicitly deleted anymore. + /// NOTE: We currently assume a centralized scheduler that owns the data. If this assumption + /// changes we need to update this function to consider to which master owns the data. + pub fn clear_exported_data(&self) { + let mut inner = self.inner.lock().unwrap(); + let cleared = inner.data.len(); + inner.data.clear(); + if cleared > 0 { + debug!("Cleared {} exported data contexts", cleared); + } + } + pub fn fetch_context(&self, data_id: u64) -> DandelionResult<(Arc, Position)> { let exported_data = self.get_exported_data(data_id)?; Ok((exported_data.context, exported_data.position)) diff --git a/server/src/config.rs b/server/src/config.rs index 9c47abfe..33ae774a 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -10,7 +10,7 @@ const DEFAULT_PORT: u16 = 8080; const DEFAULT_TIMESTAMP_COUNT: usize = 1000; const DEFAULT_MIN_SYS_CORES: usize = 1; const DEFAULT_VIRTUAL_MAX_RAM_MULTIPLIER: usize = 2; -const DEFAULT_MULTINODE_TIMEOUT: u64 = 50; +const DEFAULT_MULTINODE_RECONNECT_INTERVAL: u64 = 1000; use machine_interface::composition::DEFAULT_AUTOSHARDING_OFFLOAD_CONST; use machine_interface::function_driver::system_driver::reqwest::DEFAULT_CONCURRENCY_LIMIT; @@ -177,10 +177,10 @@ pub struct DandelionConfig { #[serde(default)] pub multinode_config: Option, - /// Timeout for how long to try to establish a connection to another node - #[arg(long, env, default_value_t = DEFAULT_MULTINODE_TIMEOUT)] + /// Multinode: how long to wait between attempts to (re-)connect to the master node. + #[arg(long, env, default_value_t = DEFAULT_MULTINODE_RECONNECT_INTERVAL)] #[serde(default)] - pub multinode_timeout_ms: u64, + pub multinode_reconnect_interval_ms: u64, /// Special modes for testing #[arg(long, env, value_enum)] @@ -241,7 +241,10 @@ impl DandelionConfig { merge_clone!(folder_path, String::from(DEFAULT_FOLDER_PATH)); merge!(node_id, 0); merge_option!(multinode_config); - merge!(multinode_timeout_ms, DEFAULT_MULTINODE_TIMEOUT); + merge!( + multinode_reconnect_interval_ms, + DEFAULT_MULTINODE_RECONNECT_INTERVAL + ); } /// Get the config from the arguments, environment and possibly config file diff --git a/server/src/main.rs b/server/src/main.rs index a9609e28..76b86df4 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -73,11 +73,56 @@ async fn remote_queue_client( dispatcher: &'static Dispatcher, export_registry: ExportRegistry, queue: WorkQueue, + reconnect_interval: std::time::Duration, ) { - let connection = tokio::net::TcpStream::connect(remote_url).await.unwrap(); - connection.set_nodelay(true).unwrap(); - connection.set_quickack(true).unwrap(); - multinode::client::remote_queue_client(connection, dispatcher, export_registry, queue).await; + loop { + // Keep retrying to (re-)establish the connection to the master node so a transient + // failure or a master restart does not permanently take this worker out of the cluster. + let connection = match tokio::net::TcpStream::connect(&remote_url).await { + Ok(connection) => connection, + Err(err) => { + debug!( + "Failed to connect to master node at {}: {}. Retrying in {:?}", + remote_url, err, reconnect_interval + ); + tokio::time::sleep(reconnect_interval).await; + continue; + } + }; + if let Err(err) = connection.set_nodelay(true) { + warn!( + "Failed to set nodelay on connection to master node: {}", + err + ); + } + if let Err(err) = connection.set_quickack(true) { + warn!( + "Failed to set quickack on connection to master node: {}", + err + ); + } + info!("Established connection to master node at {}", remote_url); + + multinode::client::remote_queue_client( + connection, + dispatcher, + export_registry.clone(), + queue.clone(), + ) + .await; + + // The connection was lost. Drop all contexts we were holding for the master node, + // since it will no longer fetch or delete them, then retry connecting. + info!( + "Lost connection to master node at {}, retrying in {:?}", + remote_url, reconnect_interval + ); + // NOTE: we currently assume a centralized scheduler that owns the data. If this assumption + // changes we need to update this function to only clear the exported data belonging + // to this node. + export_registry.clear_exported_data(); + tokio::time::sleep(reconnect_interval).await; + } } fn main() -> () { @@ -329,6 +374,7 @@ fn main() -> () { dispatcher, export_registry.clone(), work_queue, + std::time::Duration::from_millis(config.multinode_reconnect_interval_ms), )); } diff --git a/server/tests/server_tests.rs b/server/tests/server_tests.rs index 18ebadb3..82cd9b74 100644 --- a/server/tests/server_tests.rs +++ b/server/tests/server_tests.rs @@ -56,6 +56,31 @@ mod server_tests { } let _ = self.server.stdout.insert(reader.into_inner()); } + + /// Reads the server's stderr (where the logs are written) until each of `markers` + /// has appeared, in order, printing every line on the way. This is used both to + /// synchronize the test with the server's progress and to assert that the expected + /// log lines are produced. A single reader is kept for the whole scan so no buffered + /// log lines are lost between markers. Panics if the process exits before all markers + /// are seen. + fn wait_for_stderr(&mut self, markers: &[&str]) { + let mut reader = BufReader::new(self.server.stderr.take().unwrap()); + let mut next = 0; + while next < markers.len() { + let mut buf = String::new(); + let len = reader.read_line(&mut buf).unwrap(); + assert_ne!( + len, 0, + "{} exited before logging {:?}", + self.name, markers[next] + ); + print!("{} err: {}", self.name, buf); + if buf.contains(markers[next]) { + next += 1; + } + } + let _ = self.server.stderr.insert(reader.into_inner()); + } } impl Drop for ServerKiller { @@ -315,15 +340,17 @@ mod server_tests { ) } - fn start_multinode_servers() -> MultinodeServers { - let preload_path = multinode_preload_path(); - println!("Preload_path: {}", preload_path); - let multinode_config = format!( + fn multinode_config_path() -> String { + format!( "{}/tests/manifests/multinode_config.json", env!("CARGO_MANIFEST_DIR"), - ); + ) + } + + fn start_master() -> ServerKiller { + let preload_path = multinode_preload_path(); + let multinode_config = multinode_config_path(); - let remote_port = 8081; let mut master_cmd = Command::new(assert_cmd::cargo::cargo_bin!()); let master_server = master_cmd .stdout(Stdio::piped()) @@ -346,7 +373,14 @@ mod server_tests { server: master_server, }; master.check_for_start(); + master + } + fn start_worker() -> ServerKiller { + let preload_path = multinode_preload_path(); + let multinode_config = multinode_config_path(); + + let remote_port = 8081; let mut worker_cmd = Command::new(assert_cmd::cargo::cargo_bin!()); let worker_server = worker_cmd .stdout(Stdio::piped()) @@ -367,7 +401,13 @@ mod server_tests { server: worker_server, }; worker.check_for_start(); + worker + } + fn start_multinode_servers() -> MultinodeServers { + println!("Preload_path: {}", multinode_preload_path()); + let master = start_master(); + let worker = start_worker(); MultinodeServers { master, worker } } @@ -513,4 +553,46 @@ mod server_tests { servers.assert_running(); } + + /// Checks that a worker node which loses its connection to the master node notices the + /// loss, cleans up after it, and automatically re-establishes the connection once the + /// master is reachable again, all without either node panicking. + #[test] + #[serial] + fn serve_multinode_reconnect() { + let mut servers = start_multinode_servers(); + + // Make sure the cluster has actually formed before we tear it down: the worker + // connects to the master and the master registers the worker with its cores. + servers + .worker + .wait_for_stderr(&["Established connection to master node"]); + servers + .master + .wait_for_stderr(&["Established connection to worker node"]); + + // Drop the master to sever the worker's connection (also prints the master's output). + drop(servers.master); + + // Bring the master back up. The worker should notice the lost connection and keep + // retrying until it can reconnect to the new master. + let mut master = start_master(); + + // The worker logs that it lost the connection (the code path that also clears its + // ExportRegistry) and then that it re-established it, confirming the reconnect loop. + servers.worker.wait_for_stderr(&[ + "Lost connection to master node", + "Established connection to master node", + ]); + + // The (freshly started) master accepts the worker rejoining the cluster. + master.wait_for_stderr(&["Established connection to worker node"]); + + // Neither node panicked through the disconnect and reconnect. + let servers = MultinodeServers { + master, + worker: servers.worker, + }; + servers.assert_running(); + } }