use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use crate::error::Result;
use crate::future::WakerRegistry;
use crate::ring::batch::{BatchResult, OperationResult};
use crate::ring::Ring;
pub struct StandaloneBatchFuture {
operation_ids: Vec<Option<u64>>,
results: Vec<Option<OperationResult>>,
dependencies: HashMap<usize, Vec<usize>>,
fail_fast: bool,
completed: bool,
id_to_index: HashMap<u64, usize>,
waker_registry: Arc<WakerRegistry>,
}
impl StandaloneBatchFuture {
pub(crate) fn new(
operation_ids: Vec<Option<u64>>,
dependencies: HashMap<usize, Vec<usize>>,
waker_registry: Arc<WakerRegistry>,
fail_fast: bool,
) -> Self {
let operation_count = operation_ids.len();
let results = (0..operation_count).map(|_| None).collect();
let mut id_to_index = HashMap::new();
for (index, id_opt) in operation_ids.iter().enumerate() {
if let Some(id) = id_opt {
id_to_index.insert(*id, index);
}
}
Self {
operation_ids,
results,
dependencies,
fail_fast,
completed: false,
id_to_index,
waker_registry,
}
}
pub fn poll_with_ring(
&mut self,
ring: &mut Ring<'_>,
cx: &mut Context<'_>,
) -> Poll<Result<BatchResult>> {
if self.completed {
let results: Vec<OperationResult> = self
.results
.iter()
.map(|opt| opt.as_ref().cloned().unwrap_or(OperationResult::Cancelled))
.collect();
return Poll::Ready(Ok(BatchResult::new(results)));
}
match self.poll_completions_with_ring(ring, cx) {
Poll::Ready(Ok(())) => {
let results: Vec<OperationResult> = self
.results
.iter()
.map(|opt| opt.as_ref().cloned().unwrap_or(OperationResult::Cancelled))
.collect();
Poll::Ready(Ok(BatchResult::new(results)))
}
Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
Poll::Pending => Poll::Pending,
}
}
fn poll_completions_with_ring(
&mut self,
ring: &mut Ring<'_>,
cx: &mut Context<'_>,
) -> Poll<Result<()>> {
let mut any_completed = false;
let mut any_failed = false;
let mut completed_operations = Vec::new();
match ring.try_complete() {
Ok(completions) => {
for completion in completions {
let operation_id = completion.id();
if let Some(&index) = self.id_to_index.get(&operation_id) {
if self.results[index].is_some() {
continue; }
let result = match completion.result() {
Ok(bytes) => OperationResult::Success(*bytes),
Err(e) => {
let error_msg = e.to_string();
OperationResult::Error(error_msg)
}
};
let is_error = matches!(result, OperationResult::Error(_));
self.results[index] = Some(result);
any_completed = true;
if is_error {
any_failed = true;
}
completed_operations.push(index);
}
}
}
Err(e) => {
return Poll::Ready(Err(e));
}
}
for completed_index in completed_operations {
self.check_ready_operations(completed_index);
if self.fail_fast
&& matches!(
self.results[completed_index],
Some(OperationResult::Error(_))
)
{
self.cancel_dependent_operations(completed_index);
}
}
if self.fail_fast && any_failed {
self.cancel_all_remaining_operations();
return Poll::Ready(Ok(()));
}
if self.all_operations_completed() {
self.completed = true;
return Poll::Ready(Ok(()));
}
if any_completed {
cx.waker().wake_by_ref();
return Poll::Pending;
}
Poll::Pending
}
fn check_ready_operations(&mut self, completed_index: usize) {
let _newly_ready: Vec<usize> = Vec::new();
for (&_dependent_index, dependencies) in &self.dependencies {
if dependencies.contains(&completed_index) {
let _all_deps_satisfied = dependencies.iter().all(|&dep_index| {
self.results[dep_index].is_some()
&& self.results[dep_index].as_ref().unwrap().is_success()
});
}
}
}
fn cancel_dependent_operations(&mut self, failed_index: usize) {
let mut to_cancel = Vec::new();
let mut visited = std::collections::HashSet::new();
let mut stack = vec![failed_index];
while let Some(current) = stack.pop() {
if visited.contains(¤t) {
continue;
}
visited.insert(current);
for (&dependent, dependencies) in &self.dependencies {
if dependencies.contains(¤t) && !visited.contains(&dependent) {
to_cancel.push(dependent);
stack.push(dependent);
}
}
}
for &index in &to_cancel {
if self.results[index].is_none() {
self.results[index] = Some(OperationResult::Cancelled);
}
}
}
fn cancel_all_remaining_operations(&mut self) {
for result in self.results.iter_mut() {
if result.is_none() {
*result = Some(OperationResult::Cancelled);
}
}
}
fn all_operations_completed(&self) -> bool {
self.results.iter().all(|result| result.is_some())
}
}
impl Future for StandaloneBatchFuture {
type Output = Result<BatchResult>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let _ = cx;
Poll::Pending
}
}
impl Drop for StandaloneBatchFuture {
fn drop(&mut self) {
self.cancel_all_remaining_operations();
for id in self.operation_ids.iter().flatten() {
self.waker_registry.remove_waker(*id);
}
}
}