<commit_before>// UNKNOWN accessiblity
const UNKNOWN: i8 = -1;

// FOOT_FORBIDDEN that no pedestrian is allowed
const FOOT_FORBIDDEN: i8 = 0;
// FOOT_ALLOWED pedestrians are allowed in both directions
const FOOT_ALLOWED: i8 = 1;

// CAR_FORBIDDEN no car is allowed
const CAR_FORBIDDEN: i8 = 0;
// CAR_RESIDENTIAL http://wiki.openstreetmap.org/wiki/Tag:highway%3Dresidential
const CAR_RESIDENTIAL: i8 = 1;
// CAR_TERTIARY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dtertiary
const CAR_TERTIARY: i8 = 2;
// CAR_SECONDARY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dsecondary
const CAR_SECONDARY: i8 = 3;
// car_forward http://wiki.http://wiki.openstreetmap.org/wiki/Tag:highway%3Dprimary
const CAR_PRIMARY: i8 = 4;
// CAR_TRUNK http://wiki.openstreetmap.org/wiki/Tag:highway%3Dtrunk
const CAR_TRUNK: i8 = 5;
// CAR_MOTORWAY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dmotorway
const CAR_MOTORWAY: i8 = 6;

// BIKE_FORBIDDEN BIKE_ can not use this edge
const BIKE_FORBIDDEN: i8 = 0;
// BIKE_ALLOWED means that it can be used by a BIKE_, but the traffic might be shared with CAR_s
const BIKE_ALLOWED: i8 = 2;
// BIKE_LANE is a narrow lane dedicated for BIKE_, without physical separation from other traffic
const BIKE_LANE: i8 = 3;
// BIKE_BUSWAY means that BIKE_s are allowed on the bus lane
const BIKE_BUSWAY: i8 = 4;
// BIKE_TRACK is a physically separated for any other traffic
const BIKE_TRACK: i8 = 5;

// Edgeself contains what mode can use the edge in each direction
#[derive(Clone, Copy)]
pub struct EdgeProperties {
    pub foot: i8,
    pub car_forward: i8,
    pub car_backward: i8,
    pub bike_forward: i8,
    pub bike_backward: i8,
}

impl EdgeProperties {
    pub fn new() -> EdgeProperties {
        EdgeProperties {
            foot: UNKNOWN,
            car_forward: UNKNOWN,
            car_backward: UNKNOWN,
            bike_forward: UNKNOWN,
            bike_backward: UNKNOWN,
        }
    }

    // Normalize fills UNKNOWN fields
    pub fn normalize(&mut self) {
        if self.car_backward == UNKNOWN {
            self.car_backward = self.car_forward;
        }
        if self.bike_backward == UNKNOWN {
            self.bike_backward = self.bike_forward;
        }
        if self.car_forward == UNKNOWN {
            self.car_forward = CAR_FORBIDDEN;
        }
        if self.bike_forward == UNKNOWN {
            self.bike_forward = BIKE_FORBIDDEN;
        }
        if self.car_backward == UNKNOWN {
            self.car_backward = CAR_FORBIDDEN;
        }
        if self.bike_backward == UNKNOWN {
            self.bike_backward = BIKE_FORBIDDEN;
        }
        if self.foot == UNKNOWN {
            self.foot = FOOT_FORBIDDEN;
        }
    }

    // Accessible means that at least one mean of transportation can use it in one direction
    pub fn accessible(&self) -> bool {
        self.bike_forward != BIKE_FORBIDDEN || self.bike_backward != BIKE_FORBIDDEN ||
        self.car_forward != CAR_FORBIDDEN || self.car_backward != CAR_FORBIDDEN ||
        self.foot != FOOT_FORBIDDEN
    }

    pub fn update(&mut self, key_string: String, val_string: String) {
        let key = key_string.as_str();
        let val = val_string.as_str();
        self.update_with_str(key, val);
    }

    pub fn update_with_str(&mut self, key: &str, val: &str) {
        match key {
            "highway" => {
                match val {
                    "cycleway" | "path" | "foot_way" | "steps" | "pedestrian" => {
                        self.bike_forward = BIKE_TRACK;
                        self.foot = FOOT_ALLOWED;
                    }
                    "primary" | "primary_link" => {
                        self.car_forward = CAR_PRIMARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "secondary" => {
                        self.car_forward = CAR_SECONDARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "tertiary" => {
                        self.car_forward = CAR_TERTIARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "unclassified" | "residential" | "living_street" | "road" | "service" |
                    "track" => {
                        self.car_forward = CAR_RESIDENTIAL;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "motorway" | "motorway_link" => {
                        self.car_forward = CAR_MOTORWAY;
                        self.foot = CAR_FORBIDDEN;
                        self.bike_forward = BIKE_FORBIDDEN;
                    }
                    "trunk" | "trunk_link" => {
                        self.car_forward = CAR_TRUNK;
                        self.foot = CAR_FORBIDDEN;
                        self.bike_forward = BIKE_FORBIDDEN;
                    }
                    _ => {}
                }
            }
            "pedestrian" | "foot" => {
                match val {
                    "no" => self.foot = FOOT_FORBIDDEN,
                    _ => self.foot = FOOT_ALLOWED,
                }
            }

            // http://wiki.openstreetmap.org/wiki/Cycleway
            // http://wiki.openstreetmap.org/wiki/Map_Features#Cycleway
            "cycleway" => {
                match val {
                    "track" => self.bike_forward = BIKE_TRACK,
                    "opposite_track" => self.bike_backward = BIKE_TRACK,
                    "opposite" => self.bike_backward = BIKE_ALLOWED,
                    "share_busway" => self.bike_forward = BIKE_BUSWAY,
                    "lane_left" | "opposite_lane" => self.bike_backward = BIKE_LANE,
                    _ => self.bike_forward = BIKE_LANE,
                }
            }

            "bicycle" => {
                match val {
                    "no" | "false" => self.bike_forward = BIKE_FORBIDDEN,
                    _ => self.bike_forward = BIKE_ALLOWED,
                }
            }
            "busway" => {
                match val {
                    "opposite_lane" | "opposite_track" => self.bike_backward = BIKE_BUSWAY,
                    _ => self.bike_forward = BIKE_BUSWAY,
                }
            }
            "oneway" => {
                match val {
                    "yes" | "true" | "1" => {
                        self.car_backward = CAR_FORBIDDEN;
                        if self.bike_backward == UNKNOWN {
                            self.bike_backward = BIKE_FORBIDDEN;
                        }
                    }
                    _ => {}
                }
            }
            "junction" => {
                match val {
                    "roundabout" => {
                        self.car_backward = CAR_FORBIDDEN;
                        if self.bike_backward == UNKNOWN {
                            self.bike_backward = BIKE_FORBIDDEN;
                        }
                    }
                    _ => {}
                }
            }
            _ => {}
        }
    }
}

#[test]
fn test_accessible() {
    let mut p = EdgeProperties::new();
    p.normalize();
    assert!(!p.accessible());

    p.foot = FOOT_ALLOWED;
    assert!(p.accessible())
}

#[test]
fn test_normalize() {
    let mut p = EdgeProperties::new();
    p.bike_forward = BIKE_LANE;
    p.normalize();
    assert_eq!(BIKE_LANE, p.bike_backward);
    p.bike_forward = BIKE_ALLOWED;
    p.normalize();
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.car_forward = CAR_SECONDARY;
    p.car_backward = UNKNOWN;
    p.normalize();
    assert_eq!(CAR_SECONDARY, p.car_backward)
}

#[test]
fn test_update() {
    let mut p = EdgeProperties::new();
    p.update_with_str("highway", "secondary");
    assert_eq!(CAR_SECONDARY, p.car_forward);

    p.update_with_str("highway", "primary_link");
    assert_eq!(CAR_PRIMARY, p.car_forward);

    p.update_with_str("highway", "motorway");
    assert_eq!(CAR_MOTORWAY, p.car_forward);

    p.update_with_str("highway", "residential");
    assert_eq!(CAR_RESIDENTIAL, p.car_forward);

    p.update_with_str("highway", "tertiary");
    assert_eq!(CAR_TERTIARY, p.car_forward);

    p.update_with_str("highway", "trunk");
    assert_eq!(CAR_TRUNK, p.car_forward);

    p.update_with_str("highway", "cycleway");
    assert_eq!(BIKE_TRACK, p.bike_forward);
    assert_eq!(FOOT_ALLOWED, p.foot);

    p.update_with_str("foot", "designated");
    assert_eq!(FOOT_ALLOWED, p.foot);

    p.update_with_str("foot", "no");
    assert_eq!(FOOT_FORBIDDEN, p.foot);

    p.update_with_str("cycleway", "lane");
    assert_eq!(BIKE_LANE, p.bike_forward);

    p.update_with_str("cycleway", "track");
    assert_eq!(BIKE_TRACK, p.bike_forward);

    p.update_with_str("cycleway", "opposite_lane");
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.update_with_str("cycleway", "opposite_track");
    assert_eq!(BIKE_TRACK, p.bike_backward);

    p.update_with_str("cycleway", "opposite");
    assert_eq!(BIKE_ALLOWED, p.bike_backward);

    p.update_with_str("cycleway", "share_busway");
    assert_eq!(BIKE_BUSWAY, p.bike_forward);

    p.update_with_str("cycleway", "lane_left");
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.update_with_str("bicycle", "yes");
    assert_eq!(BIKE_ALLOWED, p.bike_forward);

    p.update_with_str("bicycle", "no");
    assert_eq!(BIKE_FORBIDDEN, p.bike_forward);

    p.update_with_str("busway", "yes");
    assert_eq!(BIKE_BUSWAY, p.bike_forward);

    p.update_with_str("busway", "opposite_track");
    assert_eq!(BIKE_BUSWAY, p.bike_backward);

    p.update_with_str("oneway", "yes");
    assert_eq!(BIKE_FORBIDDEN, p.car_backward);
    assert!(p.bike_backward != BIKE_FORBIDDEN);

    p.bike_backward = UNKNOWN;
    p.update_with_str("oneway", "yes");
    assert_eq!(BIKE_FORBIDDEN, p.bike_backward);

    p.update_with_str("junction", "roundabout");
    assert_eq!(BIKE_FORBIDDEN, p.car_backward);

    p.bike_backward = UNKNOWN;
    p.update_with_str("junction", "roundabout");
    assert_eq!(BIKE_FORBIDDEN, p.bike_backward);
}
<commit_msg>Change misspelled tag: 'foot_way'->'footway'<commit_after>// UNKNOWN accessiblity
const UNKNOWN: i8 = -1;

// FOOT_FORBIDDEN that no pedestrian is allowed
const FOOT_FORBIDDEN: i8 = 0;
// FOOT_ALLOWED pedestrians are allowed in both directions
const FOOT_ALLOWED: i8 = 1;

// CAR_FORBIDDEN no car is allowed
const CAR_FORBIDDEN: i8 = 0;
// CAR_RESIDENTIAL http://wiki.openstreetmap.org/wiki/Tag:highway%3Dresidential
const CAR_RESIDENTIAL: i8 = 1;
// CAR_TERTIARY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dtertiary
const CAR_TERTIARY: i8 = 2;
// CAR_SECONDARY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dsecondary
const CAR_SECONDARY: i8 = 3;
// car_forward http://wiki.http://wiki.openstreetmap.org/wiki/Tag:highway%3Dprimary
const CAR_PRIMARY: i8 = 4;
// CAR_TRUNK http://wiki.openstreetmap.org/wiki/Tag:highway%3Dtrunk
const CAR_TRUNK: i8 = 5;
// CAR_MOTORWAY http://wiki.openstreetmap.org/wiki/Tag:highway%3Dmotorway
const CAR_MOTORWAY: i8 = 6;

// BIKE_FORBIDDEN BIKE_ can not use this edge
const BIKE_FORBIDDEN: i8 = 0;
// BIKE_ALLOWED means that it can be used by a BIKE_, but the traffic might be shared with CAR_s
const BIKE_ALLOWED: i8 = 2;
// BIKE_LANE is a narrow lane dedicated for BIKE_, without physical separation from other traffic
const BIKE_LANE: i8 = 3;
// BIKE_BUSWAY means that BIKE_s are allowed on the bus lane
const BIKE_BUSWAY: i8 = 4;
// BIKE_TRACK is a physically separated for any other traffic
const BIKE_TRACK: i8 = 5;

// Edgeself contains what mode can use the edge in each direction
#[derive(Clone, Copy)]
pub struct EdgeProperties {
    pub foot: i8,
    pub car_forward: i8,
    pub car_backward: i8,
    pub bike_forward: i8,
    pub bike_backward: i8,
}

impl EdgeProperties {
    pub fn new() -> EdgeProperties {
        EdgeProperties {
            foot: UNKNOWN,
            car_forward: UNKNOWN,
            car_backward: UNKNOWN,
            bike_forward: UNKNOWN,
            bike_backward: UNKNOWN,
        }
    }

    // Normalize fills UNKNOWN fields
    pub fn normalize(&mut self) {
        if self.car_backward == UNKNOWN {
            self.car_backward = self.car_forward;
        }
        if self.bike_backward == UNKNOWN {
            self.bike_backward = self.bike_forward;
        }
        if self.car_forward == UNKNOWN {
            self.car_forward = CAR_FORBIDDEN;
        }
        if self.bike_forward == UNKNOWN {
            self.bike_forward = BIKE_FORBIDDEN;
        }
        if self.car_backward == UNKNOWN {
            self.car_backward = CAR_FORBIDDEN;
        }
        if self.bike_backward == UNKNOWN {
            self.bike_backward = BIKE_FORBIDDEN;
        }
        if self.foot == UNKNOWN {
            self.foot = FOOT_FORBIDDEN;
        }
    }

    // Accessible means that at least one mean of transportation can use it in one direction
    pub fn accessible(&self) -> bool {
        self.bike_forward != BIKE_FORBIDDEN || self.bike_backward != BIKE_FORBIDDEN ||
        self.car_forward != CAR_FORBIDDEN || self.car_backward != CAR_FORBIDDEN ||
        self.foot != FOOT_FORBIDDEN
    }

    pub fn update(&mut self, key_string: String, val_string: String) {
        let key = key_string.as_str();
        let val = val_string.as_str();
        self.update_with_str(key, val);
    }

    pub fn update_with_str(&mut self, key: &str, val: &str) {
        match key {
            "highway" => {
                match val {
                    "cycleway" | "path" | "footway" | "steps" | "pedestrian" => {
                        self.bike_forward = BIKE_TRACK;
                        self.foot = FOOT_ALLOWED;
                    }
                    "primary" | "primary_link" => {
                        self.car_forward = CAR_PRIMARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "secondary" => {
                        self.car_forward = CAR_SECONDARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "tertiary" => {
                        self.car_forward = CAR_TERTIARY;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "unclassified" | "residential" | "living_street" | "road" | "service" |
                    "track" => {
                        self.car_forward = CAR_RESIDENTIAL;
                        self.foot = FOOT_ALLOWED;
                        self.bike_forward = BIKE_ALLOWED;
                    }
                    "motorway" | "motorway_link" => {
                        self.car_forward = CAR_MOTORWAY;
                        self.foot = CAR_FORBIDDEN;
                        self.bike_forward = BIKE_FORBIDDEN;
                    }
                    "trunk" | "trunk_link" => {
                        self.car_forward = CAR_TRUNK;
                        self.foot = CAR_FORBIDDEN;
                        self.bike_forward = BIKE_FORBIDDEN;
                    }
                    _ => {}
                }
            }
            "pedestrian" | "foot" => {
                match val {
                    "no" => self.foot = FOOT_FORBIDDEN,
                    _ => self.foot = FOOT_ALLOWED,
                }
            }

            // http://wiki.openstreetmap.org/wiki/Cycleway
            // http://wiki.openstreetmap.org/wiki/Map_Features#Cycleway
            "cycleway" => {
                match val {
                    "track" => self.bike_forward = BIKE_TRACK,
                    "opposite_track" => self.bike_backward = BIKE_TRACK,
                    "opposite" => self.bike_backward = BIKE_ALLOWED,
                    "share_busway" => self.bike_forward = BIKE_BUSWAY,
                    "lane_left" | "opposite_lane" => self.bike_backward = BIKE_LANE,
                    _ => self.bike_forward = BIKE_LANE,
                }
            }

            "bicycle" => {
                match val {
                    "no" | "false" => self.bike_forward = BIKE_FORBIDDEN,
                    _ => self.bike_forward = BIKE_ALLOWED,
                }
            }
            "busway" => {
                match val {
                    "opposite_lane" | "opposite_track" => self.bike_backward = BIKE_BUSWAY,
                    _ => self.bike_forward = BIKE_BUSWAY,
                }
            }
            "oneway" => {
                match val {
                    "yes" | "true" | "1" => {
                        self.car_backward = CAR_FORBIDDEN;
                        if self.bike_backward == UNKNOWN {
                            self.bike_backward = BIKE_FORBIDDEN;
                        }
                    }
                    _ => {}
                }
            }
            "junction" => {
                match val {
                    "roundabout" => {
                        self.car_backward = CAR_FORBIDDEN;
                        if self.bike_backward == UNKNOWN {
                            self.bike_backward = BIKE_FORBIDDEN;
                        }
                    }
                    _ => {}
                }
            }
            _ => {}
        }
    }
}

#[test]
fn test_accessible() {
    let mut p = EdgeProperties::new();
    p.normalize();
    assert!(!p.accessible());

    p.foot = FOOT_ALLOWED;
    assert!(p.accessible())
}

#[test]
fn test_normalize() {
    let mut p = EdgeProperties::new();
    p.bike_forward = BIKE_LANE;
    p.normalize();
    assert_eq!(BIKE_LANE, p.bike_backward);
    p.bike_forward = BIKE_ALLOWED;
    p.normalize();
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.car_forward = CAR_SECONDARY;
    p.car_backward = UNKNOWN;
    p.normalize();
    assert_eq!(CAR_SECONDARY, p.car_backward)
}

#[test]
fn test_update() {
    let mut p = EdgeProperties::new();
    p.update_with_str("highway", "secondary");
    assert_eq!(CAR_SECONDARY, p.car_forward);

    p.update_with_str("highway", "primary_link");
    assert_eq!(CAR_PRIMARY, p.car_forward);

    p.update_with_str("highway", "motorway");
    assert_eq!(CAR_MOTORWAY, p.car_forward);

    p.update_with_str("highway", "residential");
    assert_eq!(CAR_RESIDENTIAL, p.car_forward);

    p.update_with_str("highway", "tertiary");
    assert_eq!(CAR_TERTIARY, p.car_forward);

    p.update_with_str("highway", "trunk");
    assert_eq!(CAR_TRUNK, p.car_forward);

    p.update_with_str("highway", "cycleway");
    assert_eq!(BIKE_TRACK, p.bike_forward);
    assert_eq!(FOOT_ALLOWED, p.foot);

    p.update_with_str("foot", "designated");
    assert_eq!(FOOT_ALLOWED, p.foot);

    p.update_with_str("foot", "no");
    assert_eq!(FOOT_FORBIDDEN, p.foot);

    p.update_with_str("cycleway", "lane");
    assert_eq!(BIKE_LANE, p.bike_forward);

    p.update_with_str("cycleway", "track");
    assert_eq!(BIKE_TRACK, p.bike_forward);

    p.update_with_str("cycleway", "opposite_lane");
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.update_with_str("cycleway", "opposite_track");
    assert_eq!(BIKE_TRACK, p.bike_backward);

    p.update_with_str("cycleway", "opposite");
    assert_eq!(BIKE_ALLOWED, p.bike_backward);

    p.update_with_str("cycleway", "share_busway");
    assert_eq!(BIKE_BUSWAY, p.bike_forward);

    p.update_with_str("cycleway", "lane_left");
    assert_eq!(BIKE_LANE, p.bike_backward);

    p.update_with_str("bicycle", "yes");
    assert_eq!(BIKE_ALLOWED, p.bike_forward);

    p.update_with_str("bicycle", "no");
    assert_eq!(BIKE_FORBIDDEN, p.bike_forward);

    p.update_with_str("busway", "yes");
    assert_eq!(BIKE_BUSWAY, p.bike_forward);

    p.update_with_str("busway", "opposite_track");
    assert_eq!(BIKE_BUSWAY, p.bike_backward);

    p.update_with_str("oneway", "yes");
    assert_eq!(BIKE_FORBIDDEN, p.car_backward);
    assert!(p.bike_backward != BIKE_FORBIDDEN);

    p.bike_backward = UNKNOWN;
    p.update_with_str("oneway", "yes");
    assert_eq!(BIKE_FORBIDDEN, p.bike_backward);

    p.update_with_str("junction", "roundabout");
    assert_eq!(BIKE_FORBIDDEN, p.car_backward);

    p.bike_backward = UNKNOWN;
    p.update_with_str("junction", "roundabout");
    assert_eq!(BIKE_FORBIDDEN, p.bike_backward);
}
<|endoftext|>


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

use std::env;
use std::ffi::OsString;

use common::ui::UI;
use hcore::crypto::CACHE_KEY_PATH_ENV_VAR;
use hcore::env as henv;
use hcore::fs::CACHE_KEY_PATH;
use hcore::os::users;

use config;
use error::Result;

const STUDIO_CMD: &'static str = "hab-studio";
const STUDIO_CMD_ENVVAR: &'static str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &'static str = "core/hab-studio";

pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
    try!(inner::rerun_with_sudo_if_needed(ui));

    // If the `$HAB_ORIGIN` environment variable is not present, then see if a default is set in
    // the CLI config. If so, set it as the `$HAB_ORIGIN` environment variable for the `hab-studio`
    // or `docker` execv call.
    if henv::var("HAB_ORIGIN").is_err() {
        let config = try!(config::load_with_sudo_user());
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }

    // If the `$HAB_CACHE_KEY_PATH` environment variable is not present, check if we are running
    // under a `sudo` invocation. If so, determine the non-root user that issued the command in
    // order to set their key cache location in the environment variable. This is done so that the
    // `hab-studio` command will find the correct key cache or so that the correct directory will
    // be volume mounted when used with Docker.
    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        if let Some(sudo_user) = henv::sudo_user() {
            if let Some(home) = users::get_home_for_user(&sudo_user) {
                let cache_key_path = home.join(format!(".{}", CACHE_KEY_PATH));
                debug!("Setting cache_key_path for SUDO_USER={} to: {}",
                       &sudo_user,
                       cache_key_path.display());
                env::set_var(CACHE_KEY_PATH_ENV_VAR, cache_key_path);
                // Prevent any inner `hab` invocations from triggering similar logic: we will be
                // operating in the context `hab-studio` which is running with root like privileges.
                env::remove_var("SUDO_USER");
            }
        }
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use std::env;
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::str::FromStr;

    use common::ui::UI;
    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::{am_i_root, find_command};
    use hcore::os::process;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;
    use VERSION;

    const SUDO_CMD: &'static str = "sudo";

    pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split("/").collect();
                let ident = try!(PackageIdent::from_str(&format!("{}/{}",
                                                                 super::STUDIO_PACKAGE_IDENT,
                                                                 version[0])));
                try!(exec::command_from_min_pkg(ui,
                                                super::STUDIO_CMD,
                                                &ident,
                                                &default_cache_key_path(None),
                                                0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(process::become_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }

    pub fn rerun_with_sudo_if_needed(ui: &mut UI) -> Result<()> {
        // If I have root permissions, early return, we are done.
        if am_i_root() {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into()];
                args.append(&mut env::args_os().collect());
                Ok(try!(process::become_command(sudo_prog, args)))
            }
            None => {
                try!(ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                     SUDO_CMD)));
                try!(ui.warn("Running Habitat Studio requires root or administrator privileges. \
                              Please retry this command as a super user or use a \
                              privilege-granting facility such as sudo."));
                try!(ui.br());
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use std::env;
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::str::FromStr;

    use common::ui::UI;
    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::{CACHE_KEY_PATH, find_command};
    use hcore::os::process;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;
    use VERSION;

    const DOCKER_CMD: &'static str = "docker";
    const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
    const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
    const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
    const DOCKER_OPTS: &'static str = "HAB_DOCKER_OPTS";

    const HAB_WINDOWS_STUDIO: &'static str = "HAB_WINDOWS_STUDIO";

    pub fn start(_ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        //TODO: Remove HAB_WINDOWS_STUDIO check after windows studio support is official
        if cfg!(target_os = "windows") && henv::var(HAB_WINDOWS_STUDIO).is_ok() {
            start_windows_studio(_ui, args)
        } else {
            start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split("/").collect();
                let ident = try!(PackageIdent::from_str(&format!("{}/{}",
                                                                 super::STUDIO_PACKAGE_IDENT,
                                                                 version[0])));
                try!(exec::command_from_min_pkg(ui,
                                                super::STUDIO_CMD,
                                                &ident,
                                                &default_cache_key_path(None),
                                                0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(process::become_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }

    pub fn start_docker_studio(_ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let docker = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

        let cmd = match find_command(&docker) {
            Some(cmd) => cmd,
            None => return Err(Error::ExecCommandNotFound(docker.to_string())),
        };

        let output = Command::new(&cmd)
            .arg("images")
            .arg(&image_identifier())
            .arg("-q")
            .output()
            .expect("docker failed to start");

        let stdout = String::from_utf8(output.stdout).unwrap();
        if stdout.is_empty() {
            debug!("Failed to find studio image locally.");

            let child = Command::new(&cmd)
                .arg("pull")
                .arg(&image_identifier())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit())
                .spawn()
                .expect("docker failed to start");

            let output = child
                .wait_with_output()
                .expect("failed to wait on child");

            if output.status.success() {
                debug!("Docker image is reachable. Proceeding with launching docker.");
            } else {
                debug!("Docker image is unreachable. Exit code = {:?}",
                       output.status);

                let err_output = String::from_utf8(output.stderr).unwrap();

                if err_output.contains("image") && err_output.contains("not found") {
                    return Err(Error::DockerImageNotFound(image_identifier().to_string()));
                } else if err_output.contains("Cannot connect to the Docker daemon") {
                    return Err(Error::DockerDaemonDown);
                } else {
                    return Err(Error::DockerNetworkDown(image_identifier().to_string()));
                }
            }
        } else {
            debug!("Found studio image locally.");
            debug!("Local image id: {}", stdout);
        }

        // We need to ensure that filesystem sharing has been enabled, otherwise the user will
        // be greeted with a horrible error message that's difficult to make sense of. To
        // mitigate this, we check the studio version. This will cause Docker to go through the
        // mounting steps, so we can watch stderr for failure, but has the advantage of not
        // requiring a TTY.

        let current_dir = format!("{}:/src", env::current_dir().unwrap().to_string_lossy());
        let key_cache_path = format!("{}:/{}",
                                     default_cache_key_path(None).to_string_lossy(),
                                     CACHE_KEY_PATH);
        let version_output = Command::new(&cmd)
            .arg("run")
            .arg("--rm")
            .arg("--privileged")
            .arg("--volume")
            .arg(&key_cache_path)
            .arg("--volume")
            .arg(&current_dir)
            .arg(image_identifier())
            .arg("-V")
            .output()
            .expect("docker failed to start");

        let stderr = String::from_utf8(version_output.stderr).unwrap();
        if !stderr.is_empty() &&
           (stderr.as_str().contains("Mounts denied") ||
            stderr.as_str().contains("drive is not shared")) {
            return Err(Error::DockerFileSharingNotEnabled);
        }

        // If we make it here, filesystem sharing has been setup correctly, so let's proceed like
        // normal.

        let mut cmd_args: Vec<OsString> = vec!["run".into(),
                                               "--rm".into(),
                                               "--tty".into(),
                                               "--interactive".into(),
                                               "--privileged".into()];

        // All the args already placed in `cmd_args` are things that we don't want to insert again.
        // Later args such as `--env` will overwrite any options (potentially) set mistakenly here.
        if let Ok(opts) = henv::var(DOCKER_OPTS) {
            let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
            debug!("Docker opts originating from DOCKER_OPTS = {:?}", opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }

        let env_vars = vec!["HAB_DEPOT_URL",
                            "HAB_ORIGIN",
                            "HAB_STUDIO_SUP",
                            "http_proxy",
                            "https_proxy"];
        for var in env_vars {
            if let Ok(val) = henv::var(var) {
                debug!("Propagating environment variable into container: {}={}",
                       var,
                       val);
                cmd_args.push("--env".into());
                cmd_args.push(format!("{}={}", var, val).into());
            }
        }

        cmd_args.push("--volume".into());
        cmd_args.push("/var/run/docker.sock:/var/run/docker.sock".into());
        cmd_args.push("--volume".into());
        cmd_args.push(key_cache_path.into());
        cmd_args.push("--volume".into());
        cmd_args.push(current_dir.into());

        cmd_args.push(image_identifier().into());
        cmd_args.extend_from_slice(args.as_slice());

        for var in vec!["http_proxy", "https_proxy"] {
            if let Ok(_) = henv::var(var) {
                debug!("Unsetting proxy environment variable '{}' before calling `{}'",
                       var,
                       docker);
                env::remove_var(var);
            }
        }

        debug!("Docker arguments = {:?}", cmd_args);
        Ok(try!(process::become_command(cmd, cmd_args)))
    }

    pub fn rerun_with_sudo_if_needed(_ui: &mut UI) -> Result<()> {
        // No sudo calls necessary here--we are calling `docker` commands instead
        Ok(())
    }

    /// Returns the Docker Studio image with tag for the desired version which corresponds to the
    /// same version (minus release) as this program.
    fn image_identifier() -> String {
        let version: Vec<&str> = VERSION.split("/").collect();
        henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", DOCKER_IMAGE, version[0]))
    }

    #[cfg(test)]
    mod tests {
        use super::{image_identifier, DOCKER_IMAGE};
        use VERSION;

        #[test]
        fn retrieve_image_identifier() {
            assert_eq!(image_identifier(), format!("{}:{}", DOCKER_IMAGE, VERSION));
        }
    }
}
<commit_msg>Add HAB_UPDATE_STRATEGY_FREQUENCY_MS to studio propagated vars<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::OsString;

use common::ui::UI;
use hcore::crypto::CACHE_KEY_PATH_ENV_VAR;
use hcore::env as henv;
use hcore::fs::CACHE_KEY_PATH;
use hcore::os::users;

use config;
use error::Result;

const STUDIO_CMD: &'static str = "hab-studio";
const STUDIO_CMD_ENVVAR: &'static str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &'static str = "core/hab-studio";

pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
    try!(inner::rerun_with_sudo_if_needed(ui));

    // If the `$HAB_ORIGIN` environment variable is not present, then see if a default is set in
    // the CLI config. If so, set it as the `$HAB_ORIGIN` environment variable for the `hab-studio`
    // or `docker` execv call.
    if henv::var("HAB_ORIGIN").is_err() {
        let config = try!(config::load_with_sudo_user());
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }

    // If the `$HAB_CACHE_KEY_PATH` environment variable is not present, check if we are running
    // under a `sudo` invocation. If so, determine the non-root user that issued the command in
    // order to set their key cache location in the environment variable. This is done so that the
    // `hab-studio` command will find the correct key cache or so that the correct directory will
    // be volume mounted when used with Docker.
    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        if let Some(sudo_user) = henv::sudo_user() {
            if let Some(home) = users::get_home_for_user(&sudo_user) {
                let cache_key_path = home.join(format!(".{}", CACHE_KEY_PATH));
                debug!("Setting cache_key_path for SUDO_USER={} to: {}",
                       &sudo_user,
                       cache_key_path.display());
                env::set_var(CACHE_KEY_PATH_ENV_VAR, cache_key_path);
                // Prevent any inner `hab` invocations from triggering similar logic: we will be
                // operating in the context `hab-studio` which is running with root like privileges.
                env::remove_var("SUDO_USER");
            }
        }
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use std::env;
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::str::FromStr;

    use common::ui::UI;
    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::{am_i_root, find_command};
    use hcore::os::process;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;
    use VERSION;

    const SUDO_CMD: &'static str = "sudo";

    pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split("/").collect();
                let ident = try!(PackageIdent::from_str(&format!("{}/{}",
                                                                 super::STUDIO_PACKAGE_IDENT,
                                                                 version[0])));
                try!(exec::command_from_min_pkg(ui,
                                                super::STUDIO_CMD,
                                                &ident,
                                                &default_cache_key_path(None),
                                                0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(process::become_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }

    pub fn rerun_with_sudo_if_needed(ui: &mut UI) -> Result<()> {
        // If I have root permissions, early return, we are done.
        if am_i_root() {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into()];
                args.append(&mut env::args_os().collect());
                Ok(try!(process::become_command(sudo_prog, args)))
            }
            None => {
                try!(ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                     SUDO_CMD)));
                try!(ui.warn("Running Habitat Studio requires root or administrator privileges. \
                              Please retry this command as a super user or use a \
                              privilege-granting facility such as sudo."));
                try!(ui.br());
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use std::env;
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::str::FromStr;

    use common::ui::UI;
    use hcore::crypto::{init, default_cache_key_path};
    use hcore::env as henv;
    use hcore::fs::{CACHE_KEY_PATH, find_command};
    use hcore::os::process;
    use hcore::package::PackageIdent;

    use error::{Error, Result};
    use exec;
    use VERSION;

    const DOCKER_CMD: &'static str = "docker";
    const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
    const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
    const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
    const DOCKER_OPTS: &'static str = "HAB_DOCKER_OPTS";

    const HAB_WINDOWS_STUDIO: &'static str = "HAB_WINDOWS_STUDIO";

    pub fn start(_ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        //TODO: Remove HAB_WINDOWS_STUDIO check after windows studio support is official
        if cfg!(target_os = "windows") && henv::var(HAB_WINDOWS_STUDIO).is_ok() {
            start_windows_studio(_ui, args)
        } else {
            start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split("/").collect();
                let ident = try!(PackageIdent::from_str(&format!("{}/{}",
                                                                 super::STUDIO_PACKAGE_IDENT,
                                                                 version[0])));
                try!(exec::command_from_min_pkg(ui,
                                                super::STUDIO_CMD,
                                                &ident,
                                                &default_cache_key_path(None),
                                                0))
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            try!(process::become_command(cmd, args));
        } else {
            return Err(Error::ExecCommandNotFound(command.to_string_lossy().into_owned()));
        }
        Ok(())
    }

    pub fn start_docker_studio(_ui: &mut UI, args: Vec<OsString>) -> Result<()> {
        let docker = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

        let cmd = match find_command(&docker) {
            Some(cmd) => cmd,
            None => return Err(Error::ExecCommandNotFound(docker.to_string())),
        };

        let output = Command::new(&cmd)
            .arg("images")
            .arg(&image_identifier())
            .arg("-q")
            .output()
            .expect("docker failed to start");

        let stdout = String::from_utf8(output.stdout).unwrap();
        if stdout.is_empty() {
            debug!("Failed to find studio image locally.");

            let child = Command::new(&cmd)
                .arg("pull")
                .arg(&image_identifier())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit())
                .spawn()
                .expect("docker failed to start");

            let output = child
                .wait_with_output()
                .expect("failed to wait on child");

            if output.status.success() {
                debug!("Docker image is reachable. Proceeding with launching docker.");
            } else {
                debug!("Docker image is unreachable. Exit code = {:?}",
                       output.status);

                let err_output = String::from_utf8(output.stderr).unwrap();

                if err_output.contains("image") && err_output.contains("not found") {
                    return Err(Error::DockerImageNotFound(image_identifier().to_string()));
                } else if err_output.contains("Cannot connect to the Docker daemon") {
                    return Err(Error::DockerDaemonDown);
                } else {
                    return Err(Error::DockerNetworkDown(image_identifier().to_string()));
                }
            }
        } else {
            debug!("Found studio image locally.");
            debug!("Local image id: {}", stdout);
        }

        // We need to ensure that filesystem sharing has been enabled, otherwise the user will
        // be greeted with a horrible error message that's difficult to make sense of. To
        // mitigate this, we check the studio version. This will cause Docker to go through the
        // mounting steps, so we can watch stderr for failure, but has the advantage of not
        // requiring a TTY.

        let current_dir = format!("{}:/src", env::current_dir().unwrap().to_string_lossy());
        let key_cache_path = format!("{}:/{}",
                                     default_cache_key_path(None).to_string_lossy(),
                                     CACHE_KEY_PATH);
        let version_output = Command::new(&cmd)
            .arg("run")
            .arg("--rm")
            .arg("--privileged")
            .arg("--volume")
            .arg(&key_cache_path)
            .arg("--volume")
            .arg(&current_dir)
            .arg(image_identifier())
            .arg("-V")
            .output()
            .expect("docker failed to start");

        let stderr = String::from_utf8(version_output.stderr).unwrap();
        if !stderr.is_empty() &&
           (stderr.as_str().contains("Mounts denied") ||
            stderr.as_str().contains("drive is not shared")) {
            return Err(Error::DockerFileSharingNotEnabled);
        }

        // If we make it here, filesystem sharing has been setup correctly, so let's proceed like
        // normal.

        let mut cmd_args: Vec<OsString> = vec!["run".into(),
                                               "--rm".into(),
                                               "--tty".into(),
                                               "--interactive".into(),
                                               "--privileged".into()];

        // All the args already placed in `cmd_args` are things that we don't want to insert again.
        // Later args such as `--env` will overwrite any options (potentially) set mistakenly here.
        if let Ok(opts) = henv::var(DOCKER_OPTS) {
            let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
            debug!("Docker opts originating from DOCKER_OPTS = {:?}", opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }

        let env_vars = vec!["HAB_DEPOT_URL",
                            "HAB_ORIGIN",
                            "HAB_STUDIO_SUP",
                            "HAB_UPDATE_STRATEGY_FREQUENCY_MS",
                            "http_proxy",
                            "https_proxy"];
        for var in env_vars {
            if let Ok(val) = henv::var(var) {
                debug!("Propagating environment variable into container: {}={}",
                       var,
                       val);
                cmd_args.push("--env".into());
                cmd_args.push(format!("{}={}", var, val).into());
            }
        }

        cmd_args.push("--volume".into());
        cmd_args.push("/var/run/docker.sock:/var/run/docker.sock".into());
        cmd_args.push("--volume".into());
        cmd_args.push(key_cache_path.into());
        cmd_args.push("--volume".into());
        cmd_args.push(current_dir.into());

        cmd_args.push(image_identifier().into());
        cmd_args.extend_from_slice(args.as_slice());

        for var in vec!["http_proxy", "https_proxy"] {
            if let Ok(_) = henv::var(var) {
                debug!("Unsetting proxy environment variable '{}' before calling `{}'",
                       var,
                       docker);
                env::remove_var(var);
            }
        }

        debug!("Docker arguments = {:?}", cmd_args);
        Ok(try!(process::become_command(cmd, cmd_args)))
    }

    pub fn rerun_with_sudo_if_needed(_ui: &mut UI) -> Result<()> {
        // No sudo calls necessary here--we are calling `docker` commands instead
        Ok(())
    }

    /// Returns the Docker Studio image with tag for the desired version which corresponds to the
    /// same version (minus release) as this program.
    fn image_identifier() -> String {
        let version: Vec<&str> = VERSION.split("/").collect();
        henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", DOCKER_IMAGE, version[0]))
    }

    #[cfg(test)]
    mod tests {
        use super::{image_identifier, DOCKER_IMAGE};
        use VERSION;

        #[test]
        fn retrieve_image_identifier() {
            assert_eq!(image_identifier(), format!("{}:{}", DOCKER_IMAGE, VERSION));
        }
    }
}
<|endoftext|>


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

use bluetooth_traits::{BluetoothCharacteristicMsg, BluetoothDescriptorMsg, BluetoothServiceMsg};
use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::BluetoothDeviceBinding;
use dom::bindings::codegen::Bindings::BluetoothDeviceBinding::BluetoothDeviceMethods;
use dom::bindings::codegen::Bindings::BluetoothRemoteGATTServerBinding::BluetoothRemoteGATTServerMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::js::{JS, Root, MutHeap, MutNullableHeap};
use dom::bindings::reflector::{Reflectable, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::bluetooth::Bluetooth;
use dom::bluetoothadvertisingdata::BluetoothAdvertisingData;
use dom::bluetoothcharacteristicproperties::BluetoothCharacteristicProperties;
use dom::bluetoothremotegattcharacteristic::BluetoothRemoteGATTCharacteristic;
use dom::bluetoothremotegattdescriptor::BluetoothRemoteGATTDescriptor;
use dom::bluetoothremotegattserver::BluetoothRemoteGATTServer;
use dom::bluetoothremotegattservice::BluetoothRemoteGATTService;
use dom::eventtarget::EventTarget;
use dom::globalscope::GlobalScope;
use std::collections::HashMap;

// https://webbluetoothcg.github.io/web-bluetooth/#bluetoothdevice
#[dom_struct]
pub struct BluetoothDevice {
    eventtarget: EventTarget,
    id: DOMString,
    name: Option<DOMString>,
    ad_data: MutHeap<JS<BluetoothAdvertisingData>>,
    gatt: MutNullableHeap<JS<BluetoothRemoteGATTServer>>,
    context: MutHeap<JS<Bluetooth>>,
    attribute_instance_map: (DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTService>>>>,
                             DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTCharacteristic>>>>,
                             DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTDescriptor>>>>),
}

impl BluetoothDevice {
    pub fn new_inherited(id: DOMString,
                         name: Option<DOMString>,
                         ad_data: &BluetoothAdvertisingData,
                         context: &Bluetooth)
                         -> BluetoothDevice {
        BluetoothDevice {
            eventtarget: EventTarget::new_inherited(),
            id: id,
            name: name,
            ad_data: MutHeap::new(ad_data),
            gatt: Default::default(),
            context: MutHeap::new(context),
            attribute_instance_map: (DOMRefCell::new(HashMap::new()),
                                     DOMRefCell::new(HashMap::new()),
                                     DOMRefCell::new(HashMap::new())),
        }
    }

    pub fn new(global: &GlobalScope,
               id: DOMString,
               name: Option<DOMString>,
               adData: &BluetoothAdvertisingData,
               context: &Bluetooth)
               -> Root<BluetoothDevice> {
        reflect_dom_object(box BluetoothDevice::new_inherited(id,
                                                              name,
                                                              adData,
                                                              context),
                           global,
                           BluetoothDeviceBinding::Wrap)
    }

    pub fn get_context(&self) -> Root<Bluetooth> {
        self.context.get()
    }

    pub fn get_or_create_service(&self,
                                 service: &BluetoothServiceMsg,
                                 server: &BluetoothRemoteGATTServer)
                                 -> Root<BluetoothRemoteGATTService> {
        let (ref service_map_ref, _, _) = self.attribute_instance_map;
        let mut service_map = service_map_ref.borrow_mut();
        if let Some(existing_service) = service_map.get(&service.instance_id) {
            return existing_service.get();
        }
        let bt_service = BluetoothRemoteGATTService::new(&server.global(),
                                                         &server.Device(),
                                                         DOMString::from(service.uuid.clone()),
                                                         service.is_primary,
                                                         service.instance_id.clone());
        service_map.insert(service.instance_id.clone(), MutHeap::new(&bt_service));
        return bt_service;
    }

    pub fn get_or_create_characteristic(&self,
                                        characteristic: &BluetoothCharacteristicMsg,
                                        service: &BluetoothRemoteGATTService)
                                        -> Root<BluetoothRemoteGATTCharacteristic> {
        let (_, ref characteristic_map_ref, _) = self.attribute_instance_map;
        let mut characteristic_map = characteristic_map_ref.borrow_mut();
        if let Some(existing_characteristic) = characteristic_map.get(&characteristic.instance_id) {
            return existing_characteristic.get();
        }
        let properties =
            BluetoothCharacteristicProperties::new(&service.global(),
                                                   characteristic.broadcast,
                                                   characteristic.read,
                                                   characteristic.write_without_response,
                                                   characteristic.write,
                                                   characteristic.notify,
                                                   characteristic.indicate,
                                                   characteristic.authenticated_signed_writes,
                                                   characteristic.reliable_write,
                                                   characteristic.writable_auxiliaries);
        let bt_characteristic = BluetoothRemoteGATTCharacteristic::new(&service.global(),
                                                                       service,
                                                                       DOMString::from(characteristic.uuid.clone()),
                                                                       &properties,
                                                                       characteristic.instance_id.clone());
        characteristic_map.insert(characteristic.instance_id.clone(), MutHeap::new(&bt_characteristic));
        return bt_characteristic;
    }

    pub fn get_or_create_descriptor(&self,
                                    descriptor: &BluetoothDescriptorMsg,
                                    characteristic: &BluetoothRemoteGATTCharacteristic)
                                    -> Root<BluetoothRemoteGATTDescriptor> {
        let (_, _, ref descriptor_map_ref) = self.attribute_instance_map;
        let mut descriptor_map = descriptor_map_ref.borrow_mut();
        if let Some(existing_descriptor) = descriptor_map.get(&descriptor.instance_id) {
            return existing_descriptor.get();
        }
        let bt_descriptor = BluetoothRemoteGATTDescriptor::new(&characteristic.global(),
                                                               characteristic,
                                                               DOMString::from(descriptor.uuid.clone()),
                                                               descriptor.instance_id.clone());
        descriptor_map.insert(descriptor.instance_id.clone(), MutHeap::new(&bt_descriptor));
        return bt_descriptor;
    }
}

impl BluetoothDeviceMethods for BluetoothDevice {
     // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-id
    fn Id(&self) -> DOMString {
        self.id.clone()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-name
    fn GetName(&self) -> Option<DOMString> {
        self.name.clone()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-addata
    fn AdData(&self) -> Root<BluetoothAdvertisingData> {
        self.ad_data.get()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-gatt
    fn Gatt(&self) -> Root<BluetoothRemoteGATTServer> {
        // TODO: Step 1 - 2: Implement the Permission API.
        self.gatt.or_init(|| {
            BluetoothRemoteGATTServer::new(&self.global(), self)
        })
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdeviceeventhandlers-ongattserverdisconnected
    event_handler!(gattserverdisconnected, GetOngattserverdisconnected, SetOngattserverdisconnected);
}
<commit_msg>Remove the unused get_context() function<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use bluetooth_traits::{BluetoothCharacteristicMsg, BluetoothDescriptorMsg, BluetoothServiceMsg};
use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::BluetoothDeviceBinding;
use dom::bindings::codegen::Bindings::BluetoothDeviceBinding::BluetoothDeviceMethods;
use dom::bindings::codegen::Bindings::BluetoothRemoteGATTServerBinding::BluetoothRemoteGATTServerMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::js::{JS, Root, MutHeap, MutNullableHeap};
use dom::bindings::reflector::{Reflectable, reflect_dom_object};
use dom::bindings::str::DOMString;
use dom::bluetooth::Bluetooth;
use dom::bluetoothadvertisingdata::BluetoothAdvertisingData;
use dom::bluetoothcharacteristicproperties::BluetoothCharacteristicProperties;
use dom::bluetoothremotegattcharacteristic::BluetoothRemoteGATTCharacteristic;
use dom::bluetoothremotegattdescriptor::BluetoothRemoteGATTDescriptor;
use dom::bluetoothremotegattserver::BluetoothRemoteGATTServer;
use dom::bluetoothremotegattservice::BluetoothRemoteGATTService;
use dom::eventtarget::EventTarget;
use dom::globalscope::GlobalScope;
use std::collections::HashMap;

// https://webbluetoothcg.github.io/web-bluetooth/#bluetoothdevice
#[dom_struct]
pub struct BluetoothDevice {
    eventtarget: EventTarget,
    id: DOMString,
    name: Option<DOMString>,
    ad_data: MutHeap<JS<BluetoothAdvertisingData>>,
    gatt: MutNullableHeap<JS<BluetoothRemoteGATTServer>>,
    context: MutHeap<JS<Bluetooth>>,
    attribute_instance_map: (DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTService>>>>,
                             DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTCharacteristic>>>>,
                             DOMRefCell<HashMap<String, MutHeap<JS<BluetoothRemoteGATTDescriptor>>>>),
}

impl BluetoothDevice {
    pub fn new_inherited(id: DOMString,
                         name: Option<DOMString>,
                         ad_data: &BluetoothAdvertisingData,
                         context: &Bluetooth)
                         -> BluetoothDevice {
        BluetoothDevice {
            eventtarget: EventTarget::new_inherited(),
            id: id,
            name: name,
            ad_data: MutHeap::new(ad_data),
            gatt: Default::default(),
            context: MutHeap::new(context),
            attribute_instance_map: (DOMRefCell::new(HashMap::new()),
                                     DOMRefCell::new(HashMap::new()),
                                     DOMRefCell::new(HashMap::new())),
        }
    }

    pub fn new(global: &GlobalScope,
               id: DOMString,
               name: Option<DOMString>,
               adData: &BluetoothAdvertisingData,
               context: &Bluetooth)
               -> Root<BluetoothDevice> {
        reflect_dom_object(box BluetoothDevice::new_inherited(id,
                                                              name,
                                                              adData,
                                                              context),
                           global,
                           BluetoothDeviceBinding::Wrap)
    }

    pub fn get_or_create_service(&self,
                                 service: &BluetoothServiceMsg,
                                 server: &BluetoothRemoteGATTServer)
                                 -> Root<BluetoothRemoteGATTService> {
        let (ref service_map_ref, _, _) = self.attribute_instance_map;
        let mut service_map = service_map_ref.borrow_mut();
        if let Some(existing_service) = service_map.get(&service.instance_id) {
            return existing_service.get();
        }
        let bt_service = BluetoothRemoteGATTService::new(&server.global(),
                                                         &server.Device(),
                                                         DOMString::from(service.uuid.clone()),
                                                         service.is_primary,
                                                         service.instance_id.clone());
        service_map.insert(service.instance_id.clone(), MutHeap::new(&bt_service));
        return bt_service;
    }

    pub fn get_or_create_characteristic(&self,
                                        characteristic: &BluetoothCharacteristicMsg,
                                        service: &BluetoothRemoteGATTService)
                                        -> Root<BluetoothRemoteGATTCharacteristic> {
        let (_, ref characteristic_map_ref, _) = self.attribute_instance_map;
        let mut characteristic_map = characteristic_map_ref.borrow_mut();
        if let Some(existing_characteristic) = characteristic_map.get(&characteristic.instance_id) {
            return existing_characteristic.get();
        }
        let properties =
            BluetoothCharacteristicProperties::new(&service.global(),
                                                   characteristic.broadcast,
                                                   characteristic.read,
                                                   characteristic.write_without_response,
                                                   characteristic.write,
                                                   characteristic.notify,
                                                   characteristic.indicate,
                                                   characteristic.authenticated_signed_writes,
                                                   characteristic.reliable_write,
                                                   characteristic.writable_auxiliaries);
        let bt_characteristic = BluetoothRemoteGATTCharacteristic::new(&service.global(),
                                                                       service,
                                                                       DOMString::from(characteristic.uuid.clone()),
                                                                       &properties,
                                                                       characteristic.instance_id.clone());
        characteristic_map.insert(characteristic.instance_id.clone(), MutHeap::new(&bt_characteristic));
        return bt_characteristic;
    }

    pub fn get_or_create_descriptor(&self,
                                    descriptor: &BluetoothDescriptorMsg,
                                    characteristic: &BluetoothRemoteGATTCharacteristic)
                                    -> Root<BluetoothRemoteGATTDescriptor> {
        let (_, _, ref descriptor_map_ref) = self.attribute_instance_map;
        let mut descriptor_map = descriptor_map_ref.borrow_mut();
        if let Some(existing_descriptor) = descriptor_map.get(&descriptor.instance_id) {
            return existing_descriptor.get();
        }
        let bt_descriptor = BluetoothRemoteGATTDescriptor::new(&characteristic.global(),
                                                               characteristic,
                                                               DOMString::from(descriptor.uuid.clone()),
                                                               descriptor.instance_id.clone());
        descriptor_map.insert(descriptor.instance_id.clone(), MutHeap::new(&bt_descriptor));
        return bt_descriptor;
    }
}

impl BluetoothDeviceMethods for BluetoothDevice {
     // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-id
    fn Id(&self) -> DOMString {
        self.id.clone()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-name
    fn GetName(&self) -> Option<DOMString> {
        self.name.clone()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-addata
    fn AdData(&self) -> Root<BluetoothAdvertisingData> {
        self.ad_data.get()
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdevice-gatt
    fn Gatt(&self) -> Root<BluetoothRemoteGATTServer> {
        // TODO: Step 1 - 2: Implement the Permission API.
        self.gatt.or_init(|| {
            BluetoothRemoteGATTServer::new(&self.global(), self)
        })
    }

    // https://webbluetoothcg.github.io/web-bluetooth/#dom-bluetoothdeviceeventhandlers-ongattserverdisconnected
    event_handler!(gattserverdisconnected, GetOngattserverdisconnected, SetOngattserverdisconnected);
}
<|endoftext|>


<commit_before>use reqwest::Response;
use regex::Regex;

pub fn get_query_part(words: &[String]) -> String {
    let mut ret = String::new();

    let words_length = words.len();

    for (i, word) in words.iter().enumerate() {
        if i == words_length - 1 {
            ret += &format!("{} {}*", word, word);
        } else {
            ret += &format!("{} ", word);
        }
    }

    ret
}

lazy_static! {
    static ref WORD_RE: Regex = Regex::new(r"([\w\d\-]+)").unwrap();
    static ref SHORT_NUMBERS_RE: Regex = Regex::new(r"^[\d]+$").unwrap();
}

// remove non-alphanumeric chars and then split on spaces
pub fn clean_words(q: &str) -> Vec<String> {
    let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
    let lower_q = substring(&q.to_lowercase(), 200);

    WORD_RE.captures_iter(&lower_q)
        .map(|cap| cap.get(1).unwrap().as_str().to_owned())
        // fixes the "02" and "2" from gene IDs like "SPCC1620.02.2" matching
        // lots of titles and abstracts:
        .filter(|s| !SHORT_NUMBERS_RE.is_match(s))
        .collect()
}

pub fn do_solr_request(url: &str) -> Result<Response, String> {
    print!("do_solr_request({:?})\n", url);
    match reqwest::get(url) {
        Ok(res) => {
            if res.status().is_success() {
                Ok(res)
            } else {
                if let Some(reason) = res.status().canonical_reason() {
                    Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                } else {
                    Err(format!("HTTP request to Solr failed with status code: {}",
                                res.status()))
                }
            }
        },
        Err(err) => {
            Err(format!("Error from Reqwest: {:?} for {}", err, url))
        }
    }
}

<commit_msg>Fix RE for matching sort numbers in Lucene search<commit_after>use reqwest::Response;
use regex::Regex;

pub fn get_query_part(words: &[String]) -> String {
    let mut ret = String::new();

    let words_length = words.len();

    for (i, word) in words.iter().enumerate() {
        if i == words_length - 1 {
            ret += &format!("{} {}*", word, word);
        } else {
            ret += &format!("{} ", word);
        }
    }

    ret
}

lazy_static! {
    static ref WORD_RE: Regex = Regex::new(r"([\w\d\-]+)").unwrap();
    static ref SHORT_NUMBERS_RE: Regex = Regex::new(r"^\d\d?\d?$").unwrap();
}

// remove non-alphanumeric chars and then split on spaces
pub fn clean_words(q: &str) -> Vec<String> {
    let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
    let lower_q = substring(&q.to_lowercase(), 200);

    WORD_RE.captures_iter(&lower_q)
        .map(|cap| cap.get(1).unwrap().as_str().to_owned())
        // fixes the "02" and "2" from gene IDs like "SPCC1620.02.2" matching
        // lots of titles and abstracts:
        .filter(|s| !SHORT_NUMBERS_RE.is_match(s))
        .collect()
}

pub fn do_solr_request(url: &str) -> Result<Response, String> {
    print!("do_solr_request({:?})\n", url);
    match reqwest::get(url) {
        Ok(res) => {
            if res.status().is_success() {
                Ok(res)
            } else {
                if let Some(reason) = res.status().canonical_reason() {
                    Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                } else {
                    Err(format!("HTTP request to Solr failed with status code: {}",
                                res.status()))
                }
            }
        },
        Err(err) => {
            Err(format!("Error from Reqwest: {:?} for {}", err, url))
        }
    }
}

<|endoftext|>


<commit_before>use crate::App;
use anyhow::Error;
use dashmap::{DashMap, SharedValue};
use diesel::{pg::upsert::excluded, prelude::*};
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};

/// crates.io receives a lot of download requests, and we can't execute a write query to the
/// database during each connection for performance reasons. To reduce the write load, this struct
/// collects the pending updates from the current process and writes in batch.
///
/// To avoid locking the whole data structure behind a RwLock, which could potentially delay
/// requests, this uses the dashmap crate. A DashMap has the same public API as an HashMap, but
/// stores the items into `num_cpus()*4` individually locked shards. This approach reduces the
/// likelyhood of a request encountering a locked shard.
///
/// Persisting the download counts in the database also takes advantage of the inner sharding of
/// DashMaps: to avoid locking all the download requests at the same time each iteration only
/// persists a single shard at the time.
///
/// The disadvantage of this approach is that download counts are stored in memory until they're
/// persisted, so it's possible to lose some of them if the process exits ungracefully. While
/// that's far from ideal, the advantage of batching database updates far outweights potentially
/// losing some download counts.
#[derive(Debug)]
pub struct DownloadsCounter {
    /// Inner storage for the download counts.
    inner: DashMap<i32, AtomicUsize>,
    /// Index of the next shard that should be persisted by `persist_next_shard`.
    shard_idx: AtomicUsize,
    /// Number of downloads that are not yet persisted on the database. This is just used as a
    /// metric included in log lines, and it's not guaranteed to be accurate.
    pending_count: AtomicI64,
}

impl DownloadsCounter {
    pub(crate) fn new() -> Self {
        Self {
            inner: DashMap::new(),
            shard_idx: AtomicUsize::new(0),
            pending_count: AtomicI64::new(0),
        }
    }

    pub(crate) fn increment(&self, version_id: i32) {
        self.pending_count.fetch_add(1, Ordering::SeqCst);

        if let Some(counter) = self.inner.get(&version_id) {
            // The version is already recorded in the DashMap, so we don't need to lock the whole
            // shard in write mode. The shard is instead locked in read mode, which allows an
            // unbounded number of readers as long as there are no write locks.
            counter.value().fetch_add(1, Ordering::SeqCst);
        } else {
            // The version is not in the DashMap, so we need to lock the whole shard in write mode
            // and insert the version into it. This has worse performance than the above case.
            self.inner
                .entry(version_id)
                .and_modify(|counter| {
                    // Handle the version being inserted by another thread while we were waiting
                    // for the write lock on the shard.
                    counter.fetch_add(1, Ordering::SeqCst);
                })
                .or_insert_with(|| AtomicUsize::new(1));
        }
    }

    pub fn persist_all_shards(&self, app: &App) -> Result<(), Error> {
        let conn = app.primary_database.get()?;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut pending_downloads = 0;
        for shard in self.inner.shards() {
            let shard = std::mem::take(&mut *shard.write());
            let stats = self.persist_shard(&conn, shard)?;

            counted_downloads += stats.counted_downloads;
            counted_versions += stats.counted_versions;
            pending_downloads = stats.pending_downloads;
        }

        println!(
            "download_counter all_shards counted_versions={} counted_downloads={} pending_downloads={}",
            counted_versions,
            counted_downloads,
            pending_downloads,
        );

        Ok(())
    }

    pub fn persist_next_shard(&self, app: &App) -> Result<(), Error> {
        let conn = app.primary_database.get()?;

        // Replace the next shard in the ring with an empty HashMap (clearing it), and return the
        // previous contents for processing. The fetch_add method wraps around on overflow, so it's
        // fine to keep incrementing it without resetting.
        let shards = self.inner.shards();
        let idx = self.shard_idx.fetch_add(1, Ordering::SeqCst) % shards.len();
        let shard = std::mem::take(&mut *shards[idx].write());

        let stats = self.persist_shard(&conn, shard)?;
        println!(
            "download_counter shard={} counted_versions={} counted_downloads={} pending_downloads={}",
            idx,
            stats.counted_versions,
            stats.counted_downloads,
            stats.pending_downloads,
        );

        Ok(())
    }

    fn persist_shard(
        &self,
        conn: &PgConnection,
        shard: HashMap<i32, SharedValue<AtomicUsize>>,
    ) -> Result<PersistStats, Error> {
        use crate::schema::version_downloads::dsl::*;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut to_insert = Vec::new();
        for (key, atomic) in shard.iter() {
            let count = atomic.get().load(Ordering::SeqCst);
            counted_downloads += count;
            counted_versions += 1;

            to_insert.push((version_id.eq(*key), downloads.eq(count as i32)));
        }

        if !to_insert.is_empty() {
            diesel::insert_into(version_downloads)
                .values(&to_insert)
                .on_conflict((version_id, date))
                .do_update()
                .set(downloads.eq(downloads + excluded(downloads)))
                .execute(conn)?;
        }

        let old_pending = self
            .pending_count
            .fetch_sub(counted_downloads as i64, Ordering::SeqCst);

        Ok(PersistStats {
            counted_downloads,
            counted_versions,
            pending_downloads: old_pending - counted_downloads as i64,
        })
    }

    pub fn shards_count(&self) -> usize {
        self.inner.shards().len()
    }
}

struct PersistStats {
    counted_downloads: usize,
    counted_versions: usize,
    pending_downloads: i64,
}
<commit_msg>downloads_counter: prevent deadlocks with multiple crates.io instances<commit_after>use crate::App;
use anyhow::Error;
use dashmap::{DashMap, SharedValue};
use diesel::{pg::upsert::excluded, prelude::*};
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};

/// crates.io receives a lot of download requests, and we can't execute a write query to the
/// database during each connection for performance reasons. To reduce the write load, this struct
/// collects the pending updates from the current process and writes in batch.
///
/// To avoid locking the whole data structure behind a RwLock, which could potentially delay
/// requests, this uses the dashmap crate. A DashMap has the same public API as an HashMap, but
/// stores the items into `num_cpus()*4` individually locked shards. This approach reduces the
/// likelyhood of a request encountering a locked shard.
///
/// Persisting the download counts in the database also takes advantage of the inner sharding of
/// DashMaps: to avoid locking all the download requests at the same time each iteration only
/// persists a single shard at the time.
///
/// The disadvantage of this approach is that download counts are stored in memory until they're
/// persisted, so it's possible to lose some of them if the process exits ungracefully. While
/// that's far from ideal, the advantage of batching database updates far outweights potentially
/// losing some download counts.
#[derive(Debug)]
pub struct DownloadsCounter {
    /// Inner storage for the download counts.
    inner: DashMap<i32, AtomicUsize>,
    /// Index of the next shard that should be persisted by `persist_next_shard`.
    shard_idx: AtomicUsize,
    /// Number of downloads that are not yet persisted on the database. This is just used as a
    /// metric included in log lines, and it's not guaranteed to be accurate.
    pending_count: AtomicI64,
}

impl DownloadsCounter {
    pub(crate) fn new() -> Self {
        Self {
            inner: DashMap::new(),
            shard_idx: AtomicUsize::new(0),
            pending_count: AtomicI64::new(0),
        }
    }

    pub(crate) fn increment(&self, version_id: i32) {
        self.pending_count.fetch_add(1, Ordering::SeqCst);

        if let Some(counter) = self.inner.get(&version_id) {
            // The version is already recorded in the DashMap, so we don't need to lock the whole
            // shard in write mode. The shard is instead locked in read mode, which allows an
            // unbounded number of readers as long as there are no write locks.
            counter.value().fetch_add(1, Ordering::SeqCst);
        } else {
            // The version is not in the DashMap, so we need to lock the whole shard in write mode
            // and insert the version into it. This has worse performance than the above case.
            self.inner
                .entry(version_id)
                .and_modify(|counter| {
                    // Handle the version being inserted by another thread while we were waiting
                    // for the write lock on the shard.
                    counter.fetch_add(1, Ordering::SeqCst);
                })
                .or_insert_with(|| AtomicUsize::new(1));
        }
    }

    pub fn persist_all_shards(&self, app: &App) -> Result<(), Error> {
        let conn = app.primary_database.get()?;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut pending_downloads = 0;
        for shard in self.inner.shards() {
            let shard = std::mem::take(&mut *shard.write());
            let stats = self.persist_shard(&conn, shard)?;

            counted_downloads += stats.counted_downloads;
            counted_versions += stats.counted_versions;
            pending_downloads = stats.pending_downloads;
        }

        println!(
            "download_counter all_shards counted_versions={} counted_downloads={} pending_downloads={}",
            counted_versions,
            counted_downloads,
            pending_downloads,
        );

        Ok(())
    }

    pub fn persist_next_shard(&self, app: &App) -> Result<(), Error> {
        let conn = app.primary_database.get()?;

        // Replace the next shard in the ring with an empty HashMap (clearing it), and return the
        // previous contents for processing. The fetch_add method wraps around on overflow, so it's
        // fine to keep incrementing it without resetting.
        let shards = self.inner.shards();
        let idx = self.shard_idx.fetch_add(1, Ordering::SeqCst) % shards.len();
        let shard = std::mem::take(&mut *shards[idx].write());

        let stats = self.persist_shard(&conn, shard)?;
        println!(
            "download_counter shard={} counted_versions={} counted_downloads={} pending_downloads={}",
            idx,
            stats.counted_versions,
            stats.counted_downloads,
            stats.pending_downloads,
        );

        Ok(())
    }

    fn persist_shard(
        &self,
        conn: &PgConnection,
        shard: HashMap<i32, SharedValue<AtomicUsize>>,
    ) -> Result<PersistStats, Error> {
        use crate::schema::version_downloads::dsl::*;

        let mut counted_downloads = 0;
        let mut counted_versions = 0;
        let mut to_insert = Vec::new();
        for (key, atomic) in shard.iter() {
            let count = atomic.get().load(Ordering::SeqCst);
            counted_downloads += count;
            counted_versions += 1;

            to_insert.push((*key, count));
        }

        if !to_insert.is_empty() {
            // The rows we're about to insert need to be sorted to avoid deadlocks when multiple
            // instances of crates.io are running at the same time.
            //
            // In PostgreSQL a transaction modifying a row locks that row until the transaction is
            // committed. Multiple transactions inserting rows into a table could end up
            // deadlocking each other though: PostgreSQL will detect that deadlock, abort one of
            // the transactions and allow the other one to continue. We don't want that to happen,
            // as we'd lose the downloads from the aborted transaction.
            //
            // Ensuring the rows are inserted in a consistent order (in our case by sorting them by
            // the version ID) will prevent deadlocks from occuring. For more information:
            //
            //     https://www.postgresql.org/docs/11/explicit-locking.html#LOCKING-DEADLOCKS
            //
            to_insert.sort_by_key(|(key, _)| *key);

            let to_insert = to_insert
                .into_iter()
                .map(|(key, count)| (version_id.eq(key), downloads.eq(count as i32)))
                .collect::<Vec<_>>();

            diesel::insert_into(version_downloads)
                .values(&to_insert)
                .on_conflict((version_id, date))
                .do_update()
                .set(downloads.eq(downloads + excluded(downloads)))
                .execute(conn)?;
        }

        let old_pending = self
            .pending_count
            .fetch_sub(counted_downloads as i64, Ordering::SeqCst);

        Ok(PersistStats {
            counted_downloads,
            counted_versions,
            pending_downloads: old_pending - counted_downloads as i64,
        })
    }

    pub fn shards_count(&self) -> usize {
        self.inner.shards().len()
    }
}

struct PersistStats {
    counted_downloads: usize,
    counted_versions: usize,
    pending_downloads: i64,
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::error::Error;

use crate::error::{ConfigError, Unexpected};
use crate::value::{Value, ValueKind};

#[derive(serde::Deserialize, Debug)]
#[serde(untagged)]
pub enum Val {
    Null,
    Boolean(bool),
    Integer(i64),
    Float(f64),
    String(String),
    Array(Vec<Val>),
    Object(HashMap<String, Val>),
}

pub fn parse(
    uri: Option<&String>,
    text: &str,
) -> Result<HashMap<String, Value>, Box<dyn Error + Send + Sync>> {
    let root = json5_rs::from_str::<Val>(&text)?;
    if let Some(err) = match root {
        Val::String(ref value) => Some(Unexpected::Str(value.clone())),
        Val::Integer(value) => Some(Unexpected::Integer(value)),
        Val::Float(value) => Some(Unexpected::Float(value)),
        Val::Boolean(value) => Some(Unexpected::Bool(value)),
        Val::Object(_) => None,
        Val::Array(_) => Some(Unexpected::Seq),
        Val::Null => Some(Unexpected::Unit),
    } {
        return Err(ConfigError::invalid_root(uri, err));
    }

    let value = from_json5_value(uri, root);
    match value.kind {
        ValueKind::Table(map) => Ok(map),

        _ => Ok(HashMap::new()),
    }
}

fn from_json5_value(uri: Option<&String>, value: Val) -> Value {
    match value {
        Val::String(v) => Value::new(uri, ValueKind::String(v)),

        Val::Integer(v) => Value::new(uri, ValueKind::Integer(v)),

        Val::Float(v) => Value::new(uri, ValueKind::Float(v)),

        Val::Boolean(v) => Value::new(uri, ValueKind::Boolean(v)),

        Val::Object(table) => {
            let mut m = HashMap::new();

            for (key, value) in table {
                m.insert(key, from_json5_value(uri, value));
            }

            Value::new(uri, ValueKind::Table(m))
        }

        Val::Array(array) => {
            let mut l = Vec::new();

            for value in array {
                l.push(from_json5_value(uri, value));
            }

            Value::new(uri, ValueKind::Array(l))
        }

        Val::Null => Value::new(uri, ValueKind::Nil),
    }
}
<commit_msg>Simplify parse() impl of json5 code<commit_after>use std::collections::HashMap;
use std::error::Error;

use crate::error::{ConfigError, Unexpected};
use crate::value::{Value, ValueKind};

#[derive(serde::Deserialize, Debug)]
#[serde(untagged)]
pub enum Val {
    Null,
    Boolean(bool),
    Integer(i64),
    Float(f64),
    String(String),
    Array(Vec<Val>),
    Object(HashMap<String, Val>),
}

pub fn parse(
    uri: Option<&String>,
    text: &str,
) -> Result<HashMap<String, Value>, Box<dyn Error + Send + Sync>> {
    match json5_rs::from_str::<Val>(&text)? {
        Val::String(ref value) => Err(Unexpected::Str(value.clone())),
        Val::Integer(value) => Err(Unexpected::Integer(value)),
        Val::Float(value) => Err(Unexpected::Float(value)),
        Val::Boolean(value) => Err(Unexpected::Bool(value)),
        Val::Array(_) => Err(Unexpected::Seq),
        Val::Null => Err(Unexpected::Unit),
        Val::Object(o) => match from_json5_value(uri, Val::Object(o)).kind {
            ValueKind::Table(map) => Ok(map),
            _ => Ok(HashMap::new()),
        },
    }
    .map_err(|err| ConfigError::invalid_root(uri, err))
    .map_err(|err| Box::new(err) as Box<dyn Error + Send + Sync>)
}

fn from_json5_value(uri: Option<&String>, value: Val) -> Value {
    match value {
        Val::String(v) => Value::new(uri, ValueKind::String(v)),

        Val::Integer(v) => Value::new(uri, ValueKind::Integer(v)),

        Val::Float(v) => Value::new(uri, ValueKind::Float(v)),

        Val::Boolean(v) => Value::new(uri, ValueKind::Boolean(v)),

        Val::Object(table) => {
            let mut m = HashMap::new();

            for (key, value) in table {
                m.insert(key, from_json5_value(uri, value));
            }

            Value::new(uri, ValueKind::Table(m))
        }

        Val::Array(array) => {
            let mut l = Vec::new();

            for value in array {
                l.push(from_json5_value(uri, value));
            }

            Value::new(uri, ValueKind::Array(l))
        }

        Val::Null => Value::new(uri, ValueKind::Nil),
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Binary file to return the number of rows found from Parquet file(s).
//!
//! # Install
//!
//! `parquet-rowcount` can be installed using `cargo`:
//! ```
//! cargo install parquet
//! ```
//! After this `parquet-rowcount` should be globally available:
//! ```
//! parquet-rowcount XYZ.parquet
//! ```
//!
//! The binary can also be built from the source code and run as follows:
//! ```
//! cargo run --bin parquet-rowcount XYZ.parquet ABC.parquet ZXC.parquet
//! ```
//!
//! # Usage
//!
//! ```
//! parquet-rowcount <file-path> ...
//! ```
//! where `file-path` is the path to a Parquet file and `...` is any additional number of
//! parquet files to count the number of rows from.
//!
//! Note that `parquet-rowcount` reads full file schema, no projection or filtering is
//! applied.

extern crate parquet;

use std::{env, fs::File, path::Path, process};

use parquet::file::reader::{FileReader, SerializedFileReader};

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Usage: parquet-rowcount <file-path> ...");
        process::exit(1);
    }

    for filename in args {
        let path = Path::new(&filename);
        let file = File::open(&path).unwrap();
        let parquet_reader = SerializedFileReader::new(file).unwrap();
        let row_group_metadata = parquet_reader.metadata().row_groups();
        let mut total_num_rows = 0;

        for group_metadata in row_group_metadata {
            total_num_rows += group_metadata.num_rows();
        }

        eprintln!("File {}: rowcount={}", filename, total_num_rows);
    }
}
<commit_msg>ARROW-11305 Skip first argument (which is the program name) in parquet-rowcount binary<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Binary file to return the number of rows found from Parquet file(s).
//!
//! # Install
//!
//! `parquet-rowcount` can be installed using `cargo`:
//! ```
//! cargo install parquet
//! ```
//! After this `parquet-rowcount` should be globally available:
//! ```
//! parquet-rowcount XYZ.parquet
//! ```
//!
//! The binary can also be built from the source code and run as follows:
//! ```
//! cargo run --bin parquet-rowcount XYZ.parquet ABC.parquet ZXC.parquet
//! ```
//!
//! # Usage
//!
//! ```
//! parquet-rowcount <file-path> ...
//! ```
//! where `file-path` is the path to a Parquet file and `...` is any additional number of
//! parquet files to count the number of rows from.
//!
//! Note that `parquet-rowcount` reads full file schema, no projection or filtering is
//! applied.

extern crate parquet;

use std::{env, fs::File, path::Path, process};

use parquet::file::reader::{FileReader, SerializedFileReader};

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Usage: parquet-rowcount <file-path> ...");
        process::exit(1);
    }

    for filename in &args[1..] {
        let path = Path::new(filename);
        let file = File::open(path).unwrap();
        let parquet_reader = SerializedFileReader::new(file).unwrap();
        let row_group_metadata = parquet_reader.metadata().row_groups();
        let mut total_num_rows = 0;

        for group_metadata in row_group_metadata {
            total_num_rows += group_metadata.num_rows();
        }

        eprintln!("File {}: rowcount={}", filename, total_num_rows);
    }
}
<|endoftext|>