bare-metrics-reader 0.1.0

use anyhow::bail;
use bare_metrics_core::structures::{
    get_supported_version, Frame, LogHeader, UnixTimestampMilliseconds,
};
use std::io::{Read, Seek, SeekFrom};

/// Token that is known to be usable for seeking to a frame in a metrics log.
/// TODO identify which reader is applicable? Or don't bother?
#[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct SeekToken {
    pub start_ts: UnixTimestampMilliseconds,
    pub offset: u64,
}

/// A streaming reader for metric logs.
/// If the underlying reader supports seeks, it is possible to get SeekTokens which may be used
/// to seek to previous frames in the stream.
pub struct MetricsLogReader<R: Read> {
    reader: R,
    pub header: LogHeader,
    last_read_ts: UnixTimestampMilliseconds,
}

impl<R: Read> MetricsLogReader<R> {
    /// Constructs a bare metrics log reader.
    pub fn new(mut reader: R) -> anyhow::Result<Self> {
        let header: LogHeader = serde_bare::from_reader(&mut reader)?;
        if header.bare_metrics_version != get_supported_version() {
            bail!("Wrong version. Expected {:?} got {:?}. Later versions of Bare Metrics may use a stable format.", get_supported_version(), header.bare_metrics_version);
        }
        let last_read_ts = header.start_time;
        Ok(MetricsLogReader {
            reader,
            header,
            last_read_ts,
        })
    }

    /// Reads a frame.
    /// Returns the start time of the frame and the frame itself.
    pub fn read_frame(&mut self) -> anyhow::Result<Option<(UnixTimestampMilliseconds, Frame)>> {
        let mut interceptor = EofTrackingReadInterceptor::new(&mut self.reader);
        match serde_bare::from_reader::<_, Frame>(&mut interceptor) {
            Ok(frame) => {
                let start_ts = self.last_read_ts;
                self.last_read_ts = frame.end_time;
                Ok(Some((start_ts, frame)))
            }
            // This doesn't seem to work properly for some reason...
            Err(err) if err.classify().is_eof() => Ok(None),
            Err(other_err) => {
                let eof_flag = interceptor.was_eof();
                if eof_flag == Some(true) {
                    Ok(None)
                } else {
                    bail!(
                        "Failed to read frame: {:?} class {:?}, intercepted eof flag {:?}",
                        other_err,
                        other_err.classify(),
                        eof_flag
                    );
                }
            }
        }
    }
}

struct EofTrackingReadInterceptor<R: Read> {
    inner: R,
    /// None if no reads have taken place yet.
    /// Some(true) if the first read that took place and was EOF
    /// Some(false) if the first read that took place and was not EOF
    was_eof_flag: Option<bool>,
}

impl<R: Read> EofTrackingReadInterceptor<R> {
    pub fn new(inner: R) -> EofTrackingReadInterceptor<R> {
        EofTrackingReadInterceptor {
            inner,
            was_eof_flag: None,
        }
    }

    pub fn was_eof(self) -> Option<bool> {
        self.was_eof_flag
    }
}

impl<R: Read> Read for EofTrackingReadInterceptor<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.was_eof_flag.is_none() {
            let count = self.inner.read(buf)?;
            if count == 0 {
                self.was_eof_flag = Some(true);
            } else {
                self.was_eof_flag = Some(false);
            }
            Ok(count)
        } else {
            self.inner.read(buf)
        }
    }
}

impl<R: Read + Seek> MetricsLogReader<R> {
    /// Reads a new frame, and returns a seek token that can be used to rewind such that you can
    /// 'undo' the read.
    pub fn read_frame_rewindable(
        &mut self,
    ) -> anyhow::Result<Option<(SeekToken, UnixTimestampMilliseconds, Frame)>> {
        let current_pos_in_file = self.reader.stream_position()?;
        // TODO should we rewind in case of error?
        if let Some((timestamp, frame)) = self.read_frame()? {
            Ok(Some((
                SeekToken {
                    start_ts: timestamp,
                    offset: current_pos_in_file,
                },
                timestamp,
                frame,
            )))
        } else {
            // EOF (no more things to read here).
            // Rewind to where we were before.
            self.reader.seek(SeekFrom::Start(current_pos_in_file))?;
            Ok(None)
        }
    }

    /// Seeks to a position in the stream.
    /// The given seek token MUST have come from this instance's `read_from_rewindable` function.
    /// Otherwise, corrupt frames may be read.
    /// The old position is returned as a seek token.
    pub fn seek(&mut self, seek_token: SeekToken) -> anyhow::Result<SeekToken> {
        let SeekToken {
            start_ts: seek_timestamp,
            offset: seek_pos,
        } = seek_token;
        let old_pos_in_file = self.reader.stream_position()?;
        let old_timestamp = self.last_read_ts;

        // TODO should we rewind in case of error?
        self.reader.seek(SeekFrom::Start(seek_pos))?;
        self.last_read_ts = seek_timestamp;

        Ok(SeekToken {
            start_ts: old_timestamp,
            offset: old_pos_in_file,
        })
    }
}