Skip to content

Commit

Permalink
introduce --checkpoints arg
Browse files Browse the repository at this point in the history
  • Loading branch information
shekhirin committed Nov 17, 2023
1 parent 4c15a76 commit fa793aa
Showing 1 changed file with 10 additions and 6 deletions.
16 changes: 10 additions & 6 deletions bin/reth/src/stage/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,10 @@ pub struct Command {
// e.g. query the DB size, or any table data.
#[arg(long, short)]
commit: bool,

/// Save stage checkpoints
#[arg(long)]
checkpoints: bool,
}

impl Command {
Expand Down Expand Up @@ -245,7 +249,9 @@ impl Command {
let UnwindOutput { checkpoint } = unwind_stage.unwind(&provider_rw, unwind).await?;
unwind.checkpoint = checkpoint;

provider_rw.save_stage_checkpoint(unwind_stage.id(), checkpoint)?;
if self.checkpoints {
provider_rw.save_stage_checkpoint(unwind_stage.id(), checkpoint)?;
}

if self.commit {
provider_rw.commit()?;
Expand All @@ -263,7 +269,9 @@ impl Command {
let ExecOutput { checkpoint, done } = exec_stage.execute(&provider_rw, input).await?;
input.checkpoint = Some(checkpoint);

provider_rw.save_stage_checkpoint(exec_stage.id(), checkpoint)?;
if self.checkpoints {
provider_rw.save_stage_checkpoint(exec_stage.id(), checkpoint)?;
}
if self.commit {
provider_rw.commit()?;
provider_rw = factory.provider_rw().map_err(PipelineError::Interface)?;
Expand All @@ -274,10 +282,6 @@ impl Command {
}
}

if self.commit {
provider_rw.commit()?;
}

Ok(())
}
}

0 comments on commit fa793aa

Please sign in to comment.