<commit_before>use crate::error::{ParseError, ParseErrorKind::*};
use std::fmt;
use std::iter;
use std::str::{self, FromStr};

/// A cfg expression.
#[derive(Eq, PartialEq, Hash, Ord, PartialOrd, Clone, Debug)]
pub enum CfgExpr {
    Not(Box<CfgExpr>),
    All(Vec<CfgExpr>),
    Any(Vec<CfgExpr>),
    Value(Cfg),
}

/// A cfg value.
#[derive(Eq, PartialEq, Hash, Ord, PartialOrd, Clone, Debug)]
pub enum Cfg {
    /// A named cfg value, like `unix`.
    Name(String),
    /// A key/value cfg pair, like `target_os = "linux"`.
    KeyPair(String, String),
}

#[derive(PartialEq)]
enum Token<'a> {
    LeftParen,
    RightParen,
    Ident(&'a str),
    Comma,
    Equals,
    String(&'a str),
}

#[derive(Clone)]
struct Tokenizer<'a> {
    s: iter::Peekable<str::CharIndices<'a>>,
    orig: &'a str,
}

struct Parser<'a> {
    t: Tokenizer<'a>,
}

impl FromStr for Cfg {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Cfg, Self::Err> {
        let mut p = Parser::new(s);
        let e = p.cfg()?;
        if let Some(rest) = p.rest() {
            return Err(ParseError::new(
                p.t.orig,
                UnterminatedExpression(rest.to_string()),
            ));
        }
        Ok(e)
    }
}

impl fmt::Display for Cfg {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Cfg::Name(ref s) => s.fmt(f),
            Cfg::KeyPair(ref k, ref v) => write!(f, "{} = \"{}\"", k, v),
        }
    }
}

impl CfgExpr {
    /// Utility function to check if the key, "cfg(..)" matches the `target_cfg`
    pub fn matches_key(key: &str, target_cfg: &[Cfg]) -> bool {
        if key.starts_with("cfg(") && key.ends_with(')') {
            let cfg = &key[4..key.len() - 1];

            CfgExpr::from_str(cfg)
                .ok()
                .map(|ce| ce.matches(target_cfg))
                .unwrap_or(false)
        } else {
            false
        }
    }

    pub fn matches(&self, cfg: &[Cfg]) -> bool {
        match *self {
            CfgExpr::Not(ref e) => !e.matches(cfg),
            CfgExpr::All(ref e) => e.iter().all(|e| e.matches(cfg)),
            CfgExpr::Any(ref e) => e.iter().any(|e| e.matches(cfg)),
            CfgExpr::Value(ref e) => cfg.contains(e),
        }
    }
}

impl FromStr for CfgExpr {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<CfgExpr, Self::Err> {
        let mut p = Parser::new(s);
        let e = p.expr()?;
        if let Some(rest) = p.rest() {
            return Err(ParseError::new(
                p.t.orig,
                UnterminatedExpression(rest.to_string()),
            ));
        }
        Ok(e)
    }
}

impl fmt::Display for CfgExpr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            CfgExpr::Not(ref e) => write!(f, "not({})", e),
            CfgExpr::All(ref e) => write!(f, "all({})", CommaSep(e)),
            CfgExpr::Any(ref e) => write!(f, "any({})", CommaSep(e)),
            CfgExpr::Value(ref e) => write!(f, "{}", e),
        }
    }
}

struct CommaSep<'a, T>(&'a [T]);

impl<'a, T: fmt::Display> fmt::Display for CommaSep<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, v) in self.0.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}", v)?;
        }
        Ok(())
    }
}

impl<'a> Parser<'a> {
    fn new(s: &'a str) -> Parser<'a> {
        Parser {
            t: Tokenizer {
                s: s.char_indices().peekable(),
                orig: s,
            },
        }
    }

    fn expr(&mut self) -> Result<CfgExpr, ParseError> {
        match self.peek() {
            Some(Ok(Token::Ident(op @ "all"))) | Some(Ok(Token::Ident(op @ "any"))) => {
                self.t.next();
                let mut e = Vec::new();
                self.eat(&Token::LeftParen)?;
                while !self.r#try(&Token::RightParen) {
                    e.push(self.expr()?);
                    if !self.r#try(&Token::Comma) {
                        self.eat(&Token::RightParen)?;
                        break;
                    }
                }
                if op == "all" {
                    Ok(CfgExpr::All(e))
                } else {
                    Ok(CfgExpr::Any(e))
                }
            }
            Some(Ok(Token::Ident("not"))) => {
                self.t.next();
                self.eat(&Token::LeftParen)?;
                let e = self.expr()?;
                self.eat(&Token::RightParen)?;
                Ok(CfgExpr::Not(Box::new(e)))
            }
            Some(Ok(..)) => self.cfg().map(CfgExpr::Value),
            Some(Err(..)) => Err(self.t.next().unwrap().err().unwrap()),
            None => Err(ParseError::new(
                self.t.orig,
                IncompleteExpr("start of a cfg expression"),
            )),
        }
    }

    fn cfg(&mut self) -> Result<Cfg, ParseError> {
        match self.t.next() {
            Some(Ok(Token::Ident(name))) => {
                let e = if self.r#try(&Token::Equals) {
                    let val = match self.t.next() {
                        Some(Ok(Token::String(s))) => s,
                        Some(Ok(t)) => {
                            return Err(ParseError::new(
                                self.t.orig,
                                UnexpectedToken {
                                    expected: "a string",
                                    found: t.classify(),
                                },
                            ))
                        }
                        Some(Err(e)) => return Err(e),
                        None => {
                            return Err(ParseError::new(self.t.orig, IncompleteExpr("a string")))
                        }
                    };
                    Cfg::KeyPair(name.to_string(), val.to_string())
                } else {
                    Cfg::Name(name.to_string())
                };
                Ok(e)
            }
            Some(Ok(t)) => Err(ParseError::new(
                self.t.orig,
                UnexpectedToken {
                    expected: "identifier",
                    found: t.classify(),
                },
            )),
            Some(Err(e)) => Err(e),
            None => Err(ParseError::new(self.t.orig, IncompleteExpr("identifier"))),
        }
    }

    fn peek(&mut self) -> Option<Result<Token<'a>, ParseError>> {
        self.t.clone().next()
    }

    fn r#try(&mut self, token: &Token<'a>) -> bool {
        match self.peek() {
            Some(Ok(ref t)) if token == t => {}
            _ => return false,
        }
        self.t.next();
        true
    }

    fn eat(&mut self, token: &Token<'a>) -> Result<(), ParseError> {
        match self.t.next() {
            Some(Ok(ref t)) if token == t => Ok(()),
            Some(Ok(t)) => Err(ParseError::new(
                self.t.orig,
                UnexpectedToken {
                    expected: token.classify(),
                    found: t.classify(),
                },
            )),
            Some(Err(e)) => Err(e),
            None => Err(ParseError::new(
                self.t.orig,
                IncompleteExpr(token.classify()),
            )),
        }
    }

    /// Returns the rest of the input from the current location.
    fn rest(&self) -> Option<&str> {
        let mut s = self.t.s.clone();
        loop {
            match s.next() {
                Some((_, ' ')) => {}
                Some((start, _ch)) => return Some(&self.t.orig[start..]),
                None => return None,
            }
        }
    }
}

impl<'a> Iterator for Tokenizer<'a> {
    type Item = Result<Token<'a>, ParseError>;

    fn next(&mut self) -> Option<Result<Token<'a>, ParseError>> {
        loop {
            match self.s.next() {
                Some((_, ' ')) => {}
                Some((_, '(')) => return Some(Ok(Token::LeftParen)),
                Some((_, ')')) => return Some(Ok(Token::RightParen)),
                Some((_, ',')) => return Some(Ok(Token::Comma)),
                Some((_, '=')) => return Some(Ok(Token::Equals)),
                Some((start, '"')) => {
                    while let Some((end, ch)) = self.s.next() {
                        if ch == '"' {
                            return Some(Ok(Token::String(&self.orig[start + 1..end])));
                        }
                    }
                    return Some(Err(ParseError::new(self.orig, UnterminatedString)));
                }
                Some((start, ch)) if is_ident_start(ch) => {
                    while let Some(&(end, ch)) = self.s.peek() {
                        if !is_ident_rest(ch) {
                            return Some(Ok(Token::Ident(&self.orig[start..end])));
                        } else {
                            self.s.next();
                        }
                    }
                    return Some(Ok(Token::Ident(&self.orig[start..])));
                }
                Some((_, ch)) => {
                    return Some(Err(ParseError::new(self.orig, UnexpectedChar(ch))));
                }
                None => return None,
            }
        }
    }
}

fn is_ident_start(ch: char) -> bool {
    ch == '_' || ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z')
}

fn is_ident_rest(ch: char) -> bool {
    is_ident_start(ch) || ('0' <= ch && ch <= '9')
}

impl<'a> Token<'a> {
    fn classify(&self) -> &'static str {
        match *self {
            Token::LeftParen => "`(`",
            Token::RightParen => "`)`",
            Token::Ident(..) => "an identifier",
            Token::Comma => "`,`",
            Token::Equals => "`=`",
            Token::String(..) => "a string",
        }
    }
}
<commit_msg>Auto merge of #9110 - klensy:simple-fix, r=alexcrichton<commit_after>use crate::error::{ParseError, ParseErrorKind::*};
use std::fmt;
use std::iter;
use std::str::{self, FromStr};

/// A cfg expression.
#[derive(Eq, PartialEq, Hash, Ord, PartialOrd, Clone, Debug)]
pub enum CfgExpr {
    Not(Box<CfgExpr>),
    All(Vec<CfgExpr>),
    Any(Vec<CfgExpr>),
    Value(Cfg),
}

/// A cfg value.
#[derive(Eq, PartialEq, Hash, Ord, PartialOrd, Clone, Debug)]
pub enum Cfg {
    /// A named cfg value, like `unix`.
    Name(String),
    /// A key/value cfg pair, like `target_os = "linux"`.
    KeyPair(String, String),
}

#[derive(PartialEq)]
enum Token<'a> {
    LeftParen,
    RightParen,
    Ident(&'a str),
    Comma,
    Equals,
    String(&'a str),
}

#[derive(Clone)]
struct Tokenizer<'a> {
    s: iter::Peekable<str::CharIndices<'a>>,
    orig: &'a str,
}

struct Parser<'a> {
    t: Tokenizer<'a>,
}

impl FromStr for Cfg {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Cfg, Self::Err> {
        let mut p = Parser::new(s);
        let e = p.cfg()?;
        if let Some(rest) = p.rest() {
            return Err(ParseError::new(
                p.t.orig,
                UnterminatedExpression(rest.to_string()),
            ));
        }
        Ok(e)
    }
}

impl fmt::Display for Cfg {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Cfg::Name(ref s) => s.fmt(f),
            Cfg::KeyPair(ref k, ref v) => write!(f, "{} = \"{}\"", k, v),
        }
    }
}

impl CfgExpr {
    /// Utility function to check if the key, "cfg(..)" matches the `target_cfg`
    pub fn matches_key(key: &str, target_cfg: &[Cfg]) -> bool {
        if key.starts_with("cfg(") && key.ends_with(')') {
            let cfg = &key[4..key.len() - 1];

            CfgExpr::from_str(cfg)
                .ok()
                .map(|ce| ce.matches(target_cfg))
                .unwrap_or(false)
        } else {
            false
        }
    }

    pub fn matches(&self, cfg: &[Cfg]) -> bool {
        match *self {
            CfgExpr::Not(ref e) => !e.matches(cfg),
            CfgExpr::All(ref e) => e.iter().all(|e| e.matches(cfg)),
            CfgExpr::Any(ref e) => e.iter().any(|e| e.matches(cfg)),
            CfgExpr::Value(ref e) => cfg.contains(e),
        }
    }
}

impl FromStr for CfgExpr {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<CfgExpr, Self::Err> {
        let mut p = Parser::new(s);
        let e = p.expr()?;
        if let Some(rest) = p.rest() {
            return Err(ParseError::new(
                p.t.orig,
                UnterminatedExpression(rest.to_string()),
            ));
        }
        Ok(e)
    }
}

impl fmt::Display for CfgExpr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            CfgExpr::Not(ref e) => write!(f, "not({})", e),
            CfgExpr::All(ref e) => write!(f, "all({})", CommaSep(e)),
            CfgExpr::Any(ref e) => write!(f, "any({})", CommaSep(e)),
            CfgExpr::Value(ref e) => write!(f, "{}", e),
        }
    }
}

struct CommaSep<'a, T>(&'a [T]);

impl<'a, T: fmt::Display> fmt::Display for CommaSep<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, v) in self.0.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}", v)?;
        }
        Ok(())
    }
}

impl<'a> Parser<'a> {
    fn new(s: &'a str) -> Parser<'a> {
        Parser {
            t: Tokenizer {
                s: s.char_indices().peekable(),
                orig: s,
            },
        }
    }

    fn expr(&mut self) -> Result<CfgExpr, ParseError> {
        match self.peek() {
            Some(Ok(Token::Ident(op @ "all"))) | Some(Ok(Token::Ident(op @ "any"))) => {
                self.t.next();
                let mut e = Vec::new();
                self.eat(&Token::LeftParen)?;
                while !self.r#try(&Token::RightParen) {
                    e.push(self.expr()?);
                    if !self.r#try(&Token::Comma) {
                        self.eat(&Token::RightParen)?;
                        break;
                    }
                }
                if op == "all" {
                    Ok(CfgExpr::All(e))
                } else {
                    Ok(CfgExpr::Any(e))
                }
            }
            Some(Ok(Token::Ident("not"))) => {
                self.t.next();
                self.eat(&Token::LeftParen)?;
                let e = self.expr()?;
                self.eat(&Token::RightParen)?;
                Ok(CfgExpr::Not(Box::new(e)))
            }
            Some(Ok(..)) => self.cfg().map(CfgExpr::Value),
            Some(Err(..)) => Err(self.t.next().unwrap().err().unwrap()),
            None => Err(ParseError::new(
                self.t.orig,
                IncompleteExpr("start of a cfg expression"),
            )),
        }
    }

    fn cfg(&mut self) -> Result<Cfg, ParseError> {
        match self.t.next() {
            Some(Ok(Token::Ident(name))) => {
                let e = if self.r#try(&Token::Equals) {
                    let val = match self.t.next() {
                        Some(Ok(Token::String(s))) => s,
                        Some(Ok(t)) => {
                            return Err(ParseError::new(
                                self.t.orig,
                                UnexpectedToken {
                                    expected: "a string",
                                    found: t.classify(),
                                },
                            ))
                        }
                        Some(Err(e)) => return Err(e),
                        None => {
                            return Err(ParseError::new(self.t.orig, IncompleteExpr("a string")))
                        }
                    };
                    Cfg::KeyPair(name.to_string(), val.to_string())
                } else {
                    Cfg::Name(name.to_string())
                };
                Ok(e)
            }
            Some(Ok(t)) => Err(ParseError::new(
                self.t.orig,
                UnexpectedToken {
                    expected: "identifier",
                    found: t.classify(),
                },
            )),
            Some(Err(e)) => Err(e),
            None => Err(ParseError::new(self.t.orig, IncompleteExpr("identifier"))),
        }
    }

    fn peek(&mut self) -> Option<Result<Token<'a>, ParseError>> {
        self.t.clone().next()
    }

    fn r#try(&mut self, token: &Token<'a>) -> bool {
        match self.peek() {
            Some(Ok(ref t)) if token == t => {}
            _ => return false,
        }
        self.t.next();
        true
    }

    fn eat(&mut self, token: &Token<'a>) -> Result<(), ParseError> {
        match self.t.next() {
            Some(Ok(ref t)) if token == t => Ok(()),
            Some(Ok(t)) => Err(ParseError::new(
                self.t.orig,
                UnexpectedToken {
                    expected: token.classify(),
                    found: t.classify(),
                },
            )),
            Some(Err(e)) => Err(e),
            None => Err(ParseError::new(
                self.t.orig,
                IncompleteExpr(token.classify()),
            )),
        }
    }

    /// Returns the rest of the input from the current location.
    fn rest(&self) -> Option<&str> {
        let mut s = self.t.s.clone();
        loop {
            match s.next() {
                Some((_, ' ')) => {}
                Some((start, _ch)) => return Some(&self.t.orig[start..]),
                None => return None,
            }
        }
    }
}

impl<'a> Iterator for Tokenizer<'a> {
    type Item = Result<Token<'a>, ParseError>;

    fn next(&mut self) -> Option<Result<Token<'a>, ParseError>> {
        loop {
            match self.s.next() {
                Some((_, ' ')) => {}
                Some((_, '(')) => return Some(Ok(Token::LeftParen)),
                Some((_, ')')) => return Some(Ok(Token::RightParen)),
                Some((_, ',')) => return Some(Ok(Token::Comma)),
                Some((_, '=')) => return Some(Ok(Token::Equals)),
                Some((start, '"')) => {
                    while let Some((end, ch)) = self.s.next() {
                        if ch == '"' {
                            return Some(Ok(Token::String(&self.orig[start + 1..end])));
                        }
                    }
                    return Some(Err(ParseError::new(self.orig, UnterminatedString)));
                }
                Some((start, ch)) if is_ident_start(ch) => {
                    while let Some(&(end, ch)) = self.s.peek() {
                        if !is_ident_rest(ch) {
                            return Some(Ok(Token::Ident(&self.orig[start..end])));
                        } else {
                            self.s.next();
                        }
                    }
                    return Some(Ok(Token::Ident(&self.orig[start..])));
                }
                Some((_, ch)) => {
                    return Some(Err(ParseError::new(self.orig, UnexpectedChar(ch))));
                }
                None => return None,
            }
        }
    }
}

fn is_ident_start(ch: char) -> bool {
    ch == '_' || ch.is_ascii_alphabetic()
}

fn is_ident_rest(ch: char) -> bool {
    is_ident_start(ch) || ch.is_ascii_digit()
}

impl<'a> Token<'a> {
    fn classify(&self) -> &'static str {
        match *self {
            Token::LeftParen => "`(`",
            Token::RightParen => "`)`",
            Token::Ident(..) => "an identifier",
            Token::Comma => "`,`",
            Token::Equals => "`=`",
            Token::String(..) => "a string",
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Snapshot {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogId,
    /// the highest lid observed while generating the `Snapshot`
    pub max_lid: LogId,
    /// the highest allocated pid
    pub max_pid: PageId,
    /// the mapping from pages to (lsn, lid)
    pub pt: FastMap8<PageId, PageState>,
    /// replaced pages per segment index
    pub replacements:
        FastMap8<SegmentId, (Lsn, FastSet8<(PageId, SegmentId)>)>,
    /// the free pids
    pub free: FastSet8<PageId>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr)>),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr)) {
        match *self {
            PageState::Present(ref mut items) => items.push(item),
            PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> Box<dyn Iterator<Item = (Lsn, DiskPtr)>> {
        match *self {
            PageState::Present(ref items) => {
                Box::new(items.clone().into_iter())
            }
            PageState::Free(lsn, ptr) => {
                Box::new(vec![(lsn, ptr)].into_iter())
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Default for Snapshot {
    fn default() -> Snapshot {
        Snapshot {
            max_lsn: 0,
            max_lid: 0,
            last_lid: 0,
            max_pid: 0,
            pt: FastMap8::default(),
            replacements: FastMap8::default(),
            free: FastSet8::default(),
        }
    }
}

impl Snapshot {
    fn apply<P>(
        &mut self,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        bytes: &[u8],
        config: &Config,
    ) -> Result<()>
    where
        P: 'static
            + Debug
            + Clone
            + Serialize
            + DeserializeOwned
            + Send
            + Sync,
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!(
            "trying to deserialize buf for ptr {} lsn {}",
            disk_ptr,
            lsn
        );
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    ptr {}: {:?}",
                lsn,
                disk_ptr,
                e
            );
            return Err(Error::Corruption { at: disk_ptr });
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        let replaced_at_idx =
            disk_ptr.lid() as SegmentId / config.io_buf_size;
        let replaced_at_segment_lsn = (lsn
            / config.io_buf_size as Lsn)
            * config.io_buf_size as Lsn;

        match prepend.update {
            Update::Append(append) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}: {:?}",
                        pid,
                        disk_ptr,
                        lsn,
                        append,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return Ok(());
                    }

                    lids.push((lsn, disk_ptr));
                }
                self.free.remove(&pid);
            }
            Update::Meta(_)
            | Update::Counter(_)
            | Update::Compact(_) => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}: {:?}",
                    pid,
                    disk_ptr,
                    lsn,
                    prepend.update,
                );

                self.replace_pid(
                    pid,
                    replaced_at_segment_lsn,
                    replaced_at_idx,
                    config,
                );
                self.pt.insert(
                    pid,
                    PageState::Present(vec![(lsn, disk_ptr)]),
                );
                if prepend.update.is_compact() {
                    self.free.remove(&pid);
                } else {
                    assert!(!self.free.contains(&pid));
                }
            }
            Update::Free => {
                trace!(
                    "free of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_segment_lsn,
                    replaced_at_idx,
                    config,
                );
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
                self.free.insert(pid);
            }
        }

        Ok(())
    }

    fn replace_pid(
        &mut self,
        pid: PageId,
        replaced_at_segment_lsn: Lsn,
        replaced_at_idx: usize,
        config: &Config,
    ) {
        let replacements =
            self.replacements.entry(replaced_at_idx).or_insert((
                replaced_at_segment_lsn,
                FastSet8::default(),
            ));

        assert_eq!(replaced_at_segment_lsn, replacements.0);

        match self.pt.remove(&pid) {
            Some(PageState::Present(coords)) => {
                for (_lsn, ptr) in &coords {
                    let old_idx =
                        ptr.lid() as SegmentId / config.io_buf_size;
                    if replaced_at_idx == old_idx {
                        continue;
                    }
                    replacements.1.insert((pid, old_idx));
                }

                // re-run any blob removals in case
                // they were not completed. blobs are
                // not rewritten if they are the only
                // frag for a page during rewrite.
                if coords.len() > 1 {
                    let blob_ptrs = coords
                        .iter()
                        .filter(|(_, ptr)| ptr.is_blob())
                        .map(|(_, ptr)| ptr.blob().1);

                    for blob_ptr in blob_ptrs {
                        trace!(
                            "removing blob while advancing \
                             snapshot: {}",
                            blob_ptr,
                        );

                        // we don't care if this actually works
                        // because it's possible that a previous
                        // snapshot has run over this log and
                        // removed the blob already.
                        let _ = remove_blob(blob_ptr, config);
                    }
                }
            }
            Some(PageState::Free(_lsn, ptr)) => {
                let old_idx =
                    ptr.lid() as SegmentId / config.io_buf_size;
                if replaced_at_idx != old_idx {
                    replacements.1.insert((pid, old_idx));
                }
            }
            None => {
                // we just encountered a replace without it
                // existing in the pt. this means we've
                // relocated it to a later segment than
                // the one we're currently at during recovery.
                // so we can skip it here.
            }
        }
    }
}

pub(super) fn advance_snapshot<PM, P>(
    iter: LogIter,
    mut snapshot: Snapshot,
    config: &Config,
) -> Result<Snapshot>
where
    PM: Materializer<PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
{
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let io_buf_size = config.io_buf_size;

    let mut last_seg_lsn = snapshot.max_lsn / io_buf_size as Lsn;

    for (lsn, ptr, bytes) in iter {
        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} max_lsn {}",
                lsn,
                ptr,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = ptr.lid();
        if ptr.lid() > snapshot.max_lid {
            snapshot.max_lid = ptr.lid();
        }

        if lsn / (io_buf_size as Lsn) > last_seg_lsn {
            // invalidate any removed pids
            let segment_idx = ptr.lid() as usize / io_buf_size;

            snapshot.replacements.remove(&segment_idx);
            last_seg_lsn = lsn / (io_buf_size as Lsn);
        }

        if !PM::is_null() {
            if let Err(e) =
                snapshot.apply::<P>(lsn, ptr, &*bytes, config)
            {
                error!(
                    "encountered error while reading log message: {}",
                    e
                );
                break;
            }
        }
    }

    write_snapshot(config, &snapshot)?;

    trace!("generated new snapshot: {:?}", snapshot);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<PM, P>(
    config: &Config,
) -> Result<Snapshot>
where
    PM: Materializer<PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
{
    let last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.max_lsn, config)?;

    advance_snapshot::<PM, P>(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(
    config: &Config,
) -> std::io::Result<Option<Snapshot>> {
    let mut candidates = config.get_snapshot_files()?;
    if candidates.is_empty() {
        debug!("no previous snapshot found");
        return Ok(None);
    }

    candidates.sort();

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path)?;
    if f.metadata()?.len() <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    let mut len_expected_bytes = [0u8; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0u8; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 = arr_to_u64(&len_expected_bytes);
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot>(&*bytes).ok())
}

pub(crate) fn write_snapshot(
    config: &Config,
    snapshot: &Snapshot,
) -> Result<()> {
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix =
        format!("snap.{:016X}.in___motion", snapshot.max_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.max_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    std::fs::create_dir_all(path_1.parent().unwrap())?;
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc32)?;
    f.sync_all()?;
    maybe_fail!("snap write post");

    // This is important to avoid a race condition encountered on macos APFS
    // where the rename below fails due to the metadata for path_1 not being
    // present yet.
    drop(f);

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<commit_msg>Add fsync on the parent directory to avoid the APFS race condition<commit_after>use std::io::{Read, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
/// TODO consider splitting `Snapshot` into separate
/// snapshots for PC, SA, Materializer.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Snapshot {
    /// the last lsn included in the `Snapshot`
    pub max_lsn: Lsn,
    /// the last lid included in the `Snapshot`
    pub last_lid: LogId,
    /// the highest lid observed while generating the `Snapshot`
    pub max_lid: LogId,
    /// the highest allocated pid
    pub max_pid: PageId,
    /// the mapping from pages to (lsn, lid)
    pub pt: FastMap8<PageId, PageState>,
    /// replaced pages per segment index
    pub replacements:
        FastMap8<SegmentId, (Lsn, FastSet8<(PageId, SegmentId)>)>,
    /// the free pids
    pub free: FastSet8<PageId>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr)>),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr)) {
        match *self {
            PageState::Present(ref mut items) => items.push(item),
            PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> Box<dyn Iterator<Item = (Lsn, DiskPtr)>> {
        match *self {
            PageState::Present(ref items) => {
                Box::new(items.clone().into_iter())
            }
            PageState::Free(lsn, ptr) => {
                Box::new(vec![(lsn, ptr)].into_iter())
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Default for Snapshot {
    fn default() -> Snapshot {
        Snapshot {
            max_lsn: 0,
            max_lid: 0,
            last_lid: 0,
            max_pid: 0,
            pt: FastMap8::default(),
            replacements: FastMap8::default(),
            free: FastSet8::default(),
        }
    }
}

impl Snapshot {
    fn apply<P>(
        &mut self,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        bytes: &[u8],
        config: &Config,
    ) -> Result<()>
    where
        P: 'static
            + Debug
            + Clone
            + Serialize
            + DeserializeOwned
            + Send
            + Sync,
    {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!(
            "trying to deserialize buf for ptr {} lsn {}",
            disk_ptr,
            lsn
        );
        let deserialization = deserialize::<LoggedUpdate<P>>(&*bytes);

        if let Err(e) = deserialization {
            error!(
                "failed to deserialize buffer for item in log: lsn {} \
                    ptr {}: {:?}",
                lsn,
                disk_ptr,
                e
            );
            return Err(Error::Corruption { at: disk_ptr });
        }

        let prepend = deserialization.unwrap();
        let pid = prepend.pid;

        if pid >= self.max_pid {
            self.max_pid = pid + 1;
        }

        let replaced_at_idx =
            disk_ptr.lid() as SegmentId / config.io_buf_size;
        let replaced_at_segment_lsn = (lsn
            / config.io_buf_size as Lsn)
            * config.io_buf_size as Lsn;

        match prepend.update {
            Update::Append(append) => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}: {:?}",
                        pid,
                        disk_ptr,
                        lsn,
                        append,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return Ok(());
                    }

                    lids.push((lsn, disk_ptr));
                }
                self.free.remove(&pid);
            }
            Update::Meta(_)
            | Update::Counter(_)
            | Update::Compact(_) => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}: {:?}",
                    pid,
                    disk_ptr,
                    lsn,
                    prepend.update,
                );

                self.replace_pid(
                    pid,
                    replaced_at_segment_lsn,
                    replaced_at_idx,
                    config,
                );
                self.pt.insert(
                    pid,
                    PageState::Present(vec![(lsn, disk_ptr)]),
                );
                if prepend.update.is_compact() {
                    self.free.remove(&pid);
                } else {
                    assert!(!self.free.contains(&pid));
                }
            }
            Update::Free => {
                trace!(
                    "free of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn
                );
                self.replace_pid(
                    pid,
                    replaced_at_segment_lsn,
                    replaced_at_idx,
                    config,
                );
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
                self.free.insert(pid);
            }
        }

        Ok(())
    }

    fn replace_pid(
        &mut self,
        pid: PageId,
        replaced_at_segment_lsn: Lsn,
        replaced_at_idx: usize,
        config: &Config,
    ) {
        let replacements =
            self.replacements.entry(replaced_at_idx).or_insert((
                replaced_at_segment_lsn,
                FastSet8::default(),
            ));

        assert_eq!(replaced_at_segment_lsn, replacements.0);

        match self.pt.remove(&pid) {
            Some(PageState::Present(coords)) => {
                for (_lsn, ptr) in &coords {
                    let old_idx =
                        ptr.lid() as SegmentId / config.io_buf_size;
                    if replaced_at_idx == old_idx {
                        continue;
                    }
                    replacements.1.insert((pid, old_idx));
                }

                // re-run any blob removals in case
                // they were not completed. blobs are
                // not rewritten if they are the only
                // frag for a page during rewrite.
                if coords.len() > 1 {
                    let blob_ptrs = coords
                        .iter()
                        .filter(|(_, ptr)| ptr.is_blob())
                        .map(|(_, ptr)| ptr.blob().1);

                    for blob_ptr in blob_ptrs {
                        trace!(
                            "removing blob while advancing \
                             snapshot: {}",
                            blob_ptr,
                        );

                        // we don't care if this actually works
                        // because it's possible that a previous
                        // snapshot has run over this log and
                        // removed the blob already.
                        let _ = remove_blob(blob_ptr, config);
                    }
                }
            }
            Some(PageState::Free(_lsn, ptr)) => {
                let old_idx =
                    ptr.lid() as SegmentId / config.io_buf_size;
                if replaced_at_idx != old_idx {
                    replacements.1.insert((pid, old_idx));
                }
            }
            None => {
                // we just encountered a replace without it
                // existing in the pt. this means we've
                // relocated it to a later segment than
                // the one we're currently at during recovery.
                // so we can skip it here.
            }
        }
    }
}

pub(super) fn advance_snapshot<PM, P>(
    iter: LogIter,
    mut snapshot: Snapshot,
    config: &Config,
) -> Result<Snapshot>
where
    PM: Materializer<PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
{
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let io_buf_size = config.io_buf_size;

    let mut last_seg_lsn = snapshot.max_lsn / io_buf_size as Lsn;

    for (lsn, ptr, bytes) in iter {
        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.max_lsn {
            // don't process already-processed Lsn's. max_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} max_lsn {}",
                lsn,
                ptr,
                snapshot.max_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.max_lsn);
        snapshot.max_lsn = lsn;
        snapshot.last_lid = ptr.lid();
        if ptr.lid() > snapshot.max_lid {
            snapshot.max_lid = ptr.lid();
        }

        if lsn / (io_buf_size as Lsn) > last_seg_lsn {
            // invalidate any removed pids
            let segment_idx = ptr.lid() as usize / io_buf_size;

            snapshot.replacements.remove(&segment_idx);
            last_seg_lsn = lsn / (io_buf_size as Lsn);
        }

        if !PM::is_null() {
            if let Err(e) =
                snapshot.apply::<P>(lsn, ptr, &*bytes, config)
            {
                error!(
                    "encountered error while reading log message: {}",
                    e
                );
                break;
            }
        }
    }

    write_snapshot(config, &snapshot)?;

    trace!("generated new snapshot: {:?}", snapshot);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default<PM, P>(
    config: &Config,
) -> Result<Snapshot>
where
    PM: Materializer<PageFrag = P>,
    P: 'static
        + Debug
        + Clone
        + Serialize
        + DeserializeOwned
        + Send
        + Sync,
{
    let last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.max_lsn, config)?;

    advance_snapshot::<PM, P>(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(
    config: &Config,
) -> std::io::Result<Option<Snapshot>> {
    let mut candidates = config.get_snapshot_files()?;
    if candidates.is_empty() {
        debug!("no previous snapshot found");
        return Ok(None);
    }

    candidates.sort();

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path)?;
    if f.metadata()?.len() <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    let mut len_expected_bytes = [0u8; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0u8; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 = arr_to_u64(&len_expected_bytes);
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot>(&*bytes).ok())
}

pub(crate) fn write_snapshot(
    config: &Config,
    snapshot: &Snapshot,
) -> Result<()> {
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix =
        format!("snap.{:016X}.in___motion", snapshot.max_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.max_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc32)?;
    f.sync_all()?;
    maybe_fail!("snap write post");
    drop(f);

    // This is important to avoid a race condition encountered on macos APFS
    // where the rename below fails due to the metadata for path_1 not being
    // present yet.
    let dir = std::fs::File::open(parent)?;
    dir.sync_all()?;
    drop(dir);

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::io::{Read, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Snapshot {
    /// The last read message lsn
    pub last_lsn: Lsn,
    /// The last read message lid
    pub last_lid: LogId,
    /// the mapping from pages to (lsn, lid)
    pub pt: FastMap8<PageId, PageState>,
    /// The highest stable offset persisted
    /// into a segment header
    pub max_header_stable_lsn: Lsn,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr, usize)>),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, usize)) {
        match *self {
            PageState::Present(ref mut items) => items.push(item),
            PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> impl Iterator<Item = (Lsn, DiskPtr, usize)> {
        match *self {
            PageState::Present(ref items) => items.clone().into_iter(),
            PageState::Free(lsn, ptr) => {
                vec![(lsn, ptr, MSG_HEADER_LEN)].into_iter()
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: usize,
    ) {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for ptr {} lsn {}", disk_ptr, lsn);

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt
                    .insert(pid, PageState::Present(vec![(lsn, disk_ptr, sz)]));
            }
            LogKind::Append => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return;
                    }

                    lids.push((lsn, disk_ptr, sz));
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

pub(super) fn advance_snapshot(
    iter: LogIter,
    mut snapshot: Snapshot,
    config: &Config,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let old_lsn = snapshot.last_lsn;

    for (log_kind, pid, lsn, ptr, sz) in iter {
        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.last_lsn {
            // don't process already-processed Lsn's. last_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} last_lsn {}",
                lsn,
                ptr,
                snapshot.last_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.last_lsn);
        snapshot.last_lsn = lsn;
        snapshot.last_lid = ptr.lid();

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    if snapshot.last_lsn != old_lsn {
        write_snapshot(config, &snapshot)?;
    }

    trace!("generated new snapshot: {:?}", snapshot);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default(config: &Config) -> Result<Snapshot> {
    let mut last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let (log_iter, max_header_stable_lsn) =
        raw_segment_iter_from(last_snap.last_lsn, config)?;

    last_snap.max_header_stable_lsn = max_header_stable_lsn;

    advance_snapshot(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &Config) -> std::io::Result<Option<Snapshot>> {
    let mut candidates = config.get_snapshot_files()?;
    if candidates.is_empty() {
        debug!("no previous snapshot found");
        return Ok(None);
    }

    candidates.sort();

    let path = candidates.pop().unwrap();

    let mut f = std::fs::OpenOptions::new().read(true).open(&path)?;
    if f.metadata()?.len() <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    let mut len_expected_bytes = [0u8; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0u8; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 = arr_to_u64(&len_expected_bytes);
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot>(&*bytes).ok())
}

fn write_snapshot(config: &Config, snapshot: &Snapshot) -> Result<()> {
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix = format!("snap.{:016X}.generating", snapshot.last_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.last_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc32)?;
    maybe_fail!("snap write post");

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<commit_msg>retry snapshot search if a file is deleted before we can read it<commit_after>use std::io::{Read, Write};

#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use super::*;

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Snapshot {
    /// The last read message lsn
    pub last_lsn: Lsn,
    /// The last read message lid
    pub last_lid: LogId,
    /// the mapping from pages to (lsn, lid)
    pub pt: FastMap8<PageId, PageState>,
    /// The highest stable offset persisted
    /// into a segment header
    pub max_header_stable_lsn: Lsn,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PageState {
    Present(Vec<(Lsn, DiskPtr, usize)>),
    Free(Lsn, DiskPtr),
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, usize)) {
        match *self {
            PageState::Present(ref mut items) => items.push(item),
            PageState::Free(_, _) => {
                panic!("pushed items to a PageState::Free")
            }
        }
    }

    /// Iterate over the (lsn, lid) pairs that hold this page's state.
    pub fn iter(&self) -> impl Iterator<Item = (Lsn, DiskPtr, usize)> {
        match *self {
            PageState::Present(ref items) => items.clone().into_iter(),
            PageState::Free(lsn, ptr) => {
                vec![(lsn, ptr, MSG_HEADER_LEN)].into_iter()
            }
        }
    }

    fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: usize,
    ) {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for ptr {} lsn {}", disk_ptr, lsn);

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt
                    .insert(pid, PageState::Present(vec![(lsn, disk_ptr, sz)]));
            }
            LogKind::Append => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids) = self.pt.get_mut(&pid) {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return;
                    }

                    lids.push((lsn, disk_ptr, sz));
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt.insert(pid, PageState::Free(lsn, disk_ptr));
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

pub(super) fn advance_snapshot(
    iter: LogIter,
    mut snapshot: Snapshot,
    config: &Config,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let old_lsn = snapshot.last_lsn;

    for (log_kind, pid, lsn, ptr, sz) in iter {
        trace!(
            "in advance_snapshot looking at item with lsn {} ptr {}",
            lsn,
            ptr
        );

        if lsn <= snapshot.last_lsn {
            // don't process already-processed Lsn's. last_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} last_lsn {}",
                lsn,
                ptr,
                snapshot.last_lsn
            );
            continue;
        }

        assert!(lsn > snapshot.last_lsn);
        snapshot.last_lsn = lsn;
        snapshot.last_lid = ptr.lid();

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    if snapshot.last_lsn != old_lsn {
        write_snapshot(config, &snapshot)?;
    }

    trace!("generated new snapshot: {:?}", snapshot);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default(config: &Config) -> Result<Snapshot> {
    let mut last_snap =
        read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let (log_iter, max_header_stable_lsn) =
        raw_segment_iter_from(last_snap.last_lsn, config)?;

    last_snap.max_header_stable_lsn = max_header_stable_lsn;

    advance_snapshot(log_iter, last_snap, config)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &Config) -> std::io::Result<Option<Snapshot>> {
    let mut f = loop {
        let mut candidates = config.get_snapshot_files()?;
        if candidates.is_empty() {
            debug!("no previous snapshot found");
            return Ok(None);
        }

        candidates.sort();

        let path = candidates.pop().unwrap();

        match std::fs::OpenOptions::new().read(true).open(&path) {
            Ok(f) => break f,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // this can happen if there's a race
                continue;
            }
            Err(other) => return Err(other.into()),
        }
    };

    if f.metadata()?.len() <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut buf = vec![];
    f.read_to_end(&mut buf)?;
    let len = buf.len();
    let mut len_expected_bytes = [0u8; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0u8; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        let len_expected: u64 = arr_to_u64(&len_expected_bytes);
        decompress(&*buf, len_expected as usize).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(deserialize::<Snapshot>(&*bytes).ok())
}

fn write_snapshot(config: &Config, snapshot: &Snapshot) -> Result<()> {
    let raw_bytes = serialize(&snapshot).unwrap();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix = format!("snap.{:016X}.generating", snapshot.last_lsn);

    let mut path_1 = config.snapshot_prefix();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.last_lsn);

    let mut path_2 = config.snapshot_prefix();
    path_2.push(path_2_suffix);

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    let mut f = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    maybe_fail!("snap write");
    f.write_all(&*bytes)?;
    maybe_fail!("snap write len");
    f.write_all(&len_bytes)?;
    maybe_fail!("snap write crc");
    f.write_all(&crc32)?;
    maybe_fail!("snap write post");

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    maybe_fail!("snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    maybe_fail!("snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            maybe_fail!("snap write rm old");

            if let Err(_e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    _e
                );
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019 Jeehoon Kang
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use core::mem::{self, MaybeUninit};
use core::ptr;

use crate::addr::*;
use crate::arch::*;
use crate::mm::*;
use crate::page::*;
use crate::spinlock::*;
use crate::std::*;
use crate::types::*;
use crate::vm::*;

/// From inc/hf/arch/cpu.h.
extern "C" {
    /// Disables interrupts.
    fn arch_irq_enable();

    /// Enables interrupts.
    fn arch_irq_disable();

    /// Reset the register values other than the PC and argument which are set with
    /// `arch_regs_set_pc_arg()`.
    fn arch_regs_reset(
        r: *mut ArchRegs,
        is_primary: bool,
        vm_id: spci_vm_id_t,
        vcpu_id: cpu_id_t,
        table: paddr_t,
    );

    /// Updates the given registers so that when a vcpu runs, it starts off at the
    /// given address (pc) with the given argument.
    ///
    /// This function must only be called on an arch_regs that is known not be in use
    /// by any other physical CPU.
    fn arch_regs_set_pc_arg(r: *mut ArchRegs, pc: ipaddr_t, arg: uintreg_t);

    /// Updates the register holding the return value of a function.
    ///
    /// This function must only be called on an arch_regs that is known not be in use
    /// by any other physical CPU.
    fn arch_regs_set_retval(r: *mut ArchRegs, v: uintreg_t);
}

const STACK_SIZE: usize = PAGE_SIZE;

/// The number of bits in each element of the interrupt bitfields.
const INTERRUPT_REGISTER_BITS: usize = 32;

#[repr(C)]
pub enum VCpuStatus {
    /// The vcpu is switched off.
    Off,

    /// The vcpu is ready to be run.
    Ready,

    /// The vcpu is currently running.
    Running,

    /// The vcpu is waiting for a message.
    BlockedMailbox,

    /// The vcpu is waiting for an interrupt.
    BlockedInterrupt,

    /// The vcpu has aborted.
    Aborted,
}

#[repr(C)]
pub struct Interrupts {
    /// Bitfield keeping track of which interrupts are enabled.
    enabled: [u32; HF_NUM_INTIDS / INTERRUPT_REGISTER_BITS],

    /// Bitfield keeping track of which interrupts are pending.
    pending: [u32; HF_NUM_INTIDS / INTERRUPT_REGISTER_BITS],

    /// The number of interrupts which are currently both enabled and pending. i.e. the number of
    /// bits set in enabled & pending.
    enabled_and_pending_count: u32,
}

#[repr(C)]
pub struct VCpuFaultInfo {
    ipaddr: ipaddr_t,
    vaddr: vaddr_t,
    pc: vaddr_t,
    mode: c_int,
}

#[repr(C)]
pub struct VCpu {
    lock: RawSpinLock,

    /// The state is only changed in the context of the vCPU being run. This ensures the scheduler
    /// can easily keep track of the vCPU state as transitions are indicated by the return code from
    /// the run call.
    state: VCpuStatus,

    cpu: *mut Cpu,
    vm: *mut Vm,
    regs: ArchRegs,
    interrupts: Interrupts,

    /// Determine whether the 'regs' field is available for use. This is set
    /// to false when a vCPU is about to run on a physical CPU, and is set
    /// back to true when it is descheduled.
    regs_available: bool,
}

/// Encapsulates a vCPU whose lock is held.
#[repr(C)]
#[derive(Copy, Clone)]
pub struct VCpuLocked {
    vcpu: *mut VCpu,
}

// TODO: Update alignment such that cpus are in different cache lines.
#[repr(C)]
pub struct Cpu {
    /// CPU identifier. Doesn't have to be contiguous.
    id: cpu_id_t,

    /// Pointer to bottom of the stack.
    stack_bottom: *mut c_void,

    /// Enabling/disabling irqs are counted per-cpu. They are enabled when the count is zero, and
    /// disabled when it's non-zero.
    irq_disable_count: u32,

    /// See api.c for the partial ordering on locks.
    lock: RawSpinLock,

    /// Determines whether or not the cpu is currently on.
    is_on: bool,
}

// TODO: alignas(2 * sizeof(uintreg_t))
// #[repr(align(16))]
static mut callstacks: MaybeUninit<[[u8; STACK_SIZE]; MAX_CPUS]> = MaybeUninit::uninit();

/// State of all supported CPUs. The stack of the first one is initialized.
/// struct cpu cpus[MAX_CPUS] = {
///     {
// 	        .is_on = 1,
///         .stack_bottom = &callstacks[0][STACK_SIZE],
///     },
/// };
static mut cpus: MaybeUninit<[Cpu; MAX_CPUS]> = MaybeUninit::uninit();
static mut cpu_count: u32 = 1;

#[no_mangle]
pub unsafe extern "C" fn cpu_init(c: *mut Cpu) {
    // TODO: Assumes that c is zeroed out already.
    sl_init(&mut (*c).lock);
}

#[no_mangle]
pub unsafe extern "C" fn cpu_module_init(cpu_ids: *mut cpu_id_t, count: usize) {
    let mut j: u32;
    let boot_cpu_id: cpu_id_t = cpus.get_ref()[0].id;
    let mut found_boot_cpu: bool = false;

    cpu_count = count as u32;

    // Initialize CPUs with the IDs from the configuration passed in. The
    // CPUs after the boot CPU are initialized in reverse order. The boot
    // CPU is initialized when it is found or in place of the last CPU if it
    // is not found.
    j = cpu_count;
    for i in 0..cpu_count {
        let c: *mut Cpu;
        let id: cpu_id_t = *cpu_ids.offset(i as isize);

        if found_boot_cpu || id != boot_cpu_id {
            j -= 1;
            c = &mut cpus.get_mut()[j as usize];
        } else {
            found_boot_cpu = true;
            c = &mut cpus.get_mut()[0];
        }

        cpu_init(c);
        (*c).id = id;
        {
            let callstacks_i = callstacks.get_mut()[i as usize].as_mut_ptr();
            (*c).stack_bottom = &mut *callstacks_i.offset(STACK_SIZE as isize) as *mut _ as _;

            // Note: referencing callstacks.get_mut()[i as usize][STACK_SIZE] directly
            // causes 'index out of bounds' error on compile time.
        }
    }

    if !found_boot_cpu {
        // Boot CPU was initialized but with wrong ID.
        dlog!("Boot CPU's ID not found in config.");
        cpus.get_mut()[0].id = boot_cpu_id;
    }
}

#[no_mangle]
pub unsafe extern "C" fn cpu_index(c: *mut Cpu) -> usize {
    c.offset_from(cpus.get_ref().as_ptr()) as usize
}

/// Turns CPU on and returns the previous state.
#[no_mangle]
pub unsafe extern "C" fn cpu_on(c: *mut Cpu, entry: ipaddr_t, arg: uintreg_t) -> bool {
    let prev: bool;

    sl_lock(&(*c).lock);
    prev = (*c).is_on;
    (*c).is_on = true;
    sl_unlock(&(*c).lock);

    if !prev {
        let vm = vm_find(HF_PRIMARY_VM_ID);
        let vcpu = vm_get_vcpu(vm, cpu_index(c) as u16);
        let mut vcpu_locked = vcpu_lock(vcpu);

        vcpu_on(vcpu_locked, entry, arg);
        vcpu_unlock(&mut vcpu_locked);
    }

    prev
}

/// Prepares the CPU for turning itself off.
#[no_mangle]
pub unsafe extern "C" fn cpu_off(c: *mut Cpu) {
    sl_lock(&(*c).lock);
    (*c).is_on = true;
    sl_unlock(&(*c).lock);
}

/// Searches for a CPU based on its id.
#[no_mangle]
pub unsafe extern "C" fn cpu_find(id: cpu_id_t) -> *mut Cpu {
    for i in 0usize..cpu_count as usize {
        if cpus.get_ref()[i].id == id {
            return &mut cpus.get_mut()[i];
        }
    }

    ptr::null_mut()
}

/// Locks the given vCPU and updates `locked` to hold the newly locked vCPU.
#[no_mangle]
pub unsafe extern "C" fn vcpu_lock(vcpu: *mut VCpu) -> VCpuLocked {
    sl_lock(&(*vcpu).lock);

    VCpuLocked { vcpu }
}

/// Unlocks a vCPU previously locked with vcpu_lock, and updates `locked` to
/// reflect the fact that the vCPU is no longer locked.
#[no_mangle]
pub unsafe extern "C" fn vcpu_unlock(locked: *mut VCpuLocked) {
    sl_unlock(&(*(*locked).vcpu).lock);
    (*locked).vcpu = ptr::null_mut();
}

#[no_mangle]
pub unsafe extern "C" fn vcpu_init(vcpu: *mut VCpu, vm: *mut Vm) {
    memset_s(vcpu as _, mem::size_of::<VCpu>(), 0, mem::size_of::<VCpu>());
    sl_lock(&(*vcpu).lock);
    (*vcpu).regs_available = true;
    (*vcpu).vm = vm;
    (*vcpu).state = VCpuStatus::Off;
}

/// Initialise the registers for the given vCPU and set the state to
/// VCPU_STATE_READY. The caller must hold the vCPU lock while calling this.
#[no_mangle]
pub unsafe extern "C" fn vcpu_on(vcpu: VCpuLocked, entry: ipaddr_t, arg: uintreg_t) {
    arch_regs_set_pc_arg(&mut (*vcpu.vcpu).regs, entry, arg);
    (*vcpu.vcpu).state = VCpuStatus::Ready;
}

#[no_mangle]
pub unsafe extern "C" fn vcpu_index(vcpu: *const VCpu) -> spci_vcpu_index_t {
    let vcpus = (*(*vcpu).vm).vcpus.as_ptr();
    let index = vcpu.offset_from(vcpus);
    assert!(index < core::u16::MAX as isize);
    index as u16
}

/// Check whether the given vcpu_state is an off state, for the purpose of
/// turning vCPUs on and off. Note that aborted still counts as on in this
/// context.
#[no_mangle]
pub unsafe extern "C" fn vcpu_is_off(vcpu: VCpuLocked) -> bool {
    match (*vcpu.vcpu).state {
        VCpuStatus::Off => true,
        _ =>
        // Aborted still counts as ON for the purposes of PSCI,
        // because according to the PSCI specification (section
        // 5.7.1) a core is only considered to be off if it has
        // been turned off with a CPU_OFF call or hasn't yet
        // been turned on with a CPU_ON call.
        {
            false
        }
    }
}

/// Starts a vCPU of a secondary VM.
///
/// Returns true if the secondary was reset and started, or false if it was
/// already on and so nothing was done.
#[no_mangle]
pub unsafe extern "C" fn vcpu_secondary_reset_and_start(
    vcpu: *mut VCpu,
    entry: ipaddr_t,
    arg: uintreg_t,
) -> bool {
    let mut vcpu_locked;
    let vm = (*vcpu).vm;
    let vcpu_was_off;

    assert!((*vm).id != HF_PRIMARY_VM_ID);

    vcpu_locked = vcpu_lock(vcpu);
    vcpu_was_off = vcpu_is_off(vcpu_locked);
    if vcpu_was_off {
        // Set vCPU registers to a clean state ready for boot. As this
        // is a secondary which can migrate between pCPUs, the ID of the
        // vCPU is defined as the index and does not match the ID of the
        // pCPU it is running on.
        arch_regs_reset(
            &mut (*vcpu).regs,
            false,
            (*vm).id,
            vcpu_index(vcpu) as u32,
            pa_init((*vm).ptable.root),
        );
        vcpu_on(vcpu_locked, entry, arg);
    }

    vcpu_unlock(&mut vcpu_locked);
    vcpu_was_off
}

/// Handles a page fault. It does so by determining if it's a legitimate or
/// spurious fault, and recovering from the latter.
///
/// Returns true if the caller should resume the current vcpu, or false if its
/// VM should be aborted.
#[no_mangle]
pub unsafe extern "C" fn vcpu_handle_page_fault(
    current: *const VCpu,
    f: *mut VCpuFaultInfo,
) -> bool {
    let vm = (*current).vm;
    let mut mode = mem::uninitialized(); // to avoid use-of-uninitialized error
    let mask = (*f).mode | Mode::INVALID.bits() as i32;
    let resume;

    sl_lock(&(*vm).lock);

    // Check if this is a legitimate fault, i.e., if the page table doesn't
    // allow the access attemped by the VM.
    //
    // Otherwise, this is a spurious fault, likely because another CPU is
    // updating the page table. It is responsible for issuing global TLB
    // invalidations while holding the VM lock, so we don't need to do
    // anything else to recover from it. (Acquiring/releasing the lock
    // ensured that the invalidations have completed.)
    resume = mm_vm_get_mode(
        &mut (*vm).ptable,
        (*f).ipaddr,
        ipa_add((*f).ipaddr, 1),
        &mut mode,
    ) && (mode & mask) == (*f).mode;

    sl_unlock(&(*vm).lock);

    if !resume {
        dlog!(
            "Stage-2 page fault: pc=0x{}, vmid={}, vcpu={}, vaddr=0x{}, ipaddr=0x{}, mode=0x{}\n",
            (*f).pc,
            (*vm).id,
            vcpu_index(current),
            (*f).vaddr,
            (*f).ipaddr,
            (*f).mode
        );
    }

    resume
}
<commit_msg>Type casting respect to cross compilation.<commit_after>/*
 * Copyright 2019 Jeehoon Kang
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use core::mem::{self, MaybeUninit};
use core::ptr;

use crate::addr::*;
use crate::arch::*;
use crate::mm::*;
use crate::page::*;
use crate::spinlock::*;
use crate::std::*;
use crate::types::*;
use crate::vm::*;

/// From inc/hf/arch/cpu.h.
extern "C" {
    /// Disables interrupts.
    fn arch_irq_enable();

    /// Enables interrupts.
    fn arch_irq_disable();

    /// Reset the register values other than the PC and argument which are set with
    /// `arch_regs_set_pc_arg()`.
    fn arch_regs_reset(
        r: *mut ArchRegs,
        is_primary: bool,
        vm_id: spci_vm_id_t,
        vcpu_id: cpu_id_t,
        table: paddr_t,
    );

    /// Updates the given registers so that when a vcpu runs, it starts off at the
    /// given address (pc) with the given argument.
    ///
    /// This function must only be called on an arch_regs that is known not be in use
    /// by any other physical CPU.
    fn arch_regs_set_pc_arg(r: *mut ArchRegs, pc: ipaddr_t, arg: uintreg_t);

    /// Updates the register holding the return value of a function.
    ///
    /// This function must only be called on an arch_regs that is known not be in use
    /// by any other physical CPU.
    fn arch_regs_set_retval(r: *mut ArchRegs, v: uintreg_t);
}

const STACK_SIZE: usize = PAGE_SIZE;

/// The number of bits in each element of the interrupt bitfields.
const INTERRUPT_REGISTER_BITS: usize = 32;

#[repr(C)]
pub enum VCpuStatus {
    /// The vcpu is switched off.
    Off,

    /// The vcpu is ready to be run.
    Ready,

    /// The vcpu is currently running.
    Running,

    /// The vcpu is waiting for a message.
    BlockedMailbox,

    /// The vcpu is waiting for an interrupt.
    BlockedInterrupt,

    /// The vcpu has aborted.
    Aborted,
}

#[repr(C)]
pub struct Interrupts {
    /// Bitfield keeping track of which interrupts are enabled.
    enabled: [u32; HF_NUM_INTIDS / INTERRUPT_REGISTER_BITS],

    /// Bitfield keeping track of which interrupts are pending.
    pending: [u32; HF_NUM_INTIDS / INTERRUPT_REGISTER_BITS],

    /// The number of interrupts which are currently both enabled and pending. i.e. the number of
    /// bits set in enabled & pending.
    enabled_and_pending_count: u32,
}

#[repr(C)]
pub struct VCpuFaultInfo {
    ipaddr: ipaddr_t,
    vaddr: vaddr_t,
    pc: vaddr_t,
    mode: c_int,
}

#[repr(C)]
pub struct VCpu {
    lock: RawSpinLock,

    /// The state is only changed in the context of the vCPU being run. This ensures the scheduler
    /// can easily keep track of the vCPU state as transitions are indicated by the return code from
    /// the run call.
    state: VCpuStatus,

    cpu: *mut Cpu,
    vm: *mut Vm,
    regs: ArchRegs,
    interrupts: Interrupts,

    /// Determine whether the 'regs' field is available for use. This is set
    /// to false when a vCPU is about to run on a physical CPU, and is set
    /// back to true when it is descheduled.
    regs_available: bool,
}

/// Encapsulates a vCPU whose lock is held.
#[repr(C)]
#[derive(Copy, Clone)]
pub struct VCpuLocked {
    vcpu: *mut VCpu,
}

// TODO: Update alignment such that cpus are in different cache lines.
#[repr(C)]
pub struct Cpu {
    /// CPU identifier. Doesn't have to be contiguous.
    id: cpu_id_t,

    /// Pointer to bottom of the stack.
    stack_bottom: *mut c_void,

    /// Enabling/disabling irqs are counted per-cpu. They are enabled when the count is zero, and
    /// disabled when it's non-zero.
    irq_disable_count: u32,

    /// See api.c for the partial ordering on locks.
    lock: RawSpinLock,

    /// Determines whether or not the cpu is currently on.
    is_on: bool,
}

// TODO: alignas(2 * sizeof(uintreg_t))
// #[repr(align(16))]
static mut callstacks: MaybeUninit<[[u8; STACK_SIZE]; MAX_CPUS]> = MaybeUninit::uninit();

/// State of all supported CPUs. The stack of the first one is initialized.
/// struct cpu cpus[MAX_CPUS] = {
///     {
// 	        .is_on = 1,
///         .stack_bottom = &callstacks[0][STACK_SIZE],
///     },
/// };
static mut cpus: MaybeUninit<[Cpu; MAX_CPUS]> = MaybeUninit::uninit();
static mut cpu_count: u32 = 1;

#[no_mangle]
pub unsafe extern "C" fn cpu_init(c: *mut Cpu) {
    // TODO: Assumes that c is zeroed out already.
    sl_init(&mut (*c).lock);
}

#[no_mangle]
pub unsafe extern "C" fn cpu_module_init(cpu_ids: *mut cpu_id_t, count: usize) {
    let mut j: u32;
    let boot_cpu_id: cpu_id_t = cpus.get_ref()[0].id;
    let mut found_boot_cpu: bool = false;

    cpu_count = count as u32;

    // Initialize CPUs with the IDs from the configuration passed in. The
    // CPUs after the boot CPU are initialized in reverse order. The boot
    // CPU is initialized when it is found or in place of the last CPU if it
    // is not found.
    j = cpu_count;
    for i in 0..cpu_count {
        let c: *mut Cpu;
        let id: cpu_id_t = *cpu_ids.offset(i as isize);

        if found_boot_cpu || id != boot_cpu_id {
            j -= 1;
            c = &mut cpus.get_mut()[j as usize];
        } else {
            found_boot_cpu = true;
            c = &mut cpus.get_mut()[0];
        }

        cpu_init(c);
        (*c).id = id;
        {
            let callstacks_i = callstacks.get_mut()[i as usize].as_mut_ptr();
            (*c).stack_bottom = &mut *callstacks_i.offset(STACK_SIZE as isize) as *mut _ as _;

            // Note: referencing callstacks.get_mut()[i as usize][STACK_SIZE] directly
            // causes 'index out of bounds' error on compile time.
        }
    }

    if !found_boot_cpu {
        // Boot CPU was initialized but with wrong ID.
        dlog!("Boot CPU's ID not found in config.");
        cpus.get_mut()[0].id = boot_cpu_id;
    }
}

#[no_mangle]
pub unsafe extern "C" fn cpu_index(c: *mut Cpu) -> usize {
    c.offset_from(cpus.get_ref().as_ptr()) as usize
}

/// Turns CPU on and returns the previous state.
#[no_mangle]
pub unsafe extern "C" fn cpu_on(c: *mut Cpu, entry: ipaddr_t, arg: uintreg_t) -> bool {
    let prev: bool;

    sl_lock(&(*c).lock);
    prev = (*c).is_on;
    (*c).is_on = true;
    sl_unlock(&(*c).lock);

    if !prev {
        let vm = vm_find(HF_PRIMARY_VM_ID);
        let vcpu = vm_get_vcpu(vm, cpu_index(c) as u16);
        let mut vcpu_locked = vcpu_lock(vcpu);

        vcpu_on(vcpu_locked, entry, arg);
        vcpu_unlock(&mut vcpu_locked);
    }

    prev
}

/// Prepares the CPU for turning itself off.
#[no_mangle]
pub unsafe extern "C" fn cpu_off(c: *mut Cpu) {
    sl_lock(&(*c).lock);
    (*c).is_on = true;
    sl_unlock(&(*c).lock);
}

/// Searches for a CPU based on its id.
#[no_mangle]
pub unsafe extern "C" fn cpu_find(id: cpu_id_t) -> *mut Cpu {
    for i in 0usize..cpu_count as usize {
        if cpus.get_ref()[i].id == id {
            return &mut cpus.get_mut()[i];
        }
    }

    ptr::null_mut()
}

/// Locks the given vCPU and updates `locked` to hold the newly locked vCPU.
#[no_mangle]
pub unsafe extern "C" fn vcpu_lock(vcpu: *mut VCpu) -> VCpuLocked {
    sl_lock(&(*vcpu).lock);

    VCpuLocked { vcpu }
}

/// Unlocks a vCPU previously locked with vcpu_lock, and updates `locked` to
/// reflect the fact that the vCPU is no longer locked.
#[no_mangle]
pub unsafe extern "C" fn vcpu_unlock(locked: *mut VCpuLocked) {
    sl_unlock(&(*(*locked).vcpu).lock);
    (*locked).vcpu = ptr::null_mut();
}

#[no_mangle]
pub unsafe extern "C" fn vcpu_init(vcpu: *mut VCpu, vm: *mut Vm) {
    memset_s(vcpu as _, mem::size_of::<VCpu>(), 0, mem::size_of::<VCpu>());
    sl_lock(&(*vcpu).lock);
    (*vcpu).regs_available = true;
    (*vcpu).vm = vm;
    (*vcpu).state = VCpuStatus::Off;
}

/// Initialise the registers for the given vCPU and set the state to
/// VCPU_STATE_READY. The caller must hold the vCPU lock while calling this.
#[no_mangle]
pub unsafe extern "C" fn vcpu_on(vcpu: VCpuLocked, entry: ipaddr_t, arg: uintreg_t) {
    arch_regs_set_pc_arg(&mut (*vcpu.vcpu).regs, entry, arg);
    (*vcpu.vcpu).state = VCpuStatus::Ready;
}

#[no_mangle]
pub unsafe extern "C" fn vcpu_index(vcpu: *const VCpu) -> spci_vcpu_index_t {
    let vcpus = (*(*vcpu).vm).vcpus.as_ptr();
    let index = vcpu.offset_from(vcpus);
    assert!(index < core::u16::MAX as isize);
    index as u16
}

/// Check whether the given vcpu_state is an off state, for the purpose of
/// turning vCPUs on and off. Note that aborted still counts as on in this
/// context.
#[no_mangle]
pub unsafe extern "C" fn vcpu_is_off(vcpu: VCpuLocked) -> bool {
    match (*vcpu.vcpu).state {
        VCpuStatus::Off => true,
        _ =>
        // Aborted still counts as ON for the purposes of PSCI,
        // because according to the PSCI specification (section
        // 5.7.1) a core is only considered to be off if it has
        // been turned off with a CPU_OFF call or hasn't yet
        // been turned on with a CPU_ON call.
        {
            false
        }
    }
}

/// Starts a vCPU of a secondary VM.
///
/// Returns true if the secondary was reset and started, or false if it was
/// already on and so nothing was done.
#[no_mangle]
pub unsafe extern "C" fn vcpu_secondary_reset_and_start(
    vcpu: *mut VCpu,
    entry: ipaddr_t,
    arg: uintreg_t,
) -> bool {
    let mut vcpu_locked;
    let vm = (*vcpu).vm;
    let vcpu_was_off;

    assert!((*vm).id != HF_PRIMARY_VM_ID);

    vcpu_locked = vcpu_lock(vcpu);
    vcpu_was_off = vcpu_is_off(vcpu_locked);
    if vcpu_was_off {
        // Set vCPU registers to a clean state ready for boot. As this
        // is a secondary which can migrate between pCPUs, the ID of the
        // vCPU is defined as the index and does not match the ID of the
        // pCPU it is running on.
        arch_regs_reset(
            &mut (*vcpu).regs,
            false,
            (*vm).id,
            vcpu_index(vcpu) as cpu_id_t,
            pa_init((*vm).ptable.root),
        );
        vcpu_on(vcpu_locked, entry, arg);
    }

    vcpu_unlock(&mut vcpu_locked);
    vcpu_was_off
}

/// Handles a page fault. It does so by determining if it's a legitimate or
/// spurious fault, and recovering from the latter.
///
/// Returns true if the caller should resume the current vcpu, or false if its
/// VM should be aborted.
#[no_mangle]
pub unsafe extern "C" fn vcpu_handle_page_fault(
    current: *const VCpu,
    f: *mut VCpuFaultInfo,
) -> bool {
    let vm = (*current).vm;
    let mut mode = mem::uninitialized(); // to avoid use-of-uninitialized error
    let mask = (*f).mode | Mode::INVALID.bits() as i32;
    let resume;

    sl_lock(&(*vm).lock);

    // Check if this is a legitimate fault, i.e., if the page table doesn't
    // allow the access attemped by the VM.
    //
    // Otherwise, this is a spurious fault, likely because another CPU is
    // updating the page table. It is responsible for issuing global TLB
    // invalidations while holding the VM lock, so we don't need to do
    // anything else to recover from it. (Acquiring/releasing the lock
    // ensured that the invalidations have completed.)
    resume = mm_vm_get_mode(
        &mut (*vm).ptable,
        (*f).ipaddr,
        ipa_add((*f).ipaddr, 1),
        &mut mode,
    ) && (mode & mask) == (*f).mode;

    sl_unlock(&(*vm).lock);

    if !resume {
        dlog!(
            "Stage-2 page fault: pc=0x{}, vmid={}, vcpu={}, vaddr=0x{}, ipaddr=0x{}, mode=0x{}\n",
            (*f).pc,
            (*vm).id,
            vcpu_index(current),
            (*f).vaddr,
            (*f).ipaddr,
            (*f).mode
        );
    }

    resume
}
<|endoftext|>


<commit_before>use crate::syntax::report::Errors;
use crate::syntax::Api;
use serde::Deserialize;

type Node = clang_ast::Node<Clang>;

#[derive(Deserialize)]
enum Clang {
    Unknown,
}

pub fn load(cx: &mut Errors, apis: &mut [Api]) {
    let _ = cx;
    let _ = apis;
    unimplemented!()
}
<commit_msg>Determine whether AST dump loading is necessary<commit_after>use crate::syntax::report::Errors;
use crate::syntax::Api;
use serde::Deserialize;

type Node = clang_ast::Node<Clang>;

#[derive(Deserialize)]
enum Clang {
    Unknown,
}

pub fn load(cx: &mut Errors, apis: &mut [Api]) {
    let mut variants_from_header = Vec::new();
    for api in apis {
        if let Api::Enum(enm) = api {
            if enm.variants_from_header {
                variants_from_header.push(enm);
            }
        }
    }

    if variants_from_header.is_empty() {
        return;
    }

    let _ = cx;
    unimplemented!()
}
<|endoftext|>


<commit_before>///Represents a type that can have strum metadata associated with it.
pub trait HasMetadata {
    /// Get all the metadata associated with a specific "tag".
    /// All of strum's metadata is nested inside a path such as
    /// #[strum(...)] so this let's us quickly filter down to only our metadata.
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta>;
}

fn get_metadata_inner<'a>(
    ident: &str,
    it: impl IntoIterator<Item = &'a syn::Attribute>,
) -> Vec<syn::Meta> {
    it.into_iter()
        .filter_map(|attr| attr.parse_meta().ok())
        .filter_map(|meta| match meta {
            syn::Meta::List(syn::MetaList { path, nested, .. }) => {
                if path.is_ident(ident) {
                    Some(nested)
                } else {
                    None
                }
            }
            _ => None,
        })
        .flat_map(|id| id)
        .map(|nested| match nested {
            syn::NestedMeta::Meta(meta) => meta,
            _ => panic!("unexpected literal parsing strum attributes"),
        })
        .collect()
}

impl HasMetadata for syn::Variant {
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta> {
        get_metadata_inner(ident, &self.attrs)
    }
}

impl HasMetadata for syn::DeriveInput {
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta> {
        get_metadata_inner(ident, &self.attrs)
    }
}
<commit_msg>Raise an error un malformed strum attributes (#119)<commit_after>///Represents a type that can have strum metadata associated with it.
pub trait HasMetadata {
    /// Get all the metadata associated with a specific "tag".
    /// All of strum's metadata is nested inside a path such as
    /// #[strum(...)] so this let's us quickly filter down to only our metadata.
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta>;
}

fn get_metadata_inner<'a>(
    ident: &str,
    it: impl IntoIterator<Item = &'a syn::Attribute>,
) -> Vec<syn::Meta> {
    it.into_iter()
        .filter(|attr| attr.path.is_ident(ident))
        .map(|attr| attr.parse_meta().unwrap())
        .filter_map(|meta| match meta {
            syn::Meta::List(syn::MetaList { path, nested, .. }) => {
                if path.is_ident(ident) {
                    Some(nested)
                } else {
                    None
                }
            }
            _ => None,
        })
        .flat_map(|id| id)
        .map(|nested| match nested {
            syn::NestedMeta::Meta(meta) => meta,
            _ => panic!("unexpected literal parsing strum attributes"),
        })
        .collect()
}

impl HasMetadata for syn::Variant {
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta> {
        get_metadata_inner(ident, &self.attrs)
    }
}

impl HasMetadata for syn::DeriveInput {
    fn get_metadata(&self, ident: &str) -> Vec<syn::Meta> {
        get_metadata_inner(ident, &self.attrs)
    }
}
<|endoftext|>