<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 dom::bindings::codegen::Bindings::LocationBinding;
use dom::bindings::codegen::Bindings::LocationBinding::LocationMethods;
use dom::bindings::global::Window;
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::window::Window;
use page::Page;

use servo_util::str::DOMString;

use std::rc::Rc;

#[deriving(Encodable)]
pub struct Location {
    reflector_: Reflector, //XXXjdm cycle: window->Location->window
    page: Rc<Page>,
}

impl Location {
    pub fn new_inherited(page: Rc<Page>) -> Location {
        Location {
            reflector_: Reflector::new(),
            page: page
        }
    }

    pub fn new(window: &JSRef<Window>, page: Rc<Page>) -> Temporary<Location> {
        reflect_dom_object(box Location::new_inherited(page),
                           &Window(*window),
                           LocationBinding::Wrap)
    }
}

impl<'a> LocationMethods for JSRef<'a, Location> {
    fn Href(&self) -> DOMString {
        self.page.get_url().serialize()
    }

    fn Search(&self) -> DOMString {
        match self.page.get_url().query {
            None => "".to_string(),
            Some(ref query) => "?".to_string().append(query.as_slice())
        }
    }

    fn Hash(&self) -> DOMString {
        match self.page.get_url().fragment {
            None => "".to_string(),
            Some(ref query) => "#".to_string().append(query.as_slice())
        }
    }
}

impl Reflectable for Location {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}
<commit_msg>Return the empty string as appropriate for location.{search,hash}.<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 dom::bindings::codegen::Bindings::LocationBinding;
use dom::bindings::codegen::Bindings::LocationBinding::LocationMethods;
use dom::bindings::global::Window;
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::window::Window;
use page::Page;

use servo_util::str::DOMString;

use std::rc::Rc;

#[deriving(Encodable)]
pub struct Location {
    reflector_: Reflector, //XXXjdm cycle: window->Location->window
    page: Rc<Page>,
}

impl Location {
    pub fn new_inherited(page: Rc<Page>) -> Location {
        Location {
            reflector_: Reflector::new(),
            page: page
        }
    }

    pub fn new(window: &JSRef<Window>, page: Rc<Page>) -> Temporary<Location> {
        reflect_dom_object(box Location::new_inherited(page),
                           &Window(*window),
                           LocationBinding::Wrap)
    }
}

impl<'a> LocationMethods for JSRef<'a, Location> {
    fn Href(&self) -> DOMString {
        self.page.get_url().serialize()
    }

    fn Search(&self) -> DOMString {
        match self.page.get_url().query {
            None => "".to_string(),
            Some(ref query) if query.as_slice() == "" => "".to_string(),
            Some(ref query) => "?".to_string().append(query.as_slice())
        }
    }

    fn Hash(&self) -> DOMString {
        match self.page.get_url().fragment {
            None => "".to_string(),
            Some(ref hash) if hash.as_slice() == "" => "".to_string(),
            Some(ref hash) => "#".to_string().append(hash.as_slice())
        }
    }
}

impl Reflectable for Location {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}
<|endoftext|>


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

use std::mem;

use rand::{thread_rng, Rng, ThreadRng};
use protobuf::{Message, RepeatedField};
use kvproto::coprocessor::{KeyRange, Response};
use tipb::analyze::{self, AnalyzeColumnsReq, AnalyzeIndexReq, AnalyzeReq, AnalyzeType};
use tipb::schema::ColumnInfo;
use tipb::executor::TableScan;

use coprocessor::dag::executor::{Executor, ExecutorMetrics, IndexScanExecutor, TableScanExecutor};
use coprocessor::endpoint::ReqContext;
use coprocessor::codec::datum;
use coprocessor::{Error, Result};
use storage::{Snapshot, SnapshotStore};
use super::fmsketch::FMSketch;
use super::cmsketch::CMSketch;
use super::histogram::Histogram;

// `AnalyzeContext` is used to handle `AnalyzeReq`
pub struct AnalyzeContext {
    req: AnalyzeReq,
    snap: Option<SnapshotStore>,
    ranges: Vec<KeyRange>,
}

impl AnalyzeContext {
    pub fn new(
        req: AnalyzeReq,
        ranges: Vec<KeyRange>,
        snap: Box<Snapshot>,
        req_ctx: &ReqContext,
    ) -> AnalyzeContext {
        let snap = SnapshotStore::new(
            snap,
            req.get_start_ts(),
            req_ctx.isolation_level,
            req_ctx.fill_cache,
        );
        AnalyzeContext {
            req: req,
            snap: Some(snap),
            ranges: ranges,
        }
    }

    pub fn handle_request(mut self, stats: &mut ExecutorMetrics) -> Result<Response> {
        let ret = match self.req.get_tp() {
            AnalyzeType::TypeIndex => {
                let req = self.req.take_idx_req();
                let mut scanner = IndexScanExecutor::new_with_cols_len(
                    i64::from(req.get_num_columns()),
                    mem::replace(&mut self.ranges, Vec::new()),
                    self.snap.take().unwrap(),
                )?;
                let res = AnalyzeContext::handle_index(req, &mut scanner);
                scanner.collect_metrics_into(stats);
                res
            }

            AnalyzeType::TypeColumn => {
                let col_req = self.req.take_col_req();
                let snap = self.snap.take().unwrap();
                let ranges = mem::replace(&mut self.ranges, Vec::new());
                let mut builder = SampleBuilder::new(col_req, snap, ranges)?;
                let res = AnalyzeContext::handle_column(&mut builder);
                builder.data.collect_metrics_into(stats);
                res
            }
        };
        match ret {
            Ok(data) => {
                let mut resp = Response::new();
                resp.set_data(data);
                Ok(resp)
            }
            Err(Error::Other(e)) => {
                let mut resp = Response::new();
                resp.set_other_error(format!("{}", e));
                Ok(resp)
            }
            Err(e) => Err(e),
        }
    }

    // handle_column is used to process `AnalyzeColumnsReq`
    // it would build a histogram for the primary key(if needed) and
    // collectors for each column value.
    fn handle_column(builder: &mut SampleBuilder) -> Result<Vec<u8>> {
        let (collectors, pk_builder) = builder.collect_columns_stats()?;

        let pk_hist = pk_builder.into_proto();
        let cols: Vec<analyze::SampleCollector> =
            collectors.into_iter().map(|col| col.into_proto()).collect();

        let res_data = {
            let mut res = analyze::AnalyzeColumnsResp::new();
            res.set_collectors(RepeatedField::from_vec(cols));
            res.set_pk_hist(pk_hist);
            box_try!(res.write_to_bytes())
        };
        Ok(res_data)
    }

    // handle_index is used to handle `AnalyzeIndexReq`,
    // it would build a histogram and count-min sketch of index values.
    fn handle_index(req: AnalyzeIndexReq, scanner: &mut IndexScanExecutor) -> Result<Vec<u8>> {
        let mut hist = Histogram::new(req.get_bucket_size() as usize);
        let mut cms = CMSketch::new(
            req.get_cmsketch_depth() as usize,
            req.get_cmsketch_width() as usize,
        );
        while let Some(row) = scanner.next()? {
            let bytes = row.data.get_column_values();
            hist.append(bytes);
            if let Some(c) = cms.as_mut() {
                c.insert(bytes)
            }
        }
        let mut res = analyze::AnalyzeIndexResp::new();
        res.set_hist(hist.into_proto());
        if let Some(c) = cms {
            res.set_cms(c.into_proto());
        }
        let dt = box_try!(res.write_to_bytes());
        Ok(dt)
    }
}

struct SampleBuilder {
    data: TableScanExecutor,
    cols: Vec<ColumnInfo>,
    // the number of columns need to be sampled. It equals to cols.len()
    // if cols[0] is not pk handle, or it should be cols.len() - 1.
    col_len: usize,
    max_bucket_size: usize,
    max_sample_size: usize,
    max_fm_sketch_size: usize,
    cm_sketch_depth: usize,
    cm_sketch_width: usize,
}

/// `SampleBuilder` is used to analyze columns. It collects sample from
/// the result set using Reservoir Sampling algorithm, estimates NDVs
/// using FM Sketch during the collecting process, and builds count-min sketch.
impl SampleBuilder {
    fn new(
        mut req: AnalyzeColumnsReq,
        snap: SnapshotStore,
        ranges: Vec<KeyRange>,
    ) -> Result<SampleBuilder> {
        let cols_info = req.take_columns_info();
        if cols_info.is_empty() {
            return Err(box_err!("empty columns_info"));
        }

        let mut col_len = cols_info.len();
        if cols_info[0].get_pk_handle() {
            col_len -= 1;
        }

        let mut meta = TableScan::new();
        meta.set_columns(cols_info);
        let table_scanner = TableScanExecutor::new(&meta, ranges, snap, false)?;
        Ok(SampleBuilder {
            data: table_scanner,
            cols: meta.take_columns().to_vec(),
            col_len: col_len,
            max_bucket_size: req.get_bucket_size() as usize,
            max_fm_sketch_size: req.get_sketch_size() as usize,
            max_sample_size: req.get_sample_size() as usize,
            cm_sketch_depth: req.get_cmsketch_depth() as usize,
            cm_sketch_width: req.get_cmsketch_width() as usize,
        })
    }

    // `collect_columns_stats` returns the sample collectors which contain total count,
    // null count, distinct values count and count-min sketch. And it also returns the statistic
    // builder for PK which contains the histogram.
    // See https://en.wikipedia.org/wiki/Reservoir_sampling
    fn collect_columns_stats(&mut self) -> Result<(Vec<SampleCollector>, Histogram)> {
        let mut pk_builder = Histogram::new(self.max_bucket_size);
        let mut collectors = vec![
            SampleCollector::new(
                self.max_sample_size,
                self.max_fm_sketch_size,
                self.cm_sketch_depth,
                self.cm_sketch_width
            );
            self.col_len
        ];
        while let Some(row) = self.data.next()? {
            let cols = row.get_binary_cols(&self.cols)?;
            let retrieve_len = cols.len();
            let mut cols_iter = cols.into_iter();
            if self.col_len != retrieve_len {
                if let Some(v) = cols_iter.next() {
                    pk_builder.append(&v);
                }
            }
            for (collector, val) in collectors.iter_mut().zip(cols_iter) {
                collector.collect(val);
            }
        }
        Ok((collectors, pk_builder))
    }
}

/// `SampleCollector` will collect Samples and calculate the count and ndv of an attribute.
#[derive(Clone)]
struct SampleCollector {
    samples: Vec<Vec<u8>>,
    null_count: u64,
    count: u64,
    max_sample_size: usize,
    fm_sketch: FMSketch,
    cm_sketch: Option<CMSketch>,
    rng: ThreadRng,
}

impl SampleCollector {
    fn new(
        max_sample_size: usize,
        max_fm_sketch_size: usize,
        cm_sketch_depth: usize,
        cm_sketch_width: usize,
    ) -> SampleCollector {
        SampleCollector {
            samples: Default::default(),
            null_count: 0,
            count: 0,
            max_sample_size,
            fm_sketch: FMSketch::new(max_fm_sketch_size),
            cm_sketch: CMSketch::new(cm_sketch_depth, cm_sketch_width),
            rng: thread_rng(),
        }
    }

    fn into_proto(self) -> analyze::SampleCollector {
        let mut s = analyze::SampleCollector::new();
        s.set_null_count(self.null_count as i64);
        s.set_count(self.count as i64);
        s.set_fm_sketch(self.fm_sketch.into_proto());
        s.set_samples(RepeatedField::from_vec(self.samples));
        if let Some(c) = self.cm_sketch {
            s.set_cm_sketch(c.into_proto())
        }
        s
    }

    pub fn collect(&mut self, data: Vec<u8>) {
        if data[0] == datum::NIL_FLAG {
            self.null_count += 1;
            return;
        }
        self.count += 1;
        self.fm_sketch.insert(&data);
        if let Some(c) = self.cm_sketch.as_mut() {
            c.insert(&data)
        }
        if self.samples.len() < self.max_sample_size {
            self.samples.push(data);
            return;
        }
        if self.rng.gen_range(0, self.count) < self.max_sample_size as u64 {
            let idx = self.rng.gen_range(0, self.max_sample_size);
            self.samples[idx] = data;
        }
    }
}

#[cfg(test)]
mod test {
    use coprocessor::codec::datum;
    use coprocessor::codec::datum::Datum;
    use super::*;

    #[test]
    fn test_sample_collector() {
        let max_sample_size = 3;
        let max_fm_sketch_size = 10;
        let cm_sketch_depth = 2;
        let cm_sketch_width = 16;
        let mut sample = SampleCollector::new(
            max_sample_size,
            max_fm_sketch_size,
            cm_sketch_depth,
            cm_sketch_width,
        );
        let cases = vec![Datum::I64(1), Datum::Null, Datum::I64(2), Datum::I64(5)];

        for data in cases {
            sample.collect(datum::encode_value(&[data]).unwrap());
        }
        assert_eq!(sample.samples.len(), max_sample_size);
        assert_eq!(sample.null_count, 1);
        assert_eq!(sample.count, 3);
        assert_eq!(sample.cm_sketch.unwrap().count(), 3)
    }
}
<commit_msg>cop/stats: add total column size (#2950)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::mem;

use rand::{thread_rng, Rng, ThreadRng};
use protobuf::{Message, RepeatedField};
use kvproto::coprocessor::{KeyRange, Response};
use tipb::analyze::{self, AnalyzeColumnsReq, AnalyzeIndexReq, AnalyzeReq, AnalyzeType};
use tipb::schema::ColumnInfo;
use tipb::executor::TableScan;

use coprocessor::dag::executor::{Executor, ExecutorMetrics, IndexScanExecutor, TableScanExecutor};
use coprocessor::endpoint::ReqContext;
use coprocessor::codec::datum;
use coprocessor::{Error, Result};
use storage::{Snapshot, SnapshotStore};
use super::fmsketch::FMSketch;
use super::cmsketch::CMSketch;
use super::histogram::Histogram;

// `AnalyzeContext` is used to handle `AnalyzeReq`
pub struct AnalyzeContext {
    req: AnalyzeReq,
    snap: Option<SnapshotStore>,
    ranges: Vec<KeyRange>,
}

impl AnalyzeContext {
    pub fn new(
        req: AnalyzeReq,
        ranges: Vec<KeyRange>,
        snap: Box<Snapshot>,
        req_ctx: &ReqContext,
    ) -> AnalyzeContext {
        let snap = SnapshotStore::new(
            snap,
            req.get_start_ts(),
            req_ctx.isolation_level,
            req_ctx.fill_cache,
        );
        AnalyzeContext {
            req: req,
            snap: Some(snap),
            ranges: ranges,
        }
    }

    pub fn handle_request(mut self, stats: &mut ExecutorMetrics) -> Result<Response> {
        let ret = match self.req.get_tp() {
            AnalyzeType::TypeIndex => {
                let req = self.req.take_idx_req();
                let mut scanner = IndexScanExecutor::new_with_cols_len(
                    i64::from(req.get_num_columns()),
                    mem::replace(&mut self.ranges, Vec::new()),
                    self.snap.take().unwrap(),
                )?;
                let res = AnalyzeContext::handle_index(req, &mut scanner);
                scanner.collect_metrics_into(stats);
                res
            }

            AnalyzeType::TypeColumn => {
                let col_req = self.req.take_col_req();
                let snap = self.snap.take().unwrap();
                let ranges = mem::replace(&mut self.ranges, Vec::new());
                let mut builder = SampleBuilder::new(col_req, snap, ranges)?;
                let res = AnalyzeContext::handle_column(&mut builder);
                builder.data.collect_metrics_into(stats);
                res
            }
        };
        match ret {
            Ok(data) => {
                let mut resp = Response::new();
                resp.set_data(data);
                Ok(resp)
            }
            Err(Error::Other(e)) => {
                let mut resp = Response::new();
                resp.set_other_error(format!("{}", e));
                Ok(resp)
            }
            Err(e) => Err(e),
        }
    }

    // handle_column is used to process `AnalyzeColumnsReq`
    // it would build a histogram for the primary key(if needed) and
    // collectors for each column value.
    fn handle_column(builder: &mut SampleBuilder) -> Result<Vec<u8>> {
        let (collectors, pk_builder) = builder.collect_columns_stats()?;

        let pk_hist = pk_builder.into_proto();
        let cols: Vec<analyze::SampleCollector> =
            collectors.into_iter().map(|col| col.into_proto()).collect();

        let res_data = {
            let mut res = analyze::AnalyzeColumnsResp::new();
            res.set_collectors(RepeatedField::from_vec(cols));
            res.set_pk_hist(pk_hist);
            box_try!(res.write_to_bytes())
        };
        Ok(res_data)
    }

    // handle_index is used to handle `AnalyzeIndexReq`,
    // it would build a histogram and count-min sketch of index values.
    fn handle_index(req: AnalyzeIndexReq, scanner: &mut IndexScanExecutor) -> Result<Vec<u8>> {
        let mut hist = Histogram::new(req.get_bucket_size() as usize);
        let mut cms = CMSketch::new(
            req.get_cmsketch_depth() as usize,
            req.get_cmsketch_width() as usize,
        );
        while let Some(row) = scanner.next()? {
            let bytes = row.data.get_column_values();
            hist.append(bytes);
            if let Some(c) = cms.as_mut() {
                c.insert(bytes)
            }
        }
        let mut res = analyze::AnalyzeIndexResp::new();
        res.set_hist(hist.into_proto());
        if let Some(c) = cms {
            res.set_cms(c.into_proto());
        }
        let dt = box_try!(res.write_to_bytes());
        Ok(dt)
    }
}

struct SampleBuilder {
    data: TableScanExecutor,
    cols: Vec<ColumnInfo>,
    // the number of columns need to be sampled. It equals to cols.len()
    // if cols[0] is not pk handle, or it should be cols.len() - 1.
    col_len: usize,
    max_bucket_size: usize,
    max_sample_size: usize,
    max_fm_sketch_size: usize,
    cm_sketch_depth: usize,
    cm_sketch_width: usize,
}

/// `SampleBuilder` is used to analyze columns. It collects sample from
/// the result set using Reservoir Sampling algorithm, estimates NDVs
/// using FM Sketch during the collecting process, and builds count-min sketch.
impl SampleBuilder {
    fn new(
        mut req: AnalyzeColumnsReq,
        snap: SnapshotStore,
        ranges: Vec<KeyRange>,
    ) -> Result<SampleBuilder> {
        let cols_info = req.take_columns_info();
        if cols_info.is_empty() {
            return Err(box_err!("empty columns_info"));
        }

        let mut col_len = cols_info.len();
        if cols_info[0].get_pk_handle() {
            col_len -= 1;
        }

        let mut meta = TableScan::new();
        meta.set_columns(cols_info);
        let table_scanner = TableScanExecutor::new(&meta, ranges, snap, false)?;
        Ok(SampleBuilder {
            data: table_scanner,
            cols: meta.take_columns().to_vec(),
            col_len: col_len,
            max_bucket_size: req.get_bucket_size() as usize,
            max_fm_sketch_size: req.get_sketch_size() as usize,
            max_sample_size: req.get_sample_size() as usize,
            cm_sketch_depth: req.get_cmsketch_depth() as usize,
            cm_sketch_width: req.get_cmsketch_width() as usize,
        })
    }

    // `collect_columns_stats` returns the sample collectors which contain total count,
    // null count, distinct values count and count-min sketch. And it also returns the statistic
    // builder for PK which contains the histogram.
    // See https://en.wikipedia.org/wiki/Reservoir_sampling
    fn collect_columns_stats(&mut self) -> Result<(Vec<SampleCollector>, Histogram)> {
        let mut pk_builder = Histogram::new(self.max_bucket_size);
        let mut collectors = vec![
            SampleCollector::new(
                self.max_sample_size,
                self.max_fm_sketch_size,
                self.cm_sketch_depth,
                self.cm_sketch_width
            );
            self.col_len
        ];
        while let Some(row) = self.data.next()? {
            let cols = row.get_binary_cols(&self.cols)?;
            let retrieve_len = cols.len();
            let mut cols_iter = cols.into_iter();
            if self.col_len != retrieve_len {
                if let Some(v) = cols_iter.next() {
                    pk_builder.append(&v);
                }
            }
            for (collector, val) in collectors.iter_mut().zip(cols_iter) {
                collector.collect(val);
            }
        }
        Ok((collectors, pk_builder))
    }
}

/// `SampleCollector` will collect Samples and calculate the count, ndv and total size of an attribute.
#[derive(Clone)]
struct SampleCollector {
    samples: Vec<Vec<u8>>,
    null_count: u64,
    count: u64,
    max_sample_size: usize,
    fm_sketch: FMSketch,
    cm_sketch: Option<CMSketch>,
    rng: ThreadRng,
    total_size: u64,
}

impl SampleCollector {
    fn new(
        max_sample_size: usize,
        max_fm_sketch_size: usize,
        cm_sketch_depth: usize,
        cm_sketch_width: usize,
    ) -> SampleCollector {
        SampleCollector {
            samples: Default::default(),
            null_count: 0,
            count: 0,
            max_sample_size,
            fm_sketch: FMSketch::new(max_fm_sketch_size),
            cm_sketch: CMSketch::new(cm_sketch_depth, cm_sketch_width),
            rng: thread_rng(),
            total_size: 0,
        }
    }

    fn into_proto(self) -> analyze::SampleCollector {
        let mut s = analyze::SampleCollector::new();
        s.set_null_count(self.null_count as i64);
        s.set_count(self.count as i64);
        s.set_fm_sketch(self.fm_sketch.into_proto());
        s.set_samples(RepeatedField::from_vec(self.samples));
        if let Some(c) = self.cm_sketch {
            s.set_cm_sketch(c.into_proto())
        }
        s.set_total_size(self.total_size as i64);
        s
    }

    pub fn collect(&mut self, data: Vec<u8>) {
        if data[0] == datum::NIL_FLAG {
            self.null_count += 1;
            return;
        }
        self.count += 1;
        self.fm_sketch.insert(&data);
        if let Some(c) = self.cm_sketch.as_mut() {
            c.insert(&data)
        }
        self.total_size += data.len() as u64;
        if self.samples.len() < self.max_sample_size {
            self.samples.push(data);
            return;
        }
        if self.rng.gen_range(0, self.count) < self.max_sample_size as u64 {
            let idx = self.rng.gen_range(0, self.max_sample_size);
            self.samples[idx] = data;
        }
    }
}

#[cfg(test)]
mod test {
    use coprocessor::codec::datum;
    use coprocessor::codec::datum::Datum;
    use super::*;

    #[test]
    fn test_sample_collector() {
        let max_sample_size = 3;
        let max_fm_sketch_size = 10;
        let cm_sketch_depth = 2;
        let cm_sketch_width = 16;
        let mut sample = SampleCollector::new(
            max_sample_size,
            max_fm_sketch_size,
            cm_sketch_depth,
            cm_sketch_width,
        );
        let cases = vec![Datum::I64(1), Datum::Null, Datum::I64(2), Datum::I64(5)];

        for data in cases {
            sample.collect(datum::encode_value(&[data]).unwrap());
        }
        assert_eq!(sample.samples.len(), max_sample_size);
        assert_eq!(sample.null_count, 1);
        assert_eq!(sample.count, 3);
        assert_eq!(sample.cm_sketch.unwrap().count(), 3);
        assert_eq!(sample.total_size, 6)
    }
}
<|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 dom::comment::Comment;
use dom::bindings::codegen::DocumentBinding;
use dom::bindings::utils::{Reflectable, Reflector, Traceable, reflect_dom_object};
use dom::bindings::utils::{ErrorResult, Fallible, NotSupported, InvalidCharacter};
use dom::bindings::utils::DOMString;
use dom::bindings::utils::{xml_name_type, InvalidXMLName};
use dom::documentfragment::DocumentFragment;
use dom::element::{Element};
use dom::element::{HTMLHeadElementTypeId, HTMLTitleElementTypeId};
use dom::event::{AbstractEvent, Event};
use dom::htmlcollection::HTMLCollection;
use dom::htmldocument::HTMLDocument;
use dom::mouseevent::MouseEvent;
use dom::node::{AbstractNode, ScriptView, Node, ElementNodeTypeId, DocumentNodeTypeId};
use dom::text::Text;
use dom::uievent::UIEvent;
use dom::window::Window;
use dom::htmltitleelement::HTMLTitleElement;
use html::hubbub_html_parser::build_element_from_tag;
use js::jsapi::{JSObject, JSContext, JSTracer};
use servo_util::tree::{TreeNodeRef, ElementLike};

use std::hashmap::HashMap;

use std::cast;
use std::str::eq_slice;
use std::ascii::StrAsciiExt;
use std::unstable::raw::Box;

#[deriving(Eq)]
pub enum DocumentTypeId {
    PlainDocumentTypeId,
    HTMLDocumentTypeId
}

#[deriving(Eq)]
pub struct AbstractDocument {
    document: *mut Box<Document>
}

impl AbstractDocument {
    pub fn document<'a>(&'a self) -> &'a Document {
        unsafe {
            &(*self.document).data
        }
    }

    pub fn mut_document<'a>(&'a self) -> &'a mut Document {
        unsafe {
            &mut (*self.document).data
        }
    }

    unsafe fn transmute<T, R>(&self, f: &fn(&T) -> R) -> R {
        let box: *Box<T> = cast::transmute(self.document);
        f(&(*box).data)
    }

    unsafe fn transmute_mut<T, R>(&self, f: &fn(&mut T) -> R) -> R {
        let box: *mut Box<T> = cast::transmute(self.document);
        f(&mut (*box).data)
    }

    pub fn with_html<R>(&self, callback: &fn(&HTMLDocument) -> R) -> R {
        match self.document().doctype {
            HTML => unsafe { self.transmute(callback) },
            _ => fail!("attempt to downcast a non-HTMLDocument to HTMLDocument")
        }
    }

    pub fn from_box<T>(ptr: *mut Box<T>) -> AbstractDocument {
        AbstractDocument {
            document: ptr as *mut Box<Document>
        }
    }
}

#[deriving(Eq)]
pub enum DocumentType {
    HTML,
    SVG,
    XML
}

pub struct Document {
    node: Node<ScriptView>,
    reflector_: Reflector,
    window: @mut Window,
    doctype: DocumentType,
    title: ~str,
    idmap: HashMap<~str, AbstractNode<ScriptView>>
}

impl Document {
    pub fn reflect_document<D: Reflectable>
            (document:  @mut D,
             window:    @mut Window,
             wrap_fn:   extern "Rust" fn(*JSContext, *JSObject, @mut D) -> *JSObject)
             -> AbstractDocument {
        assert!(document.reflector().get_jsobject().is_null());
        let document = reflect_dom_object(document, window, wrap_fn);
        assert!(document.reflector().get_jsobject().is_not_null());

        // This surrenders memory management of the document!
        let abstract = AbstractDocument {
            document: unsafe { cast::transmute(document) }
        };
        abstract.mut_document().node.set_owner_doc(abstract);
        abstract
    }

    pub fn new_inherited(window: @mut Window, doctype: DocumentType) -> Document {
        let node_type = match doctype {
            HTML => HTMLDocumentTypeId,
            SVG | XML => PlainDocumentTypeId
        };
        Document {
            node: Node::new_without_doc(DocumentNodeTypeId(node_type)),
            reflector_: Reflector::new(),
            window: window,
            doctype: doctype,
            title: ~"",
            idmap: HashMap::new()
        }
    }

    pub fn new(window: @mut Window, doctype: DocumentType) -> AbstractDocument {
        let document = Document::new_inherited(window, doctype);
        Document::reflect_document(@mut document, window, DocumentBinding::Wrap)
    }
}

impl Document {
    pub fn Constructor(owner: @mut Window) -> Fallible<AbstractDocument> {
        Ok(Document::new(owner, XML))
    }
}

impl Reflectable for AbstractDocument {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        self.document().reflector()
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        self.mut_document().mut_reflector()
    }
}

impl Reflectable for Document {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        self.node.reflector()
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        self.node.mut_reflector()
    }
}

impl Document {
    pub fn GetDocumentElement(&self) -> Option<AbstractNode<ScriptView>> {
        do self.node.children().find |c| {
            c.is_element()
        }
    }

    fn get_cx(&self) -> *JSContext {
        self.window.get_cx()
    }

    pub fn GetElementsByTagName(&self, tag: DOMString) -> @mut HTMLCollection {
        self.createHTMLCollection(|elem| eq_slice(elem.tag_name, tag))
    }

    pub fn GetElementsByTagNameNS(&self, _ns: Option<DOMString>, _tag: DOMString) -> @mut HTMLCollection {
        HTMLCollection::new(self.window, ~[])
    }

    pub fn GetElementsByClassName(&self, _class: DOMString) -> @mut HTMLCollection {
        HTMLCollection::new(self.window, ~[])
    }

    pub fn GetElementById(&self, id: DOMString) -> Option<AbstractNode<ScriptView>> {
        // TODO: "in tree order, within the context object's tree"
        // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
        match self.idmap.find_equiv(&id) {
            None => None,
            Some(node) => Some(*node),
        }
    }

    pub fn CreateElement(&self, abstract_self: AbstractDocument, local_name: DOMString) -> Fallible<AbstractNode<ScriptView>> {
        if xml_name_type(local_name) == InvalidXMLName {
            return Err(InvalidCharacter);
        }
        let local_name = local_name.to_ascii_lower();
        Ok(build_element_from_tag(local_name, abstract_self))
    }

    pub fn CreateDocumentFragment(&self, abstract_self: AbstractDocument) -> AbstractNode<ScriptView> {
        DocumentFragment::new(abstract_self)
    }

    pub fn CreateTextNode(&self, abstract_self: AbstractDocument, data: DOMString) -> AbstractNode<ScriptView> {
        Text::new(data, abstract_self)
    }

    pub fn CreateComment(&self, abstract_self: AbstractDocument, data: DOMString) -> AbstractNode<ScriptView> {
        Comment::new(data, abstract_self)
    }

    pub fn CreateEvent(&self, interface: DOMString) -> Fallible<AbstractEvent> {
        match interface.as_slice() {
            "UIEvents" => Ok(UIEvent::new(self.window)),
            "MouseEvents" => Ok(MouseEvent::new(self.window)),
            "HTMLEvents" => Ok(Event::new(self.window)),
            _ => Err(NotSupported)
        }
    }

    pub fn Title(&self, _: AbstractDocument) -> DOMString {
        let mut title = ~"";
        match self.doctype {
            SVG => {
                fail!("no SVG document yet")
            },
            _ => {
                match self.GetDocumentElement() {
                    None => {},
                    Some(root) => {
                        for node in root.traverse_preorder() {
                            if node.type_id() != ElementNodeTypeId(HTMLTitleElementTypeId) {
                                continue;
                            }
                            for child in node.children() {
                                if child.is_text() {
                                    do child.with_imm_text() |text| {
                                        title = title + text.element.Data();
                                    }
                                }
                            }
                            break;
                        }
                    }
                }
            }
        }
        let v: ~[&str] = title.word_iter().collect();
        title = v.connect(" ");
        title = title.trim().to_owned();
        title
    }

    pub fn SetTitle(&self, abstract_self: AbstractDocument, title: DOMString) -> ErrorResult {
        match self.doctype {
            SVG => {
                fail!("no SVG document yet")
            },
            _ => {
                match self.GetDocumentElement() {
                    None => {},
                    Some(root) => {
                        for node in root.traverse_preorder() {
                            if node.type_id() != ElementNodeTypeId(HTMLHeadElementTypeId) {
                                continue;
                            }
                            let mut has_title = false;
                            for child in node.children() {
                                if child.type_id() != ElementNodeTypeId(HTMLTitleElementTypeId) {
                                    continue;
                                }
                                has_title = true;
                                for title_child in child.children() {
                                    child.remove_child(title_child);
                                }
                                child.AppendChild(self.CreateTextNode(abstract_self, title.clone()));
                                break;
                            }
                            if !has_title {
                                let new_title = HTMLTitleElement::new(~"title", abstract_self);
                                new_title.AppendChild(self.CreateTextNode(abstract_self, title.clone()));
                                node.AppendChild(new_title);
                            }
                            break;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub fn GetElementsByName(&self, name: DOMString) -> @mut HTMLCollection {
        self.createHTMLCollection(|elem|
            elem.get_attr("name").is_some() && eq_slice(elem.get_attr("name").unwrap(), name))
    }

    pub fn createHTMLCollection(&self, callback: &fn(elem: &Element) -> bool) -> @mut HTMLCollection {
        let mut elements = ~[];
        match self.GetDocumentElement() {
            None => {},
            Some(root) => {
                for child in root.traverse_preorder() {
                    if child.is_element() {
                        do child.with_imm_element |elem| {
                            if callback(elem) {
                                elements.push(child);
                            }
                        }
                    }
                }
            }
        }
        HTMLCollection::new(self.window, elements)
    }

    pub fn content_changed(&self) {
        self.window.content_changed();
    }

    pub fn wait_until_safe_to_modify_dom(&self) {
        self.window.wait_until_safe_to_modify_dom();
    }

    pub fn register_nodes_with_id(&mut self, root: &AbstractNode<ScriptView>) {
        foreach_ided_elements(root, |id: &~str, abstract_node: &AbstractNode<ScriptView>| {
            // TODO: "in tree order, within the context object's tree"
            // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
            self.idmap.find_or_insert(id.clone(), *abstract_node);
        });
    }

    pub fn unregister_nodes_with_id(&mut self, root: &AbstractNode<ScriptView>) {
        foreach_ided_elements(root, |id: &~str, _| {
            // TODO: "in tree order, within the context object's tree"
            // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
            self.idmap.pop(id);
        });
    }

    pub fn update_idmap(&mut self,
                        abstract_self: AbstractNode<ScriptView>,
                        new_id: DOMString,
                        old_id: Option<DOMString>) {
        // remove old ids if the old ones are not same as the new one.
        match old_id {
            Some(ref old_id) if new_id != *old_id => {
                self.idmap.remove(old_id);
            }
            _ => ()
        }

        // TODO: support the case if multiple elements which haves same id are in the same document.
        self.idmap.mangle(new_id, abstract_self,
                         |_, new_node: AbstractNode<ScriptView>| -> AbstractNode<ScriptView> {
                             new_node
                         },
                         |_, old_node: &mut AbstractNode<ScriptView>, new_node: AbstractNode<ScriptView>| {
                             *old_node = new_node;
                         });
    }
}

#[inline(always)]
fn foreach_ided_elements(root: &AbstractNode<ScriptView>,
                         callback: &fn(&~str, &AbstractNode<ScriptView>)) {
    for node in root.traverse_preorder() {
        if !node.is_element() {
            continue;
        }

        do node.with_imm_element |element| {
            match element.get_attr("id") {
                Some(id) => {
                    callback(&id.to_str(), &node);
                }
                None => ()
            }
        }
    }
}

impl Traceable for Document {
    #[fixed_stack_segment]
    fn trace(&self, tracer: *mut JSTracer) {
        self.node.trace(tracer);
    }
}
<commit_msg>Use 'DOMString' to define the type insteand of '~str'.<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 dom::comment::Comment;
use dom::bindings::codegen::DocumentBinding;
use dom::bindings::utils::{Reflectable, Reflector, Traceable, reflect_dom_object};
use dom::bindings::utils::{ErrorResult, Fallible, NotSupported, InvalidCharacter};
use dom::bindings::utils::DOMString;
use dom::bindings::utils::{xml_name_type, InvalidXMLName};
use dom::documentfragment::DocumentFragment;
use dom::element::{Element};
use dom::element::{HTMLHeadElementTypeId, HTMLTitleElementTypeId};
use dom::event::{AbstractEvent, Event};
use dom::htmlcollection::HTMLCollection;
use dom::htmldocument::HTMLDocument;
use dom::mouseevent::MouseEvent;
use dom::node::{AbstractNode, ScriptView, Node, ElementNodeTypeId, DocumentNodeTypeId};
use dom::text::Text;
use dom::uievent::UIEvent;
use dom::window::Window;
use dom::htmltitleelement::HTMLTitleElement;
use html::hubbub_html_parser::build_element_from_tag;
use js::jsapi::{JSObject, JSContext, JSTracer};
use servo_util::tree::{TreeNodeRef, ElementLike};

use std::hashmap::HashMap;

use std::cast;
use std::str::eq_slice;
use std::ascii::StrAsciiExt;
use std::unstable::raw::Box;

#[deriving(Eq)]
pub enum DocumentTypeId {
    PlainDocumentTypeId,
    HTMLDocumentTypeId
}

#[deriving(Eq)]
pub struct AbstractDocument {
    document: *mut Box<Document>
}

impl AbstractDocument {
    pub fn document<'a>(&'a self) -> &'a Document {
        unsafe {
            &(*self.document).data
        }
    }

    pub fn mut_document<'a>(&'a self) -> &'a mut Document {
        unsafe {
            &mut (*self.document).data
        }
    }

    unsafe fn transmute<T, R>(&self, f: &fn(&T) -> R) -> R {
        let box: *Box<T> = cast::transmute(self.document);
        f(&(*box).data)
    }

    unsafe fn transmute_mut<T, R>(&self, f: &fn(&mut T) -> R) -> R {
        let box: *mut Box<T> = cast::transmute(self.document);
        f(&mut (*box).data)
    }

    pub fn with_html<R>(&self, callback: &fn(&HTMLDocument) -> R) -> R {
        match self.document().doctype {
            HTML => unsafe { self.transmute(callback) },
            _ => fail!("attempt to downcast a non-HTMLDocument to HTMLDocument")
        }
    }

    pub fn from_box<T>(ptr: *mut Box<T>) -> AbstractDocument {
        AbstractDocument {
            document: ptr as *mut Box<Document>
        }
    }
}

#[deriving(Eq)]
pub enum DocumentType {
    HTML,
    SVG,
    XML
}

pub struct Document {
    node: Node<ScriptView>,
    reflector_: Reflector,
    window: @mut Window,
    doctype: DocumentType,
    title: ~str,
    idmap: HashMap<DOMString, AbstractNode<ScriptView>>
}

impl Document {
    pub fn reflect_document<D: Reflectable>
            (document:  @mut D,
             window:    @mut Window,
             wrap_fn:   extern "Rust" fn(*JSContext, *JSObject, @mut D) -> *JSObject)
             -> AbstractDocument {
        assert!(document.reflector().get_jsobject().is_null());
        let document = reflect_dom_object(document, window, wrap_fn);
        assert!(document.reflector().get_jsobject().is_not_null());

        // This surrenders memory management of the document!
        let abstract = AbstractDocument {
            document: unsafe { cast::transmute(document) }
        };
        abstract.mut_document().node.set_owner_doc(abstract);
        abstract
    }

    pub fn new_inherited(window: @mut Window, doctype: DocumentType) -> Document {
        let node_type = match doctype {
            HTML => HTMLDocumentTypeId,
            SVG | XML => PlainDocumentTypeId
        };
        Document {
            node: Node::new_without_doc(DocumentNodeTypeId(node_type)),
            reflector_: Reflector::new(),
            window: window,
            doctype: doctype,
            title: ~"",
            idmap: HashMap::new()
        }
    }

    pub fn new(window: @mut Window, doctype: DocumentType) -> AbstractDocument {
        let document = Document::new_inherited(window, doctype);
        Document::reflect_document(@mut document, window, DocumentBinding::Wrap)
    }
}

impl Document {
    pub fn Constructor(owner: @mut Window) -> Fallible<AbstractDocument> {
        Ok(Document::new(owner, XML))
    }
}

impl Reflectable for AbstractDocument {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        self.document().reflector()
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        self.mut_document().mut_reflector()
    }
}

impl Reflectable for Document {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        self.node.reflector()
    }

    fn mut_reflector<'a>(&'a mut self) -> &'a mut Reflector {
        self.node.mut_reflector()
    }
}

impl Document {
    pub fn GetDocumentElement(&self) -> Option<AbstractNode<ScriptView>> {
        do self.node.children().find |c| {
            c.is_element()
        }
    }

    fn get_cx(&self) -> *JSContext {
        self.window.get_cx()
    }

    pub fn GetElementsByTagName(&self, tag: DOMString) -> @mut HTMLCollection {
        self.createHTMLCollection(|elem| eq_slice(elem.tag_name, tag))
    }

    pub fn GetElementsByTagNameNS(&self, _ns: Option<DOMString>, _tag: DOMString) -> @mut HTMLCollection {
        HTMLCollection::new(self.window, ~[])
    }

    pub fn GetElementsByClassName(&self, _class: DOMString) -> @mut HTMLCollection {
        HTMLCollection::new(self.window, ~[])
    }

    pub fn GetElementById(&self, id: DOMString) -> Option<AbstractNode<ScriptView>> {
        // TODO: "in tree order, within the context object's tree"
        // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
        match self.idmap.find_equiv(&id) {
            None => None,
            Some(node) => Some(*node),
        }
    }

    pub fn CreateElement(&self, abstract_self: AbstractDocument, local_name: DOMString) -> Fallible<AbstractNode<ScriptView>> {
        if xml_name_type(local_name) == InvalidXMLName {
            return Err(InvalidCharacter);
        }
        let local_name = local_name.to_ascii_lower();
        Ok(build_element_from_tag(local_name, abstract_self))
    }

    pub fn CreateDocumentFragment(&self, abstract_self: AbstractDocument) -> AbstractNode<ScriptView> {
        DocumentFragment::new(abstract_self)
    }

    pub fn CreateTextNode(&self, abstract_self: AbstractDocument, data: DOMString) -> AbstractNode<ScriptView> {
        Text::new(data, abstract_self)
    }

    pub fn CreateComment(&self, abstract_self: AbstractDocument, data: DOMString) -> AbstractNode<ScriptView> {
        Comment::new(data, abstract_self)
    }

    pub fn CreateEvent(&self, interface: DOMString) -> Fallible<AbstractEvent> {
        match interface.as_slice() {
            "UIEvents" => Ok(UIEvent::new(self.window)),
            "MouseEvents" => Ok(MouseEvent::new(self.window)),
            "HTMLEvents" => Ok(Event::new(self.window)),
            _ => Err(NotSupported)
        }
    }

    pub fn Title(&self, _: AbstractDocument) -> DOMString {
        let mut title = ~"";
        match self.doctype {
            SVG => {
                fail!("no SVG document yet")
            },
            _ => {
                match self.GetDocumentElement() {
                    None => {},
                    Some(root) => {
                        for node in root.traverse_preorder() {
                            if node.type_id() != ElementNodeTypeId(HTMLTitleElementTypeId) {
                                continue;
                            }
                            for child in node.children() {
                                if child.is_text() {
                                    do child.with_imm_text() |text| {
                                        title = title + text.element.Data();
                                    }
                                }
                            }
                            break;
                        }
                    }
                }
            }
        }
        let v: ~[&str] = title.word_iter().collect();
        title = v.connect(" ");
        title = title.trim().to_owned();
        title
    }

    pub fn SetTitle(&self, abstract_self: AbstractDocument, title: DOMString) -> ErrorResult {
        match self.doctype {
            SVG => {
                fail!("no SVG document yet")
            },
            _ => {
                match self.GetDocumentElement() {
                    None => {},
                    Some(root) => {
                        for node in root.traverse_preorder() {
                            if node.type_id() != ElementNodeTypeId(HTMLHeadElementTypeId) {
                                continue;
                            }
                            let mut has_title = false;
                            for child in node.children() {
                                if child.type_id() != ElementNodeTypeId(HTMLTitleElementTypeId) {
                                    continue;
                                }
                                has_title = true;
                                for title_child in child.children() {
                                    child.remove_child(title_child);
                                }
                                child.AppendChild(self.CreateTextNode(abstract_self, title.clone()));
                                break;
                            }
                            if !has_title {
                                let new_title = HTMLTitleElement::new(~"title", abstract_self);
                                new_title.AppendChild(self.CreateTextNode(abstract_self, title.clone()));
                                node.AppendChild(new_title);
                            }
                            break;
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub fn GetElementsByName(&self, name: DOMString) -> @mut HTMLCollection {
        self.createHTMLCollection(|elem|
            elem.get_attr("name").is_some() && eq_slice(elem.get_attr("name").unwrap(), name))
    }

    pub fn createHTMLCollection(&self, callback: &fn(elem: &Element) -> bool) -> @mut HTMLCollection {
        let mut elements = ~[];
        match self.GetDocumentElement() {
            None => {},
            Some(root) => {
                for child in root.traverse_preorder() {
                    if child.is_element() {
                        do child.with_imm_element |elem| {
                            if callback(elem) {
                                elements.push(child);
                            }
                        }
                    }
                }
            }
        }
        HTMLCollection::new(self.window, elements)
    }

    pub fn content_changed(&self) {
        self.window.content_changed();
    }

    pub fn wait_until_safe_to_modify_dom(&self) {
        self.window.wait_until_safe_to_modify_dom();
    }

    pub fn register_nodes_with_id(&mut self, root: &AbstractNode<ScriptView>) {
        foreach_ided_elements(root, |id: &DOMString, abstract_node: &AbstractNode<ScriptView>| {
            // TODO: "in tree order, within the context object's tree"
            // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
            self.idmap.find_or_insert(id.clone(), *abstract_node);
        });
    }

    pub fn unregister_nodes_with_id(&mut self, root: &AbstractNode<ScriptView>) {
        foreach_ided_elements(root, |id: &DOMString, _| {
            // TODO: "in tree order, within the context object's tree"
            // http://dom.spec.whatwg.org/#dom-document-getelementbyid.
            self.idmap.pop(id);
        });
    }

    pub fn update_idmap(&mut self,
                        abstract_self: AbstractNode<ScriptView>,
                        new_id: DOMString,
                        old_id: Option<DOMString>) {
        // remove old ids if the old ones are not same as the new one.
        match old_id {
            Some(ref old_id) if new_id != *old_id => {
                self.idmap.remove(old_id);
            }
            _ => ()
        }

        // TODO: support the case if multiple elements which haves same id are in the same document.
        self.idmap.mangle(new_id, abstract_self,
                         |_, new_node: AbstractNode<ScriptView>| -> AbstractNode<ScriptView> {
                             new_node
                         },
                         |_, old_node: &mut AbstractNode<ScriptView>, new_node: AbstractNode<ScriptView>| {
                             *old_node = new_node;
                         });
    }
}

#[inline(always)]
fn foreach_ided_elements(root: &AbstractNode<ScriptView>,
                         callback: &fn(&DOMString, &AbstractNode<ScriptView>)) {
    for node in root.traverse_preorder() {
        if !node.is_element() {
            continue;
        }

        do node.with_imm_element |element| {
            match element.get_attr("id") {
                Some(id) => {
                    callback(&id.to_str(), &node);
                }
                None => ()
            }
        }
    }
}

impl Traceable for Document {
    #[fixed_stack_segment]
    fn trace(&self, tracer: *mut JSTracer) {
        self.node.trace(tracer);
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::os::unix::fs::symlink;
use std::io::{stdout, stderr};
use std::path::Path;

use argparse::{ArgumentParser};
use unshare::{Command};

use config::command::CapsuleInfo;

use super::setup;
use super::Wrapper;
use super::util::{find_cmd};
use file_util::Dir;
use process_util::{run_and_wait, convert_status, copy_env_vars};

const BUSYBOX_COMMANDS: &'static [&'static str] = &[
    "[", "[[", "acpid", "add-shell", "addgroup", "adduser", "adjtimex", "arp",
    "arping", "ash", "awk", "base64", "basename", "bbconfig", "beep",
    "blkdiscard", "blkid", "blockdev", "brctl", "bunzip2", "bzcat", "bzip2",
    "cal", "cat", "catv", "chgrp", "chmod", "chown", "chpasswd", "chroot",
    "chvt", "cksum", "clear", "cmp", "comm", "conspy", "cp", "cpio", "crond",
    "crontab", "cryptpw", "cut", "date", "dc", "dd", "deallocvt", "delgroup",
    "deluser", "depmod", "df", "diff", "dirname", "dmesg", "dnsd",
    "dnsdomainname", "dos2unix", "du", "dumpkmap", "dumpleases", "echo", "ed",
    "egrep", "eject", "env", "ether-wake", "expand", "expr", "fakeidentd",
    "false", "fatattr", "fbset", "fbsplash", "fdflush", "fdformat", "fdisk",
    "fgrep", "find", "findfs", "flock", "fold", "free", "fsck", "fstrim",
    "fsync", "ftpd", "ftpget", "ftpput", "fuser", "getopt", "getty", "grep",
    "groups", "gunzip", "gzip", "halt", "hd", "hdparm", "head", "hexdump",
    "hostid", "hostname", "httpd", "hwclock", "id", "ifconfig", "ifdown",
    "ifenslave", "ifup", "inetd", "init", "inotifyd", "insmod", "install",
    "ionice", "iostat", "ip", "ipaddr", "ipcalc", "ipcrm", "ipcs", "iplink",
    "iproute", "iprule", "iptunnel", "kbd_mode", "kill", "killall",
    "killall5", "klogd", "less", "ln", "loadfont", "loadkmap", "logger",
    "login", "logread", "losetup", "ls", "lsmod", "lsof", "lspci", "lsusb",
    "lzcat", "lzma", "lzop", "lzopcat", "makemime", "md5sum", "mdev", "mesg",
    "microcom", "mkdir", "mkdosfs", "mkfifo", "mkfs.vfat", "mknod",
    "mkpasswd", "mkswap", "mktemp", "modinfo", "modprobe", "more", "mount",
    "mountpoint", "mpstat", "mv", "nameif", "nanddump", "nandwrite",
    "nbd-client", "nc", "netstat", "nice", "nmeter", "nohup", "nologin",
    "nsenter", "nslookup", "ntpd", "od", "openvt", "passwd", "patch", "pgrep",
    "pidof", "ping", "ping6", "pipe_progress", "pkill", "pmap", "poweroff",
    "powertop", "printenv", "printf", "ps", "pscan", "pstree", "pwd", "pwdx",
    "raidautorun", "rdate", "rdev", "readahead", "readlink", "readprofile",
    "realpath", "reboot", "reformime", "remove-shell", "renice", "reset",
    "resize", "rev", "rfkill", "rm", "rmdir", "rmmod", "route", "run-parts",
    "sed", "sendmail", "seq", "setconsole", "setfont", "setkeycodes",
    "setlogcons", "setserial", "setsid", "sh", "sha1sum", "sha256sum",
    "sha3sum", "sha512sum", "showkey", "shuf", "slattach", "sleep", "smemcap",
    "sort", "split", "stat", "strings", "stty", "su", "sum", "swapoff",
    "swapon", "switch_root", "sync", "sysctl", "syslogd", "tac", "tail",
    "tar", "tee", "telnet", "test", "tftp", "time", "timeout", "top", "touch",
    "tr", "traceroute", "traceroute6", "true", "truncate", "tty", "ttysize",
    "tunctl", "udhcpc", "udhcpc6", "udhcpd", "umount", "uname", "unexpand",
    "uniq", "unix2dos", "unlink", "unlzma", "unlzop", "unshare", "unxz",
    "unzip", "uptime", "usleep", "uudecode", "uuencode", "vconfig", "vi",
    "vlock", "volname", "watch", "watchdog", "wc", "wget", "which", "whoami",
    "whois", "xargs", "xzcat", "yes", "zcat"
];

pub fn symlink_busybox_commands() -> Result<(), String> {
    Dir::new("/bin").create()
        .map_err(|e| format!("Can't create /bin: {}", e))?;
    for cmd in BUSYBOX_COMMANDS {
        symlink("/vagga/bin/busybox", format!("/bin/{}", cmd))
            .map_err(|e| format!("Error symlinking {:?}: {}", cmd, e))?;
    }
    Ok(())
}

pub fn commandline_cmd(_cmd_name: &str, command: &CapsuleInfo,
    wrapper: &Wrapper, mut cmdline: Vec<String>)
    -> Result<i32, String>
{
    if command.run.len() == 0 {
        return Err(format!(
            r#"Command has empty "run" parameter. Nothing to run."#));
    }
    // TODO(tailhook) detect other shells too
    let has_args = command.accepts_arguments
            .unwrap_or(&command.run[0][..] != "/bin/sh");
    let mut args = Vec::new();
    if !has_args {
        let mut ap = ArgumentParser::new();
        ap.set_description(command.description.as_ref()
            .map(|x| &x[..]).unwrap_or(""));
        ap.stop_on_first_argument(true);
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    } else {
        cmdline.remove(0);
        args.extend(cmdline.into_iter());
    }
    let mut cmdline = command.run.clone();
    cmdline.extend(args.into_iter());

    setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings)?;

    symlink_busybox_commands()?;

    let env = setup::get_capsule_environment(&wrapper.settings, command)?;
    let cpath = find_cmd(&cmdline.remove(0), &env)?;

    let mut cmd = Command::new(&cpath);
    cmd.args(&cmdline);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &wrapper.settings);
    if let Some(ref wd) = command.work_dir {
        cmd.current_dir(Path::new("/work").join(&wd));
    } else {
        cmd.current_dir(env::var("_VAGGA_WORKDIR")
                        .unwrap_or("/work".to_string()));
    }
    for (ref k, ref v) in env.iter() {
        cmd.env(k, v);
    }

    return run_and_wait(&mut cmd).map(convert_status);
}
<commit_msg>Propagate `VAGGA_SETTINGS` into a capsule<commit_after>use std::collections::BTreeMap;
use std::env;
use std::io::{stdout, stderr};
use std::os::unix::fs::symlink;
use std::path::Path;

use argparse::{ArgumentParser};
use unshare::{Command};
use rustc_serialize::json;

use config::command::CapsuleInfo;

use super::setup;
use super::Wrapper;
use super::util::{find_cmd};
use file_util::Dir;
use process_util::{run_and_wait, convert_status, copy_env_vars};

const BUSYBOX_COMMANDS: &'static [&'static str] = &[
    "[", "[[", "acpid", "add-shell", "addgroup", "adduser", "adjtimex", "arp",
    "arping", "ash", "awk", "base64", "basename", "bbconfig", "beep",
    "blkdiscard", "blkid", "blockdev", "brctl", "bunzip2", "bzcat", "bzip2",
    "cal", "cat", "catv", "chgrp", "chmod", "chown", "chpasswd", "chroot",
    "chvt", "cksum", "clear", "cmp", "comm", "conspy", "cp", "cpio", "crond",
    "crontab", "cryptpw", "cut", "date", "dc", "dd", "deallocvt", "delgroup",
    "deluser", "depmod", "df", "diff", "dirname", "dmesg", "dnsd",
    "dnsdomainname", "dos2unix", "du", "dumpkmap", "dumpleases", "echo", "ed",
    "egrep", "eject", "env", "ether-wake", "expand", "expr", "fakeidentd",
    "false", "fatattr", "fbset", "fbsplash", "fdflush", "fdformat", "fdisk",
    "fgrep", "find", "findfs", "flock", "fold", "free", "fsck", "fstrim",
    "fsync", "ftpd", "ftpget", "ftpput", "fuser", "getopt", "getty", "grep",
    "groups", "gunzip", "gzip", "halt", "hd", "hdparm", "head", "hexdump",
    "hostid", "hostname", "httpd", "hwclock", "id", "ifconfig", "ifdown",
    "ifenslave", "ifup", "inetd", "init", "inotifyd", "insmod", "install",
    "ionice", "iostat", "ip", "ipaddr", "ipcalc", "ipcrm", "ipcs", "iplink",
    "iproute", "iprule", "iptunnel", "kbd_mode", "kill", "killall",
    "killall5", "klogd", "less", "ln", "loadfont", "loadkmap", "logger",
    "login", "logread", "losetup", "ls", "lsmod", "lsof", "lspci", "lsusb",
    "lzcat", "lzma", "lzop", "lzopcat", "makemime", "md5sum", "mdev", "mesg",
    "microcom", "mkdir", "mkdosfs", "mkfifo", "mkfs.vfat", "mknod",
    "mkpasswd", "mkswap", "mktemp", "modinfo", "modprobe", "more", "mount",
    "mountpoint", "mpstat", "mv", "nameif", "nanddump", "nandwrite",
    "nbd-client", "nc", "netstat", "nice", "nmeter", "nohup", "nologin",
    "nsenter", "nslookup", "ntpd", "od", "openvt", "passwd", "patch", "pgrep",
    "pidof", "ping", "ping6", "pipe_progress", "pkill", "pmap", "poweroff",
    "powertop", "printenv", "printf", "ps", "pscan", "pstree", "pwd", "pwdx",
    "raidautorun", "rdate", "rdev", "readahead", "readlink", "readprofile",
    "realpath", "reboot", "reformime", "remove-shell", "renice", "reset",
    "resize", "rev", "rfkill", "rm", "rmdir", "rmmod", "route", "run-parts",
    "sed", "sendmail", "seq", "setconsole", "setfont", "setkeycodes",
    "setlogcons", "setserial", "setsid", "sh", "sha1sum", "sha256sum",
    "sha3sum", "sha512sum", "showkey", "shuf", "slattach", "sleep", "smemcap",
    "sort", "split", "stat", "strings", "stty", "su", "sum", "swapoff",
    "swapon", "switch_root", "sync", "sysctl", "syslogd", "tac", "tail",
    "tar", "tee", "telnet", "test", "tftp", "time", "timeout", "top", "touch",
    "tr", "traceroute", "traceroute6", "true", "truncate", "tty", "ttysize",
    "tunctl", "udhcpc", "udhcpc6", "udhcpd", "umount", "uname", "unexpand",
    "uniq", "unix2dos", "unlink", "unlzma", "unlzop", "unshare", "unxz",
    "unzip", "uptime", "usleep", "uudecode", "uuencode", "vconfig", "vi",
    "vlock", "volname", "watch", "watchdog", "wc", "wget", "which", "whoami",
    "whois", "xargs", "xzcat", "yes", "zcat"
];

#[derive(RustcEncodable, Clone, Debug)]
pub struct Settings<'a> {
    pub version_check: bool,
    pub proxy_env_vars: bool,
    pub ubuntu_mirror: &'a Option<String>,
    pub alpine_mirror: &'a Option<String>,
    pub build_lock_wait: bool,
    pub environ: &'a BTreeMap<String, String>,
    pub index_all_images: bool,
    pub run_symlinks_as_commands: bool,
}

pub fn symlink_busybox_commands() -> Result<(), String> {
    Dir::new("/bin").create()
        .map_err(|e| format!("Can't create /bin: {}", e))?;
    for cmd in BUSYBOX_COMMANDS {
        symlink("/vagga/bin/busybox", format!("/bin/{}", cmd))
            .map_err(|e| format!("Error symlinking {:?}: {}", cmd, e))?;
    }
    Ok(())
}

pub fn commandline_cmd(_cmd_name: &str, command: &CapsuleInfo,
    wrapper: &Wrapper, mut cmdline: Vec<String>)
    -> Result<i32, String>
{
    if command.run.len() == 0 {
        return Err(format!(
            r#"Command has empty "run" parameter. Nothing to run."#));
    }
    // TODO(tailhook) detect other shells too
    let has_args = command.accepts_arguments
            .unwrap_or(&command.run[0][..] != "/bin/sh");
    let mut args = Vec::new();
    if !has_args {
        let mut ap = ArgumentParser::new();
        ap.set_description(command.description.as_ref()
            .map(|x| &x[..]).unwrap_or(""));
        ap.stop_on_first_argument(true);
        match ap.parse(cmdline, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    } else {
        cmdline.remove(0);
        args.extend(cmdline.into_iter());
    }
    let mut cmdline = command.run.clone();
    cmdline.extend(args.into_iter());

    setup::setup_base_filesystem(
        wrapper.project_root, wrapper.ext_settings)?;

    symlink_busybox_commands()?;

    let env = setup::get_capsule_environment(&wrapper.settings, command)?;
    let cpath = find_cmd(&cmdline.remove(0), &env)?;

    let mut cmd = Command::new(&cpath);
    cmd.args(&cmdline);
    cmd.env_clear();
    copy_env_vars(&mut cmd, &wrapper.settings);
    let ref s = wrapper.settings;
    cmd.env("VAGGA_SETTINGS", json::encode(&Settings {
            version_check: s.version_check,
            proxy_env_vars: s.proxy_env_vars,
            ubuntu_mirror: &s.ubuntu_mirror,
            alpine_mirror: &s.alpine_mirror,
            build_lock_wait: s.build_lock_wait,
            environ: &s.environ,
            index_all_images: s.index_all_images,
            run_symlinks_as_commands: s.run_symlinks_as_commands,
        }).unwrap());
    if let Some(ref wd) = command.work_dir {
        cmd.current_dir(Path::new("/work").join(&wd));
    } else {
        cmd.current_dir(env::var("_VAGGA_WORKDIR")
                        .unwrap_or("/work".to_string()));
    }
    for (ref k, ref v) in env.iter() {
        cmd.env(k, v);
    }

    return run_and_wait(&mut cmd).map(convert_status);
}
<|endoftext|>


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

/// Code to handle management of a pool's thinpool device.

use std::borrow::BorrowMut;
use std::process::Command;

use uuid::Uuid;

use devicemapper as dm;
use devicemapper::{DM, DmDevice, DataBlocks, DmError, LinearDev, MetaBlocks, Sectors, Segment,
                   ThinDev, ThinDevId, ThinPoolDev};
use devicemapper::ErrorEnum::CheckFailed;

use super::super::consts::IEC;
use super::super::engine::{Filesystem, HasName};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::structures::Table;
use super::super::types::{PoolUuid, FilesystemUuid, RenameAction};

use super::dmdevice::{FlexRole, ThinDevIdPool, ThinPoolRole, ThinRole, format_flex_name,
                      format_thinpool_name, format_thin_name};
use super::filesystem::{StratFilesystem, FilesystemStatus};
use super::mdv::MetadataVol;
use super::serde_structs::{FilesystemSave, FlexDevsSave, Recordable, ThinPoolDevSave};

pub const DATA_BLOCK_SIZE: Sectors = Sectors(2048);
pub const DATA_LOWATER: DataBlocks = DataBlocks(512);
pub const META_LOWATER: MetaBlocks = MetaBlocks(512);

const DEFAULT_THIN_DEV_SIZE: Sectors = Sectors(2 * IEC::Gi); // 1 TiB

/// A ThinPool struct contains the thinpool itself, but also the spare
/// segments for its metadata device.
#[derive(Debug)]
pub struct ThinPool {
    thin_pool: ThinPoolDev,
    meta_spare: Vec<Segment>,
    id_gen: ThinDevIdPool,
    filesystems: Table<StratFilesystem>,
    mdv: MetadataVol,
}

/// A struct returning the status of the distinct parts of the
/// thinpool.
pub struct ThinPoolStatus {
    /// The status of the thinpool itself.
    pub thinpool: dm::ThinPoolStatus,
    /// The status of the filesystems within the thinpool.
    pub filesystems: Vec<FilesystemStatus>,
}

impl ThinPool {
    /// Make a new thin pool.
    #[allow(too_many_arguments)]
    pub fn new(pool_uuid: PoolUuid,
               dm: &DM,
               data_block_size: Sectors,
               low_water_mark: DataBlocks,
               spare_segments: Vec<Segment>,
               meta_dev: LinearDev,
               data_dev: LinearDev,
               mdv: MetadataVol)
               -> EngineResult<ThinPool> {
        let name = format_thinpool_name(&pool_uuid, ThinPoolRole::Pool);
        let thinpool_dev = ThinPoolDev::new(name.as_ref(),
                                            dm,
                                            data_block_size,
                                            low_water_mark,
                                            meta_dev,
                                            data_dev)?;
        Ok(ThinPool {
               thin_pool: thinpool_dev,
               meta_spare: spare_segments,
               id_gen: ThinDevIdPool::new_from_ids(&[]),
               filesystems: Table::default(),
               mdv: mdv,
           })
    }

    /// Set up an "existing" thin pool.
    /// A thin pool must store the metadata for its thin devices, regardless of
    /// whether it has an existing device node. An existing thin pool device
    /// is a device where the metadata is already stored on its meta device.
    /// If initial setup fails due to a thin_check failure, attempt to fix
    /// the problem by running thin_repair. If failure recurs, return an
    /// error.
    #[allow(too_many_arguments)]
    pub fn setup(pool_uuid: PoolUuid,
                 dm: &DM,
                 data_block_size: Sectors,
                 low_water_mark: DataBlocks,
                 thin_ids: &[ThinDevId],
                 spare_segments: Vec<Segment>,
                 meta_dev: LinearDev,
                 data_dev: LinearDev,
                 mdv: MetadataVol,
                 fs_save: Vec<FilesystemSave>)
                 -> EngineResult<ThinPool> {
        let name = format_thinpool_name(&pool_uuid, ThinPoolRole::Pool);

        let res = match ThinPoolDev::setup(name.as_ref(),
                                           dm,
                                           data_block_size,
                                           low_water_mark,
                                           meta_dev,
                                           data_dev) {
            Ok(dev) => Ok((dev, spare_segments)),
            Err(DmError::Dm(CheckFailed(meta_dev, data_dev), _)) => {
                attempt_thin_repair(pool_uuid, dm, meta_dev, spare_segments)
                    .and_then(|(new_meta_dev, new_spare_segments)| {
                        ThinPoolDev::setup(name.as_ref(),
                                           dm,
                                           data_block_size,
                                           low_water_mark,
                                           new_meta_dev,
                                           data_dev)
                                .map(|dev| (dev, new_spare_segments))
                                .map_err(|e| e.into())
                    })
            }
            Err(e) => Err(e.into()),
        };
        let (thinpool_dev, spare_segments) = res?;

        // TODO: not fail completely if one filesystem setup fails?
        let filesystems = {
            // Set up a filesystem from its metadata.
            let get_filesystem = |fssave: &FilesystemSave| -> EngineResult<StratFilesystem> {
                let device_name = format_thin_name(&pool_uuid, ThinRole::Filesystem(fssave.uuid));
                let thin_dev = ThinDev::setup(device_name.as_ref(),
                                              dm,
                                              &thinpool_dev,
                                              fssave.thin_id,
                                              fssave.size)?;
                Ok(StratFilesystem::setup(fssave.uuid, &fssave.name, thin_dev))
            };

            fs_save
                .iter()
                .map(get_filesystem)
                .collect::<EngineResult<Vec<_>>>()?
        };

        let mut fs_table = Table::default();
        for fs in filesystems {
            let evicted = fs_table.insert(fs);
            if !evicted.is_empty() {
                let err_msg = "filesystems with duplicate UUID or name specified in metadata";
                return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
            }
        }

        Ok(ThinPool {
               thin_pool: thinpool_dev,
               meta_spare: spare_segments,
               id_gen: ThinDevIdPool::new_from_ids(thin_ids),
               filesystems: fs_table,
               mdv: mdv,
           })
    }

    /// The status of the thin pool as calculated by DM.
    pub fn check(&mut self, dm: &DM) -> EngineResult<ThinPoolStatus> {
        let thinpool = self.thin_pool.status(dm)?;
        self.mdv.check()?;

        let filesystems = self.filesystems
            .borrow_mut()
            .into_iter()
            .map(|fs| fs.check(dm))
            .collect::<EngineResult<Vec<_>>>()?;

        Ok(ThinPoolStatus {
               thinpool,
               filesystems,
           })
    }

    /// Tear down the components managed here: filesystems, the MDV,
    /// and the actual thinpool device itself.
    pub fn teardown(self, dm: &DM) -> EngineResult<()> {
        // Must succeed in tearing down all filesystems before the
        // thinpool..
        for fs in self.filesystems.empty() {
            fs.teardown(dm)?;
        }
        self.thin_pool.teardown(dm)?;

        // ..but MDV has no DM dependencies with the above
        self.mdv.teardown(dm)?;

        Ok(())
    }

    /// Return the FlexDevsSave data structure for variable length metadata.
    /// May return an error if mapper can not locate the UUID corresponding
    /// to a device node.
    pub fn flexdevssave<F>(&self, mapper: F) -> EngineResult<FlexDevsSave>
        where F: Fn(&Segment) -> EngineResult<(Uuid, Sectors, Sectors)>
    {
        Ok(FlexDevsSave {
               meta_dev: self.mdv
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_meta_dev: self.thin_pool
                   .meta_dev()
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_data_dev: self.thin_pool
                   .data_dev()
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_meta_dev_spare: self.meta_spare
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
           })
    }

    /// Get the devicemapper::ThinPoolDev for this pool. Used for testing.
    pub fn thinpooldev(&self) -> &ThinPoolDev {
        &self.thin_pool
    }

    /// Extend the thinpool with new data regions.
    pub fn extend_data(&mut self, dm: &DM, segs: Vec<Segment>) -> EngineResult<()> {
        Ok(self.thin_pool.extend_data(dm, segs)?)
    }

    /// The number of physical sectors in use, that is, unavailable for storage
    /// of additional user data, by this pool.
    // This includes all the sectors being held as spares for the meta device,
    // all the sectors allocated to the meta data device, and all the sectors
    // in use on the data device.
    pub fn total_physical_used(&self) -> EngineResult<Sectors> {
        let data_dev_used = match self.thin_pool.status(&DM::new()?)? {
            dm::ThinPoolStatus::Good(_, usage) => *usage.used_data * DATA_BLOCK_SIZE,
            _ => {
                let err_msg = "thin pool failed, could not obtain usage";
                return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
            }
        };

        let spare_total = self.meta_spare.iter().map(|s| s.length).sum();
        let meta_dev_total = self.thin_pool
            .meta_dev()
            .segments()
            .iter()
            .map(|s| s.length)
            .sum();

        let mdv_total = self.mdv.segments().iter().map(|s| s.length).sum();

        Ok(data_dev_used + spare_total + meta_dev_total + mdv_total)
    }

    pub fn get_filesystem_by_uuid(&self, uuid: &FilesystemUuid) -> Option<&StratFilesystem> {
        self.filesystems.get_by_uuid(uuid)
    }

    pub fn get_mut_filesystem_by_uuid(&mut self,
                                      uuid: &FilesystemUuid)
                                      -> Option<&mut StratFilesystem> {
        self.filesystems.get_mut_by_uuid(uuid)
    }

    #[allow(dead_code)]
    pub fn get_filesystem_by_name(&self, name: &str) -> Option<&StratFilesystem> {
        self.filesystems.get_by_name(name)
    }

    pub fn get_mut_filesystem_by_name(&mut self, name: &str) -> Option<&mut StratFilesystem> {
        self.filesystems.get_mut_by_name(name)
    }

    pub fn has_filesystems(&self) -> bool {
        !self.filesystems.is_empty()
    }

    pub fn filesystems(&self) -> Vec<&Filesystem> {
        self.filesystems
            .into_iter()
            .map(|x| x as &Filesystem)
            .collect()
    }

    /// Create a filesystem within the thin pool. Given name must not
    /// already be in use.
    pub fn create_filesystem(&mut self,
                             pool_uuid: &Uuid,
                             name: &str,
                             dm: &DM,
                             size: Option<Sectors>)
                             -> EngineResult<FilesystemUuid> {
        let fs_uuid = Uuid::new_v4();
        let device_name = format_thin_name(pool_uuid, ThinRole::Filesystem(fs_uuid));
        let thin_dev = ThinDev::new(device_name.as_ref(),
                                    dm,
                                    &self.thin_pool,
                                    self.id_gen.new_id()?,
                                    size.unwrap_or(DEFAULT_THIN_DEV_SIZE))?;

        let new_filesystem = StratFilesystem::initialize(fs_uuid, name, thin_dev)?;
        self.mdv.save_fs(&new_filesystem)?;
        self.filesystems.insert(new_filesystem);

        Ok(fs_uuid)
    }

    /// Destroy a filesystem within the thin pool.
    pub fn destroy_filesystem(&mut self, dm: &DM, uuid: &FilesystemUuid) -> EngineResult<()> {
        if let Some(fs) = self.filesystems.remove_by_uuid(uuid) {
            fs.destroy(dm, &self.thin_pool)?;
            self.mdv.rm_fs(uuid)?;
        }
        Ok(())
    }

    /// Rename a filesystem within the thin pool.
    pub fn rename_filesystem(&mut self,
                             uuid: &FilesystemUuid,
                             new_name: &str)
                             -> EngineResult<RenameAction> {

        let old_name = rename_filesystem_pre!(self; uuid; new_name);

        let mut filesystem =
            self.filesystems
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.filesystems.get_by_uuid() returned a value");

        filesystem.rename(new_name);
        if let Err(err) = self.mdv.save_fs(&filesystem) {
            filesystem.rename(&old_name);
            self.filesystems.insert(filesystem);
            Err(err)
        } else {
            self.filesystems.insert(filesystem);
            Ok(RenameAction::Renamed)
        }
    }
}

impl Recordable<ThinPoolDevSave> for ThinPool {
    fn record(&self) -> EngineResult<ThinPoolDevSave> {
        Ok(ThinPoolDevSave { data_block_size: self.thin_pool.data_block_size() })
    }
}

/// Attempt a thin repair operation on the meta device.
/// If the operation succeeds, teardown the old meta device,
/// and return the new meta device and the new spare segments.
fn attempt_thin_repair(pool_uuid: PoolUuid,
                       dm: &DM,
                       meta_dev: LinearDev,
                       mut spare_segments: Vec<Segment>)
                       -> EngineResult<(LinearDev, Vec<Segment>)> {
    let mut new_meta_dev = LinearDev::new(format_flex_name(&pool_uuid, FlexRole::ThinMetaSpare)
                                              .as_ref(),
                                          dm,
                                          spare_segments.drain(..).collect())?;

    if !Command::new("thin_repair")
            .arg("-i")
            .arg(&meta_dev.devnode())
            .arg("-o")
            .arg(&new_meta_dev.devnode())
            .status()?
            .success() {
        return Err(EngineError::Engine(ErrorEnum::Error,
                                       "thin_repair failed, pool unusable".into()));
    }

    let name = meta_dev.name().to_owned();
    let new_spare_segments = meta_dev
        .segments()
        .iter()
        .map(|x| {
                 Segment {
                     start: x.start,
                     length: x.length,
                     device: x.device,
                 }
             })
        .collect();
    meta_dev.teardown(dm)?;
    new_meta_dev.set_name(dm, name.as_ref())?;

    Ok((new_meta_dev, new_spare_segments))
}
<commit_msg>Update the comment about the ThinPool struct<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/// Code to handle management of a pool's thinpool device.

use std::borrow::BorrowMut;
use std::process::Command;

use uuid::Uuid;

use devicemapper as dm;
use devicemapper::{DM, DmDevice, DataBlocks, DmError, LinearDev, MetaBlocks, Sectors, Segment,
                   ThinDev, ThinDevId, ThinPoolDev};
use devicemapper::ErrorEnum::CheckFailed;

use super::super::consts::IEC;
use super::super::engine::{Filesystem, HasName};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::structures::Table;
use super::super::types::{PoolUuid, FilesystemUuid, RenameAction};

use super::dmdevice::{FlexRole, ThinDevIdPool, ThinPoolRole, ThinRole, format_flex_name,
                      format_thinpool_name, format_thin_name};
use super::filesystem::{StratFilesystem, FilesystemStatus};
use super::mdv::MetadataVol;
use super::serde_structs::{FilesystemSave, FlexDevsSave, Recordable, ThinPoolDevSave};

pub const DATA_BLOCK_SIZE: Sectors = Sectors(2048);
pub const DATA_LOWATER: DataBlocks = DataBlocks(512);
pub const META_LOWATER: MetaBlocks = MetaBlocks(512);

const DEFAULT_THIN_DEV_SIZE: Sectors = Sectors(2 * IEC::Gi); // 1 TiB

/// A ThinPool struct contains the thinpool itself, the spare
/// segments for its metadata device, and the filesystems and filesystem
/// metadata associated with it.
#[derive(Debug)]
pub struct ThinPool {
    thin_pool: ThinPoolDev,
    meta_spare: Vec<Segment>,
    id_gen: ThinDevIdPool,
    filesystems: Table<StratFilesystem>,
    mdv: MetadataVol,
}

/// A struct returning the status of the distinct parts of the
/// thinpool.
pub struct ThinPoolStatus {
    /// The status of the thinpool itself.
    pub thinpool: dm::ThinPoolStatus,
    /// The status of the filesystems within the thinpool.
    pub filesystems: Vec<FilesystemStatus>,
}

impl ThinPool {
    /// Make a new thin pool.
    #[allow(too_many_arguments)]
    pub fn new(pool_uuid: PoolUuid,
               dm: &DM,
               data_block_size: Sectors,
               low_water_mark: DataBlocks,
               spare_segments: Vec<Segment>,
               meta_dev: LinearDev,
               data_dev: LinearDev,
               mdv: MetadataVol)
               -> EngineResult<ThinPool> {
        let name = format_thinpool_name(&pool_uuid, ThinPoolRole::Pool);
        let thinpool_dev = ThinPoolDev::new(name.as_ref(),
                                            dm,
                                            data_block_size,
                                            low_water_mark,
                                            meta_dev,
                                            data_dev)?;
        Ok(ThinPool {
               thin_pool: thinpool_dev,
               meta_spare: spare_segments,
               id_gen: ThinDevIdPool::new_from_ids(&[]),
               filesystems: Table::default(),
               mdv: mdv,
           })
    }

    /// Set up an "existing" thin pool.
    /// A thin pool must store the metadata for its thin devices, regardless of
    /// whether it has an existing device node. An existing thin pool device
    /// is a device where the metadata is already stored on its meta device.
    /// If initial setup fails due to a thin_check failure, attempt to fix
    /// the problem by running thin_repair. If failure recurs, return an
    /// error.
    #[allow(too_many_arguments)]
    pub fn setup(pool_uuid: PoolUuid,
                 dm: &DM,
                 data_block_size: Sectors,
                 low_water_mark: DataBlocks,
                 thin_ids: &[ThinDevId],
                 spare_segments: Vec<Segment>,
                 meta_dev: LinearDev,
                 data_dev: LinearDev,
                 mdv: MetadataVol,
                 fs_save: Vec<FilesystemSave>)
                 -> EngineResult<ThinPool> {
        let name = format_thinpool_name(&pool_uuid, ThinPoolRole::Pool);

        let res = match ThinPoolDev::setup(name.as_ref(),
                                           dm,
                                           data_block_size,
                                           low_water_mark,
                                           meta_dev,
                                           data_dev) {
            Ok(dev) => Ok((dev, spare_segments)),
            Err(DmError::Dm(CheckFailed(meta_dev, data_dev), _)) => {
                attempt_thin_repair(pool_uuid, dm, meta_dev, spare_segments)
                    .and_then(|(new_meta_dev, new_spare_segments)| {
                        ThinPoolDev::setup(name.as_ref(),
                                           dm,
                                           data_block_size,
                                           low_water_mark,
                                           new_meta_dev,
                                           data_dev)
                                .map(|dev| (dev, new_spare_segments))
                                .map_err(|e| e.into())
                    })
            }
            Err(e) => Err(e.into()),
        };
        let (thinpool_dev, spare_segments) = res?;

        // TODO: not fail completely if one filesystem setup fails?
        let filesystems = {
            // Set up a filesystem from its metadata.
            let get_filesystem = |fssave: &FilesystemSave| -> EngineResult<StratFilesystem> {
                let device_name = format_thin_name(&pool_uuid, ThinRole::Filesystem(fssave.uuid));
                let thin_dev = ThinDev::setup(device_name.as_ref(),
                                              dm,
                                              &thinpool_dev,
                                              fssave.thin_id,
                                              fssave.size)?;
                Ok(StratFilesystem::setup(fssave.uuid, &fssave.name, thin_dev))
            };

            fs_save
                .iter()
                .map(get_filesystem)
                .collect::<EngineResult<Vec<_>>>()?
        };

        let mut fs_table = Table::default();
        for fs in filesystems {
            let evicted = fs_table.insert(fs);
            if !evicted.is_empty() {
                let err_msg = "filesystems with duplicate UUID or name specified in metadata";
                return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
            }
        }

        Ok(ThinPool {
               thin_pool: thinpool_dev,
               meta_spare: spare_segments,
               id_gen: ThinDevIdPool::new_from_ids(thin_ids),
               filesystems: fs_table,
               mdv: mdv,
           })
    }

    /// The status of the thin pool as calculated by DM.
    pub fn check(&mut self, dm: &DM) -> EngineResult<ThinPoolStatus> {
        let thinpool = self.thin_pool.status(dm)?;
        self.mdv.check()?;

        let filesystems = self.filesystems
            .borrow_mut()
            .into_iter()
            .map(|fs| fs.check(dm))
            .collect::<EngineResult<Vec<_>>>()?;

        Ok(ThinPoolStatus {
               thinpool,
               filesystems,
           })
    }

    /// Tear down the components managed here: filesystems, the MDV,
    /// and the actual thinpool device itself.
    pub fn teardown(self, dm: &DM) -> EngineResult<()> {
        // Must succeed in tearing down all filesystems before the
        // thinpool..
        for fs in self.filesystems.empty() {
            fs.teardown(dm)?;
        }
        self.thin_pool.teardown(dm)?;

        // ..but MDV has no DM dependencies with the above
        self.mdv.teardown(dm)?;

        Ok(())
    }

    /// Return the FlexDevsSave data structure for variable length metadata.
    /// May return an error if mapper can not locate the UUID corresponding
    /// to a device node.
    pub fn flexdevssave<F>(&self, mapper: F) -> EngineResult<FlexDevsSave>
        where F: Fn(&Segment) -> EngineResult<(Uuid, Sectors, Sectors)>
    {
        Ok(FlexDevsSave {
               meta_dev: self.mdv
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_meta_dev: self.thin_pool
                   .meta_dev()
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_data_dev: self.thin_pool
                   .data_dev()
                   .segments()
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
               thin_meta_dev_spare: self.meta_spare
                   .iter()
                   .map(&mapper)
                   .collect::<EngineResult<Vec<_>>>()?,
           })
    }

    /// Get the devicemapper::ThinPoolDev for this pool. Used for testing.
    pub fn thinpooldev(&self) -> &ThinPoolDev {
        &self.thin_pool
    }

    /// Extend the thinpool with new data regions.
    pub fn extend_data(&mut self, dm: &DM, segs: Vec<Segment>) -> EngineResult<()> {
        Ok(self.thin_pool.extend_data(dm, segs)?)
    }

    /// The number of physical sectors in use, that is, unavailable for storage
    /// of additional user data, by this pool.
    // This includes all the sectors being held as spares for the meta device,
    // all the sectors allocated to the meta data device, and all the sectors
    // in use on the data device.
    pub fn total_physical_used(&self) -> EngineResult<Sectors> {
        let data_dev_used = match self.thin_pool.status(&DM::new()?)? {
            dm::ThinPoolStatus::Good(_, usage) => *usage.used_data * DATA_BLOCK_SIZE,
            _ => {
                let err_msg = "thin pool failed, could not obtain usage";
                return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
            }
        };

        let spare_total = self.meta_spare.iter().map(|s| s.length).sum();
        let meta_dev_total = self.thin_pool
            .meta_dev()
            .segments()
            .iter()
            .map(|s| s.length)
            .sum();

        let mdv_total = self.mdv.segments().iter().map(|s| s.length).sum();

        Ok(data_dev_used + spare_total + meta_dev_total + mdv_total)
    }

    pub fn get_filesystem_by_uuid(&self, uuid: &FilesystemUuid) -> Option<&StratFilesystem> {
        self.filesystems.get_by_uuid(uuid)
    }

    pub fn get_mut_filesystem_by_uuid(&mut self,
                                      uuid: &FilesystemUuid)
                                      -> Option<&mut StratFilesystem> {
        self.filesystems.get_mut_by_uuid(uuid)
    }

    #[allow(dead_code)]
    pub fn get_filesystem_by_name(&self, name: &str) -> Option<&StratFilesystem> {
        self.filesystems.get_by_name(name)
    }

    pub fn get_mut_filesystem_by_name(&mut self, name: &str) -> Option<&mut StratFilesystem> {
        self.filesystems.get_mut_by_name(name)
    }

    pub fn has_filesystems(&self) -> bool {
        !self.filesystems.is_empty()
    }

    pub fn filesystems(&self) -> Vec<&Filesystem> {
        self.filesystems
            .into_iter()
            .map(|x| x as &Filesystem)
            .collect()
    }

    /// Create a filesystem within the thin pool. Given name must not
    /// already be in use.
    pub fn create_filesystem(&mut self,
                             pool_uuid: &Uuid,
                             name: &str,
                             dm: &DM,
                             size: Option<Sectors>)
                             -> EngineResult<FilesystemUuid> {
        let fs_uuid = Uuid::new_v4();
        let device_name = format_thin_name(pool_uuid, ThinRole::Filesystem(fs_uuid));
        let thin_dev = ThinDev::new(device_name.as_ref(),
                                    dm,
                                    &self.thin_pool,
                                    self.id_gen.new_id()?,
                                    size.unwrap_or(DEFAULT_THIN_DEV_SIZE))?;

        let new_filesystem = StratFilesystem::initialize(fs_uuid, name, thin_dev)?;
        self.mdv.save_fs(&new_filesystem)?;
        self.filesystems.insert(new_filesystem);

        Ok(fs_uuid)
    }

    /// Destroy a filesystem within the thin pool.
    pub fn destroy_filesystem(&mut self, dm: &DM, uuid: &FilesystemUuid) -> EngineResult<()> {
        if let Some(fs) = self.filesystems.remove_by_uuid(uuid) {
            fs.destroy(dm, &self.thin_pool)?;
            self.mdv.rm_fs(uuid)?;
        }
        Ok(())
    }

    /// Rename a filesystem within the thin pool.
    pub fn rename_filesystem(&mut self,
                             uuid: &FilesystemUuid,
                             new_name: &str)
                             -> EngineResult<RenameAction> {

        let old_name = rename_filesystem_pre!(self; uuid; new_name);

        let mut filesystem =
            self.filesystems
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.filesystems.get_by_uuid() returned a value");

        filesystem.rename(new_name);
        if let Err(err) = self.mdv.save_fs(&filesystem) {
            filesystem.rename(&old_name);
            self.filesystems.insert(filesystem);
            Err(err)
        } else {
            self.filesystems.insert(filesystem);
            Ok(RenameAction::Renamed)
        }
    }
}

impl Recordable<ThinPoolDevSave> for ThinPool {
    fn record(&self) -> EngineResult<ThinPoolDevSave> {
        Ok(ThinPoolDevSave { data_block_size: self.thin_pool.data_block_size() })
    }
}

/// Attempt a thin repair operation on the meta device.
/// If the operation succeeds, teardown the old meta device,
/// and return the new meta device and the new spare segments.
fn attempt_thin_repair(pool_uuid: PoolUuid,
                       dm: &DM,
                       meta_dev: LinearDev,
                       mut spare_segments: Vec<Segment>)
                       -> EngineResult<(LinearDev, Vec<Segment>)> {
    let mut new_meta_dev = LinearDev::new(format_flex_name(&pool_uuid, FlexRole::ThinMetaSpare)
                                              .as_ref(),
                                          dm,
                                          spare_segments.drain(..).collect())?;

    if !Command::new("thin_repair")
            .arg("-i")
            .arg(&meta_dev.devnode())
            .arg("-o")
            .arg(&new_meta_dev.devnode())
            .status()?
            .success() {
        return Err(EngineError::Engine(ErrorEnum::Error,
                                       "thin_repair failed, pool unusable".into()));
    }

    let name = meta_dev.name().to_owned();
    let new_spare_segments = meta_dev
        .segments()
        .iter()
        .map(|x| {
                 Segment {
                     start: x.start,
                     length: x.length,
                     device: x.device,
                 }
             })
        .collect();
    meta_dev.teardown(dm)?;
    new_meta_dev.set_name(dm, name.as_ref())?;

    Ok((new_meta_dev, new_spare_segments))
}
<|endoftext|>