use apalis_core::backend::Backend;
use apalis_core::backend::BackendConfig;
use apalis_core::backend::WireFormatBackend;
use apalis_core::backend::codec::Codec;
use apalis_core::task::builder::TaskBuilder;
use apalis_core::task::metadata::Meta;
use apalis_core::task::status::Status;
use apalis_core::task::task_id::GenerateId;
use apalis_core::worker::service::IntoWorkerService;
use apalis_core::worker::service::WorkerService;
use apalis_core::{
backend::WaitForCompletion,
error::BoxDynError,
task::{Task, task_id::TaskId},
};
use futures_util::future::BoxFuture;
use futures_util::future::try_join_all;
use futures_util::{FutureExt, Sink, SinkExt, StreamExt};
use petgraph::Direction;
use petgraph::graph::DiGraph;
use petgraph::graph::NodeIndex;
use serde_json::Value;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::fmt::{Debug, Display};
use std::str::FromStr;
use std::task::Poll;
use tower::Service;
use crate::GraphFlow;
use crate::NodeService;
use crate::graph::NodeInput;
use crate::graph::context::GraphFlowContext;
use crate::graph::error::{GraphFlowError, GraphServiceError};
use crate::graph::response::GraphNodeResponse;
pub struct RootGraphService<B>
where
B: Backend + WireFormatBackend,
{
pub(super) graph: DiGraph<NodeService<B::Compact>, ()>,
pub(super) node_mapping: HashMap<String, NodeIndex>,
pub(super) topological_order: Vec<NodeIndex>,
pub(super) start_nodes: Vec<NodeIndex>,
pub(super) end_nodes: Vec<NodeIndex>,
pub(super) not_ready: VecDeque<NodeIndex>,
pub(super) backend: B,
}
impl<B> std::fmt::Debug for RootGraphService<B>
where
B: Backend + WireFormatBackend,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RootGraphService")
.field("executor", &"<GraphExecutor>")
.field("backend", &"<Backend>")
.finish()
}
}
impl<B> RootGraphService<B>
where
B: Backend + WireFormatBackend,
{
pub fn get_node_by_name_mut(&mut self, name: &str) -> Option<&mut NodeService<B::Compact>> {
self.node_mapping
.get(name)
.and_then(|&idx| self.graph.node_weight_mut(idx))
}
async fn handle_task(
graph: &mut DiGraph<NodeService<B::Compact>, ()>,
task: Task<B::Compact>,
) -> Result<(B::Compact, Value), GraphFlowError>
where
B::Compact: Send + Sync,
{
let context = task
.extract::<Meta<GraphFlowContext>>()
.await
.map_err(|e| GraphFlowError::Metadata(e.into()))?
.0;
let service = graph
.node_weight_mut(context.current_node)
.ok_or_else(|| GraphFlowError::MissingService(context.current_node))?;
let result = service
.call(task.map_args(NodeInput::Single))
.await
.map_err(GraphFlowError::NodeExecutionError)?;
Ok(result)
}
async fn handle_fan_in(
graph: &mut DiGraph<NodeService<B::Compact>, ()>,
task: Task<Vec<Value>>,
) -> Result<(B::Compact, Value), GraphFlowError>
where
B::Compact: Send + Sync,
{
let context = task
.extract::<Meta<GraphFlowContext>>()
.await
.map_err(|e| GraphFlowError::Metadata(e.into()))?
.0;
let service = graph
.node_weight_mut(context.current_node)
.ok_or_else(|| GraphFlowError::MissingService(context.current_node))?;
let result = service
.call(task.map_args(NodeInput::FanIn))
.await
.map_err(GraphFlowError::NodeExecutionError)?;
Ok(result)
}
}
impl<B> Clone for RootGraphService<B>
where
B: Backend + WireFormatBackend + Clone,
{
fn clone(&self) -> Self {
Self {
graph: self.graph.clone(),
node_mapping: self.node_mapping.clone(),
topological_order: self.topological_order.clone(),
start_nodes: self.start_nodes.clone(),
end_nodes: self.end_nodes.clone(),
not_ready: self.not_ready.clone(),
backend: self.backend.clone(),
}
}
}
fn find_designated_fan_in_handler(
incoming_nodes: &[NodeIndex],
) -> Result<&NodeIndex, GraphFlowError> {
let designated_handler = incoming_nodes.iter().max_by_key(|n| n.index());
designated_handler.ok_or(GraphFlowError::Service(
GraphServiceError::MissingFaninHandler,
))
}
impl<B, Err, CdcErr, Id, Compact> Service<Task<Compact>> for RootGraphService<B>
where
B: Backend<Error = Err>
+ BackendConfig<Id = Id>
+ WireFormatBackend<Compact = Compact>
+ Send
+ Sync
+ 'static
+ Clone
+ WaitForCompletion<GraphNodeResponse>,
Id: GenerateId + Send + Sync + 'static + PartialEq + Debug + FromStr + Display,
Compact: Send + Sync + 'static + Clone,
Err: std::error::Error + Send + Sync + 'static,
B: Sink<Task<B::Compact>, Error = Err> + Unpin,
B::Codec:
Codec<Vec<Compact>, Compact = Compact, Error = CdcErr> + Send + Sync + Clone + 'static,
CdcErr: Into<BoxDynError>,
<Id as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
type Response = GraphNodeResponse;
type Error = GraphFlowError;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
loop {
if self.not_ready.is_empty() {
return Poll::Ready(Ok(()));
} else {
if self
.graph
.node_weight_mut(self.not_ready[0])
.ok_or(GraphFlowError::MissingService(self.not_ready[0]))?
.poll_ready(cx)
.map_err(GraphServiceError::PollError)
.map_err(GraphFlowError::Service)?
.is_pending()
{
return Poll::Pending;
}
self.not_ready.pop_front();
}
}
}
fn call(&mut self, mut req: Task<B::Compact>) -> Self::Future {
let backend = self.backend.clone();
let codec = backend.codec().clone();
let start_nodes = self.start_nodes.clone();
let end_nodes = self.end_nodes.clone();
let mut graph = self.graph.clone();
let mut backend = self.backend.clone();
req.inject_data(codec);
async move {
let ctx = req.extract::<Meta<GraphFlowContext>>().await;
let ((compact, response), context) = if let Ok(Meta(context)) = ctx {
#[cfg(feature = "tracing")]
tracing::debug!(
node = ?context.current_node,
"Extracted GraphFlowContext for task"
);
let incoming_nodes = graph
.neighbors_directed(context.current_node, Direction::Incoming)
.collect::<Vec<_>>();
match incoming_nodes.len() {
0 if start_nodes.len() == 1 => {
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Found a single entry node"
);
let response = Self::handle_task(&mut graph, req).await?;
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Found a single entry node: done"
);
(response, context)
}
0 if start_nodes.len() > 1 => {
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Entry node with multiple start node"
);
let response = Self::handle_task(&mut graph, req).await?;
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Entry node with multiple start node: done"
);
(response, context)
}
1 => {
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Single incoming node"
);
let response = Self::handle_task(&mut graph, req).await?;
#[cfg(feature = "tracing")]
tracing::trace!(
node = ?context.current_node,
"Single incoming node complete"
);
(response, context)
}
_ => {
#[cfg(feature = "tracing")]
tracing::debug!(
node = ?context.current_node,
"Multiple incoming nodes, fan-in scenario"
);
let dependency_task_ids = context.get_dependency_task_ids(&incoming_nodes);
let prev_node = context.prev_node.ok_or(GraphFlowError::Service(
GraphServiceError::MissingPreviousNode,
))?;
let fan_in_node = *find_designated_fan_in_handler(&incoming_nodes)?;
#[cfg(feature = "tracing")]
tracing::debug!(
prev_node = ?prev_node,
node = ?context.current_node,
deps = ?dependency_task_ids.len(),
fan_in = ?fan_in_node,
"Fanning in from multiple dependencies",
);
if fan_in_node != prev_node {
return Ok(GraphNodeResponse::WaitingForDependencies {
pending_dependencies: dependency_task_ids,
});
}
#[cfg(feature = "tracing")]
tracing::debug!(
prev_node = ?prev_node,
node = ?context.current_node,
deps = ?dependency_task_ids.len(),
fan_in = ?fan_in_node,
"designated_fan_in_handler: Waiting for other dependencies",
);
let results = backend
.wait_for(dependency_task_ids.values().cloned().collect::<Vec<_>>())
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<Result<Vec<_>, _>>()
.map_err(|e| GraphFlowError::Backend(e.into()))?;
#[cfg(feature = "tracing")]
tracing::debug!(
prev_node = ?prev_node,
node = ?context.current_node,
deps = ?dependency_task_ids.len(),
results = ?results.len(),
fan_in = ?fan_in_node,
"Found results",
);
if results.iter().all(|s| matches!(s.status, Status::Done)) {
let sorted_results = {
let res = incoming_nodes
.iter()
.rev()
.map(|node_index| {
let task_id = context
.node_task_ids
.iter()
.find(|(n, _)| *n == node_index)
.map(|(_, task_id)| task_id)
.ok_or(GraphFlowError::Service(
GraphServiceError::MissingIncomingTaskId,
))?;
let task_result = results
.iter()
.find(|r| &r.task_id == task_id)
.ok_or(GraphFlowError::Service(
GraphServiceError::MissingTaskIdResult(format!(
"{task_id:?}"
)),
))?;
Ok(task_result)
})
.collect::<Result<Vec<_>, GraphFlowError>>();
match res {
Ok(v) => v,
Err(e) => {
#[cfg(feature = "tracing")]
tracing::error!(
node = ?context.current_node,
error = ?e,
"Encountered an error resolving result",
);
return Ok(GraphNodeResponse::WaitingForDependencies {
pending_dependencies: dependency_task_ids,
});
}
}
};
let res = sorted_results
.iter()
.map(|s| match &s.result {
Ok(val) => match val {
GraphNodeResponse::FanOut { response, .. } => {
Ok(response.clone())
}
GraphNodeResponse::EnqueuedNext { result }
| GraphNodeResponse::Complete { result } => {
Ok(result.clone())
}
_ => Err(GraphFlowError::Service(
GraphServiceError::InvalidFanInDependencyResult,
)),
},
Err(e) => Err(GraphFlowError::Service(
GraphServiceError::DependencyTaskFailed(e.as_str().into()),
)),
})
.collect::<Result<Vec<_>, _>>()?;
let req = req.map_args(|_| res); let response = Self::handle_fan_in(&mut graph, req).await?;
(response, context)
} else {
return Err(GraphFlowError::Service(
GraphServiceError::DependencyTaskFailed(
"An adjacent node failed. Terminating".into(),
),
));
}
}
}
} else {
#[cfg(feature = "tracing")]
tracing::debug!("Extracting GraphFlowContext for task without meta");
if start_nodes.len() == 1 {
#[cfg(feature = "tracing")]
tracing::debug!("Single start node detected, proceeding with execution");
let context = GraphFlowContext::new(req.task_id().cloned());
req.inject_metadata(&context)?;
let response = Self::handle_task(&mut graph, req).await?;
#[cfg(feature = "tracing")]
tracing::debug!(node = ?context.current_node, "Execution complete at node");
(response, context)
} else {
#[cfg(feature = "tracing")]
tracing::debug!("Multiple nodes detected, proceeding with fan_out_entry_nodes");
let new_node_task_ids = fan_out_entry_nodes(
&backend,
&start_nodes,
&GraphFlowContext::new(req.task_id().cloned()),
&req.args,
)
.await?;
return Ok(GraphNodeResponse::EntryFanOut {
node_task_ids: new_node_task_ids,
});
}
};
let current_node = context.current_node;
let outgoing_nodes = graph
.neighbors_directed(current_node, Direction::Outgoing)
.collect::<Vec<_>>();
match outgoing_nodes.len() {
0 => {
assert!(
end_nodes.contains(¤t_node),
"Current node is not an end node"
);
return Ok(GraphNodeResponse::Complete { result: response });
}
1 => {
let next_node = outgoing_nodes[0];
let mut new_context = context.clone();
new_context.prev_node = Some(current_node);
new_context.current_node = next_node;
new_context.current_position += 1;
new_context.is_initial = false;
let task = TaskBuilder::new(compact)
.task_id(B::Id::generate())
.metadata(&new_context)
.build();
backend
.send(task)
.await
.map_err(|e| GraphFlowError::Backend(e.into()))?;
}
_ => {
let mut new_context = context.clone();
new_context.prev_node = Some(current_node);
new_context.current_position += 1;
new_context.is_initial = false;
let next_task_ids =
fan_out_next_nodes(&backend, outgoing_nodes, &new_context, &compact)
.await?;
return Ok(GraphNodeResponse::FanOut {
response,
node_task_ids: next_task_ids,
});
}
}
Ok(GraphNodeResponse::EnqueuedNext { result: response })
}
.boxed()
}
}
async fn fan_out_next_nodes<B, Err, CdcErr>(
backend: &B,
outgoing_nodes: Vec<NodeIndex>,
context: &GraphFlowContext,
input: &B::Compact,
) -> Result<HashMap<NodeIndex, TaskId>, GraphFlowError>
where
B::Id: GenerateId + Send + Sync + 'static + PartialEq,
B::Compact: Send + Sync + 'static + Clone,
B: Sink<Task<B::Compact>, Error = Err> + BackendConfig + Unpin,
Err: std::error::Error + Send + Sync + 'static,
B: Backend<Error = Err> + WireFormatBackend + Send + Sync + 'static + Clone,
B::Codec: Codec<Vec<B::Compact>, Compact = B::Compact, Error = CdcErr>,
CdcErr: Into<BoxDynError>,
B::Id: FromStr + Display,
<B::Id as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
let mut enqueue_futures = vec![];
let next_nodes = outgoing_nodes
.iter()
.map(|node| (*node, B::Id::generate()))
.collect::<HashMap<NodeIndex, TaskId>>();
let mut node_task_ids = next_nodes.clone();
node_task_ids.extend(context.node_task_ids.clone());
for outgoing_node in outgoing_nodes.into_iter() {
let task_id = next_nodes
.get(&outgoing_node)
.ok_or(GraphFlowError::Service(GraphServiceError::MissingNextNode))?
.clone();
let task = TaskBuilder::new(input.clone())
.task_id(task_id)
.metadata(&GraphFlowContext {
prev_node: context.prev_node,
current_node: outgoing_node,
completed_nodes: context.completed_nodes.clone(),
node_task_ids: node_task_ids.clone(),
current_position: context.current_position + 1,
is_initial: context.is_initial,
root_task_id: context.root_task_id.clone(),
})
.build();
let mut b = backend.clone();
enqueue_futures.push(
async move {
b.send(task)
.await
.map_err(|e| GraphFlowError::Backend(e.into()))?;
Ok::<(), GraphFlowError>(())
}
.boxed(),
);
}
try_join_all(enqueue_futures).await?;
Ok(next_nodes)
}
async fn fan_out_entry_nodes<B, Err, CdcErr>(
backend: &B,
start_nodes: &[NodeIndex],
context: &GraphFlowContext,
input: &B::Compact,
) -> Result<HashMap<NodeIndex, TaskId>, GraphFlowError>
where
B::Id: GenerateId + Send + Sync + 'static + PartialEq + Debug,
B::Compact: Send + Sync + 'static + Clone,
B: Sink<Task<B::Compact>, Error = Err> + Unpin,
Err: std::error::Error + Send + Sync + 'static,
B: Backend<Error = Err> + WireFormatBackend + BackendConfig + Send + Sync + 'static + Clone,
B::Codec: Codec<Vec<B::Compact>, Compact = B::Compact, Error = CdcErr> + Clone,
CdcErr: Into<BoxDynError>,
B::Id: FromStr + Display,
<B::Id as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
let codec = backend.codec().clone();
let values: Vec<B::Compact> =
B::Codec::decode(&codec, input).map_err(|e: CdcErr| GraphFlowError::Codec(e.into()))?;
if values.len() != start_nodes.len() {
return Err(GraphFlowError::InputCountMismatch {
expected: start_nodes.len(),
actual: values.len(),
});
}
let mut enqueue_futures = vec![];
let next_nodes = start_nodes
.iter()
.map(|node| (*node, B::Id::generate()))
.collect::<HashMap<NodeIndex, TaskId>>();
let mut node_task_ids = next_nodes.clone();
node_task_ids.extend(context.node_task_ids.clone());
for (outgoing_node, input) in start_nodes.iter().zip(values) {
let task_id = next_nodes
.get(outgoing_node)
.ok_or(GraphFlowError::Service(GraphServiceError::MissingNextNode))?;
let task = TaskBuilder::new(input)
.task_id(task_id.clone())
.metadata(&GraphFlowContext {
prev_node: None,
current_node: *outgoing_node,
completed_nodes: Default::default(),
node_task_ids: node_task_ids.clone(),
current_position: context.current_position,
is_initial: true,
root_task_id: context.root_task_id.clone(),
})
.build();
let mut b = backend.clone();
enqueue_futures.push(
async move {
b.send(task)
.await
.map_err(|e| GraphFlowError::Backend(BoxDynError::from(e)))?;
Ok::<(), GraphFlowError>(())
}
.boxed(),
);
}
try_join_all(enqueue_futures).await?;
Ok(next_nodes)
}
impl<B, Compact, Err> IntoWorkerService<B, RootGraphService<B>> for GraphFlow<B>
where
B: Backend<Error = Err, Task = Task<Compact>> + WireFormatBackend<Compact = Compact> + Clone,
Err: std::error::Error + Send + Sync + 'static,
B::Compact: Send + Sync + 'static + Clone,
RootGraphService<B>: Service<Task<Compact>>,
{
type Task = Task<Compact>;
type Backend = B;
fn into_service(self, b: B) -> WorkerService<B, RootGraphService<B>> {
let service = self.build(b.clone()).expect("Execution should be valid");
WorkerService {
backend: b,
service,
}
}
}