agent_tk/execution/

default.rs

//! This module defines a default execution module.

use async_broadcast::Sender;
use ccutils::sync::ArcMutex;
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, decision, delegation, execution, module, utils, uuid, Error, Result};

mod task_executor;

//  _____         _    _____                     _
// |_   _|_ _ ___| | _| ____|_  _____  ___ _   _| |_ ___  _ __
//   | |/ _` / __| |/ /  _| \ \/ / _ \/ __| | | | __/ _ \| '__|
//   | | (_| \__ \   <| |___ >  |  __/ (__| |_| | || (_) | |
//   |_|\__,_|___/_|\_\_____/_/\_\___|\___|\__,_|\__\___/|_|

/// Interface to a task executor
pub trait TaskExecutor: Sync + Send + 'static
{
  /// Execute the given task. The function is expected to block until the execution is completed.
  /// It will be called from an async context by the default executor.
  fn execute_task(
    &self,
    task: definitions_task::Task,
  ) -> futures::future::BoxFuture<'static, Result<()>>;

  //  impl std::future::Future<Output = Result<()>> + std::marker::Send + 'static;
  /// Check if the given executor can execute the task
  fn can_execute(&self, task: &definitions_task::Task) -> bool;
}

//   ____        _   _
//  / __ \      | | (_)
//  | |  | |_ __ | |_ _  ___  _ __  ___
//  | |  | | '_ \| __| |/ _ \| '_ \/ __|
//  | |__| | |_) | |_| | (_) | | | \__ \
//   \____/| .__/ \__|_|\___/|_| |_|___/
//         | |
//         |_|

/// Options for the default execution module
pub struct Options
{
  task_executor: Box<dyn TaskExecutor>,
}

impl Options
{
  /// Create options for the default execution module
  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>,
}

/// This structure implements a execution module
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(),
            )
          };
          output_sender
            .broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
              current_cost,
              final_states,
              queue_length: module_data.task_sender.len(),
            })
            .await;
        }
        module_data.task_sender.broadcast(task).await;
      }
      execution::InputMessage::CancelExecution { uuid } =>
      {
        module_data.cancelled_tasks.lock()?.push(uuid);
      }
    }
    Ok(())
  }
  fn handle_task_execution(
    task: definitions_task::Task,
    cancelled_tasks: &ccutils::sync::ArcMutex<Vec<uuid::Uuid>>,
    task_executor: &dyn TaskExecutor,
  ) -> Result<Option<futures::future::BoxFuture<'static, 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_interface, 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;
            }
            output_sender
              .broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
                current_cost: 0.0,
                final_states: agent_data.states.to_owned_states().unwrap(),
                queue_length: 0,
              })
              .await;
          }
          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,
              )
            };
            output_sender
              .broadcast(execution::OutputMessage::CurrentEstimatedIdlingState {
                current_cost,
                final_states,
                queue_length,
              })
              .await;
          }
        }
        else
        {
          return;
        }
      }
    };

    join!(msg_fut, exec_fut);
  }
}

impl module::Module for Module
{
  type InputMessage = execution::InputMessage;
  type OutputMessage = execution::OutputMessage;
  type ModulePrivateInterface = ModulePrivateInterface;
}