<commit_before>use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use std::time::{SystemTime, UNIX_EPOCH};

use domain::wallet::Tags;
use errors::prelude::*;
use services::wallet::WalletService;
use api::{WalletHandle, PoolHandle, CommandHandle};
use commands::{Command, CommandExecutor};
use commands::ledger::LedgerCommand;

use api::next_command_handle;

const CRED_DEF_CACHE: &str = "cred_def_cache";
const SCHEMA_CACHE: &str = "schema_cache";

pub enum CacheCommand {
    GetSchema(PoolHandle,
              WalletHandle,
              String, // submitter_did
              String, // id
              String, // options_json
              Box<Fn(IndyResult<String>) + Send>),
    GetSchemaContinue(
        WalletHandle,
        IndyResult<(String, String)>, // ledger_response
        GetCacheOptions,              // options
        CommandHandle,                          // cb_id
    ),
    GetCredDef(PoolHandle,
               WalletHandle,
               String, // submitter_did
               String, // id
               String, // options_json
               Box<Fn(IndyResult<String>) + Send>),
    GetCredDefContinue(
        WalletHandle,
        IndyResult<(String, String)>, // ledger_response
        GetCacheOptions,              // options
        CommandHandle,                          // cb_id
    ),
    PurgeSchemaCache(WalletHandle,
                     String, // options json
                     Box<Fn(IndyResult<()>) + Send>),
    PurgeCredDefCache(WalletHandle,
                      String, // options json
                      Box<Fn(IndyResult<()>) + Send>),
}

pub struct CacheCommandExecutor {
    wallet_service: Rc<WalletService>,

    pending_callbacks: RefCell<HashMap<CommandHandle, Box<Fn(IndyResult<String>)>>>,
}

impl CacheCommandExecutor {
    pub fn new(wallet_service: Rc<WalletService>) -> CacheCommandExecutor {
        CacheCommandExecutor {
            wallet_service,
            pending_callbacks: RefCell::new(HashMap::new()),
        }
    }

    pub fn execute(&self, command: CacheCommand) {
        match command {
            CacheCommand::GetSchema(pool_handle, wallet_handle, submitter_did, id, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "GetSchema command received");
                self.get_schema(pool_handle, wallet_handle, &submitter_did, &id, &options_json, cb);
            }
            CacheCommand::GetSchemaContinue(wallet_handle, ledger_response, options, cb_id) => {
                info!(target: "non_secrets_command_executor", "GetSchemaContinue command received");
                self._get_schema_continue(wallet_handle, ledger_response, options, cb_id);
            }
            CacheCommand::GetCredDef(pool_handle, wallet_handle, submitter_did, id, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "GetCredDef command received");
                self.get_cred_def(pool_handle, wallet_handle, &submitter_did, &id, &options_json, cb);
            }
            CacheCommand::GetCredDefContinue(wallet_handle, ledger_response, options, cb_id) => {
                info!(target: "non_secrets_command_executor", "GetCredDefContinue command received");
                self._get_cred_def_continue(wallet_handle, ledger_response, options, cb_id);
            }
            CacheCommand::PurgeSchemaCache(wallet_handle, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "PurgeSchemaCache command received");
                cb(self.purge_schema_cache(wallet_handle, &options_json));
            }
            CacheCommand::PurgeCredDefCache(wallet_handle, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "PurgeCredDefCache command received");
                cb(self.purge_cred_def_cache(wallet_handle, &options_json));
            }
        }
    }

    fn get_schema(&self,
                  pool_handle: PoolHandle,
                  wallet_handle: WalletHandle,
                  submitter_did: &str,
                  id: &str,
                  options_json: &str,
                  cb: Box<Fn(IndyResult<String>) + Send>) {
        trace!("get_schema >>> pool_handle: {:?}, wallet_handle: {:?}, submitter_did: {:?}, id: {:?}, options_json: {:?}",
               pool_handle, wallet_handle, submitter_did, id, options_json);

        let options = try_cb!(serde_json::from_str::<GetCacheOptions>(options_json).to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options"), cb);

        let cache = if !options.no_cache.unwrap_or(false) {
            let options_json = json!({
                "retrieveType": false,
                "retrieveValue": true,
                "retrieveTags": true,
            }).to_string();
            match self.wallet_service.get_record(wallet_handle, SCHEMA_CACHE, id, &options_json) {
                Ok(record) => Ok(Some(record)),
                Err(err) => if err.kind() == IndyErrorKind::WalletItemNotFound { Ok(None) } else { Err(err) }
            }
        } else { Ok(None) };
        let cache = try_cb!(cache, cb);

        if let Some(cache) = cache {
            let min_fresh = options.min_fresh.unwrap_or(-1);
            if min_fresh >= 0 {
                let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                    Ok(ts) => ts.as_secs() as i32,
                    Err(err) => {
                        error!("Cannot get time: {:?}", err);
                        return cb(Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err))))
                    }
                };
                if ts - min_fresh <= cache.get_tags().unwrap_or(&Tags::new()).get("timestamp").unwrap_or(&"-1".to_string()).parse().unwrap_or(-1) {
                    return cb(Ok(cache.get_value().unwrap_or("").to_string()))
                }
            } else {
                return cb(Ok(cache.get_value().unwrap_or("").to_string()))
            }
        }

        if options.no_update.unwrap_or(false) {
            return cb(Err(IndyError::from(IndyErrorKind::LedgerItemNotFound)));
        }

        let cb_id = next_command_handle();
        self.pending_callbacks.borrow_mut().insert(cb_id, cb);

        CommandExecutor::instance().send(
            Command::Ledger(
                LedgerCommand::GetSchema(
                    pool_handle,
                    Some(submitter_did.to_string()),
                    id.to_string(),
                    Box::new(move |ledger_response| {
                        CommandExecutor::instance().send(
                            Command::Cache(
                                CacheCommand::GetSchemaContinue(
                                    wallet_handle,
                                    ledger_response,
                                    options.clone(),
                                    cb_id,
                                )
                            )
                        ).unwrap();
                    })
                )
            )
        ).unwrap();
    }

    fn _get_schema_continue(&self, wallet_handle: WalletHandle, ledger_response: IndyResult<(String, String)>, options: GetCacheOptions, cb_id: CommandHandle) {
        let cb = self.pending_callbacks.borrow_mut().remove(&cb_id).expect("FIXME INVALID STATE");

        let (schema_id, schema_json) = try_cb!(ledger_response, cb);

        if !options.no_store.unwrap_or(false) {
            let mut tags = Tags::new();
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    warn!("Cannot get time: {:?}", err);
                    0
                }
            };
            tags.insert("timestamp".to_string(), ts.to_string());
            let _ = self.wallet_service.delete_record(wallet_handle, SCHEMA_CACHE, &schema_id);
            let _ = self.wallet_service.add_record(wallet_handle, SCHEMA_CACHE, &schema_id, &schema_json, &tags);
        }

        cb(Ok(schema_json));
    }

    fn get_cred_def(&self,
                    pool_handle: PoolHandle,
                    wallet_handle: WalletHandle,
                    submitter_did: &str,
                    id: &str,
                    options_json: &str,
                    cb: Box<Fn(IndyResult<String>) + Send>) {
        trace!("get_cred_def >>> pool_handle: {:?}, wallet_handle: {:?}, submitter_did: {:?}, id: {:?}, options_json: {:?}",
               pool_handle, wallet_handle, submitter_did, id, options_json);

        let options = try_cb!(serde_json::from_str::<GetCacheOptions>(options_json).to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options"), cb);

        let cache = if !options.no_cache.unwrap_or(false) {
            let options_json = json!({
                "retrieveType": false,
                "retrieveValue": true,
                "retrieveTags": true,
            }).to_string();
            match self.wallet_service.get_record(wallet_handle, CRED_DEF_CACHE, id, &options_json) {
                Ok(record) => Ok(Some(record)),
                Err(err) => if err.kind() == IndyErrorKind::WalletItemNotFound { Ok(None) } else { Err(err) }
            }
        } else { Ok(None) };
        let cache = try_cb!(cache, cb);

        if let Some(cache) = cache {
            let min_fresh = options.min_fresh.unwrap_or(-1);
            if min_fresh >= 0 {
                let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                    Ok(ts) => ts.as_secs() as i32,
                    Err(err) => {
                        error!("Cannot get time: {:?}", err);
                        return cb(Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err))))
                    }
                };
                if ts - min_fresh <= cache.get_tags().unwrap_or(&Tags::new()).get("timestamp").unwrap_or(&"-1".to_string()).parse().unwrap_or(-1) {
                    return cb(Ok(cache.get_value().unwrap_or("").to_string()))
                }
            } else {
                return cb(Ok(cache.get_value().unwrap_or("").to_string()))
            }
        }

        if options.no_update.unwrap_or(false) {
            return cb(Err(IndyError::from(IndyErrorKind::LedgerItemNotFound)));
        }

        let cb_id = next_command_handle();
        self.pending_callbacks.borrow_mut().insert(cb_id, cb);

        CommandExecutor::instance().send(
            Command::Ledger(
                LedgerCommand::GetCredDef(
                    pool_handle,
                    Some(submitter_did.to_string()),
                    id.to_string(),
                    Box::new(move |ledger_response| {
                        CommandExecutor::instance().send(
                            Command::Cache(
                                CacheCommand::GetCredDefContinue(
                                    wallet_handle,
                                    ledger_response,
                                    options.clone(),
                                    cb_id,
                                )
                            )
                        ).unwrap();
                    })
                )
            )
        ).unwrap();
    }

    fn _get_cred_def_continue(&self, wallet_handle: WalletHandle, ledger_response: IndyResult<(String, String)>, options: GetCacheOptions, cb_id: CommandHandle) {
        let cb = self.pending_callbacks.borrow_mut().remove(&cb_id).expect("FIXME INVALID STATE");

        let (cred_def_id, cred_def_json) = try_cb!(ledger_response, cb);

        if !options.no_store.unwrap_or(false) {
            let mut tags = Tags::new();
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    warn!("Cannot get time: {:?}", err);
                    0
                }
            };
            tags.insert("timestamp".to_string(), ts.to_string());
            let _ = self.wallet_service.delete_record(wallet_handle, CRED_DEF_CACHE, &cred_def_id);
            let _ = self.wallet_service.add_record(wallet_handle, CRED_DEF_CACHE, &cred_def_id, &cred_def_json, &tags);
        }

        cb(Ok(cred_def_json));
    }

    fn purge_schema_cache(&self,
                          wallet_handle: WalletHandle,
                          options_json: &str) -> IndyResult<()> {
        trace!("purge_schema_cache >>> wallet_handle: {:?}, options_json: {:?}", wallet_handle, options_json);

        let options = serde_json::from_str::<PurgeOptions>(options_json)
            .to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options")?;

        let max_age = options.max_age.unwrap_or(-1);
        let query_json = if max_age >= 0 {
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    error!("Cannot get time: {:?}", err);
                    return Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err)))
                }
            };
            json!({"timestamp": {"$lt": ts - max_age}}).to_string()
        } else {
            "{}".to_string()
        };

        let options_json = json!({
            "retrieveType": false,
            "retrieveValue": false,
            "retrieveTags": false,
        }).to_string();

        let mut search = self.wallet_service.search_records(
            wallet_handle,
            SCHEMA_CACHE,
            &query_json,
            &options_json,
        )?;

        while let Some(record) = search.fetch_next_record()? {
            self.wallet_service.delete_record(wallet_handle, SCHEMA_CACHE, record.get_id())?;
        }

        trace!("purge_schema_cache <<< res: ()");

        Ok(())
    }

    fn purge_cred_def_cache(&self,
                            wallet_handle: WalletHandle,
                            options_json: &str) -> IndyResult<()> {
        trace!("purge_cred_def_cache >>> wallet_handle: {:?}, options_json: {:?}", wallet_handle, options_json);

        let options = serde_json::from_str::<PurgeOptions>(options_json)
            .to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options")?;

        let max_age = options.max_age.unwrap_or(-1);
        let query_json = if max_age >= 0 {
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    error!("Cannot get time: {:?}", err);
                    return Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err)))
                }
            };
            json!({"timestamp": {"$lt": ts - max_age}}).to_string()
        } else {
            "{}".to_string()
        };

        let options_json = json!({
            "retrieveType": false,
            "retrieveValue": false,
            "retrieveTags": false,
        }).to_string();

        let mut search = self.wallet_service.search_records(
            wallet_handle,
            CRED_DEF_CACHE,
            &query_json,
            &options_json,
        )?;

        while let Some(record) = search.fetch_next_record()? {
            self.wallet_service.delete_record(wallet_handle, CRED_DEF_CACHE, record.get_id())?;
        }

        trace!("purge_cred_def_cache <<< res: ()");

        Ok(())
    }
}

#[serde(rename_all = "camelCase")]
#[derive(Debug, Deserialize, Serialize)]
struct PurgeOptions {
    pub max_age: Option<i32>,
}

#[serde(rename_all = "camelCase")]
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct GetCacheOptions {
    pub no_cache: Option<bool>,     // Skip usage of cache,
    pub no_update: Option<bool>,    // Use only cached data, do not try to update.
    pub no_store: Option<bool>,     // Skip storing fresh data if updated
    pub min_fresh: Option<i32>,     // Return cached data if not older than this many seconds. -1 means do not check age.
}
<commit_msg>remove duplicate code get_record_from_cache<commit_after>use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use std::time::{SystemTime, UNIX_EPOCH};

use domain::wallet::Tags;
use errors::prelude::*;
use services::wallet::{WalletService, WalletRecord};
use api::{WalletHandle, PoolHandle, CommandHandle};
use commands::{Command, CommandExecutor};
use commands::ledger::LedgerCommand;

use api::next_command_handle;

const CRED_DEF_CACHE: &str = "cred_def_cache";
const SCHEMA_CACHE: &str = "schema_cache";

pub enum CacheCommand {
    GetSchema(PoolHandle,
              WalletHandle,
              String, // submitter_did
              String, // id
              String, // options_json
              Box<Fn(IndyResult<String>) + Send>),
    GetSchemaContinue(
        WalletHandle,
        IndyResult<(String, String)>, // ledger_response
        GetCacheOptions,              // options
        CommandHandle,                          // cb_id
    ),
    GetCredDef(PoolHandle,
               WalletHandle,
               String, // submitter_did
               String, // id
               String, // options_json
               Box<Fn(IndyResult<String>) + Send>),
    GetCredDefContinue(
        WalletHandle,
        IndyResult<(String, String)>, // ledger_response
        GetCacheOptions,              // options
        CommandHandle,                          // cb_id
    ),
    PurgeSchemaCache(WalletHandle,
                     String, // options json
                     Box<Fn(IndyResult<()>) + Send>),
    PurgeCredDefCache(WalletHandle,
                      String, // options json
                      Box<Fn(IndyResult<()>) + Send>),
}

pub struct CacheCommandExecutor {
    wallet_service: Rc<WalletService>,

    pending_callbacks: RefCell<HashMap<CommandHandle, Box<Fn(IndyResult<String>)>>>,
}

impl CacheCommandExecutor {
    pub fn new(wallet_service: Rc<WalletService>) -> CacheCommandExecutor {
        CacheCommandExecutor {
            wallet_service,
            pending_callbacks: RefCell::new(HashMap::new()),
        }
    }

    pub fn execute(&self, command: CacheCommand) {
        match command {
            CacheCommand::GetSchema(pool_handle, wallet_handle, submitter_did, id, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "GetSchema command received");
                self.get_schema(pool_handle, wallet_handle, &submitter_did, &id, &options_json, cb);
            }
            CacheCommand::GetSchemaContinue(wallet_handle, ledger_response, options, cb_id) => {
                info!(target: "non_secrets_command_executor", "GetSchemaContinue command received");
                self._get_schema_continue(wallet_handle, ledger_response, options, cb_id);
            }
            CacheCommand::GetCredDef(pool_handle, wallet_handle, submitter_did, id, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "GetCredDef command received");
                self.get_cred_def(pool_handle, wallet_handle, &submitter_did, &id, &options_json, cb);
            }
            CacheCommand::GetCredDefContinue(wallet_handle, ledger_response, options, cb_id) => {
                info!(target: "non_secrets_command_executor", "GetCredDefContinue command received");
                self._get_cred_def_continue(wallet_handle, ledger_response, options, cb_id);
            }
            CacheCommand::PurgeSchemaCache(wallet_handle, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "PurgeSchemaCache command received");
                cb(self.purge_schema_cache(wallet_handle, &options_json));
            }
            CacheCommand::PurgeCredDefCache(wallet_handle, options_json, cb) => {
                info!(target: "non_secrets_command_executor", "PurgeCredDefCache command received");
                cb(self.purge_cred_def_cache(wallet_handle, &options_json));
            }
        }
    }

    fn get_schema(&self,
                  pool_handle: PoolHandle,
                  wallet_handle: WalletHandle,
                  submitter_did: &str,
                  id: &str,
                  options_json: &str,
                  cb: Box<Fn(IndyResult<String>) + Send>) {
        trace!("get_schema >>> pool_handle: {:?}, wallet_handle: {:?}, submitter_did: {:?}, id: {:?}, options_json: {:?}",
               pool_handle, wallet_handle, submitter_did, id, options_json);

        let options = try_cb!(serde_json::from_str::<GetCacheOptions>(options_json).to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options"), cb);

        let cache = self.get_record_from_cache(wallet_handle, id, &options, SCHEMA_CACHE);
        let cache = try_cb!(cache, cb);

        if let Some(cache) = cache {
            let min_fresh = options.min_fresh.unwrap_or(-1);
            if min_fresh >= 0 {
                let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                    Ok(ts) => ts.as_secs() as i32,
                    Err(err) => {
                        error!("Cannot get time: {:?}", err);
                        return cb(Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err))))
                    }
                };
                if ts - min_fresh <= cache.get_tags().unwrap_or(&Tags::new()).get("timestamp").unwrap_or(&"-1".to_string()).parse().unwrap_or(-1) {
                    return cb(Ok(cache.get_value().unwrap_or("").to_string()))
                }
            } else {
                return cb(Ok(cache.get_value().unwrap_or("").to_string()))
            }
        }

        if options.no_update.unwrap_or(false) {
            return cb(Err(IndyError::from(IndyErrorKind::LedgerItemNotFound)));
        }

        let cb_id = next_command_handle();
        self.pending_callbacks.borrow_mut().insert(cb_id, cb);

        CommandExecutor::instance().send(
            Command::Ledger(
                LedgerCommand::GetSchema(
                    pool_handle,
                    Some(submitter_did.to_string()),
                    id.to_string(),
                    Box::new(move |ledger_response| {
                        CommandExecutor::instance().send(
                            Command::Cache(
                                CacheCommand::GetSchemaContinue(
                                    wallet_handle,
                                    ledger_response,
                                    options.clone(),
                                    cb_id,
                                )
                            )
                        ).unwrap();
                    })
                )
            )
        ).unwrap();
    }

    fn _get_schema_continue(&self, wallet_handle: WalletHandle, ledger_response: IndyResult<(String, String)>, options: GetCacheOptions, cb_id: CommandHandle) {
        let cb = self.pending_callbacks.borrow_mut().remove(&cb_id).expect("FIXME INVALID STATE");

        let (schema_id, schema_json) = try_cb!(ledger_response, cb);

        if !options.no_store.unwrap_or(false) {
            let mut tags = Tags::new();
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    warn!("Cannot get time: {:?}", err);
                    0
                }
            };
            tags.insert("timestamp".to_string(), ts.to_string());
            let _ = self.wallet_service.delete_record(wallet_handle, SCHEMA_CACHE, &schema_id);
            let _ = self.wallet_service.add_record(wallet_handle, SCHEMA_CACHE, &schema_id, &schema_json, &tags);
        }

        cb(Ok(schema_json));
    }

    fn get_cred_def(&self,
                    pool_handle: PoolHandle,
                    wallet_handle: WalletHandle,
                    submitter_did: &str,
                    id: &str,
                    options_json: &str,
                    cb: Box<Fn(IndyResult<String>) + Send>) {
        trace!("get_cred_def >>> pool_handle: {:?}, wallet_handle: {:?}, submitter_did: {:?}, id: {:?}, options_json: {:?}",
               pool_handle, wallet_handle, submitter_did, id, options_json);

        let options = try_cb!(serde_json::from_str::<GetCacheOptions>(options_json).to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options"), cb);

        let cache = self.get_record_from_cache(wallet_handle, id, &options, CRED_DEF_CACHE);
        let cache = try_cb!(cache, cb);

        if let Some(cache) = cache {
            let min_fresh = options.min_fresh.unwrap_or(-1);
            if min_fresh >= 0 {
                let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                    Ok(ts) => ts.as_secs() as i32,
                    Err(err) => {
                        error!("Cannot get time: {:?}", err);
                        return cb(Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err))))
                    }
                };
                if ts - min_fresh <= cache.get_tags().unwrap_or(&Tags::new()).get("timestamp").unwrap_or(&"-1".to_string()).parse().unwrap_or(-1) {
                    return cb(Ok(cache.get_value().unwrap_or("").to_string()))
                }
            } else {
                return cb(Ok(cache.get_value().unwrap_or("").to_string()))
            }
        }

        if options.no_update.unwrap_or(false) {
            return cb(Err(IndyError::from(IndyErrorKind::LedgerItemNotFound)));
        }

        let cb_id = next_command_handle();
        self.pending_callbacks.borrow_mut().insert(cb_id, cb);

        CommandExecutor::instance().send(
            Command::Ledger(
                LedgerCommand::GetCredDef(
                    pool_handle,
                    Some(submitter_did.to_string()),
                    id.to_string(),
                    Box::new(move |ledger_response| {
                        CommandExecutor::instance().send(
                            Command::Cache(
                                CacheCommand::GetCredDefContinue(
                                    wallet_handle,
                                    ledger_response,
                                    options.clone(),
                                    cb_id,
                                )
                            )
                        ).unwrap();
                    })
                )
            )
        ).unwrap();
    }

    fn get_record_from_cache(&self, wallet_handle: WalletHandle, id: &str, options: &GetCacheOptions, which_cache: &str) -> Result<Option<WalletRecord>, IndyError> {
        if !options.no_cache.unwrap_or(false) {
            let options_json = json!({
                "retrieveType": false,
                "retrieveValue": true,
                "retrieveTags": true,
            }).to_string();
            match self.wallet_service.get_record(wallet_handle, which_cache, id, &options_json) {
                Ok(record) => Ok(Some(record)),
                Err(err) => if err.kind() == IndyErrorKind::WalletItemNotFound { Ok(None) } else { Err(err) }
            }
        } else { Ok(None) }
    }

    fn _get_cred_def_continue(&self, wallet_handle: WalletHandle, ledger_response: IndyResult<(String, String)>, options: GetCacheOptions, cb_id: CommandHandle) {
        let cb = self.pending_callbacks.borrow_mut().remove(&cb_id).expect("FIXME INVALID STATE");

        let (cred_def_id, cred_def_json) = try_cb!(ledger_response, cb);

        if !options.no_store.unwrap_or(false) {
            let mut tags = Tags::new();
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    warn!("Cannot get time: {:?}", err);
                    0
                }
            };
            tags.insert("timestamp".to_string(), ts.to_string());
            let _ = self.wallet_service.delete_record(wallet_handle, CRED_DEF_CACHE, &cred_def_id);
            let _ = self.wallet_service.add_record(wallet_handle, CRED_DEF_CACHE, &cred_def_id, &cred_def_json, &tags);
        }

        cb(Ok(cred_def_json));
    }

    fn purge_schema_cache(&self,
                          wallet_handle: WalletHandle,
                          options_json: &str) -> IndyResult<()> {
        trace!("purge_schema_cache >>> wallet_handle: {:?}, options_json: {:?}", wallet_handle, options_json);

        let options = serde_json::from_str::<PurgeOptions>(options_json)
            .to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options")?;

        let max_age = options.max_age.unwrap_or(-1);
        let query_json = if max_age >= 0 {
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    error!("Cannot get time: {:?}", err);
                    return Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err)))
                }
            };
            json!({"timestamp": {"$lt": ts - max_age}}).to_string()
        } else {
            "{}".to_string()
        };

        let options_json = json!({
            "retrieveType": false,
            "retrieveValue": false,
            "retrieveTags": false,
        }).to_string();

        let mut search = self.wallet_service.search_records(
            wallet_handle,
            SCHEMA_CACHE,
            &query_json,
            &options_json,
        )?;

        while let Some(record) = search.fetch_next_record()? {
            self.wallet_service.delete_record(wallet_handle, SCHEMA_CACHE, record.get_id())?;
        }

        trace!("purge_schema_cache <<< res: ()");

        Ok(())
    }

    fn purge_cred_def_cache(&self,
                            wallet_handle: WalletHandle,
                            options_json: &str) -> IndyResult<()> {
        trace!("purge_cred_def_cache >>> wallet_handle: {:?}, options_json: {:?}", wallet_handle, options_json);

        let options = serde_json::from_str::<PurgeOptions>(options_json)
            .to_indy(IndyErrorKind::InvalidStructure, "Cannot deserialize options")?;

        let max_age = options.max_age.unwrap_or(-1);
        let query_json = if max_age >= 0 {
            let ts = match SystemTime::now().duration_since(UNIX_EPOCH) {
                Ok(ts) => ts.as_secs() as i32,
                Err(err) => {
                    error!("Cannot get time: {:?}", err);
                    return Err(IndyError::from_msg(IndyErrorKind::InvalidState, format!("Cannot get time: {:?}", err)))
                }
            };
            json!({"timestamp": {"$lt": ts - max_age}}).to_string()
        } else {
            "{}".to_string()
        };

        let options_json = json!({
            "retrieveType": false,
            "retrieveValue": false,
            "retrieveTags": false,
        }).to_string();

        let mut search = self.wallet_service.search_records(
            wallet_handle,
            CRED_DEF_CACHE,
            &query_json,
            &options_json,
        )?;

        while let Some(record) = search.fetch_next_record()? {
            self.wallet_service.delete_record(wallet_handle, CRED_DEF_CACHE, record.get_id())?;
        }

        trace!("purge_cred_def_cache <<< res: ()");

        Ok(())
    }
}

#[serde(rename_all = "camelCase")]
#[derive(Debug, Deserialize, Serialize)]
struct PurgeOptions {
    pub max_age: Option<i32>,
}

#[serde(rename_all = "camelCase")]
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct GetCacheOptions {
    pub no_cache: Option<bool>,     // Skip usage of cache,
    pub no_update: Option<bool>,    // Use only cached data, do not try to update.
    pub no_store: Option<bool>,     // Skip storing fresh data if updated
    pub min_fresh: Option<i32>,     // Return cached data if not older than this many seconds. -1 means do not check age.
}
<|endoftext|>


<commit_before>use ide_db::{
    assists::{AssistId, AssistKind},
    base_db::FileId,
    defs::Definition,
    search::FileReference,
    syntax_helpers::node_ext::full_path_of_name_ref,
};
use syntax::{
    ast::{self, NameLike, NameRef},
    AstNode, SyntaxKind, TextRange,
};

use crate::{AssistContext, Assists};

// Assist: unnecessary_async
//
// Removes the `async` mark from functions which have no `.await` in their body.
// Looks for calls to the functions and removes the `.await` on the call site.
//
// ```
// pub async f$0n foo() {}
// pub async fn bar() { foo().await }
// ```
// ->
// ```
// pub fn foo() {}
// pub async fn bar() { foo() }
// ```
pub(crate) fn unnecessary_async(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let function: ast::Fn = ctx.find_node_at_offset()?;

    // Do nothing if the cursor is not on the prototype. This is so that the check does not pollute
    // when the user asks us for assists when in the middle of the function body.
    // We consider the prototype to be anything that is before the body of the function.
    let cursor_position = ctx.offset();
    if cursor_position >= function.body()?.syntax().text_range().start() {
        return None;
    }
    // Do nothing if the function isn't async.
    if let None = function.async_token() {
        return None;
    }
    // Do nothing if the function has an `await` expression in its body.
    if function.body()?.syntax().descendants().find_map(ast::AwaitExpr::cast).is_some() {
        return None;
    }

    // Remove the `async` keyword plus whitespace after it, if any.
    let async_range = {
        let async_token = function.async_token()?;
        let next_token = async_token.next_token()?;
        if matches!(next_token.kind(), SyntaxKind::WHITESPACE) {
            TextRange::new(async_token.text_range().start(), next_token.text_range().end())
        } else {
            async_token.text_range()
        }
    };

    // Otherwise, we may remove the `async` keyword.
    acc.add(
        AssistId("unnecessary_async", AssistKind::QuickFix),
        "Remove unnecessary async",
        async_range,
        |edit| {
            // Remove async on the function definition.
            edit.replace(async_range, "");

            // Remove all `.await`s from calls to the function we remove `async` from.
            if let Some(fn_def) = ctx.sema.to_def(&function) {
                for await_expr in find_all_references(ctx, &Definition::Function(fn_def))
                    // Keep only references that correspond NameRefs.
                    .filter_map(|(_, reference)| match reference.name {
                        NameLike::NameRef(nameref) => Some(nameref),
                        _ => None,
                    })
                    // Keep only references that correspond to await expressions
                    .filter_map(|nameref| find_await_expression(ctx, &nameref))
                {
                    if let Some(await_token) = &await_expr.await_token() {
                        edit.replace(await_token.text_range(), "");
                    }
                    if let Some(dot_token) = &await_expr.dot_token() {
                        edit.replace(dot_token.text_range(), "");
                    }
                }
            }
        },
    )
}

fn find_all_references(
    ctx: &AssistContext<'_>,
    def: &Definition,
) -> impl Iterator<Item = (FileId, FileReference)> {
    def.usages(&ctx.sema).all().into_iter().flat_map(|(file_id, references)| {
        references.into_iter().map(move |reference| (file_id, reference))
    })
}

/// Finds the await expression for the given `NameRef`.
/// If no await expression is found, returns None.
fn find_await_expression(ctx: &AssistContext<'_>, nameref: &NameRef) -> Option<ast::AwaitExpr> {
    // From the nameref, walk up the tree to the await expression.
    let await_expr = if let Some(path) = full_path_of_name_ref(&nameref) {
        // Function calls.
        path.syntax()
            .parent()
            .and_then(ast::PathExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::CallExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::AwaitExpr::cast)
    } else {
        // Method calls.
        nameref
            .syntax()
            .parent()
            .and_then(ast::MethodCallExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::AwaitExpr::cast)
    };

    ctx.sema.original_ast_node(await_expr?)
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::tests::{check_assist, check_assist_not_applicable};

    #[test]
    fn applies_on_empty_function() {
        check_assist(unnecessary_async, "pub async f$0n f() {}", "pub fn f() {}")
    }

    #[test]
    fn applies_and_removes_whitespace() {
        check_assist(unnecessary_async, "pub async       f$0n f() {}", "pub fn f() {}")
    }

    #[test]
    fn does_not_apply_on_non_async_function() {
        check_assist_not_applicable(unnecessary_async, "pub f$0n f() {}")
    }

    #[test]
    fn applies_on_function_with_a_non_await_expr() {
        check_assist(unnecessary_async, "pub async f$0n f() { f2() }", "pub fn f() { f2() }")
    }

    #[test]
    fn does_not_apply_on_function_with_an_await_expr() {
        check_assist_not_applicable(unnecessary_async, "pub async f$0n f() { f2().await }")
    }

    #[test]
    fn applies_and_removes_await_on_reference() {
        check_assist(
            unnecessary_async,
            r#"
pub async fn f4() { }
pub async f$0n f2() { }
pub async fn f() { f2().await }
pub async fn f3() { f2().await }"#,
            r#"
pub async fn f4() { }
pub fn f2() { }
pub async fn f() { f2() }
pub async fn f3() { f2() }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_from_within_module() {
        check_assist(
            unnecessary_async,
            r#"
pub async fn f4() { }
mod a { pub async f$0n f2() { } }
pub async fn f() { a::f2().await }
pub async fn f3() { a::f2().await }"#,
            r#"
pub async fn f4() { }
mod a { pub fn f2() { } }
pub async fn f() { a::f2() }
pub async fn f3() { a::f2() }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_on_inner_await() {
        check_assist(
            unnecessary_async,
            // Ensure that it is the first await on the 3rd line that is removed
            r#"
pub async fn f() { f2().await }
pub async f$0n f2() -> i32 { 1 }
pub async fn f3() { f4(f2().await).await }
pub async fn f4(i: i32) { }"#,
            r#"
pub async fn f() { f2() }
pub fn f2() -> i32 { 1 }
pub async fn f3() { f4(f2()).await }
pub async fn f4(i: i32) { }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_on_outer_await() {
        check_assist(
            unnecessary_async,
            // Ensure that it is the second await on the 3rd line that is removed
            r#"
pub async fn f() { f2().await }
pub async f$0n f2(i: i32) { }
pub async fn f3() { f2(f4().await).await }
pub async fn f4() -> i32 { 1 }"#,
            r#"
pub async fn f() { f2() }
pub fn f2(i: i32) { }
pub async fn f3() { f2(f4().await) }
pub async fn f4() -> i32 { 1 }"#,
        )
    }

    #[test]
    fn applies_on_method_call() {
        check_assist(
            unnecessary_async,
            r#"
pub struct S { }
impl S { pub async f$0n f2(&self) { } }
pub async fn f(s: &S) { s.f2().await }"#,
            r#"
pub struct S { }
impl S { pub fn f2(&self) { } }
pub async fn f(s: &S) { s.f2() }"#,
        )
    }

    #[test]
    fn does_not_apply_on_function_with_a_nested_await_expr() {
        check_assist_not_applicable(
            unnecessary_async,
            "async f$0n f() { if true { loop { f2().await } } }",
        )
    }

    #[test]
    fn does_not_apply_when_not_on_prototype() {
        check_assist_not_applicable(unnecessary_async, "pub async fn f() { $0f2() }")
    }
}
<commit_msg>fix: async trait method for `unnecessary_async`<commit_after>use hir::AssocItem;
use ide_db::{
    assists::{AssistId, AssistKind},
    base_db::FileId,
    defs::Definition,
    search::FileReference,
    syntax_helpers::node_ext::full_path_of_name_ref,
    traits::resolve_target_trait,
};
use syntax::{
    ast::{self, HasName, NameLike, NameRef},
    AstNode, SyntaxKind, TextRange,
};

use crate::{AssistContext, Assists};

// Assist: unnecessary_async
//
// Removes the `async` mark from functions which have no `.await` in their body.
// Looks for calls to the functions and removes the `.await` on the call site.
//
// ```
// pub async f$0n foo() {}
// pub async fn bar() { foo().await }
// ```
// ->
// ```
// pub fn foo() {}
// pub async fn bar() { foo() }
// ```
pub(crate) fn unnecessary_async(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let function: ast::Fn = ctx.find_node_at_offset()?;

    // Do nothing if the cursor is not on the prototype. This is so that the check does not pollute
    // when the user asks us for assists when in the middle of the function body.
    // We consider the prototype to be anything that is before the body of the function.
    let cursor_position = ctx.offset();
    if cursor_position >= function.body()?.syntax().text_range().start() {
        return None;
    }
    // Do nothing if the function isn't async.
    if let None = function.async_token() {
        return None;
    }
    // Do nothing if the function has an `await` expression in its body.
    if function.body()?.syntax().descendants().find_map(ast::AwaitExpr::cast).is_some() {
        return None;
    }
    // Do nothing if the method is an async member of trait.
    if let Some(fname) = function.name() {
        if let Some(trait_item) = find_corresponding_trait_member(ctx, fname.to_string()) {
            if let AssocItem::Function(method) = trait_item {
                if method.is_async(ctx.db()) {
                    return None;
                }
            }
        }
    }
    // Remove the `async` keyword plus whitespace after it, if any.
    let async_range = {
        let async_token = function.async_token()?;
        let next_token = async_token.next_token()?;
        if matches!(next_token.kind(), SyntaxKind::WHITESPACE) {
            TextRange::new(async_token.text_range().start(), next_token.text_range().end())
        } else {
            async_token.text_range()
        }
    };

    // Otherwise, we may remove the `async` keyword.
    acc.add(
        AssistId("unnecessary_async", AssistKind::QuickFix),
        "Remove unnecessary async",
        async_range,
        |edit| {
            // Remove async on the function definition.
            edit.replace(async_range, "");

            // Remove all `.await`s from calls to the function we remove `async` from.
            if let Some(fn_def) = ctx.sema.to_def(&function) {
                for await_expr in find_all_references(ctx, &Definition::Function(fn_def))
                    // Keep only references that correspond NameRefs.
                    .filter_map(|(_, reference)| match reference.name {
                        NameLike::NameRef(nameref) => Some(nameref),
                        _ => None,
                    })
                    // Keep only references that correspond to await expressions
                    .filter_map(|nameref| find_await_expression(ctx, &nameref))
                {
                    if let Some(await_token) = &await_expr.await_token() {
                        edit.replace(await_token.text_range(), "");
                    }
                    if let Some(dot_token) = &await_expr.dot_token() {
                        edit.replace(dot_token.text_range(), "");
                    }
                }
            }
        },
    )
}

fn find_corresponding_trait_member(
    ctx: &AssistContext<'_>,
    function_name: String,
) -> Option<AssocItem> {
    let impl_ = ctx.find_node_at_offset::<ast::Impl>()?;
    let trait_ = resolve_target_trait(&ctx.sema, &impl_)?;

    trait_
        .items(ctx.db())
        .iter()
        .find(|item| match item.name(ctx.db()) {
            Some(method_name) => method_name.to_string() == function_name,
            _ => false,
        })
        .cloned()
}

fn find_all_references(
    ctx: &AssistContext<'_>,
    def: &Definition,
) -> impl Iterator<Item = (FileId, FileReference)> {
    def.usages(&ctx.sema).all().into_iter().flat_map(|(file_id, references)| {
        references.into_iter().map(move |reference| (file_id, reference))
    })
}

/// Finds the await expression for the given `NameRef`.
/// If no await expression is found, returns None.
fn find_await_expression(ctx: &AssistContext<'_>, nameref: &NameRef) -> Option<ast::AwaitExpr> {
    // From the nameref, walk up the tree to the await expression.
    let await_expr = if let Some(path) = full_path_of_name_ref(&nameref) {
        // Function calls.
        path.syntax()
            .parent()
            .and_then(ast::PathExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::CallExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::AwaitExpr::cast)
    } else {
        // Method calls.
        nameref
            .syntax()
            .parent()
            .and_then(ast::MethodCallExpr::cast)?
            .syntax()
            .parent()
            .and_then(ast::AwaitExpr::cast)
    };

    ctx.sema.original_ast_node(await_expr?)
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::tests::{check_assist, check_assist_not_applicable};

    #[test]
    fn applies_on_empty_function() {
        check_assist(unnecessary_async, "pub async f$0n f() {}", "pub fn f() {}")
    }

    #[test]
    fn applies_and_removes_whitespace() {
        check_assist(unnecessary_async, "pub async       f$0n f() {}", "pub fn f() {}")
    }

    #[test]
    fn does_not_apply_on_non_async_function() {
        check_assist_not_applicable(unnecessary_async, "pub f$0n f() {}")
    }

    #[test]
    fn applies_on_function_with_a_non_await_expr() {
        check_assist(unnecessary_async, "pub async f$0n f() { f2() }", "pub fn f() { f2() }")
    }

    #[test]
    fn does_not_apply_on_function_with_an_await_expr() {
        check_assist_not_applicable(unnecessary_async, "pub async f$0n f() { f2().await }")
    }

    #[test]
    fn applies_and_removes_await_on_reference() {
        check_assist(
            unnecessary_async,
            r#"
pub async fn f4() { }
pub async f$0n f2() { }
pub async fn f() { f2().await }
pub async fn f3() { f2().await }"#,
            r#"
pub async fn f4() { }
pub fn f2() { }
pub async fn f() { f2() }
pub async fn f3() { f2() }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_from_within_module() {
        check_assist(
            unnecessary_async,
            r#"
pub async fn f4() { }
mod a { pub async f$0n f2() { } }
pub async fn f() { a::f2().await }
pub async fn f3() { a::f2().await }"#,
            r#"
pub async fn f4() { }
mod a { pub fn f2() { } }
pub async fn f() { a::f2() }
pub async fn f3() { a::f2() }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_on_inner_await() {
        check_assist(
            unnecessary_async,
            // Ensure that it is the first await on the 3rd line that is removed
            r#"
pub async fn f() { f2().await }
pub async f$0n f2() -> i32 { 1 }
pub async fn f3() { f4(f2().await).await }
pub async fn f4(i: i32) { }"#,
            r#"
pub async fn f() { f2() }
pub fn f2() -> i32 { 1 }
pub async fn f3() { f4(f2()).await }
pub async fn f4(i: i32) { }"#,
        )
    }

    #[test]
    fn applies_and_removes_await_on_outer_await() {
        check_assist(
            unnecessary_async,
            // Ensure that it is the second await on the 3rd line that is removed
            r#"
pub async fn f() { f2().await }
pub async f$0n f2(i: i32) { }
pub async fn f3() { f2(f4().await).await }
pub async fn f4() -> i32 { 1 }"#,
            r#"
pub async fn f() { f2() }
pub fn f2(i: i32) { }
pub async fn f3() { f2(f4().await) }
pub async fn f4() -> i32 { 1 }"#,
        )
    }

    #[test]
    fn applies_on_method_call() {
        check_assist(
            unnecessary_async,
            r#"
pub struct S { }
impl S { pub async f$0n f2(&self) { } }
pub async fn f(s: &S) { s.f2().await }"#,
            r#"
pub struct S { }
impl S { pub fn f2(&self) { } }
pub async fn f(s: &S) { s.f2() }"#,
        )
    }

    #[test]
    fn does_not_apply_on_function_with_a_nested_await_expr() {
        check_assist_not_applicable(
            unnecessary_async,
            "async f$0n f() { if true { loop { f2().await } } }",
        )
    }

    #[test]
    fn does_not_apply_when_not_on_prototype() {
        check_assist_not_applicable(unnecessary_async, "pub async fn f() { $0f2() }")
    }

    #[test]
    fn applies_on_unnecessary_async_on_trait_method() {
        check_assist(
            unnecessary_async,
            r#"
trait Trait {
    fn foo();
}
impl Trait for () {
    $0async fn foo() {}
}"#,
            r#"
trait Trait {
    fn foo();
}
impl Trait for () {
    fn foo() {}
}"#,
        );
    }

    #[test]
    fn does_not_apply_on_async_trait_method() {
        check_assist_not_applicable(
            unnecessary_async,
            r#"
trait Trait {
    async fn foo();
}
impl Trait for () {
    $0async fn foo() {}
}"#,
        );
    }
}
<|endoftext|>


<commit_before>use std;
use toml;

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub enum Error {
    Config(String),
    Migration(String),
    MigrationComplete(String),
    PathError(String),
    IoOpen(std::io::Error),
    IoCreate(std::io::Error),
    IoRead(std::io::Error),
    IoWrite(std::io::Error),
    IoProc(std::io::Error),
    Utf8Error(std::string::FromUtf8Error),
    TomlDe(toml::de::Error),
    TomlSe(toml::ser::Error),
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use Error::*;
        match *self {
            Config(ref s)             => write!(f, "Config Error: {}", s),
            Migration(ref s)          => write!(f, "Migration Error: {}", s),
            MigrationComplete(ref s)  => write!(f, "MigrationComplete: {}", s),
            PathError(ref s)          => write!(f, "PathError: {}", s),
            IoOpen(ref e)             => write!(f, "IoOpen Error: {}", e),
            IoCreate(ref e)           => write!(f, "IoCreate Error: {}", e),
            IoRead(ref e)             => write!(f, "IoRead Error: {}", e),
            IoWrite(ref e)            => write!(f, "IoWrite Error: {}", e),
            IoProc(ref e)             => write!(f, "IoProcess Error: {}", e),
            Utf8Error(ref e)          => write!(f, "Utf8 Error: {}", e),
            TomlDe(ref e)             => write!(f, "Toml Deserialization Error: {}", e),
            TomlSe(ref e)             => write!(f, "Toml Serialization Error: {}", e),
        }
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        "Migrant Error"
    }

    fn cause(&self) -> Option<&std::error::Error> {
        use Error::*;
        Some(match *self {
            IoOpen(ref e)     => e,
            IoCreate(ref e)   => e,
            IoRead(ref e)     => e,
            IoWrite(ref e)    => e,
            IoProc(ref e)     => e,
            Utf8Error(ref e)  => e,
            TomlDe(ref e)     => e,
            TomlSe(ref e)     => e,
            _ => return None
        })
    }
}

#[macro_export]
macro_rules! format_err {
    ($e_type:expr, $literal:expr) => {
        $e_type(format!($literal))
    };
    ($e_type:expr, $literal:expr, $($arg:expr),*) => {
        $e_type(format!($literal, $($arg),*))
    };
}

#[macro_export]
macro_rules! bail {
    (Config <- $msg:expr) => {
        return Err(format_err!(Error::Config, $msg))
    };
    (Config <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::Config, $msg, $($arg),*))
    };
    (Migration <- $msg:expr) => {
        return Err(format_err!(Error::Migration, $msg))
    };
    (Migration <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::Migration, $msg, $($arg),*))
    };
    (MigrationComplete <- $msg:expr) => {
        return Err(format_err!(Error::MigrationComplete, $msg))
    };
    (MigrationComplete <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::MigrationComplete, $msg, $($arg),*))
    };
}

<commit_msg>don't export macros<commit_after>use std;
use toml;

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub enum Error {
    Config(String),
    Migration(String),
    MigrationComplete(String),
    PathError(String),
    IoOpen(std::io::Error),
    IoCreate(std::io::Error),
    IoRead(std::io::Error),
    IoWrite(std::io::Error),
    IoProc(std::io::Error),
    Utf8Error(std::string::FromUtf8Error),
    TomlDe(toml::de::Error),
    TomlSe(toml::ser::Error),
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use Error::*;
        match *self {
            Config(ref s)             => write!(f, "Config Error: {}", s),
            Migration(ref s)          => write!(f, "Migration Error: {}", s),
            MigrationComplete(ref s)  => write!(f, "MigrationComplete: {}", s),
            PathError(ref s)          => write!(f, "PathError: {}", s),
            IoOpen(ref e)             => write!(f, "IoOpen Error: {}", e),
            IoCreate(ref e)           => write!(f, "IoCreate Error: {}", e),
            IoRead(ref e)             => write!(f, "IoRead Error: {}", e),
            IoWrite(ref e)            => write!(f, "IoWrite Error: {}", e),
            IoProc(ref e)             => write!(f, "IoProcess Error: {}", e),
            Utf8Error(ref e)          => write!(f, "Utf8 Error: {}", e),
            TomlDe(ref e)             => write!(f, "Toml Deserialization Error: {}", e),
            TomlSe(ref e)             => write!(f, "Toml Serialization Error: {}", e),
        }
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        "Migrant Error"
    }

    fn cause(&self) -> Option<&std::error::Error> {
        use Error::*;
        Some(match *self {
            IoOpen(ref e)     => e,
            IoCreate(ref e)   => e,
            IoRead(ref e)     => e,
            IoWrite(ref e)    => e,
            IoProc(ref e)     => e,
            Utf8Error(ref e)  => e,
            TomlDe(ref e)     => e,
            TomlSe(ref e)     => e,
            _ => return None
        })
    }
}

macro_rules! format_err {
    ($e_type:expr, $literal:expr) => {
        $e_type(format!($literal))
    };
    ($e_type:expr, $literal:expr, $($arg:expr),*) => {
        $e_type(format!($literal, $($arg),*))
    };
}

macro_rules! bail {
    (Config <- $msg:expr) => {
        return Err(format_err!(Error::Config, $msg))
    };
    (Config <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::Config, $msg, $($arg),*))
    };
    (Migration <- $msg:expr) => {
        return Err(format_err!(Error::Migration, $msg))
    };
    (Migration <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::Migration, $msg, $($arg),*))
    };
    (MigrationComplete <- $msg:expr) => {
        return Err(format_err!(Error::MigrationComplete, $msg))
    };
    (MigrationComplete <- $msg:expr, $($arg:expr),*) => {
        return Err(format_err!(Error::MigrationComplete, $msg, $($arg),*))
    };
}

<|endoftext|>


<commit_before>//! Asynchronous filesystem manipulation operations (and stdin, stdout, stderr).
//!
//! This module contains basic methods and types for manipulating the contents
//! of the local filesystem from within the context of the Tokio runtime.
//!
//! Tasks running on the Tokio runtime are expected to be asynchronous, i.e.,
//! they will not block the thread of execution. Filesystem operations do not
//! satisfy this requirement. In order to perform filesystem operations
//! asynchronously, this library uses the [`blocking`][blocking] annotation
//! to signal to the runtime that a blocking operation is being performed. This
//! allows the runtime to compensate.
//!
//! [blocking]: https://docs.rs/tokio-threadpool/0.1/tokio_threadpool/fn.blocking.html

#[macro_use]
extern crate futures;
extern crate tokio_io;
extern crate tokio_threadpool;

pub mod file;
mod stdin;
mod stdout;
mod stderr;

pub use file::File;
pub use stdin::{stdin, Stdin};
pub use stdout::{stdout, Stdout};
pub use stderr::{stderr, Stderr};

use futures::Poll;
use futures::Async::*;

use std::io;
use std::io::ErrorKind::{Other, WouldBlock};

fn blocking_io<F, T>(f: F) -> Poll<T, io::Error>
where F: FnOnce() -> io::Result<T>,
{
    match tokio_threadpool::blocking(f) {
        Ok(Ready(Ok(v))) => Ok(v.into()),
        Ok(Ready(Err(err))) => Err(err),
        Ok(NotReady) => Ok(NotReady),
        Err(_) => Err(blocking_err()),
    }
}

fn would_block<F, T>(f: F) -> io::Result<T>
where F: FnOnce() -> io::Result<T>,
{
    match tokio_threadpool::blocking(f) {
        Ok(Ready(Ok(v))) => Ok(v),
        Ok(Ready(Err(err))) => {
            debug_assert_ne!(err.kind(), WouldBlock);
            Err(err)
        }
        Ok(NotReady) => Err(WouldBlock.into()),
        Err(_) => Err(blocking_err()),
    }
}

fn blocking_err() -> io::Error {
    io::Error::new(Other, "tokio-fs::File::open must be called \
                   from the context of the Tokio runtime.")
}
<commit_msg>Fix description of BlockingError as io::Error (#384)<commit_after>//! Asynchronous filesystem manipulation operations (and stdin, stdout, stderr).
//!
//! This module contains basic methods and types for manipulating the contents
//! of the local filesystem from within the context of the Tokio runtime.
//!
//! Tasks running on the Tokio runtime are expected to be asynchronous, i.e.,
//! they will not block the thread of execution. Filesystem operations do not
//! satisfy this requirement. In order to perform filesystem operations
//! asynchronously, this library uses the [`blocking`][blocking] annotation
//! to signal to the runtime that a blocking operation is being performed. This
//! allows the runtime to compensate.
//!
//! [blocking]: https://docs.rs/tokio-threadpool/0.1/tokio_threadpool/fn.blocking.html

#[macro_use]
extern crate futures;
extern crate tokio_io;
extern crate tokio_threadpool;

pub mod file;
mod stdin;
mod stdout;
mod stderr;

pub use file::File;
pub use stdin::{stdin, Stdin};
pub use stdout::{stdout, Stdout};
pub use stderr::{stderr, Stderr};

use futures::Poll;
use futures::Async::*;

use std::io;
use std::io::ErrorKind::{Other, WouldBlock};

fn blocking_io<F, T>(f: F) -> Poll<T, io::Error>
where F: FnOnce() -> io::Result<T>,
{
    match tokio_threadpool::blocking(f) {
        Ok(Ready(Ok(v))) => Ok(v.into()),
        Ok(Ready(Err(err))) => Err(err),
        Ok(NotReady) => Ok(NotReady),
        Err(_) => Err(blocking_err()),
    }
}

fn would_block<F, T>(f: F) -> io::Result<T>
where F: FnOnce() -> io::Result<T>,
{
    match tokio_threadpool::blocking(f) {
        Ok(Ready(Ok(v))) => Ok(v),
        Ok(Ready(Err(err))) => {
            debug_assert_ne!(err.kind(), WouldBlock);
            Err(err)
        }
        Ok(NotReady) => Err(WouldBlock.into()),
        Err(_) => Err(blocking_err()),
    }
}

fn blocking_err() -> io::Error {
    io::Error::new(Other, "`blocking` annotated I/O must be called \
                   from the context of the Tokio runtime.")
}
<|endoftext|>


<commit_before>// Kevin Cantu
// hash function testing

use std;
use murmur;
use djb;
use str2;

fn word_of_god () -> [str]
{
   let path = "./lolcat/LOLCatBible_2012-01-04.xml";
   let bible = std::io::read_whole_file (path);
   let bible_ = str::unsafe_from_bytes (result::get (bible));

   ret str2::words(bible_) +
       vec::map(str2::lines(bible_), {|ln| ln+ln+ln+ln});
}

// given a string hash function
// (modified to provide a string return value)
// test its behavior
fn hash_bench <TT> ( label: str, hashfn: fn(&&str)->TT, data: [str] )
{
   let t0 = std::time::precise_time_s();
   let _v = vec::map(data, hashfn);
   let t1 = std::time::precise_time_s();

   std::io::println(label + #fmt("%06.3f sec", t1 - t0));
}

// main
fn main () {
   let meow = word_of_god();

   hash_bench ("Benching djb...     ", djb::djb,          meow);
   hash_bench ("Benching murmur3... ", murmur::murmur,       meow);
}

<commit_msg>comments<commit_after>// Kevin Cantu
// hash function testing

use std;
use murmur;
use djb;
use str2;

fn word_of_god () -> [str]
{
   let path = "./lolcat/LOLCatBible_2012-01-04.xml";
   let bible = std::io::read_whole_file (path);
   let bible_ = str::unsafe_from_bytes (result::get (bible));

   ret str2::words(bible_) +
       vec::map(str2::lines(bible_), {|ln| ln+ln+ln+ln});
}

// given a string hash function, test its behavior
fn hash_bench <TT> ( label: str, hashfn: fn(&&str)->TT, data: [str] )
{
   let t0 = std::time::precise_time_s();
   let _v = vec::map(data, hashfn);
   let t1 = std::time::precise_time_s();

   std::io::println(label + #fmt("%06.3f sec", t1 - t0));
}

// main
fn main () {
   let meow = word_of_god();

   hash_bench ("Benching djb...     ", djb::djb,          meow);
   hash_bench ("Benching murmur3... ", murmur::murmur,       meow);
}

<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

/// This module provides an experiment which simulates a multi-region environment.
/// It undoes the simulation in the cluster after the given duration
use crate::effects::Effect;
use crate::effects::NetworkDelay;
use crate::prometheus::Prometheus;
use crate::util::unix_timestamp_now;
use crate::{cluster::Cluster, experiments::Experiment, thread_pool_executor::ThreadPoolExecutor};
use failure;
use slog_scope::{info, warn};
use std::{collections::HashSet, fmt, thread, time::Duration};

#[derive(Default, Debug)]
struct Metrics {
    split_size: usize,
    cross_region_latency: u64,
    txn_per_sec: Option<f64>,
    blocks_per_sec: Option<f64>,
    avg_vote_time: Option<f64>,
    avg_qc_time: Option<f64>,
}

pub struct MultiRegionSimulation {
    splits: Vec<usize>,
    cross_region_latencies: Vec<u64>,
    exp_duration_per_config: Duration,
    thread_pool_executor: ThreadPoolExecutor,
    cluster: Cluster,
    prometheus: Prometheus,
}

impl MultiRegionSimulation {
    pub fn new(
        splits: Vec<usize>,
        cross_region_latencies: Vec<u64>,
        exp_duration_per_config: Duration,
        cluster: Cluster,
        thread_pool_executor: ThreadPoolExecutor,
        prometheus: Prometheus,
    ) -> Self {
        Self {
            splits,
            cross_region_latencies,
            exp_duration_per_config,
            thread_pool_executor,
            cluster,
            prometheus,
        }
    }

    fn single_run(&self, count: usize, cross_region_delay: Duration) -> failure::Result<Metrics> {
        let (cluster1, cluster2) = self.cluster.split_n_random(count);
        let (region1, region2) = (cluster1.into_instances(), cluster2.into_instances());
        let (smaller_region, larger_region);
        if region1.len() < region2.len() {
            smaller_region = &region1;
            larger_region = &region2;
        } else {
            smaller_region = &region2;
            larger_region = &region1;
        }
        let network_delays_effects: Vec<_> = smaller_region
            .iter()
            .map(|instance| {
                NetworkDelay::new(
                    instance.clone(),
                    Some(larger_region.clone()),
                    cross_region_delay,
                )
            })
            .collect();
        // Add network delay commands only to the instances of the smaller
        // region because we have to run fewer ssh commands.
        let start_effects: Vec<_> = network_delays_effects
            .iter()
            .map(|effect| {
                move || {
                    effect.activate().unwrap();
                }
            })
            .collect();
        self.thread_pool_executor.execute_jobs(start_effects);

        thread::sleep(self.exp_duration_per_config);
        let metrics = self.get_metrics(count, cross_region_delay.as_millis() as u64);
        let stop_effects: Vec<_> = network_delays_effects
            .iter()
            .map(|effect| {
                move || {
                    effect.deactivate().unwrap();
                }
            })
            .collect();
        self.thread_pool_executor.execute_jobs(stop_effects);
        Ok(metrics)
    }

    fn get_metrics(&self, split_size: usize, cross_region_latency: u64) -> Metrics {
        let mut metrics: Metrics = Default::default();
        metrics.split_size = split_size;
        metrics.cross_region_latency = cross_region_latency;
        let step = 10;
        let end = unix_timestamp_now();
        // Measure metrics 30s after the experiment started to ignore initial warmup metrics
        let start = end - self.exp_duration_per_config + Duration::from_secs(30);
        metrics.txn_per_sec = self
            .prometheus
            .query_range_avg(
                "irate(consensus_gauge{op='last_committed_version'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok();
        metrics.blocks_per_sec = self
            .prometheus
            .query_range_avg(
                "irate(consensus_gauge{op='last_committed_round'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok();
        metrics.avg_vote_time = self
            .prometheus
            .query_range_avg(
                "irate(consensus_duration_sum{op='creation_to_receival_s'}[1m])/irate(consensus_duration_count{op='creation_to_receival_s'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok().map(|x| x * 1000f64);
        metrics.avg_qc_time = self
            .prometheus
            .query_range_avg(
                "irate(consensus_duration_sum{op='creation_to_qc_s'}[1m])/irate(consensus_duration_count{op='creation_to_qc_s'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok().map(|x| x * 1000f64);
        metrics
    }
}

fn print_results(metrics: Vec<Metrics>) {
    info!("Results:");
    println!(
        "split_size, cross_region_latency, txn_per_sec, blocks_per_sec, avg_vote_time, avg_qc_time"
    );
    let to_str = |x: &Option<f64>| -> String {
        if let Some(val) = x {
            format!("{:.2}", val)
        } else {
            "".to_string()
        }
    };
    for metric in metrics {
        println!(
            "{}, {}, {}, {}, {}, {}",
            metric.split_size,
            metric.cross_region_latency,
            to_str(&metric.txn_per_sec),
            to_str(&metric.blocks_per_sec),
            to_str(&metric.avg_vote_time),
            to_str(&metric.avg_qc_time),
        );
    }
}

impl Experiment for MultiRegionSimulation {
    fn affected_validators(&self) -> HashSet<String> {
        let mut r = HashSet::new();
        for instance in self.cluster.clone().into_instances() {
            r.insert(instance.short_hash().clone());
        }
        r
    }

    fn run(&self) -> failure::Result<()> {
        let mut results = vec![];
        for split in &self.splits {
            for cross_region_latency in &self.cross_region_latencies {
                info!(
                    "Running multi region simulation: split: {}, cross region latency: {}ms",
                    split, cross_region_latency
                );
                if let Ok(metrics) = self
                    .single_run(*split, Duration::from_millis(*cross_region_latency))
                    .map_err(|e| warn!("{}", e))
                {
                    info!("metrics for this run: {:?}", metrics);
                    results.push(metrics);
                }
                thread::sleep(Duration::from_secs(60));
            }
        }
        print_results(results);
        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(24 * 60 * 60)
    }
}

impl fmt::Display for MultiRegionSimulation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Multi Region Simulation splits: {:?} cross region latencies: {:?}",
            self.splits, self.cross_region_latencies
        )
    }
}
<commit_msg>[cluster-test] Emit transactions when running network delay simulations<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

/// This module provides an experiment which simulates a multi-region environment.
/// It undoes the simulation in the cluster after the given duration
use crate::effects::Effect;
use crate::effects::NetworkDelay;
use crate::prometheus::Prometheus;
use crate::tx_emitter::{EmitJobRequest, EmitThreadParams, TxEmitter};
use crate::util::unix_timestamp_now;
use crate::{cluster::Cluster, experiments::Experiment, thread_pool_executor::ThreadPoolExecutor};
use failure;
use slog_scope::{info, warn};
use std::{collections::HashSet, fmt, thread, time::Duration};

#[derive(Default, Debug)]
struct Metrics {
    split_size: usize,
    cross_region_latency: u64,
    txn_per_sec: Option<f64>,
    blocks_per_sec: Option<f64>,
    avg_vote_time: Option<f64>,
    avg_qc_time: Option<f64>,
}

pub struct MultiRegionSimulation {
    splits: Vec<usize>,
    cross_region_latencies: Vec<u64>,
    exp_duration_per_config: Duration,
    thread_pool_executor: ThreadPoolExecutor,
    cluster: Cluster,
    prometheus: Prometheus,
}

impl MultiRegionSimulation {
    pub fn new(
        splits: Vec<usize>,
        cross_region_latencies: Vec<u64>,
        exp_duration_per_config: Duration,
        cluster: Cluster,
        thread_pool_executor: ThreadPoolExecutor,
        prometheus: Prometheus,
    ) -> Self {
        Self {
            splits,
            cross_region_latencies,
            exp_duration_per_config,
            thread_pool_executor,
            cluster,
            prometheus,
        }
    }

    fn single_run(&self, count: usize, cross_region_delay: Duration) -> failure::Result<Metrics> {
        let (cluster1, cluster2) = self.cluster.split_n_random(count);
        let (region1, region2) = (cluster1.into_instances(), cluster2.into_instances());
        let (smaller_region, larger_region);
        if region1.len() < region2.len() {
            smaller_region = &region1;
            larger_region = &region2;
        } else {
            smaller_region = &region2;
            larger_region = &region1;
        }
        let network_delays_effects: Vec<_> = smaller_region
            .iter()
            .map(|instance| {
                NetworkDelay::new(
                    instance.clone(),
                    Some(larger_region.clone()),
                    cross_region_delay,
                )
            })
            .collect();
        // Add network delay commands only to the instances of the smaller
        // region because we have to run fewer ssh commands.
        let start_effects: Vec<_> = network_delays_effects
            .iter()
            .map(|effect| {
                move || {
                    effect.activate().unwrap();
                }
            })
            .collect();
        self.thread_pool_executor.execute_jobs(start_effects);

        thread::sleep(self.exp_duration_per_config);
        let metrics = self.get_metrics(count, cross_region_delay.as_millis() as u64);
        let stop_effects: Vec<_> = network_delays_effects
            .iter()
            .map(|effect| {
                move || {
                    effect.deactivate().unwrap();
                }
            })
            .collect();
        self.thread_pool_executor.execute_jobs(stop_effects);
        Ok(metrics)
    }

    fn get_metrics(&self, split_size: usize, cross_region_latency: u64) -> Metrics {
        let mut metrics: Metrics = Default::default();
        metrics.split_size = split_size;
        metrics.cross_region_latency = cross_region_latency;
        let step = 10;
        let end = unix_timestamp_now();
        // Measure metrics 30s after the experiment started to ignore initial warmup metrics
        let start = end - self.exp_duration_per_config + Duration::from_secs(30);
        metrics.txn_per_sec = self
            .prometheus
            .query_range_avg(
                "irate(consensus_gauge{op='last_committed_version'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok();
        metrics.blocks_per_sec = self
            .prometheus
            .query_range_avg(
                "irate(consensus_gauge{op='last_committed_round'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok();
        metrics.avg_vote_time = self
            .prometheus
            .query_range_avg(
                "irate(consensus_duration_sum{op='creation_to_receival_s'}[1m])/irate(consensus_duration_count{op='creation_to_receival_s'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok().map(|x| x * 1000f64);
        metrics.avg_qc_time = self
            .prometheus
            .query_range_avg(
                "irate(consensus_duration_sum{op='creation_to_qc_s'}[1m])/irate(consensus_duration_count{op='creation_to_qc_s'}[1m])".to_string(),
                &start,
                &end,
                step,
            )
            .map_err(|e| warn!("{}", e))
            .ok().map(|x| x * 1000f64);
        metrics
    }
}

fn print_results(metrics: Vec<Metrics>) {
    info!("Results:");
    println!(
        "split_size, cross_region_latency, txn_per_sec, blocks_per_sec, avg_vote_time, avg_qc_time"
    );
    let to_str = |x: &Option<f64>| -> String {
        if let Some(val) = x {
            format!("{:.2}", val)
        } else {
            "".to_string()
        }
    };
    for metric in metrics {
        println!(
            "{}, {}, {}, {}, {}, {}",
            metric.split_size,
            metric.cross_region_latency,
            to_str(&metric.txn_per_sec),
            to_str(&metric.blocks_per_sec),
            to_str(&metric.avg_vote_time),
            to_str(&metric.avg_qc_time),
        );
    }
}

impl Experiment for MultiRegionSimulation {
    fn affected_validators(&self) -> HashSet<String> {
        let mut r = HashSet::new();
        for instance in self.cluster.clone().into_instances() {
            r.insert(instance.short_hash().clone());
        }
        r
    }

    fn run(&self) -> failure::Result<()> {
        let mut emitter = TxEmitter::new(&self.cluster);
        let mut results = vec![];
        for split in &self.splits {
            for cross_region_latency in &self.cross_region_latencies {
                let job = emitter
                    .start_job(EmitJobRequest {
                        instances: self.cluster.instances().clone(),
                        accounts_per_client: 10,
                        thread_params: EmitThreadParams::default(),
                    })
                    .expect("Failed to start emit job");
                // Wait for minting to complete and transactions to start
                thread::sleep(Duration::from_secs(30));
                info!(
                    "Running multi region simulation: split: {}, cross region latency: {}ms",
                    split, cross_region_latency
                );
                if let Ok(metrics) = self
                    .single_run(*split, Duration::from_millis(*cross_region_latency))
                    .map_err(|e| warn!("{}", e))
                {
                    info!("metrics for this run: {:?}", metrics);
                    results.push(metrics);
                }
                emitter.stop_job(job);
                // Sleep for a minute before different experiments
                thread::sleep(Duration::from_secs(60));
            }
        }
        print_results(results);
        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(24 * 60 * 60)
    }
}

impl fmt::Display for MultiRegionSimulation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Multi Region Simulation splits: {:?} cross region latencies: {:?}",
            self.splits, self.cross_region_latencies
        )
    }
}
<|endoftext|>


<commit_before>use cdchunking::{Chunker, ChunkInput, ZPAQ};
use rusqlite;
use rusqlite::{Connection, Transaction};
use rusqlite::types::ToSql;
use sha1::Sha1;
use std::fs::File;
use std::path::{Path, PathBuf};

use crate::{Error, HashDigest};

const SCHEMA: &'static str = "
    CREATE TABLE version(
        name VARCHAR(8) NOT NULL,
        version VARCHAR(16) NOT NULL
    );
    INSERT INTO version(name, version) VALUES('rs-sync', '0.1');

    CREATE TABLE files(
        file_id INTEGER NOT NULL PRIMARY KEY,
        name VARCHAR(512) NOT NULL,
        modified DATETIME NOT NULL
    );
    CREATE INDEX idx_files_name ON files(name);

    CREATE TABLE blocks(
        hash VARCHAR(40) NOT NULL,
        file_id INTEGER NOT NULL,
        offset INTEGER NOT NULL,
        PRIMARY KEY(file_id, offset)
    );
    CREATE INDEX idx_blocks_hash ON blocks(hash);
    CREATE INDEX idx_blocks_file ON blocks(file_id);
    CREATE INDEX idx_blocks_file_offset ON blocks(file_id, offset);
";

/// Index of files and blocks
pub struct Index {
    db: Connection,
}

impl Index {
    /// Open an index from a file
    pub fn open(filename: &Path) -> Result<Index, Error> {
        let exists = filename.exists();
        let db = Connection::open(filename)?;
        if !exists {
            warn!("Database doesn't exist, creating tables...");
            db.execute_batch(SCHEMA)?;
        }
        Ok(Index { db })
    }

    /// Open an in-memory index
    pub fn open_in_memory() -> Result<Index, Error> {
        let db = Connection::open_in_memory()?;
        db.execute_batch(SCHEMA)?;
        Ok(Index { db })
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize)>, Error>
    {
        let mut stmt = self.db.prepare(
            "
            SELECT files.name, blocks.offset
            FROM blocks
            INNER JOIN files ON blocks.file_id = files.file_id
            WHERE blocks.hash = ?;
            ",
        )?;
        let mut rows = stmt.query(&[&hash as &dyn ToSql])?;
        if let Some(row) = rows.next() {
            let row = row?;
            let path: String = row.get(0);
            let path: PathBuf = path.into();
            let offset: i64 = row.get(1);
            let offset = offset as usize;
            Ok(Some((path, offset)))
        } else {
            Ok(None)
        }
    }

    /// Start a transaction to update the index
    pub fn transaction<'a>(
        &'a mut self
    ) -> Result<IndexTransaction<'a>, rusqlite::Error>
    {
        let tx = self.db.transaction()?;
        Ok(IndexTransaction { tx })
    }
}

/// A transaction on the index, for safety and performance
pub struct IndexTransaction<'a> {
    tx: Transaction<'a>,
}

const ZPAQ_BITS: usize = 13; // 13 bits = 8 KiB block average
const MAX_BLOCK_SIZE: usize = 1 << 15; // 32 KiB

impl<'a> IndexTransaction<'a> {
    /// Add a file to the index
    ///
    /// This returns a tuple `(file_id, up_to_date)` where `file_id` can be
    /// used to insert blocks, and `up_to_date` indicates whether the file's
    /// modification date has changed and it should be re-indexed.
    pub fn add_file(
        &mut self,
        name: &Path,
        modified: chrono::DateTime<chrono::Utc>,
    ) -> Result<(u32, bool), Error>
    {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, modified FROM files
            WHERE name = ?;
            ",
        )?;
        let mut rows = stmt.query(&[name.to_str().expect("encoding")])?;
        if let Some(row) = rows.next() {
            let row = row?;
            let file_id: u32 = row.get(0);
            let old_modified: chrono::DateTime<chrono::Utc> = row.get(1);
            if old_modified != modified {
                info!("Resetting file {:?}, modified", name);
                // Delete blocks
                self.tx.execute(
                    "
                    DELETE FROM blocks WHERE file_id = ?;
                    ",
                    &[&file_id],
                )?;
                // Update modification time
                self.tx.execute(
                    "
                    UPDATE files SET modified = ? WHERE file_id = ?;
                    ",
                    &[&modified as &dyn ToSql, &file_id],
                )?;
                Ok((file_id, false))
            } else {
                debug!("File {:?} up to date", name);
                Ok((file_id, true))
            }
        } else {
            info!("Inserting new file {:?}", name);
            self.tx.execute(
                "
                INSERT INTO files(name, modified)
                VALUES(?, ?);
                ",
                &[&name.to_str().expect("encoding") as &dyn ToSql, &modified],
            )?;
            let file_id = self.tx.last_insert_rowid();
            Ok((file_id as u32, false))
        }
    }

    /// Remove a file and all its blocks from the index
    pub fn remove_file(
        &mut self,
        file_id: u32,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            DELETE FROM blocks WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        self.tx.execute(
            "
            DELETE FROM files WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        Ok(())
    }

    /// Get a list of all the files in the index
    pub fn list_files(&self) -> Result<Vec<(u32, PathBuf)>, Error> {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, name FROM files;
            ",
        )?;
        let mut rows = stmt.query(rusqlite::NO_PARAMS)?;
        let mut results = Vec::new();
        loop {
            match rows.next() {
                Some(Ok(row)) => {
                    let path: String = row.get(1);
                    results.push((row.get(0), path.into()))
                }
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        Ok(results)
    }

    /// Add a block to the index
    pub fn add_block(
        &mut self,
        hash: &HashDigest,
        file_id: u32,
        offset: usize,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            INSERT INTO blocks(hash, file_id, offset)
            VALUES(?, ?, ?);
            ",
            &[&hash as &dyn ToSql, &file_id, &(offset as i64)],
        )?;
        Ok(())
    }

    /// Cut up a file into blocks and add them to the index
    pub fn index_file(
        &mut self,
        name: &Path,
    ) -> Result<(), Error>
    {
        let file = File::open(name)?;
        let (file_id, up_to_date) = self.add_file(
            name,
            file.metadata()?.modified()?.into(),
        )?;
        if !up_to_date {
            // Use ZPAQ to cut the stream into blocks
            let chunker = Chunker::new(
                ZPAQ::new(ZPAQ_BITS) // 13 bits = 8 KiB block average
            ).max_size(MAX_BLOCK_SIZE);
            let mut chunk_iterator = chunker.stream(file);
            let mut start_offset = 0;
            let mut offset = 0;
            let mut sha1 = Sha1::new();
            while let Some(chunk) = chunk_iterator.read() {
                match chunk? {
                    ChunkInput::Data(d) => {
                        sha1.update(d);
                        offset += d.len();
                    }
                    ChunkInput::End => {
                        let digest = HashDigest(sha1.digest().bytes());
                        debug!(
                            "Adding block, offset={}, size={}, sha1={}",
                            start_offset, offset - start_offset, digest,
                        );
                        self.add_block(&digest, file_id, start_offset)?;
                        start_offset = offset;
                        sha1.reset();
                    }
                }
            }
        }
        Ok(())
    }

    /// Commit the transaction
    pub fn commit(self) -> Result<(), rusqlite::Error> {
        self.tx.commit()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tempfile::NamedTempFile;

    use crate::HashDigest;
    use super::{Index, MAX_BLOCK_SIZE};

    #[test]
    fn test() {
        let mut file = NamedTempFile::new().expect("tempfile");
        for i in 0..2000 {
            write!(file, "Line {}\n", i + 1).expect("tempfile");
        }
        for _ in 0..2000 {
            write!(file, "Test content\n").expect("tempfile");
        }
        file.flush().expect("tempfile");
        let mut index = Index::open_in_memory().expect("db");
        {
            let mut tx = index.transaction().expect("db");
            tx.index_file(file.path()).expect("index");
            tx.commit().expect("db");
        }
        assert!(
            index.get_block(&HashDigest(*b"12345678901234567890"))
                .expect("get")
                .is_none()
        );
        let block1 = index.get_block(&HashDigest(
            *b"\xfb\x5e\xf7\xeb\xad\xd8\x2c\x80\x85\xc5\
               \xff\x63\x82\x36\x22\xba\xe0\xe2\x63\xf6"
        )).expect("get");
        assert_eq!(
            block1,
            Some((file.path().into(), 0)),
        );
        let block2 = index.get_block(&HashDigest(
            *b"\x57\x0d\x8b\x30\xfc\xfd\x58\x5e\x41\x27\
               \xb5\x61\xf5\xec\xd3\x76\xff\x4d\x01\x01"
        )).expect("get");
        assert_eq!(
            block2,
            Some((file.path().into(), 11579)),
        );
        let block3 = index.get_block(&HashDigest(
            *b"\xb9\xa8\xc2\x64\x1a\xf2\xcf\x8f\xd8\xf3\
               \x6a\x24\x56\xa3\xea\xa9\x5c\x02\x91\x27"
        )).expect("get");
        assert_eq!(
            block3,
            Some((file.path().into(), 44347)),
        );
        assert_eq!(block3.unwrap().1 - block2.unwrap().1, MAX_BLOCK_SIZE);
    }
}
<commit_msg>Have get_block() on both Index & IndexTransaction<commit_after>use cdchunking::{Chunker, ChunkInput, ZPAQ};
use rusqlite;
use rusqlite::{Connection, Transaction};
use rusqlite::types::ToSql;
use sha1::Sha1;
use std::fs::File;
use std::path::{Path, PathBuf};

use crate::{Error, HashDigest};

const SCHEMA: &'static str = "
    CREATE TABLE version(
        name VARCHAR(8) NOT NULL,
        version VARCHAR(16) NOT NULL
    );
    INSERT INTO version(name, version) VALUES('rs-sync', '0.1');

    CREATE TABLE files(
        file_id INTEGER NOT NULL PRIMARY KEY,
        name VARCHAR(512) NOT NULL,
        modified DATETIME NOT NULL
    );
    CREATE INDEX idx_files_name ON files(name);

    CREATE TABLE blocks(
        hash VARCHAR(40) NOT NULL,
        file_id INTEGER NOT NULL,
        offset INTEGER NOT NULL,
        PRIMARY KEY(file_id, offset)
    );
    CREATE INDEX idx_blocks_hash ON blocks(hash);
    CREATE INDEX idx_blocks_file ON blocks(file_id);
    CREATE INDEX idx_blocks_file_offset ON blocks(file_id, offset);
";

fn get_block(
    db: &Connection,
    hash: &HashDigest,
) -> Result<Option<(PathBuf, usize)>, Error>
{
    let mut stmt = db.prepare(
        "
        SELECT files.name, blocks.offset
        FROM blocks
        INNER JOIN files ON blocks.file_id = files.file_id
        WHERE blocks.hash = ?;
        ",
    )?;
    let mut rows = stmt.query(&[hash as &dyn ToSql])?;
    if let Some(row) = rows.next() {
        let row = row?;
        let path: String = row.get(0);
        let path: PathBuf = path.into();
        let offset: i64 = row.get(1);
        let offset = offset as usize;
        Ok(Some((path, offset)))
    } else {
        Ok(None)
    }
}

/// Index of files and blocks
pub struct Index {
    db: Connection,
}

impl Index {
    /// Open an index from a file
    pub fn open(filename: &Path) -> Result<Index, Error> {
        let exists = filename.exists();
        let db = Connection::open(filename)?;
        if !exists {
            warn!("Database doesn't exist, creating tables...");
            db.execute_batch(SCHEMA)?;
        }
        Ok(Index { db })
    }

    /// Open an in-memory index
    pub fn open_in_memory() -> Result<Index, Error> {
        let db = Connection::open_in_memory()?;
        db.execute_batch(SCHEMA)?;
        Ok(Index { db })
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize)>, Error>
    {
        get_block(&self.db, hash)
    }

    /// Start a transaction to update the index
    pub fn transaction<'a>(
        &'a mut self
    ) -> Result<IndexTransaction<'a>, rusqlite::Error>
    {
        let tx = self.db.transaction()?;
        Ok(IndexTransaction { tx })
    }
}

/// A transaction on the index, for safety and performance
pub struct IndexTransaction<'a> {
    tx: Transaction<'a>,
}

const ZPAQ_BITS: usize = 13; // 13 bits = 8 KiB block average
const MAX_BLOCK_SIZE: usize = 1 << 15; // 32 KiB

impl<'a> IndexTransaction<'a> {
    /// Add a file to the index
    ///
    /// This returns a tuple `(file_id, up_to_date)` where `file_id` can be
    /// used to insert blocks, and `up_to_date` indicates whether the file's
    /// modification date has changed and it should be re-indexed.
    pub fn add_file(
        &mut self,
        name: &Path,
        modified: chrono::DateTime<chrono::Utc>,
    ) -> Result<(u32, bool), Error>
    {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, modified FROM files
            WHERE name = ?;
            ",
        )?;
        let mut rows = stmt.query(&[name.to_str().expect("encoding")])?;
        if let Some(row) = rows.next() {
            let row = row?;
            let file_id: u32 = row.get(0);
            let old_modified: chrono::DateTime<chrono::Utc> = row.get(1);
            if old_modified != modified {
                info!("Resetting file {:?}, modified", name);
                // Delete blocks
                self.tx.execute(
                    "
                    DELETE FROM blocks WHERE file_id = ?;
                    ",
                    &[&file_id],
                )?;
                // Update modification time
                self.tx.execute(
                    "
                    UPDATE files SET modified = ? WHERE file_id = ?;
                    ",
                    &[&modified as &dyn ToSql, &file_id],
                )?;
                Ok((file_id, false))
            } else {
                debug!("File {:?} up to date", name);
                Ok((file_id, true))
            }
        } else {
            info!("Inserting new file {:?}", name);
            self.tx.execute(
                "
                INSERT INTO files(name, modified)
                VALUES(?, ?);
                ",
                &[&name.to_str().expect("encoding") as &dyn ToSql, &modified],
            )?;
            let file_id = self.tx.last_insert_rowid();
            Ok((file_id as u32, false))
        }
    }

    /// Remove a file and all its blocks from the index
    pub fn remove_file(
        &mut self,
        file_id: u32,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            DELETE FROM blocks WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        self.tx.execute(
            "
            DELETE FROM files WHERE file_id = ?;
            ",
            &[&file_id],
        )?;
        Ok(())
    }

    /// Get a list of all the files in the index
    pub fn list_files(&self) -> Result<Vec<(u32, PathBuf)>, Error> {
        let mut stmt = self.tx.prepare(
            "
            SELECT file_id, name FROM files;
            ",
        )?;
        let mut rows = stmt.query(rusqlite::NO_PARAMS)?;
        let mut results = Vec::new();
        loop {
            match rows.next() {
                Some(Ok(row)) => {
                    let path: String = row.get(1);
                    results.push((row.get(0), path.into()))
                }
                Some(Err(e)) => return Err(e.into()),
                None => break,
            }
        }
        Ok(results)
    }

    /// Add a block to the index
    pub fn add_block(
        &mut self,
        hash: &HashDigest,
        file_id: u32,
        offset: usize,
    ) -> Result<(), Error>
    {
        self.tx.execute(
            "
            INSERT INTO blocks(hash, file_id, offset)
            VALUES(?, ?, ?);
            ",
            &[&hash as &dyn ToSql, &file_id, &(offset as i64)],
        )?;
        Ok(())
    }

    /// Try to find a block in the indexed files
    pub fn get_block(
        &self,
        hash: &HashDigest,
    ) -> Result<Option<(PathBuf, usize)>, Error>
    {
        get_block(&self.tx, hash)
    }

    /// Cut up a file into blocks and add them to the index
    pub fn index_file(
        &mut self,
        name: &Path,
    ) -> Result<(), Error>
    {
        let file = File::open(name)?;
        let (file_id, up_to_date) = self.add_file(
            name,
            file.metadata()?.modified()?.into(),
        )?;
        if !up_to_date {
            // Use ZPAQ to cut the stream into blocks
            let chunker = Chunker::new(
                ZPAQ::new(ZPAQ_BITS) // 13 bits = 8 KiB block average
            ).max_size(MAX_BLOCK_SIZE);
            let mut chunk_iterator = chunker.stream(file);
            let mut start_offset = 0;
            let mut offset = 0;
            let mut sha1 = Sha1::new();
            while let Some(chunk) = chunk_iterator.read() {
                match chunk? {
                    ChunkInput::Data(d) => {
                        sha1.update(d);
                        offset += d.len();
                    }
                    ChunkInput::End => {
                        let digest = HashDigest(sha1.digest().bytes());
                        debug!(
                            "Adding block, offset={}, size={}, sha1={}",
                            start_offset, offset - start_offset, digest,
                        );
                        self.add_block(&digest, file_id, start_offset)?;
                        start_offset = offset;
                        sha1.reset();
                    }
                }
            }
        }
        Ok(())
    }

    /// Commit the transaction
    pub fn commit(self) -> Result<(), rusqlite::Error> {
        self.tx.commit()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use tempfile::NamedTempFile;

    use crate::HashDigest;
    use super::{Index, MAX_BLOCK_SIZE};

    #[test]
    fn test() {
        let mut file = NamedTempFile::new().expect("tempfile");
        for i in 0..2000 {
            write!(file, "Line {}\n", i + 1).expect("tempfile");
        }
        for _ in 0..2000 {
            write!(file, "Test content\n").expect("tempfile");
        }
        file.flush().expect("tempfile");
        let mut index = Index::open_in_memory().expect("db");
        {
            let mut tx = index.transaction().expect("db");
            tx.index_file(file.path()).expect("index");
            tx.commit().expect("db");
        }
        assert!(
            index.get_block(&HashDigest(*b"12345678901234567890"))
                .expect("get")
                .is_none()
        );
        let block1 = index.get_block(&HashDigest(
            *b"\xfb\x5e\xf7\xeb\xad\xd8\x2c\x80\x85\xc5\
               \xff\x63\x82\x36\x22\xba\xe0\xe2\x63\xf6"
        )).expect("get");
        assert_eq!(
            block1,
            Some((file.path().into(), 0)),
        );
        let block2 = index.get_block(&HashDigest(
            *b"\x57\x0d\x8b\x30\xfc\xfd\x58\x5e\x41\x27\
               \xb5\x61\xf5\xec\xd3\x76\xff\x4d\x01\x01"
        )).expect("get");
        assert_eq!(
            block2,
            Some((file.path().into(), 11579)),
        );
        let block3 = index.get_block(&HashDigest(
            *b"\xb9\xa8\xc2\x64\x1a\xf2\xcf\x8f\xd8\xf3\
               \x6a\x24\x56\xa3\xea\xa9\x5c\x02\x91\x27"
        )).expect("get");
        assert_eq!(
            block3,
            Some((file.path().into(), 44347)),
        );
        assert_eq!(block3.unwrap().1 - block2.unwrap().1, MAX_BLOCK_SIZE);
    }
}
<|endoftext|>