From fbdb85ea7bfdaf64ad308c5cfcd6b3dfd7168a8f Mon Sep 17 00:00:00 2001 From: DaLaw2 Date: Mon, 8 Apr 2024 21:42:20 +0800 Subject: [PATCH] Management side and Agent side communication architecture. --- .../channel/control_channel_receive_thread.rs | 4 +- .../channel/control_channel_receiver.rs | 24 +- .../channel/data_channel_receive_thread.rs | 2 + .../channel/data_channel_receiver.rs | 11 + AgentLibrary/src/management/agent.rs | 399 ++++++++++++++---- AgentLibrary/src/management/manager.rs | 29 +- AgentLibrary/src/utils/config.rs | 2 - Common/src/connection/packet/mod.rs | 18 +- Common/src/management/utils/file_header.rs | 6 +- .../management/utils/file_transfer_result.rs | 6 + Common/src/management/utils/task_info.rs | 6 +- .../packet/file_transfer_end_packet.rs | 57 +++ .../src/connection/packet/mod.rs | 1 + ManagementLibrary/src/management/agent.rs | 239 +++++------ 14 files changed, 546 insertions(+), 258 deletions(-) create mode 100644 ManagementLibrary/src/connection/packet/file_transfer_end_packet.rs diff --git a/AgentLibrary/src/connection/channel/control_channel_receive_thread.rs b/AgentLibrary/src/connection/channel/control_channel_receive_thread.rs index 2365bcc..b675b50 100644 --- a/AgentLibrary/src/connection/channel/control_channel_receive_thread.rs +++ b/AgentLibrary/src/connection/channel/control_channel_receive_thread.rs @@ -30,8 +30,10 @@ impl ReceiveThread { Ok(packet) => { let packet_type = PacketType::parse_packet_type(&packet.clone_id_byte()); let result = match packet_type { - PacketType::ConfirmPacket => self.receiver_tx.confirm_packet.send(packet), + PacketType::AgentInformationAcknowledgePacket => self.receiver_tx.agent_information_acknowledge_packet.send(packet), + PacketType::ControlPacket => self.receiver_tx.control_packet.send(packet), PacketType::DataChannelPortPacket => self.receiver_tx.data_channel_port_packet.send(packet), + PacketType::ResultAcknowledgePacket => self.receiver_tx.performance_acknowledge_packet.send(packet), _ => { logging_warning!("Receive Thread: Receive unknown packet."); Ok(()) diff --git a/AgentLibrary/src/connection/channel/control_channel_receiver.rs b/AgentLibrary/src/connection/channel/control_channel_receiver.rs index 716ddc4..6662234 100644 --- a/AgentLibrary/src/connection/channel/control_channel_receiver.rs +++ b/AgentLibrary/src/connection/channel/control_channel_receiver.rs @@ -7,18 +7,24 @@ use crate::connection::channel::control_channel_receive_thread::ReceiveThread; pub struct ControlChannelReceiver { stop_signal_tx: Option>, - pub confirm_packet: mpsc::UnboundedReceiver, + pub agent_information_acknowledge_packet: mpsc::UnboundedReceiver, + pub control_packet: mpsc::UnboundedReceiver, pub data_channel_port_packet: mpsc::UnboundedReceiver, + pub performance_acknowledge_packet: mpsc::UnboundedReceiver, } impl ControlChannelReceiver { pub fn new(socket_rx: ReadHalf) -> Self { let (stop_signal_tx, stop_signal_rx) = oneshot::channel(); - let (confirm_packet_tx, confirm_packet_rx) = mpsc::unbounded_channel(); + let (agent_information_acknowledge_packet_tx, agent_information_acknowledge_packet_rx) = mpsc::unbounded_channel(); + let (control_packet_tx, control_packet_rx) = mpsc::unbounded_channel(); let (data_channel_port_packet_tx, data_channel_port_packet_rx) = mpsc::unbounded_channel(); + let (performance_acknowledge_packet_tx, performance_acknowledge_packet_rx) = mpsc::unbounded_channel(); let receiver_tx = ReceiverTX { - confirm_packet: confirm_packet_tx, + agent_information_acknowledge_packet: agent_information_acknowledge_packet_tx, + control_packet: control_packet_tx, data_channel_port_packet: data_channel_port_packet_tx, + performance_acknowledge_packet: performance_acknowledge_packet_tx, }; let mut receive_thread = ReceiveThread::new(socket_rx, receiver_tx, stop_signal_rx); tokio::spawn(async move { @@ -26,14 +32,18 @@ impl ControlChannelReceiver { }); Self { stop_signal_tx: Some(stop_signal_tx), - confirm_packet: confirm_packet_rx, + agent_information_acknowledge_packet: agent_information_acknowledge_packet_rx, + control_packet: control_packet_rx, data_channel_port_packet: data_channel_port_packet_rx, + performance_acknowledge_packet: performance_acknowledge_packet_rx, } } pub async fn disconnect(&mut self) { - self.confirm_packet.close(); + self.agent_information_acknowledge_packet.close(); + self.control_packet.close(); self.data_channel_port_packet.close(); + self.performance_acknowledge_packet.close(); match self.stop_signal_tx.take() { Some(stop_signal) => { let _ = stop_signal.send(()); @@ -45,6 +55,8 @@ impl ControlChannelReceiver { } pub struct ReceiverTX { - pub confirm_packet: mpsc::UnboundedSender, + pub agent_information_acknowledge_packet: mpsc::UnboundedSender, + pub control_packet: mpsc::UnboundedSender, pub data_channel_port_packet: mpsc::UnboundedSender, + pub performance_acknowledge_packet: mpsc::UnboundedSender, } diff --git a/AgentLibrary/src/connection/channel/data_channel_receive_thread.rs b/AgentLibrary/src/connection/channel/data_channel_receive_thread.rs index 602db71..8d05234 100644 --- a/AgentLibrary/src/connection/channel/data_channel_receive_thread.rs +++ b/AgentLibrary/src/connection/channel/data_channel_receive_thread.rs @@ -33,6 +33,8 @@ impl ReceiveThread { PacketType::AlivePacket => self.receiver_tx.alive_packet.send(packet), PacketType::FileBodyPacket => self.receiver_tx.file_body_packet.send(packet), PacketType::FileHeaderPacket => self.receiver_tx.file_header_packet.send(packet), + PacketType::FileTransferEndPacket => self.receiver_tx.file_transfer_end_packet.send(packet), + PacketType::ResultAcknowledgePacket => self.receiver_tx.result_acknowledge_packet.send(packet), PacketType::StillProcessPacket => self.receiver_tx.still_process_packet.send(packet), PacketType::TaskInfoPacket => self.receiver_tx.task_info_packet.send(packet), _ => { diff --git a/AgentLibrary/src/connection/channel/data_channel_receiver.rs b/AgentLibrary/src/connection/channel/data_channel_receiver.rs index d37f5c6..c9c7010 100644 --- a/AgentLibrary/src/connection/channel/data_channel_receiver.rs +++ b/AgentLibrary/src/connection/channel/data_channel_receiver.rs @@ -10,6 +10,8 @@ pub struct DataChannelReceiver { pub alive_packet: UnboundedReceiver, pub file_body_packet: UnboundedReceiver, pub file_header_packet: UnboundedReceiver, + pub file_transfer_end_packet: UnboundedReceiver, + pub result_acknowledge_packet: UnboundedReceiver, pub still_process_packet: UnboundedReceiver, pub task_info_packet: UnboundedReceiver, } @@ -20,12 +22,16 @@ impl DataChannelReceiver { let (alive_packet_tx, alive_packet_rx) = mpsc::unbounded_channel(); let (file_body_packet_tx, file_body_packet_rx) = mpsc::unbounded_channel(); let (file_header_packet_tx, file_header_packet_rx) = mpsc::unbounded_channel(); + let (file_transfer_end_packet_tx, file_transfer_end_packet_rx) = mpsc::unbounded_channel(); + let (result_acknowledge_packet_tx, result_acknowledge_packet_rx) = mpsc::unbounded_channel(); let (still_process_packet_tx, still_process_packet_rx) = mpsc::unbounded_channel(); let (task_info_packet_tx, task_info_packet_rx) = mpsc::unbounded_channel(); let receiver_tx = ReceiverTX { alive_packet: alive_packet_tx, file_body_packet: file_body_packet_tx, file_header_packet: file_header_packet_tx, + file_transfer_end_packet: file_transfer_end_packet_tx, + result_acknowledge_packet: result_acknowledge_packet_tx, still_process_packet: still_process_packet_tx, task_info_packet: task_info_packet_tx, }; @@ -38,6 +44,8 @@ impl DataChannelReceiver { alive_packet: alive_packet_rx, file_body_packet: file_body_packet_rx, file_header_packet: file_header_packet_rx, + file_transfer_end_packet: file_transfer_end_packet_rx, + result_acknowledge_packet: result_acknowledge_packet_rx, still_process_packet: still_process_packet_rx, task_info_packet: task_info_packet_rx, } @@ -47,6 +55,7 @@ impl DataChannelReceiver { self.alive_packet.close(); self.file_body_packet.close(); self.file_header_packet.close(); + self.file_transfer_end_packet.close(); self.still_process_packet.close(); self.task_info_packet.close(); match self.stop_signal_tx.take() { @@ -63,6 +72,8 @@ pub struct ReceiverTX { pub alive_packet: UnboundedSender, pub file_body_packet: UnboundedSender, pub file_header_packet: UnboundedSender, + pub file_transfer_end_packet: UnboundedSender, + pub result_acknowledge_packet: UnboundedSender, pub still_process_packet: UnboundedSender, pub task_info_packet: UnboundedSender, } diff --git a/AgentLibrary/src/management/agent.rs b/AgentLibrary/src/management/agent.rs index 78e6091..1c877cb 100644 --- a/AgentLibrary/src/management/agent.rs +++ b/AgentLibrary/src/management/agent.rs @@ -1,31 +1,38 @@ -use tokio::net::TcpStream; -use std::sync::Arc; -use std::time::Duration; -use tokio::select; -use tokio::sync::RwLock; -use tokio::time::{Instant, sleep}; -use crate::management::utils::performance::Performance; +use std::mem; use uuid::Uuid; -use crate::management::utils::task_info::TaskInfo; -use crate::utils::clear_unbounded_channel; +use tokio::select; +use std::sync::Arc; +use std::path::PathBuf; +use tokio::sync::RwLock; +use std::time::Duration; +use tokio::net::TcpStream; +use std::collections::HashMap; +use tokio::time::{Instant, sleep}; use crate::utils::logger::*; +use crate::utils::config::Config; use crate::connection::packet::Packet; -use crate::management::utils::confirm_type::ConfirmType; +use crate::management::manager::Manager; +use crate::management::monitor::Monitor; +use crate::utils::clear_unbounded_channel; +use crate::management::utils::task_info::TaskInfo; +use crate::management::utils::agent_state::AgentState; +use crate::management::utils::file_header::FileHeader; use crate::connection::socket::socket_stream::SocketStream; +use crate::connection::channel::{ControlChannel, DataChannel}; +use crate::connection::packet::performance_packet::PerformancePacket; use crate::connection::channel::data_channel_sender::DataChannelSender; use crate::connection::channel::data_channel_receiver::DataChannelReceiver; use crate::connection::channel::control_channel_sender::ControlChannelSender; -use crate::connection::channel::control_channel_receiver::ControlChannelReceiver; -use crate::connection::channel::{ControlChannel, DataChannel}; use crate::connection::packet::agent_information_packet::AgentInformationPacket; use crate::connection::packet::alive_acknowledge_packet::AliveAcknowledgePacket; -use crate::connection::packet::performance_packet::PerformancePacket; -use crate::management::monitor::Monitor; -use crate::utils::config::Config; +use crate::connection::channel::control_channel_receiver::ControlChannelReceiver; +use crate::connection::packet::control_acknowledge_packet::ControlAcknowledgePacket; +use crate::connection::packet::file_transfer_result_packet::FileTransferResultPacket; +use crate::connection::packet::task_info_acknowledge_packet::TaskInfoAcknowledgePacket; +use crate::connection::packet::file_header_acknowledge_packet::FileHeaderAcknowledgePacket; pub struct Agent { - terminate: bool, - previous_task: Option, + previous_task_uuid: Option, control_channel_sender: ControlChannelSender, control_channel_receiver: ControlChannelReceiver, data_channel_sender: Option, @@ -35,10 +42,10 @@ pub struct Agent { impl Agent { pub async fn new(socket_stream: SocketStream) -> Result { let config = Config::now().await; - let information = serde_json::to_vec(&Monitor::get_system_info().await) - .map_err(|_| error_entry!("Agent: Unable to serialized agent information."))?; let (mut control_channel_sender, mut control_channel_receiver) = ControlChannel::new(socket_stream); let mut information_confirm = false; + let information = serde_json::to_vec(&Monitor::get_system_info().await) + .map_err(|_| error_entry!("Agent: Unable to serialized agent information."))?; let timer = Instant::now(); let mut polling_times = 0_u32; let polling_interval = Duration::from_millis(config.polling_interval); @@ -56,28 +63,27 @@ impl Agent { } select! { biased; - reply = control_channel_receiver.confirm_packet.recv() => { - let packet = reply + packet = control_channel_receiver.agent_information_acknowledge_packet.recv() => { + let packet = packet .ok_or(info_entry!("Agent: Channel has been closed."))?; - clear_unbounded_channel(&mut control_channel_receiver.confirm_packet).await; - let confirm = serde_json::from_slice::(packet.as_data_byte()) - .map_err(|_| error_entry!("Agent: Unable to parse confirm type."))?; - match confirm { - ConfirmType::ReceivedAgentInformation => { - information_confirm = true; - continue - }, - ConfirmType::ReceivedPerformance => { - let agent = Self { - terminate: false, - control_channel_sender, - control_channel_receiver, - data_channel_sender: None, - data_channel_receiver: None, - }; - return Ok(agent) - }, + clear_unbounded_channel(&mut control_channel_receiver.agent_information_acknowledge_packet).await; + information_confirm = true; + }, + packet = control_channel_receiver.performance_acknowledge_packet.recv() => { + let packet = packet + .ok_or(info_entry!("Agent: Channel has been closed."))?; + clear_unbounded_channel(&mut control_channel_receiver.performance_acknowledge_packet).await; + if !information_confirm { + Err(error_entry!("Agent: Agent information not acknowledge."))?; } + let agent = Self { + previous_task_uuid: None, + control_channel_sender, + control_channel_receiver, + data_channel_sender: None, + data_channel_receiver: None, + }; + return Ok(agent); }, _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, } @@ -87,30 +93,15 @@ impl Agent { pub async fn run(agent: Arc>) { let for_performance = agent.clone(); - let for_management = agent.clone(); + let for_management = agent; tokio::spawn(async move { Self::performance(for_performance).await; }); tokio::spawn(async move { - Self::create_data_channel(for_management).await; + Self::management(for_management).await; }); } - pub async fn terminate(agent: Arc>) { - logging_info!("Agent: Terminating agent."); - let mut agent = agent.write().await; - agent.terminate = true; - agent.control_channel_sender.disconnect().await; - agent.control_channel_receiver.disconnect().await; - if let Some(data_channel_sender) = &mut agent.data_channel_sender { - data_channel_sender.disconnect().await; - } - if let Some(data_channel_receiver) = &mut agent.data_channel_receiver { - data_channel_receiver.disconnect().await; - } - logging_info!("Agent: Termination complete."); - } - async fn performance(agent: Arc>) { let config = Config::now().await; let mut polling_times = 0_u32; @@ -118,16 +109,18 @@ impl Agent { let polling_interval = Duration::from_millis(config.polling_interval); let mut timeout_timer = Instant::now(); let timeout_duration = Duration::from_secs(config.control_channel_timeout); - while !agent.read().await.terminate { + loop { + if Manager::get_state() == AgentState::Terminate { + return; + } if timeout_timer.elapsed() > timeout_duration { logging_warning!("Agent: Control Channel timeout."); - Agent::terminate(agent).await; - return; + break; } if polling_timer.elapsed() > polling_times * polling_interval { let performance = Monitor::get_performance().await; - if let Ok(performance) = serde_json::to_vec(&performance) { - agent.write().await.control_channel_sender.send(PerformancePacket::new(performance)).await; + if let Ok(performance_data) = serde_json::to_vec(&performance) { + agent.write().await.control_channel_sender.send(PerformancePacket::new(performance_data)).await; } else { logging_error!("Agent: Unable to serialized performance data."); } @@ -135,78 +128,273 @@ impl Agent { let mut agent = agent.write().await; select! { biased; - reply = agent.control_channel_receiver.confirm_packet.recv() => { + reply = agent.control_channel_receiver.performance_acknowledge_packet.recv() => { if let Some(packet) = reply { - clear_unbounded_channel(&mut agent.control_channel_receiver.confirm_packet).await; - if let Ok(_) = serde_json::from_slice::(packet.as_data_byte()) { - timeout_timer = Instant::now(); - } else { - logging_error!("Agent: Unable to parse confirm data."); - } + clear_unbounded_channel(&mut agent.control_channel_receiver.performance_acknowledge_packet).await; + timeout_timer = Instant::now(); } else { logging_info!("Agent: Channel has been closed."); - return; + break; } }, _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, } } + Manager::store_state(AgentState::Terminate).await; + } + + async fn management(agent: Arc>) { + loop { + Self::refresh_state(agent.clone()).await; + let state = Manager::get_state().await; + match state { + AgentState::ProcessTask => Self::process_task(agent.clone()).await, + AgentState::Idle(idle_time) => Self::idle(agent.clone(), Duration::from_secs(idle_time)).await, + AgentState::CreateDataChannel => Self::create_data_channel(agent.clone()).await, + AgentState::Terminate => { + Self::terminate(agent.clone()).await; + return; + }, + _ => {}, + } + } + } + + async fn refresh_state(agent: Arc>) { + let config = Config::now().await; + let timer = Instant::now(); + let timeout_duration = Duration::from_secs(config.control_channel_timeout); + while Manager::get_state() != AgentState::Terminate { + if timer.elapsed() > timeout_duration { + Manager::store_state(AgentState::Terminate).await; + return; + } + let agent = agent.write().await; + select! { + packet = agent.control_channel_receiver.control_packet.recv() => { + match packet { + Some(packet) => { + clear_unbounded_channel(&mut agent.control_channel_receiver.control_packet).await; + match serde_json::from_slice::(packet.as_data_byte()) { + Some(state) => Manager::store_state(state).await, + None => { + logging_error!("Agent: Unable to parse control state."); + continue; + }, + } + }, + None => { + logging_warning!("Agent: Channel has been closed."); + Manager::store_state(AgentState::Terminate).await; + }, + } + }, + _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, + } + agent.control_channel_sender.send(ControlAcknowledgePacket::new()).await; + return; + } } async fn process_task(agent: Arc>) { - while !agent.read().await.terminate { - let task_info = match Self::receive_task_info(agent.clone()).await { - Ok(task_info) => task_info, - Err(entry) => { - logging_entry!(entry); - continue; - } - }; - + if let Err(entry) = Self::receive_task(agent.clone()).await { + logging_entry!(entry); } + if let Err(entry) = Self::inference_task(agent.clone()).await { + logging_entry!(entry); + } + } + + async fn receive_task(agent: Arc>) -> Result<(), LogEntry> { + let task_info = Self::receive_task_info(agent.clone()).await?; + let previous_task_uuid = agent.read().await.previous_task_uuid; + let need_receive_model = if let Some(previous_task_uuid) = previous_task_uuid { + previous_task_uuid != task_info.uuid + } else { + true + }; + if need_receive_model { + Self::receive_file(agent.clone())?; + } + Self::receive_file(agent.clone())?; + Ok(()) } async fn receive_task_info(agent: Arc>) -> Result { let config = Config::now().await; let timer = Instant::now(); let timeout_duration = Duration::from_secs(config.data_channel_timeout); - while !agent.read().await.terminate { + let task_info = loop { + if Manager::get_state() == AgentState::Terminate { + Err(info_entry!("Agent: Terminating. Receive task info cancel."))?; + } if timer.elapsed() > timeout_duration { - Err(info_entry!("Data Channel timeout."))?; + Err(info_entry!("Agent: Data Channel timeout."))?; } if let Some(data_channel_receiver) = &mut agent.write().await.data_channel_receiver { select! { - reply = data_channel_receiver.task_info_packet.recv() => { - let packet = reply + packet = data_channel_receiver.task_info_packet.recv() => { + let packet = packet .ok_or(warning_entry!("Agent: Channel has been closed."))?; clear_unbounded_channel(&mut data_channel_receiver.task_info_packet).await; - return serde_json::from_slice::(packet.as_data_byte()) - .map_err(|_| error_entry!("Agent: Unable to parse task info.")); + break serde_json::from_slice::(packet.as_data_byte()) + .map_err(|_| error_entry!("Agent: Unable to parse task info."))?; }, _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, } } else { - sleep(Duration::from_secs(config.agent_idle_interval)).await; Err(warning_entry!("Agent: Data Channel is not available."))? } + }; + if let Some(data_channel_sender) = &mut agent.write().await.data_channel_sender { + data_channel_sender.send(TaskInfoAcknowledgePacket::new()).await; + } else { + Err(warning_entry!("Agent: Data Channel is not available."))?; } - Err(info_entry!("Agent: Terminating. Receive task info cancel.")) + return Ok(task_info); } - async fn receive_file(agent: Arc>) + async fn receive_file(agent: Arc>, save_path: &PathBuf) -> Result<(), LogEntry> { + let file_header = Self::receive_file_header(agent.clone()).await?; + let file_body = Self::receive_file_body(agent.clone(), &file_header).await?; + Self::create_file(agent, file_body, save_path).await?; + Ok(()) + } - async fn idle(agent: Arc>) { + async fn receive_file_header(agent: Arc>) -> Result { let config = Config::now().await; - while !agent.read().await.terminate { + let timer = Instant::now(); + let timeout_duration = Duration::from_secs(config.data_channel_timeout); + let file_header = loop { + if Manager::get_state() == AgentState::Terminate { + Err(info_entry!("Agent: Terminating. Receive file header cancel."))?; + } + if timer.elapsed() > timeout_duration { + Err(warning_entry!("Agent: Data Channel timeout."))?; + } if let Some(data_channel_receiver) = &mut agent.write().await.data_channel_receiver { + select! { + packet = data_channel_receiver.task_info_packet.recv() => { + let packet = packet + .ok_or(warning_entry!("Agent: Channel has been closed."))?; + clear_unbounded_channel(&mut data_channel_receiver.task_info_packet).await; + break serde_json::from_slice::(packet.as_data_byte()) + .map_err(|_| error_entry!("Agent: Unable to parse task info."))?; + }, + _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, + } + } else { + Err(warning_entry!("Agent: Data Channel is not available."))?; + } + }; + if let Some(data_channel_sender) = &mut agent.write().await.data_channel_sender { + data_channel_sender.send(FileHeaderAcknowledgePacket::new()).await; + } else { + Err(warning_entry!("Agent: Data Channel is not available."))?; + } + return Ok(file_header); + } + + async fn receive_file_body(agent: Arc>, file_header: &FileHeader) -> Result>, LogEntry> { + let config = Config::now().await; + let mut file_block: HashMap> = HashMap::new(); + let mut missing_blocks = Vec::new(); + let timer = Instant::now(); + let timeout_duration = Duration::from_secs(config.data_channel_timeout); + loop { + if Manager::get_state() == AgentState::Terminate { + Err(info_entry!("Agent: Terminating. Receive file header cancel."))?; + } + if timer.elapsed() > timeout_duration { + Err(warning_entry!("Agent: Data Channel timeout."))?; + } + if let Some(data_channel_receiver) = agent.write().await.data_channel_receiver.as_mut() { + select! { + biased; + packet = data_channel_receiver.file_body_packet.recv() => { + let packet = &packet + .ok_or(warning_entry!("Agent: Channel has been closed."))?; + clear_unbounded_channel(&mut data_channel_receiver.file_body_packet).await; + let (sequence_bytes, file_body) = packet.data.split_at(std::mem::size_of::()); + let sequence_number = usize::from_be_bytes(sequence_bytes.try_into().map_err(|_| error_entry!("Agent: Unable to parse file body."))?); + file_block.insert(sequence_number, Vec::from(file_body)); + continue; + }, + packet = data_channel_receiver.file_transfer_end_packet.recv() => { + let packet = &packet + .ok_or(warning_entry!("Agent: Channel has been closed."))?; + clear_unbounded_channel(&mut data_channel_receiver.file_transfer_end_packet).await; + for sequence_number in 0..file_header.packet_count { + if !file_block.contains_key(&sequence_number) { + missing_blocks.push(sequence_number); + } + } + }, + _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, + } + } else { + Err(warning_entry!("Agent: Data Channel is not available."))?; + } + if let Some(data_channel_sender) = agent.write().await.data_channel_sender.as_mut() { + if missing_blocks.len() != 0_usize { + let result = Some(mem::take(&mut missing_blocks)); + let result_data = serde_json::to_vec(&result) + .map_err(|_| error_entry!("Agent: Unable to serialized result data."))?; + data_channel_sender.send(FileTransferResultPacket::new(result_data)).await; + } else { + let result: Option> = None; + let result_data = serde_json::to_vec(&result) + .map_err(|_| error_entry!("Agent: Unable to serialized result data."))?; + data_channel_sender.send(FileTransferResultPacket::new(result_data)).await; + let mut sorted_blocks: Vec> = Vec::with_capacity(file_header.packet_count); + for index in 0..file_header.packet_count { + if let Some(block) = file_block.remove(&index) { + sorted_blocks.push(block); + } + } + return Ok(sorted_blocks); + }; + } else { + Err(warning_entry!("Agent: Data Channel is not available."))?; + } + } + } + + async fn create_file(agent: Arc>, file_body: Vec>, save_path: &PathBuf) -> Result<(), LogEntry> { + + } + + async fn inference_task(agent: Arc>) -> Result<(), LogEntry> { + + } + + async fn idle(agent: Arc>, idle_duration: Duration) { + let config = Config::now().await; + let timer = Instant::now(); + loop { + if Manager::get_state() == AgentState::Terminate { + logging_info!("Agent: Terminating. Stop idle."); + return; + } + if timer.elapsed() > idle_duration { + return; + } + let agent = agent.write().await; + if let Some(data_channel_receiver) = &mut agent.data_channel_receiver { select! { biased; _ = data_channel_receiver.alive_packet.recv() => clear_unbounded_channel(&mut data_channel_receiver.alive_packet).await, _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, } + } else { + logging_warning!("Agent: Data Channel is not available."); + return; } - if let Some(data_channel_sender) = &mut agent.write().await.data_channel_sender { + if let Some(data_channel_sender) = &mut agent.data_channel_sender { data_channel_sender.send(AliveAcknowledgePacket::new()).await; + } else { + logging_warning!("Agent: Data Channel is not available."); + return; } } } @@ -214,7 +402,18 @@ impl Agent { async fn create_data_channel(agent: Arc>) { let config = Config::now().await; let mut port: Option = None; - while !agent.read().await.terminate { + let timer = Instant::now(); + let timeout_duration = Duration::from_secs(config.control_channel_timeout); + loop { + if Manager::get_state() == AgentState::Terminate { + logging_info!("Agent: Terminating. Cancel create data channel."); + return; + } + if timer.elapsed() > timeout_duration { + Manager::store_state(AgentState::Terminate).await; + logging_warning!("Agent: Control channel timout."); + return; + } { let mut agent = agent.write().await; select! { @@ -230,6 +429,7 @@ impl Agent { continue; } } else { + Manager::store_state(AgentState::Terminate).await; logging_warning!("Agent: Channel has been closed."); return; } @@ -245,6 +445,9 @@ impl Agent { let mut agent = agent.write().await; agent.data_channel_sender = Some(data_channel_sender); agent.data_channel_receiver = Some(data_channel_receiver); + } else { + logging_error!("Agent: Unable to create data channel connect."); + return; } } else { logging_error!("Agent: Port data not ready."); @@ -252,4 +455,18 @@ impl Agent { } } } + + pub async fn terminate(agent: Arc>) { + logging_info!("Agent: Terminating agent."); + let mut agent = agent.write().await; + agent.control_channel_sender.disconnect().await; + agent.control_channel_receiver.disconnect().await; + if let Some(data_channel_sender) = &mut agent.data_channel_sender { + data_channel_sender.disconnect().await; + } + if let Some(data_channel_receiver) = &mut agent.data_channel_receiver { + data_channel_receiver.disconnect().await; + } + logging_info!("Agent: Termination complete."); + } } diff --git a/AgentLibrary/src/management/manager.rs b/AgentLibrary/src/management/manager.rs index 67f9332..a0f4803 100644 --- a/AgentLibrary/src/management/manager.rs +++ b/AgentLibrary/src/management/manager.rs @@ -1,6 +1,7 @@ use std::sync::Arc; -use tokio::sync::RwLock; +use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; use lazy_static::lazy_static; +use crate::management::utils::agent_state::AgentState; use crate::management::agent::Agent; lazy_static! { @@ -9,6 +10,7 @@ lazy_static! { pub struct Manager { agent: Option>>, + state: Option, terminate: bool, } @@ -16,10 +18,19 @@ impl Manager { pub fn new() -> Self { Self { agent: None, + state: None, terminate: false, } } + pub async fn instance() -> RwLockReadGuard<'static, Self> { + MANAGER.read().await + } + + pub async fn instance_mut() -> RwLockWriteGuard<'static, Self> { + MANAGER.write().await + } + pub async fn run() { } @@ -35,4 +46,20 @@ impl Manager { fn cleanup() { } + + pub async fn store_state(state: AgentState) { + let mut manager = Self::instance_mut().await; + if let Some(origin_state) = manager.state { + if origin_state != AgentState::Terminate { + manager.state = Some(state); + } + } else { + manager.state = Some(state) + } + } + + pub async fn get_state() -> AgentState { + let manager = Self::instance().await; + manager.state.unwrap_or_else(|| AgentState::None) + } } diff --git a/AgentLibrary/src/utils/config.rs b/AgentLibrary/src/utils/config.rs index bea65d0..9e9434f 100644 --- a/AgentLibrary/src/utils/config.rs +++ b/AgentLibrary/src/utils/config.rs @@ -19,7 +19,6 @@ pub struct Config { pub internal_timestamp: u64, pub management_address: String, pub management_port: u16, - pub agent_idle_interval: u64, pub polling_interval: u64, pub control_channel_timeout: u64, pub data_channel_timeout: u64, @@ -50,7 +49,6 @@ impl Config { pub fn validate(config: &Config) -> bool { Config::validate_mini_second(config.internal_timestamp) && Config::validate_full_address(&config.management_address, config.management_port) - && Config::validate_second(config.agent_idle_interval) && Config::validate_second(config.control_channel_timeout) && Config::validate_second(config.data_channel_timeout) && Config::validate_second(config.file_transfer_timeout) diff --git a/Common/src/connection/packet/mod.rs b/Common/src/connection/packet/mod.rs index 44c4e49..198617f 100644 --- a/Common/src/connection/packet/mod.rs +++ b/Common/src/connection/packet/mod.rs @@ -26,6 +26,7 @@ pub enum PacketType { FileHeaderPacket, FileHeaderAcknowledgePacket, FileTransferResultPacket, + FileTransferEndPacket, PerformancePacket, PerformanceAcknowledgePacket, ResultPacket, @@ -58,14 +59,15 @@ impl PacketType { 9 => PacketType::FileHeaderPacket, 10 => PacketType::FileHeaderAcknowledgePacket, 11 => PacketType::FileTransferResultPacket, - 12 => PacketType::PerformancePacket, - 13 => PacketType::PerformanceAcknowledgePacket, - 14 => PacketType::ResultPacket, - 15 => PacketType::ResultAcknowledgePacket, - 16 => PacketType::StillProcessPacket, - 17 => PacketType::StillProcessAcknowledgePacket, - 18 => PacketType::TaskInfoPacket, - 19 => PacketType::TaskInfoAcknowledgePacket, + 12 => PacketType::FileTransferEndPacket, + 13 => PacketType::PerformancePacket, + 14 => PacketType::PerformanceAcknowledgePacket, + 15 => PacketType::ResultPacket, + 16 => PacketType::ResultAcknowledgePacket, + 17 => PacketType::StillProcessPacket, + 18 => PacketType::StillProcessAcknowledgePacket, + 19 => PacketType::TaskInfoPacket, + 20 => PacketType::TaskInfoAcknowledgePacket, _ => PacketType::BasePacket, } } diff --git a/Common/src/management/utils/file_header.rs b/Common/src/management/utils/file_header.rs index 0cafbd2..4758677 100644 --- a/Common/src/management/utils/file_header.rs +++ b/Common/src/management/utils/file_header.rs @@ -2,9 +2,9 @@ use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize, Clone)] pub struct FileHeader { - filename: String, - filesize: usize, - packet_count: usize, + pub filename: String, + pub filesize: usize, + pub packet_count: usize, } impl FileHeader { diff --git a/Common/src/management/utils/file_transfer_result.rs b/Common/src/management/utils/file_transfer_result.rs index e73991b..ea31b96 100644 --- a/Common/src/management/utils/file_transfer_result.rs +++ b/Common/src/management/utils/file_transfer_result.rs @@ -6,6 +6,12 @@ pub struct FileTransferResult { } impl FileTransferResult { + pub fn new(result: Option>) -> Self { + Self { + result, + } + } + pub fn into(self) -> Option> { self.result } diff --git a/Common/src/management/utils/task_info.rs b/Common/src/management/utils/task_info.rs index cde1211..51ea625 100644 --- a/Common/src/management/utils/task_info.rs +++ b/Common/src/management/utils/task_info.rs @@ -4,9 +4,9 @@ use crate::management::utils::model_type::ModelType; #[derive(Serialize, Deserialize, Clone)] pub struct TaskInfo { - uuid: Uuid, - model_filename: String, - model_type: ModelType, + pub uuid: Uuid, + pub model_filename: String, + pub model_type: ModelType, } impl TaskInfo { diff --git a/ManagementLibrary/src/connection/packet/file_transfer_end_packet.rs b/ManagementLibrary/src/connection/packet/file_transfer_end_packet.rs new file mode 100644 index 0000000..ef88419 --- /dev/null +++ b/ManagementLibrary/src/connection/packet/file_transfer_end_packet.rs @@ -0,0 +1,57 @@ +use crate::connection::packet::{Packet, PacketType, length_to_byte}; + +pub struct FileTransferEndPacket { + length: Vec, + id: Vec, + data: Vec, + packet_type: PacketType, +} + +impl FileTransferEndPacket { + pub fn new() -> Self { + Self { + length: length_to_byte(16), + id: PacketType::FileTransferEndPacket.as_byte(), + data: Vec::new(), + packet_type: PacketType::FileTransferEndPacket, + } + } +} + +impl Packet for FileTransferEndPacket { + fn as_length_byte(&self) -> &[u8] { + &self.length + } + + fn as_id_byte(&self) -> &[u8] { + &self.id + } + + fn as_data_byte(&self) -> &[u8] { + &self.data + } + + fn clone_length_byte(&self) -> Vec { + self.length.clone() + } + + fn clone_id_byte(&self) -> Vec { + self.id.clone() + } + + fn clone_data_byte(&self) -> Vec { + self.data.clone() + } + + fn data_to_string(&self) -> String { + String::from_utf8_lossy(&*self.data.clone()).to_string() + } + + fn packet_type(&self) -> PacketType { + self.packet_type + } + + fn equal(&self, packet_type: PacketType) -> bool { + self.packet_type.eq(&packet_type) + } +} diff --git a/ManagementLibrary/src/connection/packet/mod.rs b/ManagementLibrary/src/connection/packet/mod.rs index 8b05321..5892f4d 100644 --- a/ManagementLibrary/src/connection/packet/mod.rs +++ b/ManagementLibrary/src/connection/packet/mod.rs @@ -4,6 +4,7 @@ pub mod control_packet; pub mod data_channel_port_packet; pub mod file_body_packet; pub mod file_header_packet; +pub mod file_transfer_end_packet; pub mod performance_acknowledge_packet; pub mod result_acknowledge_packet; pub mod still_process_packet; diff --git a/ManagementLibrary/src/management/agent.rs b/ManagementLibrary/src/management/agent.rs index eee2e89..1b0eb7a 100644 --- a/ManagementLibrary/src/management/agent.rs +++ b/ManagementLibrary/src/management/agent.rs @@ -8,7 +8,6 @@ use tokio::{fs, select}; use tokio::io::AsyncReadExt; use tokio::net::TcpListener; use std::collections::VecDeque; -use std::sync::atomic::AtomicBool; use tokio::time::{sleep, Duration, Instant}; use crate::utils::logger::*; use crate::utils::config::Config; @@ -28,7 +27,6 @@ use crate::management::utils::performance::Performance; use crate::connection::packet::alive_packet::AlivePacket; use crate::connection::socket::socket_stream::SocketStream; use crate::connection::packet::control_packet::ControlPacket; -use crate::management::utils::prevent_reenter::PreventReenter; use crate::connection::packet::task_info_packet::TaskInfoPacket; use crate::connection::packet::file_body_packet::FileBodyPacket; use crate::management::utils::agent_information::AgentInformation; @@ -39,6 +37,7 @@ use crate::connection::packet::still_process_packet::StillProcessPacket; use crate::connection::channel::data_channel_receiver::DataChannelReceiver; use crate::connection::channel::control_channel_sender::ControlChannelSender; use crate::connection::packet::data_channel_port_packet::DataChannelPortPacket; +use crate::connection::packet::file_transfer_end_packet::FileTransferEndPacket; use crate::connection::channel::control_channel_receiver::ControlChannelReceiver; use crate::connection::packet::result_acknowledge_packet::ResultAcknowledgePacket; use crate::connection::packet::performance_acknowledge_packet::PerformanceAcknowledgePacket; @@ -281,15 +280,12 @@ impl Agent { } if timer.elapsed() > polling_times * polling_interval { let mut agent = agent.write().await; - match agent.data_channel_sender.as_mut() { - Some(data_channel_sender) => { - data_channel_sender.send(TaskInfoPacket::new(task_info_data.clone())).await; - polling_times += 1; - }, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, + if let Some(data_channel_sender) = agent.data_channel_sender.as_mut() { + data_channel_sender.send(TaskInfoPacket::new(task_info_data.clone())).await; + polling_times += 1; + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))? } } let mut agent = agent.write().await; @@ -319,8 +315,8 @@ impl Agent { async fn transfer_file(agent: Arc>, filename: &String, filepath: &PathBuf) -> Result<(), LogEntry> { Self::transfer_file_header(agent.clone(), filename, filepath).await?; - let sent_packets = Self::transfer_file_body(agent.clone(), filepath).await?; - Self::retransmit_file(agent, sent_packets).await + let file_body_packets = Self::read_file(agent.clone(), filepath).await?; + Self::transfer_file_body(agent, file_body_packets).await } async fn transfer_file_header(agent: Arc>, filename: &String, filepath: &PathBuf) -> Result<(), LogEntry> { @@ -346,44 +342,40 @@ impl Agent { } if timer.elapsed() > polling_times * polling_interval { let mut agent = agent.write().await; - match agent.data_channel_sender.as_mut() { - Some(data_channel_sender) => data_channel_sender.send(FileHeaderPacket::new(file_header_data.clone())).await, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, + if let Some(data_channel_sender) = agent.data_channel_sender.as_mut() { + data_channel_sender.send(FileHeaderPacket::new(file_header_data.clone())).await + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))? } polling_times += 1; } let mut agent = agent.write().await; - match agent.data_channel_receiver.as_mut() { - Some(data_channel_receiver) => { - select! { - packet = data_channel_receiver.file_header_acknowledge_packet.recv() => { - if packet.is_some() { - clear_unbounded_channel(&mut data_channel_receiver.file_header_acknowledge_packet).await; - return Ok(()) - } else { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Channel has been closed."))?; - } - }, - _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, - } - }, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, + if let Some(data_channel_receiver) = agent.data_channel_receiver.as_mut() { + select! { + packet = data_channel_receiver.file_header_acknowledge_packet.recv() => { + if packet.is_some() { + clear_unbounded_channel(&mut data_channel_receiver.file_header_acknowledge_packet).await; + return Ok(()) + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Channel has been closed."))?; + } + }, + _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, + } + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))? } } } - async fn transfer_file_body(agent: Arc>, filepath: &PathBuf) -> Result>, LogEntry> { + async fn read_file(agent: Arc>, filepath: &PathBuf) -> Result>, LogEntry> { let uuid = agent.read().await.uuid; let mut sequence_number = 0_usize; let mut buffer = vec![0; 1_048_576]; - let mut sent_packets = Vec::new(); + let mut packets = Vec::new(); let mut file = File::open(filepath.clone()).await .map_err(|err| error_entry!(format!("Agent: Cannot read file {filepath}.\nReason: {err}", filepath = filepath.display())))?; loop { @@ -393,89 +385,67 @@ impl Agent { let bytes_read = file.read(&mut buffer).await .map_err(|_| error_entry!(format!("Agent: An error occurred while reading file {filepath}.", filepath = filepath.display())))?; if bytes_read == 0 { - break; + return Ok(packets); } let mut data = sequence_number.to_be_bytes().to_vec(); data.extend_from_slice(&buffer[..bytes_read]); - let mut agent = agent.write().await; - match agent.data_channel_sender.as_mut() { - Some(data_channel_sender) => { - data_channel_sender.send(FileBodyPacket::new(data.clone())).await; - sent_packets.push(data); - sequence_number += 1; - }, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, - } + packets.push(data); + sequence_number += 1; } - Ok(sent_packets) } - #[allow(unused_assignments)] - async fn retransmit_file(agent: Arc>, sent_packets: Vec>) -> Result<(), LogEntry> { + async fn transfer_file_body(agent: Arc>, sent_packets: Vec>) -> Result<(), LogEntry> { let uuid = agent.read().await.uuid; let config = Config::now().await; + let mut require_send: Vec = (0..sent_packets.len()).collect(); let time = Instant::now(); let timeout_duration = Duration::from_secs(config.file_transfer_timeout); - let mut require_resend = Vec::new(); loop { - { - if AgentManager::get_state(uuid).await == AgentState::Terminate { - Err(info_entry!("Agent: Terminating. File transfer cancel."))?; + if AgentManager::get_state(uuid).await == AgentState::Terminate { + Err(info_entry!("Agent: Terminating. File transfer cancel."))?; + } + if time.elapsed() > timeout_duration { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: File transfer timeout."))?; + } + if let Some(data_channel_sender) = agent.write().await.data_channel_sender.as_mut() { + for chunk in &require_send { + if let Some(data) = sent_packets.get(*chunk) { + data_channel_sender.send(FileBodyPacket::new(data.clone())).await; + } else { + Err(warning_entry!("Agent: File body packet missing."))? + } } - if time.elapsed() > timeout_duration { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: File transfer timeout."))?; - } - let mut agent = agent.write().await; - match agent.data_channel_receiver.as_mut() { - Some(data_channel_receiver) => { - select! { - biased; - packet = data_channel_receiver.file_transfer_result_packet.recv() => { - match packet { - Some(packet) => { - clear_unbounded_channel(&mut data_channel_receiver.file_transfer_result_packet).await; - let file_transfer_result = serde_json::from_slice::(packet.as_data_byte()) - .map_err(|_| error_entry!("Agent: Unable to parse file transfer result."))?; - match file_transfer_result.into() { - Some(missing_chunks) => require_resend = missing_chunks, - None => return Ok(()), - } - }, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Channel has been closed."))?; - }, + data_channel_sender.send(FileTransferEndPacket::new()).await; + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))? + } + if let Some(data_channel_receiver) = agent.write().await.data_channel_receiver.as_mut() { + select! { + biased; + packet = data_channel_receiver.file_transfer_result_packet.recv() => { + match packet { + Some(packet) => { + clear_unbounded_channel(&mut data_channel_receiver.file_transfer_result_packet).await; + let file_transfer_result = serde_json::from_slice::(packet.as_data_byte()) + .map_err(|_| error_entry!("Agent: Unable to parse file transfer result."))?; + match file_transfer_result.into() { + Some(missing_chunks) => require_send = missing_chunks, + None => return Ok(()), } }, - _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, + None => { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Channel has been closed."))?; + }, } }, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, - } - } - for missing_chunk in &require_resend { - if let Some(data) = sent_packets.get(*missing_chunk) { - if AgentManager::get_state(uuid).await == AgentState::Terminate { - Err(info_entry!("Agent: Terminating. File transfer cancel."))?; - } - let mut agent = agent.write().await; - match agent.data_channel_sender.as_mut() { - Some(data_channel_sender) => data_channel_sender.send(FileBodyPacket::new(data.clone())).await, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))? - }, - } - } else { - Err(error_entry!("Agent: File block missing."))? + _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, } + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))? } } } @@ -483,13 +453,12 @@ impl Agent { async fn waiting_result(agent: Arc>, image_task: &mut ImageTask) -> Result<(), LogEntry> { let uuid = agent.read().await.uuid; let config = Config::now().await; - let mut agent_process_error: Option = None; let mut polling_times = 0_u32; let polling_timer = Instant::now(); let polling_interval = Duration::from_millis(config.polling_interval); let mut timeout_timer = Instant::now(); let timeout_duration = Duration::from_secs(config.control_channel_timeout); - loop { + let bounding_box = loop { if AgentManager::get_state(uuid).await == AgentState::Terminate { Err(info_entry!("Agent: Terminating. Interrupt task processing."))?; } @@ -498,16 +467,15 @@ impl Agent { Err(warning_entry!("Agent: Data Channel timeout."))?; } if polling_timer.elapsed() > polling_times * polling_interval { - match &mut agent.write().await.data_channel_sender { - Some(data_channel_sender) => data_channel_sender.send(StillProcessPacket::new()).await, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))?; - } + if let Some(data_channel_sender) = agent.write().await.data_channel_sender.as_mut() { + data_channel_sender.send(StillProcessPacket::new()).await + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))?; } polling_times += 1; } - if let Some(data_channel_receiver) = &mut agent.write().await.data_channel_receiver { + if let Some(data_channel_receiver) = agent.write().await.data_channel_receiver.as_mut() { select! { biased; packet = data_channel_receiver.still_process_acknowledge_packet.recv() => { @@ -523,16 +491,7 @@ impl Agent { if let Some(packet) = &packet { clear_unbounded_channel(&mut data_channel_receiver.result_packet).await; if let Ok(task_result) = serde_json::from_slice::(packet.as_data_byte()) { - match task_result.into() { - Ok(bounding_box) => { - image_task.bounding_boxes = bounding_box; - break; - }, - Err(err) => { - agent_process_error = Some(err); - break; - }, - } + break task_result.into(); } else { Err(error_entry!("Agent: Unable to parse task result."))?; } @@ -547,17 +506,16 @@ impl Agent { AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; Err(warning_entry!("Agent: Data Channel is not available."))?; } + }; + if let Some(data_channel_sender) = agent.write().await.data_channel_sender.as_mut() { + data_channel_sender.send(ResultAcknowledgePacket::new()).await + } else { + AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; + Err(warning_entry!("Agent: Data Channel is not available."))?; } - match &mut agent.write().await.data_channel_sender { - Some(data_channel_sender) => data_channel_sender.send(ResultAcknowledgePacket::new()).await, - None => { - AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; - Err(warning_entry!("Agent: Data Channel is not available."))?; - } - } - if let Err(err) = agent_process_error { - Err(error_entry!(format!("Agent: An error occurred while processing.\nReason: {err}")))?; - } + let bounding_box = bounding_box + .map_err(|err| error_entry!(format!("Agent: An error occurred while processing.\nReason: {err}")))?; + image_task.bounding_boxes = bounding_box; Ok(()) } @@ -584,7 +542,7 @@ impl Agent { return; } if timer.elapsed() > polling_times * polling_interval { - if let Some(data_channel_sender) = &mut agent.write().await.data_channel_sender { + if let Some(data_channel_sender) = agent.write().await.data_channel_sender.as_mut() { data_channel_sender.send(AlivePacket::new()).await } else { AgentManager::store_state(uuid, AgentState::CreateDataChannel).await; @@ -593,7 +551,7 @@ impl Agent { } polling_times += 1; } - if let Some(data_channel_receiver) = &mut agent.write().await.data_channel_receiver { + if let Some(data_channel_receiver) = agent.write().await.data_channel_receiver.as_mut() { select! { biased; packet = data_channel_receiver.alive_acknowledge_packet.recv() => { @@ -674,7 +632,7 @@ impl Agent { connection = listener.accept() => { match connection { Ok(connection) => break connection, - Err(err) => Err(error_entry!(format!("Agent: Failed to establish connection.\nReason: {}", err)))? + Err(err) => Err(error_entry!(format!("Agent: Failed to establish connection.\nReason: {}", err)))?, } }, _ = sleep(Duration::from_millis(config.internal_timestamp)) => continue, @@ -690,11 +648,6 @@ impl Agent { } pub async fn terminate(agent: Arc>) { - static TERMINATING_PROCESSING: AtomicBool = AtomicBool::new(false); - let prevent_reenter = PreventReenter::new(&TERMINATING_PROCESSING); - if prevent_reenter.is_none() { - return; - } let uuid = agent.read().await.uuid; logging_info!(uuid, "Agent: Terminating agent."); let image_task = {