oura 2.2.0

The tail of Cardano
Documentation
use std::io::Write;
use std::path::PathBuf;

use file_rotate::compression::Compression;
use file_rotate::suffix::AppendTimestamp;
use file_rotate::suffix::FileLimit;
use file_rotate::ContentLimit;
use file_rotate::FileRotate;
use gasket::framework::*;
use serde::Deserialize;
use serde_json::Value as JsonValue;

use crate::framework::*;

pub struct Worker {
    writer: FileRotate<AppendTimestamp>,
}

#[async_trait::async_trait(?Send)]
impl gasket::framework::Worker<Stage> for Worker {
    async fn bootstrap(stage: &Stage) -> Result<Self, WorkerError> {
        let output_path = match &stage.config.output_path {
            Some(x) => PathBuf::from(x),
            None => stage.current_dir.clone(),
        };

        let suffix_scheme = AppendTimestamp::default(FileLimit::MaxFiles(
            stage
                .config
                .max_total_files
                .unwrap_or(DEFAULT_MAX_TOTAL_FILES),
        ));

        let content_limit = ContentLimit::BytesSurpassed(
            stage
                .config
                .max_bytes_per_file
                .unwrap_or(DEFAULT_MAX_BYTES_PER_FILE),
        );

        let compression = if let Some(true) = stage.config.compress_files {
            Compression::OnRotate(2)
        } else {
            Compression::None
        };

        #[cfg(unix)]
        let writer = FileRotate::new(output_path, suffix_scheme, content_limit, compression, None);

        #[cfg(not(unix))]
        let writer = FileRotate::new(output_path, suffix_scheme, content_limit, compression);

        Ok(Self { writer })
    }

    async fn schedule(
        &mut self,
        stage: &mut Stage,
    ) -> Result<WorkSchedule<ChainEvent>, WorkerError> {
        let msg = stage.input.recv().await.or_panic()?;
        Ok(WorkSchedule::Unit(msg.payload))
    }

    async fn execute(&mut self, unit: &ChainEvent, stage: &mut Stage) -> Result<(), WorkerError> {
        let point = unit.point();
        let json = JsonValue::from(unit.clone());

        self.writer
            .write_all(json.to_string().as_bytes())
            .and_then(|_| self.writer.write_all(b"\n"))
            .or_retry()?;

        stage.ops_count.inc(1);

        stage.latest_block.set(point.slot_or_default() as i64);
        stage.cursor.send(point.clone().into()).await.or_panic()?;

        Ok(())
    }
}

#[derive(Stage)]
#[stage(name = "sink-filerotate", unit = "ChainEvent", worker = "Worker")]
pub struct Stage {
    config: Config,
    current_dir: PathBuf,

    pub input: MapperInputPort,
    pub cursor: SinkCursorPort,

    #[metric]
    ops_count: gasket::metrics::Counter,

    #[metric]
    latest_block: gasket::metrics::Gauge,
}

#[derive(Debug, Deserialize, Clone)]
pub enum Format {
    #[serde(rename = "JSONL")]
    Jsonl,
}

#[derive(Default, Debug, Deserialize)]
pub struct Config {
    pub output_format: Option<Format>,
    pub output_path: Option<String>,
    pub max_bytes_per_file: Option<usize>,
    pub max_total_files: Option<usize>,
    pub compress_files: Option<bool>,
}

const DEFAULT_MAX_BYTES_PER_FILE: usize = 50 * 1024 * 1024;
const DEFAULT_MAX_TOTAL_FILES: usize = 200;

impl Config {
    pub fn bootstrapper(self, ctx: &Context) -> Result<Stage, Error> {
        let stage = Stage {
            config: self,
            current_dir: ctx.current_dir.clone(),
            ops_count: Default::default(),
            latest_block: Default::default(),
            input: Default::default(),
            cursor: Default::default(),
        };

        Ok(stage)
    }
}