<commit_before>use std::io::prelude::*;
use std::fs::File;
use std::string::String;
use std::path::Path;
use std::vec::Vec;
use json::{parse};

use game::data::Rectangle;

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(_) => panic!("Couldn't read file {}", filename),
  }
}

pub fn load_character() -> Vec<Rectangle> {
  let mut sprites = Vec::with_capacity(256);
  let mut sprite_names = Vec::with_capacity(256);
  let character_json = read_sprite_file("./assets/character.json");
  let character = match parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };
  for x in 0..15 {
    for y in 0..13 {
      sprite_names.push(format!("run_{}_{}", x, y));
    }
  }
  for &ref sprite in &sprite_names {
    let x = character["frames"][sprite]["frame"]["x"].as_f64().unwrap();
    let y = character["frames"][sprite]["frame"]["y"].as_f64().unwrap();
    let w = character["frames"][sprite]["frame"]["w"].as_f64().unwrap();
    let h = character["frames"][sprite]["frame"]["h"].as_f64().unwrap();
    sprites.push(Rectangle {
      w: w,
      h: h,
      x: x,
      y: y,
    });
  }
  sprites
}

<commit_msg>Load fire sprites<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::string::String;
use std::path::Path;
use std::vec::Vec;
use json::{parse};

use game::data::Rectangle;

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(_) => panic!("Couldn't read file {}", filename),
  }
}

pub fn load_character() -> Vec<Rectangle> {
  let mut sprites = Vec::with_capacity(512);
  let mut move_sprite_names = Vec::with_capacity(256);
  let mut fire_sprite_names = Vec::with_capacity(256);
  let character_json = read_sprite_file("./assets/character.json");
  let character = match parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };
  for x in 0..15 {
    for y in 0..13 {
      move_sprite_names.push(format!("run_{}_{}", x, y));
    }
  }
  for x in 0..15 {
    for y in 0..3 {
      fire_sprite_names.push(format!("fire_{}_{}", x, y));
    }
  }

  for &ref move_sprite in &move_sprite_names {
    let x = character["frames"][move_sprite]["frame"]["x"].as_f64().unwrap();
    let y = character["frames"][move_sprite]["frame"]["y"].as_f64().unwrap();
    let w = character["frames"][move_sprite]["frame"]["w"].as_f64().unwrap();
    let h = character["frames"][move_sprite]["frame"]["h"].as_f64().unwrap();
    sprites.push(Rectangle {
      w: w,
      h: h,
      x: x,
      y: y,
    });
  }

  for &ref fire_sprite in &fire_sprite_names {
    let x = character["frames"][fire_sprite]["frame"]["x"].as_f64().unwrap();
    let y = character["frames"][fire_sprite]["frame"]["y"].as_f64().unwrap();
    let w = character["frames"][fire_sprite]["frame"]["w"].as_f64().unwrap();
    let h = character["frames"][fire_sprite]["frame"]["h"].as_f64().unwrap();
    sprites.push(Rectangle {
      w: w,
      h: h,
      x: x,
      y: y,
    });
  }
  sprites
}

<|endoftext|>


<commit_before>use super::utils::LspPositionConversion;
use crate::{
    compiler::{
        ast_to_hir::AstToHir,
        hir::{self, Expression, HirDb, Lambda},
    },
    discover::{result::DiscoverResult, run::Discover, value::Value},
    input::{Input, InputDb},
    language_server::utils::TupleToPosition,
};
use itertools::Itertools;
use lsp_types::{notification::Notification, Position};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
pub struct Hint {
    kind: HintKind,
    text: String,
    position: Position,
}
#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum HintKind {
    Value,
    Panic,
}

#[derive(Serialize, Deserialize)]
pub struct HintsNotification {
    pub uri: String,
    pub hints: Vec<Hint>,
}
impl Notification for HintsNotification {
    const METHOD: &'static str = "candy/textDocument/publishHints";

    type Params = Self;
}

#[salsa::query_group(HintsDbStorage)]
pub trait HintsDb: AstToHir + Discover + HirDb + InputDb + LspPositionConversion {
    fn hints(&self, input: Input) -> Vec<Hint>;
}

fn hints(db: &dyn HintsDb, input: Input) -> Vec<Hint> {
    log::debug!("Calculating hints for {input}");

    // TODO: migrate LSP hints to use VM
    return vec![];

    let (hir, _) = db.hir(input.clone()).unwrap();
    let discover_results: HashMap<hir::Id, DiscoverResult> = HashMap::new(); // db.run_all(input.clone(), vec![]);

    collect_hir_ids_for_hints_list(db, hir.expressions.keys().cloned().collect())
        .into_iter()
        .filter_map(|id| {
            let (kind, value) = match discover_results.get(&id).unwrap() {
                DiscoverResult::Value(value) if value != &Value::nothing() => {
                    (HintKind::Value, value.to_owned())
                }
                DiscoverResult::Panic(value) => (HintKind::Panic, value.to_owned()),
                DiscoverResult::CircularImport(import_chain) => (
                    HintKind::Panic,
                    Value::Text(format!(
                        "Circular import detected: {}",
                        import_chain.iter().map(|it| format!("{it}")).join(" → ")
                    )),
                ),
                _ => return None,
            };

            let span = db.hir_id_to_display_span(id)?;

            let line = db
                .offset_to_lsp(input.clone(), span.start)
                .to_position()
                .line;
            let line_start_offsets = db.line_start_utf8_byte_offsets(input.clone());
            let last_characer_of_line = if line as usize == line_start_offsets.len() - 1 {
                db.get_input(input.clone()).unwrap().len()
            } else {
                line_start_offsets[(line + 1) as usize] - 1
            };
            let position = db
                .offset_to_lsp(input.clone(), last_characer_of_line)
                .to_position();

            None as Option<Hint>
            // TODO: Include more hints
            // Some(Hint {
            //     kind,
            //     text: format!(" # {}", db.value_to_display_string(value.to_owned())),
            //     position,
            // })
        })
        // If multiple hints are on the same line, only show the last one.
        .group_by(|hint| hint.position.line)
        .into_iter()
        .map(|(_, hints)| hints.into_iter().last().unwrap())
        .collect()
}

fn collect_hir_ids_for_hints_list(db: &dyn HintsDb, ids: Vec<hir::Id>) -> Vec<hir::Id> {
    ids.into_iter()
        .flat_map(|id| collect_hir_ids_for_hints(db, id))
        .collect()
}
fn collect_hir_ids_for_hints(db: &dyn HintsDb, id: hir::Id) -> Vec<hir::Id> {
    match db.find_expression(id.clone()).unwrap() {
        Expression::Int(_) => vec![],
        Expression::Text(_) => vec![],
        Expression::Reference(_) => vec![id],
        Expression::Symbol(_) => vec![],
        Expression::Struct(_) => vec![], // Handled separately // TODO
        Expression::Lambda(Lambda { body, .. }) => {
            collect_hir_ids_for_hints_list(db, body.expressions.keys().cloned().collect())
        }
        Expression::Call { arguments, .. } => {
            let mut ids = vec![id.to_owned()];
            for argument_id in arguments {
                let argument = match db.find_expression(argument_id.clone()) {
                    Some(argument) => argument,
                    None => continue, // Generated code
                };
                if let Expression::Reference(_) = argument {
                    continue;
                }

                ids.extend(collect_hir_ids_for_hints(db, argument_id));
            }
            ids
        }
        Expression::Builtin(_) => vec![],
        Expression::Needs { .. } => vec![],
        Expression::Error { .. } => vec![],
    }
}
<commit_msg>Add basic hints<commit_after>use super::utils::LspPositionConversion;
use crate::{
    compiler::{
        ast_to_hir::AstToHir,
        hir::{self, HirDb},
        hir_to_lir::HirToLir,
    },
    discover::run::Discover,
    input::{Input, InputDb},
    language_server::utils::TupleToPosition,
    vm::{tracer::TraceEntry, use_provider::FunctionUseProvider, value::Value, Status, Vm},
};
use itertools::Itertools;
use lsp_types::{notification::Notification, Position};
use serde::{Deserialize, Serialize};

#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
pub struct Hint {
    kind: HintKind,
    text: String,
    position: Position,
}
#[derive(PartialEq, Eq, Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum HintKind {
    Value,
    Panic,
}

#[derive(Serialize, Deserialize)]
pub struct HintsNotification {
    pub uri: String,
    pub hints: Vec<Hint>,
}
impl Notification for HintsNotification {
    const METHOD: &'static str = "candy/textDocument/publishHints";

    type Params = Self;
}

#[salsa::query_group(HintsDbStorage)]
pub trait HintsDb:
    HirToLir + AstToHir + Discover + HirDb + InputDb + LspPositionConversion
{
    fn hints(&self, input: Input) -> Vec<Hint>;
}

fn hints(db: &dyn HintsDb, input: Input) -> Vec<Hint> {
    log::debug!("Calculating hints for {input}");

    let lir = match db.lir(input.clone()) {
        Some(lir) => (*lir).clone(),
        None => return vec![],
    };
    let module_closure = Value::module_closure_of_lir(input.clone(), lir);

    let mut vm = Vm::new();
    let use_provider = FunctionUseProvider {
        use_asset: &|input| {
            db.get_input(input.clone())
                .map(|bytes| (*bytes).clone())
                .ok_or_else(|| format!("Couldn't import file '{}'.", input))
        },
        use_local_module: &|input| db.lir(input).map(|lir| (*lir).clone()),
    };

    vm.set_up_module_closure_execution(&use_provider, module_closure);
    vm.run(&use_provider, 1000);
    let panicked = match vm.status() {
        Status::Running => {
            log::info!("VM is still running.");
            None
        }
        Status::Done => {
            let return_value = vm.tear_down_module_closure_execution();
            log::info!("VM is done. Export map: {return_value}");
            None
        }
        Status::Panicked(value) => {
            log::error!("VM panicked with value {value}.");
            // log::error!("This is the stack trace:");
            // vm.tracer.dump_stack_trace(&db, input);
            Some(value)
        }
    };

    let id_of_this_module = hir::Id::new(input.clone(), vec![]);
    let values = vm
        .tracer
        .log()
        .iter()
        .filter_map(|entry| match entry {
            TraceEntry::ValueEvaluated { id, value } => {
                if id_of_this_module.is_parent_of(id) {
                    Some((id.clone(), value.clone()))
                } else {
                    None
                }
            }
            _ => None,
        })
        .collect_vec();

    let value_hints = values
        .into_iter()
        .filter_map(|(id, value)| {
            let span = db.hir_id_to_display_span(id)?;

            let line = db
                .offset_to_lsp(input.clone(), span.start)
                .to_position()
                .line;
            let line_start_offsets = db.line_start_utf8_byte_offsets(input.clone());
            let last_characer_of_line = if line as usize == line_start_offsets.len() - 1 {
                db.get_input(input.clone()).unwrap().len()
            } else {
                line_start_offsets[(line + 1) as usize] - 1
            };
            let position = db
                .offset_to_lsp(input.clone(), last_characer_of_line)
                .to_position();

            Some(Hint {
                kind: HintKind::Value,
                text: format!(" # {value}"),
                position,
            })
        })
        .collect_vec();

    // If multiple hints are on the same line, only show the last one.
    let mut hints = value_hints
        .into_iter()
        .group_by(|hint| hint.position.line)
        .into_iter()
        .map(|(_, hints)| hints.into_iter().last().unwrap())
        .collect_vec();

    if let Some(message) = panicked {
        hints.push(Hint {
            kind: HintKind::Panic,
            text: format!(
                " # The code in this module panicked because {}.",
                if let Value::Text(message) = message {
                    message
                } else {
                    format!("{message}")
                }
            ),
            position: Position {
                line: 0,
                character: 1,
            },
        });
    }

    hints
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate cargo_registry;
extern crate chrono;
#[macro_use]
extern crate diesel;

use diesel::prelude::*;
use std::env;
use std::time::Duration;

use cargo_registry::models::VersionDownload;
use cargo_registry::schema::*;

static LIMIT: i64 = 1000;

fn main() {
    let daemon = env::args().nth(1).as_ref().map(|s| &s[..]) == Some("daemon");
    let sleep = env::args().nth(2).map(|s| s.parse().unwrap());
    loop {
        let conn = cargo_registry::db::connect_now().unwrap();
        update(&conn).unwrap();
        drop(conn);
        if daemon {
            std::thread::sleep(Duration::new(sleep.unwrap(), 0));
        } else {
            break;
        }
    }
}

fn update(conn: &PgConnection) -> QueryResult<()> {
    use diesel::select;
    use version_downloads::dsl::*;

    let mut max = Some(0);
    while let Some(m) = max {
        let rows = version_downloads
            .filter(processed.eq(false))
            .filter(id.gt(m))
            .order(id)
            .limit(LIMIT)
            .load(conn)?;
        collect(conn, &rows)?;
        max = rows.last().map(|d| d.id);
    }

    no_arg_sql_function!(refresh_recent_crate_downloads, ());
    select(refresh_recent_crate_downloads).execute(conn)?;
    Ok(())
}

fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
    use diesel::{insert_into, update};

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    let now = chrono::Utc::now();
    let cutoff = now.naive_utc().date() - chrono::Duration::days(1);

    println!(
        "updating {} versions (cutoff {})",
        rows.len(),
        now.to_rfc2822()
    );

    for download in rows {
        let amt = download.downloads - download.counted;

        conn.transaction::<_, diesel::result::Error, _>(|| {
            // Flag this row as having been processed if we're passed the cutoff,
            // and unconditionally increment the number of counted downloads.
            update(version_downloads::table.find(download.id))
                .set((
                    version_downloads::processed.eq(download.date < cutoff),
                    version_downloads::counted.eq(version_downloads::counted + amt),
                ))
                .execute(conn)?;

            // Update the total number of version downloads
            let crate_id = update(versions::table.find(download.version_id))
                .set(versions::downloads.eq(versions::downloads + amt))
                .returning(versions::crate_id)
                .get_result::<i32>(conn)?;

            // Update the total number of crate downloads
            update(crates::table.find(crate_id))
                .set(crates::downloads.eq(crates::downloads + amt))
                .execute(conn)?;

            // Update the total number of crate downloads for today
            insert_into(crate_downloads::table)
                .values((
                    crate_downloads::crate_id.eq(crate_id),
                    crate_downloads::downloads.eq(amt),
                    crate_downloads::date.eq(download.date),
                ))
                .on_conflict(crate_downloads::table.primary_key())
                .do_update()
                .set(crate_downloads::downloads.eq(crate_downloads::downloads + amt))
                .execute(conn)?;

            // Now that everything else for this crate is done, update the global counter of total
            // downloads
            update(metadata::table)
                .set(metadata::total_downloads.eq(metadata::total_downloads + i64::from(amt)))
                .execute(conn)?;

            Ok(())
        })?;
    }

    Ok(())
}

#[cfg(test)]
mod test {
    extern crate semver;

    use std::collections::HashMap;

    use diesel::insert_into;
    use super::*;
    use cargo_registry::env;
    use cargo_registry::models::{Crate, NewCrate, NewUser, NewVersion, User, Version};

    fn conn() -> PgConnection {
        let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
        conn.begin_test_transaction().unwrap();
        conn
    }

    fn user(conn: &PgConnection) -> User {
        NewUser::new(2, "login", None, None, None, "access_token")
            .create_or_update(conn)
            .unwrap()
    }

    fn crate_and_version(conn: &PgConnection, user_id: i32) -> (Crate, Version) {
        let krate = NewCrate {
            name: "foo",
            ..Default::default()
        }.create_or_update(&conn, None, user_id)
            .unwrap();
        let version = NewVersion::new(
            krate.id,
            &semver::Version::parse("1.0.0").unwrap(),
            &HashMap::new(),
            None,
            None,
        ).unwrap();
        let version = version.save(&conn, &[]).unwrap();
        (krate, version)
    }

    macro_rules! crate_downloads {
        ($conn: expr, $id: expr, $expected: expr) => {
            let dl = crate_downloads::table
                .filter(crate_downloads::crate_id.eq($id))
                .select(crate_downloads::downloads)
                .first($conn);
            assert_eq!(Ok($expected as i32), dl);
        }
    }

    #[test]
    fn increment() {
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (krate, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values(version_downloads::version_id.eq(version.id))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::date.eq(date(now - 1.day())),
                version_downloads::processed.eq(true),
            ))
            .execute(&conn)
            .unwrap();

        ::update(&conn).unwrap();
        let version_downloads = versions::table
            .find(version.id)
            .select(versions::downloads)
            .first(&conn);
        assert_eq!(Ok(1), version_downloads);
        let crate_downloads = crates::table
            .find(krate.id)
            .select(crates::downloads)
            .first(&conn);
        assert_eq!(Ok(1), crate_downloads);
        crate_downloads!(&conn, krate.id, 1);
        ::update(&conn).unwrap();
        let version_downloads = versions::table
            .find(version.id)
            .select(versions::downloads)
            .first(&conn);
        assert_eq!(Ok(1), version_downloads);
    }

    #[test]
    fn set_processed_true() {
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.days())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let processed = version_downloads::table
            .filter(version_downloads::version_id.eq(version.id))
            .select(version_downloads::processed)
            .first(&conn);
        assert_eq!(Ok(true), processed);
    }

    #[test]
    fn dont_process_recent_row() {
        use diesel::dsl::*;
        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.hours())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let processed = version_downloads::table
            .filter(version_downloads::version_id.eq(version.id))
            .select(version_downloads::processed)
            .first(&conn);
        assert_eq!(Ok(false), processed);
    }

    #[test]
    fn increment_a_little() {
        use diesel::dsl::*;
        use diesel::update;

        let conn = conn();
        let user = user(&conn);
        let (krate, version) = crate_and_version(&conn, user.id);
        update(versions::table)
            .set(versions::updated_at.eq(now - 2.hours()))
            .execute(&conn)
            .unwrap();
        update(crates::table)
            .set(crates::updated_at.eq(now - 2.hours()))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(1),
                version_downloads::date.eq(date(now)),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::date.eq(date(now - 1.day())),
            ))
            .execute(&conn)
            .unwrap();

        let version_before = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        let krate_before = Crate::all()
            .filter(crates::id.eq(krate.id))
            .first::<Crate>(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let version2 = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        assert_eq!(version2.downloads, 2);
        assert_eq!(version2.updated_at, version_before.updated_at);
        let krate2 = Crate::all()
            .filter(crates::id.eq(krate.id))
            .first::<Crate>(&conn)
            .unwrap();
        assert_eq!(krate2.downloads, 2);
        assert_eq!(krate2.updated_at, krate_before.updated_at);
        crate_downloads!(&conn, krate.id, 1);
        ::update(&conn).unwrap();
        let version3 = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        assert_eq!(version3.downloads, 2);
    }

    #[test]
    fn set_processed_no_set_updated_at() {
        use diesel::update;
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        update(versions::table)
            .set(versions::updated_at.eq(now - 2.days()))
            .execute(&conn)
            .unwrap();
        update(crates::table)
            .set(crates::updated_at.eq(now - 2.days()))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.days())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();

        ::update(&conn).unwrap();
        let versions_changed = versions::table
            .select(versions::updated_at.ne(now - 2.days()))
            .get_result(&conn);
        let crates_changed = crates::table
            .select(crates::updated_at.ne(now - 2.days()))
            .get_result(&conn);
        assert_eq!(Ok(false), versions_changed);
        assert_eq!(Ok(false), crates_changed);
    }
}
<commit_msg>Do less work in update-downloads<commit_after>#![deny(warnings)]

extern crate cargo_registry;
#[macro_use]
extern crate diesel;

use diesel::prelude::*;
use std::env;
use std::time::Duration;

use cargo_registry::models::VersionDownload;
use cargo_registry::schema::*;

static LIMIT: i64 = 1000;

fn main() {
    let daemon = env::args().nth(1).as_ref().map(|s| &s[..]) == Some("daemon");
    let sleep = env::args().nth(2).map(|s| s.parse().unwrap());
    loop {
        let conn = cargo_registry::db::connect_now().unwrap();
        update(&conn).unwrap();
        drop(conn);
        if daemon {
            std::thread::sleep(Duration::new(sleep.unwrap(), 0));
        } else {
            break;
        }
    }
}

fn update(conn: &PgConnection) -> QueryResult<()> {
    use diesel::dsl::now;
    use diesel::select;
    use version_downloads::dsl::*;

    let mut max = Some(0);
    while let Some(m) = max {
        let rows = version_downloads
            .filter(processed.eq(false))
            .filter(id.gt(m))
            .filter(downloads.ne(counted))
            .order(id)
            .limit(LIMIT)
            .load(conn)?;
        collect(conn, &rows)?;
        max = rows.last().map(|d| d.id);
    }

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    diesel::update(version_downloads)
        .set(processed.eq(true))
        .filter(date.lt(diesel::dsl::date(now)))
        .filter(downloads.eq(counted))
        .filter(processed.eq(false))
        .execute(conn)?;

    no_arg_sql_function!(refresh_recent_crate_downloads, ());
    select(refresh_recent_crate_downloads).execute(conn)?;
    Ok(())
}

fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
    use diesel::{insert_into, update};

    println!("updating {} versions", rows.len());

    for download in rows {
        let amt = download.downloads - download.counted;

        conn.transaction::<_, diesel::result::Error, _>(|| {
            // increment the number of counted downloads
            update(version_downloads::table.find(download.id))
                .set(version_downloads::counted.eq(version_downloads::counted + amt))
                .execute(conn)?;

            // Update the total number of version downloads
            let crate_id = update(versions::table.find(download.version_id))
                .set(versions::downloads.eq(versions::downloads + amt))
                .returning(versions::crate_id)
                .get_result::<i32>(conn)?;

            // Update the total number of crate downloads
            update(crates::table.find(crate_id))
                .set(crates::downloads.eq(crates::downloads + amt))
                .execute(conn)?;

            // Update the total number of crate downloads for today
            insert_into(crate_downloads::table)
                .values((
                    crate_downloads::crate_id.eq(crate_id),
                    crate_downloads::downloads.eq(amt),
                    crate_downloads::date.eq(download.date),
                ))
                .on_conflict(crate_downloads::table.primary_key())
                .do_update()
                .set(crate_downloads::downloads.eq(crate_downloads::downloads + amt))
                .execute(conn)?;

            // Now that everything else for this crate is done, update the global counter of total
            // downloads
            update(metadata::table)
                .set(metadata::total_downloads.eq(metadata::total_downloads + i64::from(amt)))
                .execute(conn)?;

            Ok(())
        })?;
    }

    Ok(())
}

#[cfg(test)]
mod test {
    extern crate semver;

    use std::collections::HashMap;

    use diesel::insert_into;
    use super::*;
    use cargo_registry::env;
    use cargo_registry::models::{Crate, NewCrate, NewUser, NewVersion, User, Version};

    fn conn() -> PgConnection {
        let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
        conn.begin_test_transaction().unwrap();
        conn
    }

    fn user(conn: &PgConnection) -> User {
        NewUser::new(2, "login", None, None, None, "access_token")
            .create_or_update(conn)
            .unwrap()
    }

    fn crate_and_version(conn: &PgConnection, user_id: i32) -> (Crate, Version) {
        let krate = NewCrate {
            name: "foo",
            ..Default::default()
        }.create_or_update(&conn, None, user_id)
            .unwrap();
        let version = NewVersion::new(
            krate.id,
            &semver::Version::parse("1.0.0").unwrap(),
            &HashMap::new(),
            None,
            None,
        ).unwrap();
        let version = version.save(&conn, &[]).unwrap();
        (krate, version)
    }

    macro_rules! crate_downloads {
        ($conn: expr, $id: expr, $expected: expr) => {
            let dl = crate_downloads::table
                .filter(crate_downloads::crate_id.eq($id))
                .select(crate_downloads::downloads)
                .first($conn);
            assert_eq!(Ok($expected as i32), dl);
        }
    }

    #[test]
    fn increment() {
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (krate, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values(version_downloads::version_id.eq(version.id))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::date.eq(date(now - 1.day())),
                version_downloads::processed.eq(true),
            ))
            .execute(&conn)
            .unwrap();

        ::update(&conn).unwrap();
        let version_downloads = versions::table
            .find(version.id)
            .select(versions::downloads)
            .first(&conn);
        assert_eq!(Ok(1), version_downloads);
        let crate_downloads = crates::table
            .find(krate.id)
            .select(crates::downloads)
            .first(&conn);
        assert_eq!(Ok(1), crate_downloads);
        crate_downloads!(&conn, krate.id, 1);
        ::update(&conn).unwrap();
        let version_downloads = versions::table
            .find(version.id)
            .select(versions::downloads)
            .first(&conn);
        assert_eq!(Ok(1), version_downloads);
    }

    #[test]
    fn set_processed_true() {
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.days())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let processed = version_downloads::table
            .filter(version_downloads::version_id.eq(version.id))
            .select(version_downloads::processed)
            .first(&conn);
        assert_eq!(Ok(true), processed);
    }

    #[test]
    fn dont_process_recent_row() {
        use diesel::dsl::*;
        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.hours())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let processed = version_downloads::table
            .filter(version_downloads::version_id.eq(version.id))
            .select(version_downloads::processed)
            .first(&conn);
        assert_eq!(Ok(false), processed);
    }

    #[test]
    fn increment_a_little() {
        use diesel::dsl::*;
        use diesel::update;

        let conn = conn();
        let user = user(&conn);
        let (krate, version) = crate_and_version(&conn, user.id);
        update(versions::table)
            .set(versions::updated_at.eq(now - 2.hours()))
            .execute(&conn)
            .unwrap();
        update(crates::table)
            .set(crates::updated_at.eq(now - 2.hours()))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(1),
                version_downloads::date.eq(date(now)),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::date.eq(date(now - 1.day())),
            ))
            .execute(&conn)
            .unwrap();

        let version_before = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        let krate_before = Crate::all()
            .filter(crates::id.eq(krate.id))
            .first::<Crate>(&conn)
            .unwrap();
        ::update(&conn).unwrap();
        let version2 = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        assert_eq!(version2.downloads, 2);
        assert_eq!(version2.updated_at, version_before.updated_at);
        let krate2 = Crate::all()
            .filter(crates::id.eq(krate.id))
            .first::<Crate>(&conn)
            .unwrap();
        assert_eq!(krate2.downloads, 2);
        assert_eq!(krate2.updated_at, krate_before.updated_at);
        crate_downloads!(&conn, krate.id, 1);
        ::update(&conn).unwrap();
        let version3 = versions::table
            .find(version.id)
            .first::<Version>(&conn)
            .unwrap();
        assert_eq!(version3.downloads, 2);
    }

    #[test]
    fn set_processed_no_set_updated_at() {
        use diesel::update;
        use diesel::dsl::*;

        let conn = conn();
        let user = user(&conn);
        let (_, version) = crate_and_version(&conn, user.id);
        update(versions::table)
            .set(versions::updated_at.eq(now - 2.days()))
            .execute(&conn)
            .unwrap();
        update(crates::table)
            .set(crates::updated_at.eq(now - 2.days()))
            .execute(&conn)
            .unwrap();
        insert_into(version_downloads::table)
            .values((
                version_downloads::version_id.eq(version.id),
                version_downloads::downloads.eq(2),
                version_downloads::counted.eq(2),
                version_downloads::date.eq(date(now - 2.days())),
                version_downloads::processed.eq(false),
            ))
            .execute(&conn)
            .unwrap();

        ::update(&conn).unwrap();
        let versions_changed = versions::table
            .select(versions::updated_at.ne(now - 2.days()))
            .get_result(&conn);
        let crates_changed = crates::table
            .select(crates::updated_at.ne(now - 2.days()))
            .get_result(&conn);
        assert_eq!(Ok(false), versions_changed);
        assert_eq!(Ok(false), crates_changed);
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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
//
//     http://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,
// See the License for the specific language governing permissions and
// limitations under the License.

use storage::engine::{Snapshot, Cursor, ScanMode};
use storage::{Key, Value, CF_LOCK, CF_WRITE};
use super::{Error, Result};
use super::lock::Lock;
use super::write::{Write, WriteType};

pub struct MvccReader<'a> {
    snapshot: &'a Snapshot,
    // cursors are used for speeding up scans.
    data_cursor: Option<Cursor<'a>>,
    lock_cursor: Option<Cursor<'a>>,
    write_cursor: Option<Cursor<'a>>,

    scan_mode: Option<ScanMode>,
    key_only: bool,
}

impl<'a> MvccReader<'a> {
    pub fn new(snapshot: &Snapshot, scan_mode: Option<ScanMode>) -> MvccReader {
        MvccReader {
            snapshot: snapshot,
            data_cursor: None,
            lock_cursor: None,
            write_cursor: None,
            scan_mode: scan_mode,
            key_only: false,
        }
    }

    pub fn set_key_only(&mut self, key_only: bool) {
        self.key_only = key_only;
    }

    pub fn load_data(&mut self, key: &Key, ts: u64) -> Result<Value> {
        if self.key_only {
            return Ok(vec![]);
        }
        if self.scan_mode.is_some() && self.data_cursor.is_none() {
            self.data_cursor = Some(try!(self.snapshot.iter(None, self.scan_mode.unwrap())));
        }

        let k = key.append_ts(ts);
        if let Some(ref mut cursor) = self.data_cursor {
            cursor.get(&k).map(|x| x.unwrap().to_vec()).map_err(Error::from)
        } else {
            self.snapshot.get(&k).map(|x| x.unwrap()).map_err(Error::from)
        }
    }

    pub fn load_lock(&mut self, key: &Key) -> Result<Option<Lock>> {
        if self.scan_mode.is_some() && self.lock_cursor.is_none() {
            self.lock_cursor = Some(try!(self.snapshot
                .iter_cf(CF_LOCK, None, self.scan_mode.unwrap())));
        }

        if let Some(ref mut cursor) = self.lock_cursor {
            match try!(cursor.get(&key)) {
                Some(v) => Ok(Some(try!(Lock::parse(v)))),
                None => Ok(None),
            }
        } else {
            match try!(self.snapshot.get_cf(CF_LOCK, &key)) {
                Some(v) => Ok(Some(try!(Lock::parse(&v)))),
                None => Ok(None),
            }
        }
    }

    pub fn seek_write(&mut self, key: &Key, ts: u64) -> Result<Option<(u64, Write)>> {
        self.seek_write_impl(key, ts, false)
    }

    pub fn reverse_seek_write(&mut self, key: &Key, ts: u64) -> Result<Option<(u64, Write)>> {
        self.seek_write_impl(key, ts, true)
    }

    fn seek_write_impl(&mut self,
                       key: &Key,
                       ts: u64,
                       reverse: bool)
                       -> Result<Option<(u64, Write)>> {
        if self.scan_mode.is_some() {
            if self.write_cursor.is_none() {
                self.write_cursor = Some(try!(self.snapshot
                    .iter_cf(CF_WRITE, None, self.scan_mode.unwrap())));
            }
        } else {
            let upper_bound_key = key.append_ts(0u64);
            let upper_bound = upper_bound_key.encoded().as_slice();
            self.write_cursor = Some(try!(self.snapshot
                .iter_cf(CF_WRITE, Some(upper_bound), ScanMode::Mixed)));
        }

        let mut cursor = self.write_cursor.as_mut().unwrap();
        let ok = if reverse {
            try!(cursor.near_reverse_seek(&key.append_ts(ts)))
        } else {
            try!(cursor.near_seek(&key.append_ts(ts)))
        };
        if !ok {
            return Ok(None);
        }
        let write_key = Key::from_encoded(cursor.key().to_vec());
        let commit_ts = try!(write_key.decode_ts());
        let k = try!(write_key.truncate_ts());
        if &k != key {
            return Ok(None);
        }
        let write = try!(Write::parse(cursor.value()));
        Ok(Some((commit_ts, write)))
    }

    pub fn get(&mut self, key: &Key, mut ts: u64) -> Result<Option<Value>> {
        // Check for locks that signal concurrent writes.
        if let Some(lock) = try!(self.load_lock(key)) {
            if lock.ts <= ts {
                // There is a pending lock. Client should wait or clean it.
                return Err(Error::KeyIsLocked {
                    key: try!(key.raw()),
                    primary: lock.primary,
                    ts: lock.ts,
                });
            }
        }
        loop {
            match try!(self.seek_write(key, ts)) {
                Some((commit_ts, write)) => {
                    match write.write_type {
                        WriteType::Put => return self.load_data(key, write.start_ts).map(Some),
                        WriteType::Delete => return Ok(None),
                        WriteType::Lock | WriteType::Rollback => ts = commit_ts - 1,
                    }
                }
                None => return Ok(None),
            }
        }
    }

    pub fn get_txn_commit_ts(&mut self, key: &Key, start_ts: u64) -> Result<Option<u64>> {
        if let Some((commit_ts, write)) = try!(self.reverse_seek_write(key, start_ts)) {
            if write.start_ts == start_ts {
                match write.write_type {
                    WriteType::Put | WriteType::Delete | WriteType::Lock => {
                        return Ok(Some(commit_ts))
                    }
                    WriteType::Rollback => {}
                }
            }
        }
        Ok(None)
    }

    fn create_write_cursor(&mut self) -> Result<()> {
        if self.write_cursor.is_none() {
            self.write_cursor = Some(try!(self.snapshot
                .iter_cf(CF_WRITE, None, self.scan_mode.unwrap_or(ScanMode::Mixed))));
        }
        Ok(())
    }

    fn create_lock_cursor(&mut self) -> Result<()> {
        if self.lock_cursor.is_none() {
            self.lock_cursor = Some(try!(self.snapshot
                .iter_cf(CF_LOCK, None, self.scan_mode.unwrap_or(ScanMode::Mixed))));
        }
        Ok(())
    }

    pub fn seek(&mut self, mut key: Key, ts: u64) -> Result<Option<(Key, Value)>> {
        assert!(self.scan_mode.is_some());
        try!(self.create_write_cursor());
        try!(self.create_lock_cursor());

        let (mut write_valid, mut lock_valid) = (true, true);

        loop {
            key = {
                let mut w_cur = self.write_cursor.as_mut().unwrap();
                let mut l_cur = self.lock_cursor.as_mut().unwrap();
                let (mut w_key, mut l_key) = (None, None);
                if write_valid {
                    if try!(w_cur.near_seek(&key)) {
                        w_key = Some(w_cur.key());
                    } else {
                        w_key = None;
                        write_valid = false;
                    }
                }
                if lock_valid {
                    if try!(l_cur.near_seek(&key)) {
                        l_key = Some(l_cur.key());
                    } else {
                        l_key = None;
                        lock_valid = false;
                    }
                }
                match (w_key, l_key) {
                    (None, None) => return Ok(None),
                    (None, Some(k)) => Key::from_encoded(k.to_vec()),
                    (Some(k), None) => try!(Key::from_encoded(k.to_vec()).truncate_ts()),
                    (Some(wk), Some(lk)) => {
                        if wk < lk {
                            try!(Key::from_encoded(wk.to_vec()).truncate_ts())
                        } else {
                            Key::from_encoded(lk.to_vec())
                        }
                    }
                }
            };
            if let Some(v) = try!(self.get(&key, ts)) {
                return Ok(Some((key, v)));
            }
            key = key.append_ts(0);
        }
    }

    pub fn reverse_seek(&mut self, mut key: Key, ts: u64) -> Result<Option<(Key, Value)>> {
        assert!(self.scan_mode.is_some());
        try!(self.create_write_cursor());
        try!(self.create_lock_cursor());

        let (mut write_valid, mut lock_valid) = (true, true);

        loop {
            key = {
                let mut w_cur = self.write_cursor.as_mut().unwrap();
                let mut l_cur = self.lock_cursor.as_mut().unwrap();
                let (mut w_key, mut l_key) = (None, None);
                if write_valid {
                    if try!(w_cur.near_reverse_seek(&key)) {
                        w_key = Some(w_cur.key());
                    } else {
                        w_key = None;
                        write_valid = false;
                    }
                }
                if lock_valid {
                    if try!(l_cur.near_reverse_seek(&key)) {
                        l_key = Some(l_cur.key());
                    } else {
                        l_key = None;
                        lock_valid = false;
                    }
                }
                match (w_key, l_key) {
                    (None, None) => return Ok(None),
                    (None, Some(k)) => Key::from_encoded(k.to_vec()),
                    (Some(k), None) => try!(Key::from_encoded(k.to_vec()).truncate_ts()),
                    (Some(wk), Some(lk)) => {
                        if wk < lk {
                            Key::from_encoded(lk.to_vec())
                        } else {
                            try!(Key::from_encoded(wk.to_vec()).truncate_ts())
                        }
                    }
                }
            };
            if let Some(v) = try!(self.get(&key, ts)) {
                return Ok(Some((key, v)));
            }
        }
    }

    pub fn scan_lock<F>(&mut self, filter: F) -> Result<Vec<(Key, Lock)>>
        where F: Fn(&Lock) -> bool
    {
        try!(self.create_lock_cursor());
        let mut cursor = self.lock_cursor.as_mut().unwrap();
        cursor.seek_to_first();
        let mut locks = vec![];
        while cursor.valid() {
            let key = Key::from_encoded(cursor.key().to_vec());
            let lock = try!(Lock::parse(cursor.value()));
            if filter(&lock) {
                locks.push((key, lock));
            }
            cursor.next();
        }
        Ok(locks)
    }

    pub fn scan_keys(&mut self,
                     mut start: Option<Key>,
                     limit: usize)
                     -> Result<(Vec<Key>, Option<Key>)> {
        let mut cursor = try!(self.snapshot
            .iter_cf(CF_WRITE, None, self.scan_mode.unwrap_or(ScanMode::Mixed)));
        let mut keys = vec![];
        loop {
            let ok = match start {
                Some(ref x) => try!(cursor.near_seek(x)),
                None => cursor.seek_to_first(),
            };
            if !ok {
                return Ok((keys, None));
            }
            if keys.len() >= limit {
                return Ok((keys, start));
            }
            let key = try!(Key::from_encoded(cursor.key().to_vec()).truncate_ts());
            start = Some(key.append_ts(0));
            keys.push(key);
        }
    }
}
<commit_msg>storage: log suspicious none (#1237)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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
//
//     http://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,
// See the License for the specific language governing permissions and
// limitations under the License.

use storage::engine::{Snapshot, Cursor, ScanMode};
use storage::{Key, Value, CF_LOCK, CF_WRITE};
use super::{Error, Result};
use super::lock::Lock;
use super::write::{Write, WriteType};

pub struct MvccReader<'a> {
    snapshot: &'a Snapshot,
    // cursors are used for speeding up scans.
    data_cursor: Option<Cursor<'a>>,
    lock_cursor: Option<Cursor<'a>>,
    write_cursor: Option<Cursor<'a>>,

    scan_mode: Option<ScanMode>,
    key_only: bool,
}

impl<'a> MvccReader<'a> {
    pub fn new(snapshot: &Snapshot, scan_mode: Option<ScanMode>) -> MvccReader {
        MvccReader {
            snapshot: snapshot,
            data_cursor: None,
            lock_cursor: None,
            write_cursor: None,
            scan_mode: scan_mode,
            key_only: false,
        }
    }

    pub fn set_key_only(&mut self, key_only: bool) {
        self.key_only = key_only;
    }

    pub fn load_data(&mut self, key: &Key, ts: u64) -> Result<Value> {
        if self.key_only {
            return Ok(vec![]);
        }
        if self.scan_mode.is_some() && self.data_cursor.is_none() {
            self.data_cursor = Some(try!(self.snapshot.iter(None, self.scan_mode.unwrap())));
        }

        let k = key.append_ts(ts);
        if let Some(ref mut cursor) = self.data_cursor {
            match try!(cursor.get(&k)) {
                None => panic!("key {} not found, ts {}", key, ts),
                Some(v) => Ok(v.to_vec()),
            }
        } else {
            match try!(self.snapshot.get(&k)) {
                None => panic!("key {} not found, ts: {}", key, ts),
                Some(v) => Ok(v),
            }
        }
    }

    pub fn load_lock(&mut self, key: &Key) -> Result<Option<Lock>> {
        if self.scan_mode.is_some() && self.lock_cursor.is_none() {
            self.lock_cursor = Some(try!(self.snapshot
                .iter_cf(CF_LOCK, None, self.scan_mode.unwrap())));
        }

        if let Some(ref mut cursor) = self.lock_cursor {
            match try!(cursor.get(&key)) {
                Some(v) => Ok(Some(try!(Lock::parse(v)))),
                None => Ok(None),
            }
        } else {
            match try!(self.snapshot.get_cf(CF_LOCK, &key)) {
                Some(v) => Ok(Some(try!(Lock::parse(&v)))),
                None => Ok(None),
            }
        }
    }

    pub fn seek_write(&mut self, key: &Key, ts: u64) -> Result<Option<(u64, Write)>> {
        self.seek_write_impl(key, ts, false)
    }

    pub fn reverse_seek_write(&mut self, key: &Key, ts: u64) -> Result<Option<(u64, Write)>> {
        self.seek_write_impl(key, ts, true)
    }

    fn seek_write_impl(&mut self,
                       key: &Key,
                       ts: u64,
                       reverse: bool)
                       -> Result<Option<(u64, Write)>> {
        if self.scan_mode.is_some() {
            if self.write_cursor.is_none() {
                self.write_cursor = Some(try!(self.snapshot
                    .iter_cf(CF_WRITE, None, self.scan_mode.unwrap())));
            }
        } else {
            let upper_bound_key = key.append_ts(0u64);
            let upper_bound = upper_bound_key.encoded().as_slice();
            self.write_cursor = Some(try!(self.snapshot
                .iter_cf(CF_WRITE, Some(upper_bound), ScanMode::Mixed)));
        }

        let mut cursor = self.write_cursor.as_mut().unwrap();
        let ok = if reverse {
            try!(cursor.near_reverse_seek(&key.append_ts(ts)))
        } else {
            try!(cursor.near_seek(&key.append_ts(ts)))
        };
        if !ok {
            return Ok(None);
        }
        let write_key = Key::from_encoded(cursor.key().to_vec());
        let commit_ts = try!(write_key.decode_ts());
        let k = try!(write_key.truncate_ts());
        if &k != key {
            return Ok(None);
        }
        let write = try!(Write::parse(cursor.value()));
        Ok(Some((commit_ts, write)))
    }

    pub fn get(&mut self, key: &Key, mut ts: u64) -> Result<Option<Value>> {
        // Check for locks that signal concurrent writes.
        if let Some(lock) = try!(self.load_lock(key)) {
            if lock.ts <= ts {
                // There is a pending lock. Client should wait or clean it.
                return Err(Error::KeyIsLocked {
                    key: try!(key.raw()),
                    primary: lock.primary,
                    ts: lock.ts,
                });
            }
        }
        loop {
            match try!(self.seek_write(key, ts)) {
                Some((commit_ts, write)) => {
                    match write.write_type {
                        WriteType::Put => return self.load_data(key, write.start_ts).map(Some),
                        WriteType::Delete => return Ok(None),
                        WriteType::Lock | WriteType::Rollback => ts = commit_ts - 1,
                    }
                }
                None => return Ok(None),
            }
        }
    }

    pub fn get_txn_commit_ts(&mut self, key: &Key, start_ts: u64) -> Result<Option<u64>> {
        if let Some((commit_ts, write)) = try!(self.reverse_seek_write(key, start_ts)) {
            if write.start_ts == start_ts {
                match write.write_type {
                    WriteType::Put | WriteType::Delete | WriteType::Lock => {
                        return Ok(Some(commit_ts))
                    }
                    WriteType::Rollback => {}
                }
            }
        }
        Ok(None)
    }

    fn create_write_cursor(&mut self) -> Result<()> {
        if self.write_cursor.is_none() {
            self.write_cursor = Some(try!(self.snapshot
                .iter_cf(CF_WRITE, None, self.scan_mode.unwrap_or(ScanMode::Mixed))));
        }
        Ok(())
    }

    fn create_lock_cursor(&mut self) -> Result<()> {
        if self.lock_cursor.is_none() {
            self.lock_cursor = Some(try!(self.snapshot
                .iter_cf(CF_LOCK, None, self.scan_mode.unwrap_or(ScanMode::Mixed))));
        }
        Ok(())
    }

    pub fn seek(&mut self, mut key: Key, ts: u64) -> Result<Option<(Key, Value)>> {
        assert!(self.scan_mode.is_some());
        try!(self.create_write_cursor());
        try!(self.create_lock_cursor());

        let (mut write_valid, mut lock_valid) = (true, true);

        loop {
            key = {
                let mut w_cur = self.write_cursor.as_mut().unwrap();
                let mut l_cur = self.lock_cursor.as_mut().unwrap();
                let (mut w_key, mut l_key) = (None, None);
                if write_valid {
                    if try!(w_cur.near_seek(&key)) {
                        w_key = Some(w_cur.key());
                    } else {
                        w_key = None;
                        write_valid = false;
                    }
                }
                if lock_valid {
                    if try!(l_cur.near_seek(&key)) {
                        l_key = Some(l_cur.key());
                    } else {
                        l_key = None;
                        lock_valid = false;
                    }
                }
                match (w_key, l_key) {
                    (None, None) => return Ok(None),
                    (None, Some(k)) => Key::from_encoded(k.to_vec()),
                    (Some(k), None) => try!(Key::from_encoded(k.to_vec()).truncate_ts()),
                    (Some(wk), Some(lk)) => {
                        if wk < lk {
                            try!(Key::from_encoded(wk.to_vec()).truncate_ts())
                        } else {
                            Key::from_encoded(lk.to_vec())
                        }
                    }
                }
            };
            if let Some(v) = try!(self.get(&key, ts)) {
                return Ok(Some((key, v)));
            }
            key = key.append_ts(0);
        }
    }

    pub fn reverse_seek(&mut self, mut key: Key, ts: u64) -> Result<Option<(Key, Value)>> {
        assert!(self.scan_mode.is_some());
        try!(self.create_write_cursor());
        try!(self.create_lock_cursor());

        let (mut write_valid, mut lock_valid) = (true, true);

        loop {
            key = {
                let mut w_cur = self.write_cursor.as_mut().unwrap();
                let mut l_cur = self.lock_cursor.as_mut().unwrap();
                let (mut w_key, mut l_key) = (None, None);
                if write_valid {
                    if try!(w_cur.near_reverse_seek(&key)) {
                        w_key = Some(w_cur.key());
                    } else {
                        w_key = None;
                        write_valid = false;
                    }
                }
                if lock_valid {
                    if try!(l_cur.near_reverse_seek(&key)) {
                        l_key = Some(l_cur.key());
                    } else {
                        l_key = None;
                        lock_valid = false;
                    }
                }
                match (w_key, l_key) {
                    (None, None) => return Ok(None),
                    (None, Some(k)) => Key::from_encoded(k.to_vec()),
                    (Some(k), None) => try!(Key::from_encoded(k.to_vec()).truncate_ts()),
                    (Some(wk), Some(lk)) => {
                        if wk < lk {
                            Key::from_encoded(lk.to_vec())
                        } else {
                            try!(Key::from_encoded(wk.to_vec()).truncate_ts())
                        }
                    }
                }
            };
            if let Some(v) = try!(self.get(&key, ts)) {
                return Ok(Some((key, v)));
            }
        }
    }

    pub fn scan_lock<F>(&mut self, filter: F) -> Result<Vec<(Key, Lock)>>
        where F: Fn(&Lock) -> bool
    {
        try!(self.create_lock_cursor());
        let mut cursor = self.lock_cursor.as_mut().unwrap();
        cursor.seek_to_first();
        let mut locks = vec![];
        while cursor.valid() {
            let key = Key::from_encoded(cursor.key().to_vec());
            let lock = try!(Lock::parse(cursor.value()));
            if filter(&lock) {
                locks.push((key, lock));
            }
            cursor.next();
        }
        Ok(locks)
    }

    pub fn scan_keys(&mut self,
                     mut start: Option<Key>,
                     limit: usize)
                     -> Result<(Vec<Key>, Option<Key>)> {
        let mut cursor = try!(self.snapshot
            .iter_cf(CF_WRITE, None, self.scan_mode.unwrap_or(ScanMode::Mixed)));
        let mut keys = vec![];
        loop {
            let ok = match start {
                Some(ref x) => try!(cursor.near_seek(x)),
                None => cursor.seek_to_first(),
            };
            if !ok {
                return Ok((keys, None));
            }
            if keys.len() >= limit {
                return Ok((keys, start));
            }
            let key = try!(Key::from_encoded(cursor.key().to_vec()).truncate_ts());
            start = Some(key.append_ts(0));
            keys.push(key);
        }
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code to handle initial setup steps for a pool.
// Initial setup steps are steps that do not alter the environment.

use std::collections::{HashMap, HashSet};
use std::io::ErrorKind;
use std::fs::{OpenOptions, read_dir};
use std::path::PathBuf;

use nix::Errno;
use serde_json;

use devicemapper::{Device, devnode_to_devno};

use super::super::errors::{EngineResult, EngineError, ErrorEnum};
use super::super::types::PoolUuid;

use super::blockdev::StratBlockDev;
use super::device::blkdev_size;
use super::engine::DevOwnership;
use super::metadata::{BDA, StaticHeader};
use super::range_alloc::RangeAllocator;
use super::serde_structs::PoolSave;

/// Find all Stratis devices.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> EngineResult<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {

    let mut pool_map = HashMap::new();
    let mut devno_set = HashSet::new();
    for dir_e in read_dir("/dev")? {
        let dir_e = dir_e?;
        let devnode = dir_e.path();

        let devno = match devnode_to_devno(&devnode)? {
            None => continue,
            Some(devno) => {
                // If this device has already been processed, continue.
                if devno_set.insert(devno) {
                    devno
                } else {
                    continue;
                }
            }
        };

        let f = OpenOptions::new().read(true).open(&devnode);

        // There are some reasons for OpenOptions::open() to return an error
        // which are not reasons for this method to return an error.
        // Try to distinguish. Non-error conditions are:
        //
        // 1. ENXIO: The device does not exist anymore. This means that the device
        // was volatile for some reason; in that case it can not belong to
        // Stratis so it is safe to ignore it.
        //
        // 2. ENOMEDIUM: The device has no medium. An example of this case is an
        // empty optical drive.
        //
        // Note that it is better to be conservative and return with an
        // error in any case where failure to read the device could result
        // in bad data for Stratis. Additional exceptions may be added,
        // but only with a complete justification.
        if f.is_err() {
            let err = f.unwrap_err();
            match err.kind() {
                ErrorKind::NotFound => {
                    continue;
                }
                _ => {
                    if let Some(errno) = err.raw_os_error() {
                        match Errno::from_i32(errno) {
                            Errno::ENXIO | Errno::ENOMEDIUM => continue,
                            _ => return Err(EngineError::Io(err)),
                        };
                    } else {
                        return Err(EngineError::Io(err));
                    }
                }
            }
        }

        let mut f = f.expect("unreachable if f is err");
        if let DevOwnership::Ours(pool_uuid, _) = StaticHeader::determine_ownership(&mut f)? {
            // No value should ever be ejected, because duplicate device nodes
            // are filtered out above. Therefore, the return value of insert()
            // might as well be ignored.
            let _ = pool_map
                .entry(pool_uuid)
                .or_insert_with(HashMap::new)
                .insert(Device::from(devno), devnode);
        };
    }

    Ok(pool_map)
}

/// Get the most recent metadata from a set of Devices for a given pool UUID.
/// Returns None if no metadata found for this pool.
#[allow(implicit_hasher)]
pub fn get_metadata(pool_uuid: PoolUuid,
                    devnodes: &HashMap<Device, PathBuf>)
                    -> EngineResult<Option<PoolSave>> {

    // Get pairs of device nodes and matching BDAs
    // If no BDA, or BDA UUID does not match pool UUID, skip.
    // If there is an error reading the BDA, error. There could have been
    // vital information on that BDA, for example, it may have contained
    // the newest metadata.
    let mut bdas = Vec::new();
    for devnode in devnodes.values() {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                bdas.push((devnode, bda));
            }
        }
    }

    // Most recent time should never be None if this was a properly
    // created pool; this allows for the method to be called in other
    // circumstances.
    let most_recent_time = {
        match bdas.iter()
                  .filter_map(|&(_, ref bda)| bda.last_update_time())
                  .max() {
            Some(time) => time,
            None => return Ok(None),
        }
    };

    // Try to read from all available devnodes that could contain most
    // recent metadata. In the event of errors, continue to try until all are
    // exhausted.
    for &(devnode, ref bda) in
        bdas.iter()
            .filter(|&&(_, ref bda)| bda.last_update_time() == Some(most_recent_time)) {

        let poolsave = OpenOptions::new()
            .read(true)
            .open(devnode)
            .ok()
            .and_then(|mut f| bda.load_state(&mut f).ok())
            .and_then(|opt| opt)
            .and_then(|data| serde_json::from_slice(&data).ok());

        if poolsave.is_some() {
            return Ok(poolsave);
        }
    }

    // If no data has yet returned, we have an error. That is, we should have
    // some metadata, because we have a most recent time, but we failed to
    // get any.
    let err_str = "timestamp indicates data was written, but no data successfully read";
    Err(EngineError::Engine(ErrorEnum::NotFound, err_str.into()))
}

/// Get all the blockdevs corresponding to this pool that can be obtained from
/// the given devices.
/// Returns an error if the blockdevs obtained do not match the metadata.
#[allow(implicit_hasher)]
pub fn get_blockdevs(pool_uuid: PoolUuid,
                     pool_save: &PoolSave,
                     devnodes: &HashMap<Device, PathBuf>)
                     -> EngineResult<Vec<StratBlockDev>> {
    let segments = pool_save
        .flex_devs
        .meta_dev
        .iter()
        .chain(pool_save.flex_devs.thin_meta_dev.iter())
        .chain(pool_save.flex_devs.thin_data_dev.iter());

    let mut segment_table = HashMap::new();
    for seg in segments {
        segment_table
            .entry(seg.0)
            .or_insert_with(Vec::default)
            .push((seg.1, seg.2))
    }

    let mut blockdevs = vec![];
    for (device, devnode) in devnodes {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                let actual_size = blkdev_size(&OpenOptions::new().read(true).open(devnode)?)?
                    .sectors();

                // If size of device has changed and is less, then it is
                // possible that the segments previously allocated for this
                // blockdev no longer exist. If that is the case,
                // RangeAllocator::new() will return an error.
                let allocator =
                    RangeAllocator::new(actual_size,
                                        segment_table.get(&bda.dev_uuid()).unwrap_or(&vec![]))?;

                let bd_save = pool_save
                    .block_devs
                    .get(&bda.dev_uuid())
                    .ok_or_else(|| {
                                    let err_msg = format!("Blockdev {} not found in metadata",
                                                          bda.dev_uuid());
                                    EngineError::Engine(ErrorEnum::NotFound, err_msg)
                                })?;

                blockdevs.push(StratBlockDev::new(*device,
                                                  devnode.to_owned(),
                                                  bda,
                                                  allocator,
                                                  bd_save.user_info.clone(),
                                                  bd_save.hardware_info.clone()));
            }
        }
    }

    // Verify that blockdevs found match blockdevs recorded.
    let current_uuids: HashSet<_> = blockdevs.iter().map(|b| b.uuid()).collect();
    let recorded_uuids: HashSet<_> = pool_save.block_devs.keys().cloned().collect();

    if current_uuids != recorded_uuids {
        let err_msg = "Recorded block dev UUIDs != discovered blockdev UUIDs";
        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
    }

    if blockdevs.len() != current_uuids.len() {
        let err_msg = "Duplicate block devices found in environment";
        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
    }

    Ok(blockdevs)
}
<commit_msg>Do not forget thin_meta_dev_spare when getting segments for metadata<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code to handle initial setup steps for a pool.
// Initial setup steps are steps that do not alter the environment.

use std::collections::{HashMap, HashSet};
use std::io::ErrorKind;
use std::fs::{OpenOptions, read_dir};
use std::path::PathBuf;

use nix::Errno;
use serde_json;

use devicemapper::{Device, devnode_to_devno};

use super::super::errors::{EngineResult, EngineError, ErrorEnum};
use super::super::types::PoolUuid;

use super::blockdev::StratBlockDev;
use super::device::blkdev_size;
use super::engine::DevOwnership;
use super::metadata::{BDA, StaticHeader};
use super::range_alloc::RangeAllocator;
use super::serde_structs::PoolSave;

/// Find all Stratis devices.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> EngineResult<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {

    let mut pool_map = HashMap::new();
    let mut devno_set = HashSet::new();
    for dir_e in read_dir("/dev")? {
        let dir_e = dir_e?;
        let devnode = dir_e.path();

        let devno = match devnode_to_devno(&devnode)? {
            None => continue,
            Some(devno) => {
                // If this device has already been processed, continue.
                if devno_set.insert(devno) {
                    devno
                } else {
                    continue;
                }
            }
        };

        let f = OpenOptions::new().read(true).open(&devnode);

        // There are some reasons for OpenOptions::open() to return an error
        // which are not reasons for this method to return an error.
        // Try to distinguish. Non-error conditions are:
        //
        // 1. ENXIO: The device does not exist anymore. This means that the device
        // was volatile for some reason; in that case it can not belong to
        // Stratis so it is safe to ignore it.
        //
        // 2. ENOMEDIUM: The device has no medium. An example of this case is an
        // empty optical drive.
        //
        // Note that it is better to be conservative and return with an
        // error in any case where failure to read the device could result
        // in bad data for Stratis. Additional exceptions may be added,
        // but only with a complete justification.
        if f.is_err() {
            let err = f.unwrap_err();
            match err.kind() {
                ErrorKind::NotFound => {
                    continue;
                }
                _ => {
                    if let Some(errno) = err.raw_os_error() {
                        match Errno::from_i32(errno) {
                            Errno::ENXIO | Errno::ENOMEDIUM => continue,
                            _ => return Err(EngineError::Io(err)),
                        };
                    } else {
                        return Err(EngineError::Io(err));
                    }
                }
            }
        }

        let mut f = f.expect("unreachable if f is err");
        if let DevOwnership::Ours(pool_uuid, _) = StaticHeader::determine_ownership(&mut f)? {
            // No value should ever be ejected, because duplicate device nodes
            // are filtered out above. Therefore, the return value of insert()
            // might as well be ignored.
            let _ = pool_map
                .entry(pool_uuid)
                .or_insert_with(HashMap::new)
                .insert(Device::from(devno), devnode);
        };
    }

    Ok(pool_map)
}

/// Get the most recent metadata from a set of Devices for a given pool UUID.
/// Returns None if no metadata found for this pool.
#[allow(implicit_hasher)]
pub fn get_metadata(pool_uuid: PoolUuid,
                    devnodes: &HashMap<Device, PathBuf>)
                    -> EngineResult<Option<PoolSave>> {

    // Get pairs of device nodes and matching BDAs
    // If no BDA, or BDA UUID does not match pool UUID, skip.
    // If there is an error reading the BDA, error. There could have been
    // vital information on that BDA, for example, it may have contained
    // the newest metadata.
    let mut bdas = Vec::new();
    for devnode in devnodes.values() {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                bdas.push((devnode, bda));
            }
        }
    }

    // Most recent time should never be None if this was a properly
    // created pool; this allows for the method to be called in other
    // circumstances.
    let most_recent_time = {
        match bdas.iter()
                  .filter_map(|&(_, ref bda)| bda.last_update_time())
                  .max() {
            Some(time) => time,
            None => return Ok(None),
        }
    };

    // Try to read from all available devnodes that could contain most
    // recent metadata. In the event of errors, continue to try until all are
    // exhausted.
    for &(devnode, ref bda) in
        bdas.iter()
            .filter(|&&(_, ref bda)| bda.last_update_time() == Some(most_recent_time)) {

        let poolsave = OpenOptions::new()
            .read(true)
            .open(devnode)
            .ok()
            .and_then(|mut f| bda.load_state(&mut f).ok())
            .and_then(|opt| opt)
            .and_then(|data| serde_json::from_slice(&data).ok());

        if poolsave.is_some() {
            return Ok(poolsave);
        }
    }

    // If no data has yet returned, we have an error. That is, we should have
    // some metadata, because we have a most recent time, but we failed to
    // get any.
    let err_str = "timestamp indicates data was written, but no data successfully read";
    Err(EngineError::Engine(ErrorEnum::NotFound, err_str.into()))
}

/// Get all the blockdevs corresponding to this pool that can be obtained from
/// the given devices.
/// Returns an error if the blockdevs obtained do not match the metadata.
#[allow(implicit_hasher)]
pub fn get_blockdevs(pool_uuid: PoolUuid,
                     pool_save: &PoolSave,
                     devnodes: &HashMap<Device, PathBuf>)
                     -> EngineResult<Vec<StratBlockDev>> {
    let flex_devs = &pool_save.flex_devs;
    let segments = flex_devs
        .meta_dev
        .iter()
        .chain(flex_devs.thin_meta_dev.iter())
        .chain(flex_devs.thin_data_dev.iter())
        .chain(flex_devs.thin_meta_dev_spare.iter());

    let mut segment_table = HashMap::new();
    for seg in segments {
        segment_table
            .entry(seg.0)
            .or_insert_with(Vec::default)
            .push((seg.1, seg.2))
    }

    let mut blockdevs = vec![];
    for (device, devnode) in devnodes {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                let actual_size = blkdev_size(&OpenOptions::new().read(true).open(devnode)?)?
                    .sectors();

                // If size of device has changed and is less, then it is
                // possible that the segments previously allocated for this
                // blockdev no longer exist. If that is the case,
                // RangeAllocator::new() will return an error.
                let allocator =
                    RangeAllocator::new(actual_size,
                                        segment_table.get(&bda.dev_uuid()).unwrap_or(&vec![]))?;

                let bd_save = pool_save
                    .block_devs
                    .get(&bda.dev_uuid())
                    .ok_or_else(|| {
                                    let err_msg = format!("Blockdev {} not found in metadata",
                                                          bda.dev_uuid());
                                    EngineError::Engine(ErrorEnum::NotFound, err_msg)
                                })?;

                blockdevs.push(StratBlockDev::new(*device,
                                                  devnode.to_owned(),
                                                  bda,
                                                  allocator,
                                                  bd_save.user_info.clone(),
                                                  bd_save.hardware_info.clone()));
            }
        }
    }

    // Verify that blockdevs found match blockdevs recorded.
    let current_uuids: HashSet<_> = blockdevs.iter().map(|b| b.uuid()).collect();
    let recorded_uuids: HashSet<_> = pool_save.block_devs.keys().cloned().collect();

    if current_uuids != recorded_uuids {
        let err_msg = "Recorded block dev UUIDs != discovered blockdev UUIDs";
        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
    }

    if blockdevs.len() != current_uuids.len() {
        let err_msg = "Duplicate block devices found in environment";
        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
    }

    Ok(blockdevs)
}
<|endoftext|>


<commit_before>use std::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
pub struct ChunksDecoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,

    // data from the start of the current chunk
    buffer: Vec<u8>,
}

#[derive(Debug, Copy, Clone)]
pub struct ChunksError;

impl fmt::Display for ChunksError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for ChunksError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

impl<R> ChunksDecoder<R> where R: Read {
    pub fn new(source: R) -> ChunksDecoder<R> {
        ChunksDecoder {
            source: source,
            remaining_chunks_size: None,
            buffer: Vec::with_capacity(128),
        }
    }
}

impl<R> Read for ChunksDecoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        // first possibility: we are not in a chunk
        if self.remaining_chunks_size.is_none() {
            // trying the read the chunk size
            let mut chunk_size = Vec::new();

            loop {
                let byte = match self.source.by_ref().bytes().next() {
                    Some(b) => try!(b),
                    None => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError)),
                };

                if byte == b'\r' {
                    break;
                }

                chunk_size.push(byte);
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError)),
            }

            let chunk_size = match String::from_utf8(chunk_size) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError))
            };

            let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError))
            };

            // if the chunk size is 0, we are at EOF
            if chunk_size == 0 {
                return Ok(0);
            }

            // now that we now the current chunk size, calling ourselves recursively
            self.remaining_chunks_size = Some(chunk_size);
            return self.read(buf);
        }

        assert!(self.remaining_chunks_size.is_some());

        // second possibility: we continue reading from a chunk
        if buf.len() < *self.remaining_chunks_size.as_ref().unwrap() {
            let read = try!(self.source.read(buf));
            *self.remaining_chunks_size.as_mut().unwrap() -= read;
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= *self.remaining_chunks_size.as_ref().unwrap());

        let remaining_chunks_size = *self.remaining_chunks_size.as_ref().unwrap();

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));
        *self.remaining_chunks_size.as_mut().unwrap() -= read;

        if read == remaining_chunks_size {
            self.remaining_chunks_size = None;

            if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\r' {
                return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
            }
            if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\n' {
                return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
            }
        }

        return Ok(read);
    }
}

#[cfg(test)]
mod test {
    use super::ChunksDecoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    #[should_panic]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();
    }

    #[test]
    #[should_panic]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();
    }
}
<commit_msg>Revert "Fix chunk decoder bug. We already parse \r\n for 0"<commit_after>use std::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
pub struct ChunksDecoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,

    // data from the start of the current chunk
    buffer: Vec<u8>,
}

#[derive(Debug, Copy, Clone)]
pub struct ChunksError;

impl fmt::Display for ChunksError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for ChunksError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

impl<R> ChunksDecoder<R> where R: Read {
    pub fn new(source: R) -> ChunksDecoder<R> {
        ChunksDecoder {
            source: source,
            remaining_chunks_size: None,
            buffer: Vec::with_capacity(128),
        }
    }
}

impl<R> Read for ChunksDecoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        // first possibility: we are not in a chunk
        if self.remaining_chunks_size.is_none() {
            // trying the read the chunk size
            let mut chunk_size = Vec::new();

            loop {
                let byte = match self.source.by_ref().bytes().next() {
                    Some(b) => try!(b),
                    None => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError)),
                };

                if byte == b'\r' {
                    break;
                }

                chunk_size.push(byte);
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError)),
            }

            let chunk_size = match String::from_utf8(chunk_size) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError))
            };

            let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, ChunksError))
            };

            // if the chunk size is 0, we are at EOF
            if chunk_size == 0 {
                if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\r' {
                    return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
                }
                if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\n' {
                    return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
                }
                return Ok(0);
            }

            // now that we now the current chunk size, calling ourselves recursively
            self.remaining_chunks_size = Some(chunk_size);
            return self.read(buf);
        }

        assert!(self.remaining_chunks_size.is_some());

        // second possibility: we continue reading from a chunk
        if buf.len() < *self.remaining_chunks_size.as_ref().unwrap() {
            let read = try!(self.source.read(buf));
            *self.remaining_chunks_size.as_mut().unwrap() -= read;
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= *self.remaining_chunks_size.as_ref().unwrap());

        let remaining_chunks_size = *self.remaining_chunks_size.as_ref().unwrap();

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));
        *self.remaining_chunks_size.as_mut().unwrap() -= read;

        if read == remaining_chunks_size {
            self.remaining_chunks_size = None;

            if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\r' {
                return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
            }
            if try!(self.source.by_ref().bytes().next().unwrap_or(Ok(0))) != b'\n' {
                return Err(IoError::new(ErrorKind::InvalidInput, ChunksError));
            }
        }

        return Ok(read);
    }
}

#[cfg(test)]
mod test {
    use super::ChunksDecoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    #[should_panic]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();
    }

    #[test]
    #[should_panic]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = ChunksDecoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::firmware::hook::run_script;

use failure::bail;
use std::{fmt, result, str::FromStr};

const GET_SCRIPT: &str = "updatehub-active-get";
const SET_SCRIPT: &str = "updatehub-active-set";

#[derive(PartialEq, Debug, Copy, Clone)]
pub enum Set {
    A,
    B,
}

impl FromStr for Set {
    type Err = failure::Error;

    fn from_str(s: &str) -> result::Result<Self, Self::Err> {
        match s.parse::<u8>() {
            Ok(0) => Ok(Set::A),
            Ok(1) => Ok(Set::B),
            Ok(v) => bail!("{} is a invalid value. The only know ones are 0 or 1", v),
            Err(e) => Err(e.into()),
        }
    }
}

impl fmt::Display for Set {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Set::A => "0",
                Set::B => "1",
            }
        )
    }
}

pub fn active() -> Result<Set, failure::Error> {
    Ok(run_script(GET_SCRIPT)?.parse()?)
}

pub fn inactive() -> Result<Set, failure::Error> {
    match active()? {
        Set::A => Ok(Set::B),
        Set::B => Ok(Set::A),
    }
}

pub fn swap_active() -> Result<(), failure::Error> {
    let _ = run_script(&format!("{} {}", SET_SCRIPT, inactive()?))?;
    Ok(())
}

#[test]
fn as_str() {
    assert_eq!("0", format!("{}", Set::A));
    assert_eq!("1", format!("{}", Set::B));
}

#[test]
fn works() {
    use std::env;
    use tempfile::tempdir;

    // create the fake backend
    let tmpdir = tempdir().unwrap();
    let tmpdir = tmpdir.path();
    env::set_var("PATH", format!("{}", &tmpdir.to_string_lossy()));

    let create_fake_backend = |active: usize| {
        use std::{
            fs::{create_dir_all, metadata, File},
            io::Write,
            os::unix::fs::PermissionsExt,
        };

        create_dir_all(&tmpdir).unwrap();

        let mut file = File::create(&tmpdir.join(GET_SCRIPT)).unwrap();
        writeln!(file, "#!/bin/sh\necho {}", active).unwrap();

        let mut permissions = metadata(tmpdir).unwrap().permissions();

        permissions.set_mode(0o755);
        file.set_permissions(permissions).unwrap();

        let mut file = File::create(&tmpdir.join(SET_SCRIPT)).unwrap();
        writeln!(file, "#!/bin/sh\nexit 0").unwrap();

        let mut permissions = metadata(tmpdir).unwrap().permissions();
        permissions.set_mode(0o755);
        file.set_permissions(permissions).unwrap();
    };

    // Create a fake backend using 0 as active. It must test the
    // following:
    // - active is A
    // - inactive is B
    // - swap works
    create_fake_backend(0);
    assert_eq!(active().unwrap(), Set::A);
    assert_eq!(inactive().unwrap(), Set::B);
    assert!(swap_active().is_ok());

    // Create a fake backend using 1 as active. It must test the
    // following:
    // - active is B
    // - inactive is A
    // - swap works
    create_fake_backend(1);
    assert_eq!(active().unwrap(), Set::B);
    assert_eq!(inactive().unwrap(), Set::A);
    assert!(swap_active().is_ok());
}
<commit_msg>firmware: tests: Drop duplicated code inside installation_set<commit_after>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::firmware::hook::run_script;

use failure::bail;
use std::{fmt, result, str::FromStr};

const GET_SCRIPT: &str = "updatehub-active-get";
const SET_SCRIPT: &str = "updatehub-active-set";

#[derive(PartialEq, Debug, Copy, Clone)]
pub enum Set {
    A,
    B,
}

impl FromStr for Set {
    type Err = failure::Error;

    fn from_str(s: &str) -> result::Result<Self, Self::Err> {
        match s.parse::<u8>() {
            Ok(0) => Ok(Set::A),
            Ok(1) => Ok(Set::B),
            Ok(v) => bail!("{} is a invalid value. The only know ones are 0 or 1", v),
            Err(e) => Err(e.into()),
        }
    }
}

impl fmt::Display for Set {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Set::A => "0",
                Set::B => "1",
            }
        )
    }
}

pub fn active() -> Result<Set, failure::Error> {
    Ok(run_script(GET_SCRIPT)?.parse()?)
}

pub fn inactive() -> Result<Set, failure::Error> {
    match active()? {
        Set::A => Ok(Set::B),
        Set::B => Ok(Set::A),
    }
}

pub fn swap_active() -> Result<(), failure::Error> {
    let _ = run_script(&format!("{} {}", SET_SCRIPT, inactive()?))?;
    Ok(())
}

#[test]
fn as_str() {
    assert_eq!("0", format!("{}", Set::A));
    assert_eq!("1", format!("{}", Set::B));
}

#[test]
fn works() {
    use crate::firmware::tests::create_fake_installation_set;
    use std::env;
    use tempfile::tempdir;

    // create the fake backend
    let tmpdir = tempdir().unwrap();
    let tmpdir = tmpdir.path();
    env::set_var("PATH", format!("{}", &tmpdir.to_string_lossy()));

    // Create a fake backend using 0 as active. It must test the
    // following:
    // - active is A
    // - inactive is B
    // - swap works
    create_fake_installation_set(&tmpdir, 0);
    assert_eq!(active().unwrap(), Set::A);
    assert_eq!(inactive().unwrap(), Set::B);
    assert!(swap_active().is_ok());

    // Create a fake backend using 1 as active. It must test the
    // following:
    // - active is B
    // - inactive is A
    // - swap works
    create_fake_installation_set(&tmpdir, 1);
    assert_eq!(active().unwrap(), Set::B);
    assert_eq!(inactive().unwrap(), Set::A);
    assert!(swap_active().is_ok());
}
<|endoftext|>