use async_broadcast::Sender;
use futures::join;
use crate::definitions::task as definitions_task;
use crate::delegation::Task;
use crate::knowledge_base::KnowledgeBaseInterface;
use crate::simulation::tst::SimulationResult;
use crate::{agent, consts, execution, module, uuid, Error, Result};
mod task_executor;
pub trait TaskExecutor: Sync + Send + 'static
{
fn execute_task<'a>(
&'a self,
task: definitions_task::Task,
) -> futures::future::BoxFuture<'a, Result<()>>;
fn can_execute(&self, task: &definitions_task::Task) -> bool;
}
pub struct Options
{
task_executor: Box<dyn TaskExecutor>,
}
impl Options
{
pub fn new<TTaskExecutor>(task_executor: TTaskExecutor) -> Self
where
TTaskExecutor: TaskExecutor,
{
Self {
task_executor: Box::new(task_executor),
}
}
}
module::create_module_private_interface!(
ModulePrivateInterface,
execution::InputMessage,
execution::OutputMessage
);
struct EndOfQueue
{
total_cost: f32,
final_states: crate::states::States,
queue_length: usize,
}
struct ModuleData
{
cancelled_tasks: ccutils::sync::ArcMutex<Vec<uuid::Uuid>>,
end_of_queue: ccutils::sync::ArcMutex<EndOfQueue>,
task_sender: Sender<definitions_task::Task>,
}
pub struct Module {}
impl Module
{
async fn handle_input_message(
msg: execution::InputMessage,
agent_data: &agent::AgentData,
module_data: &ModuleData,
output_sender: &Sender<execution::OutputMessage>,
) -> Result<()>
{
match msg
{
execution::InputMessage::QueueExecution { uuid } =>
{
let sr: SimulationResult = agent_data
.knowledge_base
.retrieve("task_simulation_result", uuid.to_hex())?;
let task = agent_data
.knowledge_base
.retrieve::<definitions_task::Task>("tasks", uuid.to_hex())?;
{
let (current_cost, final_states) = {
let mut end_of_queue = module_data.end_of_queue.lock()?;
end_of_queue.total_cost += sr.get_estimated_cost();
end_of_queue.final_states = sr.get_final_states().to_owned();
(
end_of_queue.total_cost,
end_of_queue.final_states.to_owned(),
)
};
ccutils::log::log_error!(
output_sender
.broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
current_cost,
final_states,
queue_length: module_data.task_sender.len(),
})
.await,
"sending current cost"
);
}
ccutils::log::log_error!(
module_data.task_sender.broadcast(task).await,
"sending task"
);
}
execution::InputMessage::CancelExecution { uuid } =>
{
module_data.cancelled_tasks.lock()?.push(uuid);
}
}
Ok(())
}
fn handle_task_execution<'a>(
task: definitions_task::Task,
cancelled_tasks: &ccutils::sync::ArcMutex<Vec<uuid::Uuid>>,
task_executor: &'a dyn TaskExecutor,
) -> Result<Option<futures::future::BoxFuture<'a, Result<()>>>>
{
if cancelled_tasks.lock()?.contains(&task.task_id())
{
Ok(None)
}
else if task_executor.can_execute(&task)
{
Ok(Some(task_executor.execute_task(task)))
}
else
{
Err(Error::NoExecutor())
}
}
}
impl execution::Module for Module
{
type Options = Options;
async fn start(
agent_data: agent::AgentData,
module_interfaces: (
crate::module::ModuleInterface<execution::InputMessage, execution::OutputMessage>,
ModulePrivateInterface,
),
options: Options,
)
{
let (_, module_private_interface) = module_interfaces;
let mut input_receiver = module_private_interface.input_receiver.activate();
let output_sender = module_private_interface.output_sender.clone();
let (task_sender, mut task_receiver) =
async_broadcast::broadcast::<definitions_task::Task>(consts::TASK_QUEUE_LENGTH);
let cancelled_tasks: ccutils::sync::ArcMutex<Vec<uuid::Uuid>> = Default::default();
let end_of_queue: ccutils::sync::ArcMutex<_> = EndOfQueue {
total_cost: 0.0,
final_states: agent_data.states.to_owned_states().unwrap(),
queue_length: 0,
}
.into();
let msg_fut = async {
let module_data = ModuleData {
cancelled_tasks: cancelled_tasks.to_owned(),
end_of_queue: end_of_queue.to_owned(),
task_sender,
};
loop
{
let msg = input_receiver.recv().await;
if let Ok(msg) = msg
{
if let Err(e) =
Self::handle_input_message(msg, &agent_data, &module_data, &output_sender).await
{
log::error!(
"An error occured when handling input execution message: {} for agent {}",
e,
agent_data.agent_uri
);
}
}
else
{
return;
}
}
};
let output_sender = module_private_interface.output_sender.clone();
let agent_data = agent_data.clone();
let cancelled_tasks = cancelled_tasks.clone();
let end_of_queue = end_of_queue.clone();
let exec_fut = async move {
loop
{
let msg = task_receiver.recv().await;
if let Ok(msg) = msg
{
let sr = agent_data
.knowledge_base
.retrieve::<SimulationResult>("task_simulation_result", msg.task_id().to_hex());
let cost = match sr
{
Ok(sr) => sr.get_estimated_cost(),
Err(_) => 0.0,
};
match Self::handle_task_execution(msg, &cancelled_tasks, &*options.task_executor)
{
Ok(Some(fut)) =>
{
if let Err(e) = fut.await
{
log::error!(
"An error occured when executing a task: {} for agent {}",
e,
agent_data.agent_uri
);
}
}
Ok(None) =>
{}
Err(e) =>
{
log::error!(
"An error occured when handling task execution message: {} for agent {}",
e,
agent_data.agent_uri
);
}
}
if task_receiver.is_empty()
{
{
let mut end_of_queue = end_of_queue.lock().unwrap();
end_of_queue.total_cost = 0.0;
end_of_queue.final_states = agent_data.states.to_owned_states().unwrap();
end_of_queue.queue_length = 0;
}
ccutils::log::log_error!(
output_sender
.broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
current_cost: 0.0,
final_states: agent_data.states.to_owned_states().unwrap(),
queue_length: 0,
})
.await,
"sending output message"
);
}
else
{
let (current_cost, final_states, queue_length) = {
let mut end_of_queue = end_of_queue.lock().unwrap();
end_of_queue.total_cost -= cost;
end_of_queue.queue_length = task_receiver.len();
(
end_of_queue.total_cost,
end_of_queue.final_states.to_owned(),
end_of_queue.queue_length,
)
};
ccutils::log::log_error!(
output_sender
.broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
current_cost,
final_states,
queue_length,
})
.await,
"sending output message"
);
}
}
else
{
return;
}
}
};
join!(msg_fut, exec_fut);
}
}
impl module::Module for Module
{
type InputMessage = execution::InputMessage;
type OutputMessage = execution::OutputMessage;
type ModulePrivateInterface = ModulePrivateInterface;
}