[Chain Orchestrator] Retry on failure

[frisitano]

Overview

If we have a failure when processing a future in the chain orchestrator and the error is associated with some transient error, such as DatabaseError, NetworkRequestError or RpcError, then we should retry the future until completion. This prevents transient errors from impacting the chain orchestrator. We can do this by marking the error types as retryable or not e.g.

impl ChainOrchestratorError {
    fn can_retry(&self) -> bool {
        matches!(
            self,
            ChainOrchestratorError::DatabaseError(_) |
                ChainOrchestratorError::NetworkRequestError(_) |
                ChainOrchestratorError::RpcError(_)
        )
    }
}

We would then implement the retry logic in the Stream implementation of the ChainOrchestrator:

rollup-node/crates/chain-orchestrator/src/lib.rs

Lines 775 to 804 in 3d400f7

	impl<
	ChainSpec: ScrollHardforks + 'static,
	BC: BlockClient<Block = ScrollBlock> + Send + Sync + 'static,
	P: Provider<Scroll> + Send + Sync + 'static,
	> Stream for ChainOrchestrator<ChainSpec, BC, P>
	{
	type Item = Result<ChainOrchestratorEvent, ChainOrchestratorError>;
	fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
	// Register the waker such that we can wake when required.
	self.waker.register(cx.waker());
	// Remove and poll the next future in the queue
	while let Some(mut action) = self.pending_futures.pop_front() {
	match action.poll(cx) {
	Poll::Ready(result) => match result {
	Ok(None) => {}
	Ok(Some(event)) => return Poll::Ready(Some(Ok(event))),
	Err(e) => return Poll::Ready(Some(Err(e))),
	},
	Poll::Pending => {
	self.pending_futures.push_front(action);
	return Poll::Pending
	}
	};
	}
	Poll::Pending
	}
	}

[greged93]

I propose to implement a retry logic, trying to mimic the RetryBackoffLayer in alloy. I was not able to end with the same API, as their implementation involves services, which allows you to stack the various services in a pipeline manner.

In our case, we would want a way to have control over the retry attempts on any DatabaseReadOperations or DatabaseWriteOperations call. I came up with the below design, which creates a retry object, in which the retry mechanism is implemented in the Future implementation. We would need to modify the DatabaseReadOperations and DatabaseWriteOperations into factory traits, which capture the inputs and return an Fn trait object. Let me know if you think it's ok or if you have another solution:

pub type FactoryFn<OK, ERR> =
    Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<OK, ERR>> + Send>> + Send + Sync>;

trait MyTraitFactory {
    fn do_something(&self, num: u64) -> FactoryFn<u64, ()> {
        Box::new(move || {
            let num = num;
            Box::pin(async move {
                // Simulate an operation that may fail
                if num % 2 == 0 {
                    Ok(num)
                } else {
                    Err(())
                }
            })
        })
    }
}

struct MyStruct;

impl MyTraitFactory for MyStruct {}

#[tokio::main]
async fn main() {
    let retried_fut = MyStruct.do_something(2).with_retry(4, 10);
    let res = retried_fut.await.expect("retry failed");
    dbg!(res);
}

pub struct RetryTask<Factory: Fn() -> Fut, Fut: Future> {
    factory: Factory,
    timer: Option<Pin<Box<tokio::time::Sleep>>>,
    fut: Option<Pin<Box<dyn Future<Output = Fut::Output>>>>,
    retries: usize,
    initial_delay: u64,
    attempt: usize,
}

impl<Factory: Fn() -> Fut, Fut: Future> RetryTask<Factory, Fut> {
    pub fn new(factory: Factory, retries: usize, initial_delay: u64) -> Self {
        Self {
            factory,
            fut: None,
            retries,
            initial_delay,
            attempt: 0,
            timer: None,
        }
    }
}

pub trait WithRetry {
    fn with_retry<Fut>(self, retries: usize, initial_delay: u64) -> RetryTask<Self, Fut>
    where
        Self: Fn() -> Fut + Sized,
        Fut: Future,
    {
        RetryTask::new(self, retries, initial_delay)
    }
}

impl<T> WithRetry for T {}

impl<
        Factory: Fn() -> Fut + Unpin,
        Fut: Future<Output = Result<OK, ERR>> + 'static,
        OK,
        ERR,
    > Future for RetryTask<Factory, Fut>
{
    type Output = Result<OK, ()>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.attempt >= self.retries {
            return Poll::Ready(Err(()));
        }

        let this = self.get_mut();
        if let Some(timer) = &mut this.timer {
            ready!(timer.as_mut().poll(cx));
            this.timer.take();
            cx.waker().wake_by_ref();
            Poll::Pending
        } else if let Some(fut) = &mut this.fut {
            let res = ready!(fut.as_mut().poll(cx));
            if res.is_err() {
                let delay = this.initial_delay * 2u64.pow(this.attempt as u32);
                this.attempt += 1;
                this.timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(delay))));
                this.fut.take();
                cx.waker().wake_by_ref();
                Poll::Pending
            } else {
                Poll::Ready(res.map_err(|_| ()))
            }
        } else {
            let fut = (this.factory)();
            this.fut = Some(Box::pin(fut));
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

[frisitano]

closed by #322

[greged93]

@frisitano let's discuss here

[frisitano]

I'm thinking we can implement something like this using the tower stack. This isn't compiling at the moment due to using references for the transaction. We may need to pass an owned TX / TXMut instead. Let me know what you think.

pub enum DatabaseQuery<T> {
    Read(
        Box<
            dyn FnOnce(&TX) -> Pin<Box<dyn Future<Output = Result<T, DatabaseError>> + Send>>
                + Send,
        >,
    ),
    Write(
        Box<
            dyn FnOnce(&mut TXMut) -> Pin<Box<dyn Future<Output = Result<T, DatabaseError>> + Send>>
                + Send,
        >,
    ),
}

impl<T> DatabaseQuery<T>
where
    T: Send + 'static,
{
    pub fn read<F, Fut>(f: F) -> Self
    where
        F: FnOnce(&TX) -> Fut + Send + 'static,
        Fut: Future<Output = Result<T, DatabaseError>> + Send + 'static,
    {
        DatabaseQuery::Read(Box::new(move |tx| Box::pin(f(tx))))
    }

    pub fn write<F, Fut>(f: F) -> Self
    where
        F: FnOnce(&mut TXMut) -> Fut + Send + 'static,
        Fut: Future<Output = Result<T, DatabaseError>> + Send + 'static,
    {
        DatabaseQuery::Write(Box::new(move |tx| Box::pin(f(tx))))
    }
}

impl<T> Service<DatabaseQuery<T>> for std::sync::Arc<Database>
where
    T: Send + 'static,
{
    type Response = T;
    type Error = DatabaseError;
    type Future =
        std::pin::Pin<Box<dyn std::future::Future<Output = Result<T, DatabaseError>> + Send>>;

    fn poll_ready(
        &mut self,
        _: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        std::task::Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: DatabaseQuery<T>) -> Self::Future {
        let db = self.clone();
        match req {
            DatabaseQuery::Read(f) => Box::pin(async move {
                let tx = db.tx().await?;
                f(&tx).await
            }),
            DatabaseQuery::Write(f) => Box::pin(async move {
                let mut tx = db.tx_mut().await?;
                let res = f(&mut tx).await;
                if res.is_ok() {
                    tx.commit().await?;
                } else {
                    tx.rollback().await?;
                }
                res
            }),
        }
    }
}

async fn example(db: std::sync::Arc<Database>) -> Result<(), DatabaseError> {
    db.call(DatabaseQuery::read(|tx| async {
        let batch = tx.get_batch_by_index(5).await?;
        Ok(batch)
    }))
    .await?;

    Ok(())
}

[greged93]

I think that is a nice improvement but I'm not sure it improves the interface of the retry (but I might be misunderstanding). Does this mean we would implement the retry as below?

impl<T> Service<DatabaseQuery<T>> for Retry
{
    ....
}

async fn example_retry(retry: Retry) -> Result<(), DatabaseError> {
    retry.call(DatabaseQuery::read(|tx| async {
        let batch = tx.get_batch_by_index(5).await?;
        Ok(batch)
    }))
    .await?;

    Ok(())
}

I was trying to mimic the alloy_provider::Provider trait, where you can directly call the method on the provider, which holds the transport abstracted behind a trait object. Awaiting the return value from the provider allows you to start the polling loop which in turns poll the transport (which is made up of the various layers of transport).

My idea for the interface was to get close to something where performing the call (and any intermediate layered service calls) would just be a matter of calling:

db_client.get_batch_by_index(5).await?;

I think with your implementation and changing the DatabaseRead(Write)Operations traits as below we can have a better interfacing.

#[must_use = "futures do nothing unless you `.await` or poll them"]
#[pin_project::pin_project(project = DatabaseCallProj)]
pub struct DatabaseCall<T> {
    #[pin]
    pub state: DatabaseCallState<T>,
}

#[must_use = "futures do nothing unless you `.await` or poll them"]
#[warn(clippy::large_enum_variant)]
#[pin_project::pin_project(project = DatabaseCallStateProj)]
pub enum DatabaseCallState<T> {
    Preparing {
        database_connection: BoxedService<T>,
        request: Option<DatabaseRequest<T>>,
    },
    Awaiting {
        #[pin]
        fut: Box<dyn Future<Output = Result<T, DatabaseError>> + Send + Unpin>,
    },
}

impl<T> Future for DatabaseCall<T> {
    type Output = Result<T, DatabaseError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.as_mut().project();
        Poll::Ready(ready!(this.state.poll_unpin(cx)))
    }
}

impl<T> Future for DatabaseCallState<T> {
    type Output = Result<T, DatabaseError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.as_mut().project();
        match this {
            DatabaseCallStateProj::Preparing {
                database_connection,
                request,
            } => {
                let fut = database_connection.call(request.take().expect("first call"));
                self.set(Self::Awaiting { fut });
                Poll::Pending
            }
            DatabaseCallStateProj::Awaiting { mut fut } => {
                let res = ready!(fut.poll_unpin(cx));
                Poll::Ready(res)
            }
        }
    }
}

pub type BoxedService<T> = Box<
    dyn Service<
        DatabaseRequest<T>,
        Response = T,
        Error = DatabaseError,
        Future = Box<dyn Future<Output = Result<T, DatabaseError>> + Send + Unpin>,
    >,
>;

pub trait IntoBoxedService {
    fn service<T>(&self) -> BoxedService<T>;
}

pub trait DatabaseReadOperations: IntoBoxedService + Sync + Sized {
    /// Get a [`BatchCommitData`] from the database by its batch index.
    fn get_batch_by_index(
        &self,
        batch_index: u64,
    ) -> DatabaseCall<Option<BatchCommitData>> {
        let database = self.service();
        let query = DatabaseRequest::Read(Box::new(move |tx: Arc<TX>| {
            let batch_index = batch_index;
            Box::pin(async move {
                let conn = tx.get_connection();
                Ok(models::Entity::find_by_id(
                    TryInto::<i64>::try_into(batch_index).expect("index should fit in i64"),
                )
                .one(tx.get_connection())
                .await
                .map(|x| x.map(Into::into))?)
            })
        }));
        DatabaseCall {
            state: DatabaseCallState::Preparing {
                request: Some(query),
                database_connection: database,
            },
        }
    }
}