<commit_before>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::defer_stream::DEFER_STREAM_CONSTANTS;
use crate::{DeferDirective, NoInlineFragmentSpreadMetadata, StreamDirective};
use common::{Diagnostic, DiagnosticsResult, NamedItem};
use graphql_ir::{
    transform_list_multi, Condition, ConditionValue, ConstantValue, FragmentDefinition,
    FragmentSpread, InlineFragment, LinkedField, Program, Selection, Transformed, TransformedMulti,
    TransformedValue, Transformer, Value,
};
use intern::string_key::{StringKey, StringKeyMap};
use std::sync::Arc;
use thiserror::Error;

enum ValidationMode {
    Strict,
    Loose,
}

pub fn skip_unreachable_node_strict(program: &Program) -> DiagnosticsResult<Program> {
    skip_unreachable_node(program, ValidationMode::Strict)
}

pub fn skip_unreachable_node_loose(program: &Program) -> Program {
    skip_unreachable_node(program, ValidationMode::Loose).unwrap_or_else(|error| {
        panic!(
            "Unexpected error in `skip_unreachable_node_loose`: {:?}",
            error
        )
    })
}

fn skip_unreachable_node(
    program: &Program,
    validation_mode: ValidationMode,
) -> DiagnosticsResult<Program> {
    let mut skip_unreachable_node_transform =
        SkipUnreachableNodeTransform::new(program, validation_mode);
    let transformed = skip_unreachable_node_transform.transform_program(program);
    if skip_unreachable_node_transform.errors.is_empty() {
        Ok(transformed.replace_or_else(|| program.clone()))
    } else {
        Err(skip_unreachable_node_transform.errors)
    }
}

type VisitedFragments = StringKeyMap<(Arc<FragmentDefinition>, Transformed<FragmentDefinition>)>;

pub struct SkipUnreachableNodeTransform<'s> {
    errors: Vec<Diagnostic>,
    visited_fragments: VisitedFragments,
    program: &'s Program,
    validation_mode: ValidationMode,
}

impl<'s> Transformer for SkipUnreachableNodeTransform<'s> {
    const NAME: &'static str = "SkipUnreachableNodeTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_program(&mut self, program: &Program) -> TransformedValue<Program> {
        // Iterate over the operation ASTs. Whenever we encounter a FragmentSpread (by way of
        // fn transform_fragment_spread) look up that the associated FragmentDefinition.
        // Crawl that FragmentDefinition for other FragmentSpreads and repeat.
        // Delete that FragmentSpread if the associated FragmentDefinition was deleted.
        //
        // Throughout, when we encounter a condition node with a constant predicate, either remove
        // it or replace it with its contents.
        //
        // @include(if: false)  => remove
        // @include(if: true)   => replace with contents
        // @skip(if: true)      => remove
        // @skip(if: false)     => replace with contents
        // @defer(if: false)    => replace with contents
        //
        // Removal of a condition or spread can result in a FragmentDefinition being deleted.

        let mut next_program = Program::new(Arc::clone(&program.schema));
        let mut has_changes = false;

        for operation in program.operations() {
            match self.transform_operation(operation) {
                Transformed::Delete => {
                    if matches!(self.validation_mode, ValidationMode::Strict) {
                        self.errors.push(Diagnostic::error(
                            ValidationMessage::EmptySelectionsInDocument {
                                document: "query",
                                name: operation.name.item,
                            },
                            operation.name.location,
                        ));
                    }
                    has_changes = true;
                }
                Transformed::Keep => next_program.insert_operation(Arc::clone(operation)),
                Transformed::Replace(replacement) => {
                    has_changes = true;
                    next_program.insert_operation(Arc::new(replacement))
                }
            }
        }
        for fragment in program.fragments() {
            if let Some(visited_fragment) = self.visited_fragments.get(&fragment.name.item) {
                match visited_fragment {
                    (_, Transformed::Delete) => {
                        has_changes = true;
                    }
                    (fragment, Transformed::Keep) => {
                        next_program.insert_fragment(Arc::clone(fragment));
                    }
                    (_, Transformed::Replace(replacement)) => {
                        next_program.insert_fragment(Arc::new(replacement.clone()));
                        has_changes = true;
                    }
                }
            } else {
                match self.transform_fragment(fragment) {
                    Transformed::Delete => {
                        if matches!(self.validation_mode, ValidationMode::Strict) {
                            self.errors.push(Diagnostic::error(
                                ValidationMessage::EmptySelectionsInDocument {
                                    document: "fragment",
                                    name: fragment.name.item,
                                },
                                fragment.name.location,
                            ));
                        }
                        has_changes = true;
                    }
                    Transformed::Keep => {
                        next_program.insert_fragment(Arc::clone(fragment));
                    }
                    Transformed::Replace(replacement) => {
                        next_program.insert_fragment(Arc::new(replacement.clone()));
                        has_changes = true;
                    }
                }
            }
        }

        if has_changes {
            TransformedValue::Replace(next_program)
        } else {
            TransformedValue::Keep
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        // Remove the fragment with empty selections
        let selections = self.transform_selections(&fragment.selections);
        if let TransformedValue::Replace(selections) = &selections {
            if selections.is_empty() {
                return Transformed::Delete;
            }
        }
        let directives = self.transform_directives(&fragment.directives);
        if selections.should_keep() && directives.should_keep() {
            return Transformed::Keep;
        }
        Transformed::Replace(FragmentDefinition {
            directives: directives.replace_or_else(|| fragment.directives.clone()),
            selections: selections.replace_or_else(|| fragment.selections.clone()),
            ..fragment.clone()
        })
    }

    fn transform_selections(
        &mut self,
        selections: &[Selection],
    ) -> TransformedValue<Vec<Selection>> {
        transform_list_multi(selections, |selection| self.map_selection_multi(selection))
    }

    fn transform_fragment_spread(&mut self, spread: &FragmentSpread) -> Transformed<Selection> {
        if spread
            .directives
            .named(NoInlineFragmentSpreadMetadata::directive_name())
            .is_some()
        {
            return Transformed::Keep;
        }
        if self.should_delete_fragment_definition(spread.fragment.item) {
            Transformed::Delete
        } else {
            Transformed::Keep
        }
    }

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Selection> {
        let transformed_field = self.default_transform_linked_field(field);
        if let Some(directive) = field.directives.named(DEFER_STREAM_CONSTANTS.stream_name) {
            if let Some(if_arg) = StreamDirective::from(directive).if_arg {
                if let Value::Constant(ConstantValue::Boolean(false)) = &if_arg.value.item {
                    let mut next_field = match transformed_field {
                        Transformed::Delete => return Transformed::Delete,
                        Transformed::Keep => Arc::new(field.clone()),
                        Transformed::Replace(Selection::LinkedField(replacement)) => replacement,
                        Transformed::Replace(other) => {
                            panic!("unexpected replacement: {:?}", other)
                        }
                    };
                    let previous_directive_len = next_field.directives.len();
                    Arc::make_mut(&mut next_field)
                        .directives
                        .retain(|directive| {
                            directive.name.item != DEFER_STREAM_CONSTANTS.stream_name
                        });
                    assert_eq!(
                        previous_directive_len,
                        next_field.directives.len() + 1,
                        "should have removed exactly one directive"
                    );
                    return Transformed::Replace(Selection::LinkedField(next_field));
                }
            }
        }
        transformed_field
    }
}

impl<'s> SkipUnreachableNodeTransform<'s> {
    fn new(program: &'s Program, validation_mode: ValidationMode) -> Self {
        Self {
            errors: Vec::new(),
            visited_fragments: Default::default(),
            program,
            validation_mode,
        }
    }

    // Visit the fragment definition and return true if it should be deleted
    fn should_delete_fragment_definition(&mut self, key: StringKey) -> bool {
        if let Some((_, transformed)) = self.visited_fragments.get(&key) {
            return matches!(transformed, Transformed::Delete);
        }
        if let Some(fragment) = self.program.fragment(key) {
            self.visited_fragments
                .insert(key, (Arc::clone(fragment), Transformed::Keep));
            let transformed = self.transform_fragment(fragment);
            let should_delete = matches!(transformed, Transformed::Delete);

            // N.B. we must call self.visited_fragments.get* twice, because we cannot have
            // a reference to visited_opt and call transform_segment, which requires an
            // exclusive reference to self.
            let (_, visited_opt) = self.visited_fragments.get_mut(&key).unwrap();
            *visited_opt = transformed;
            should_delete
        } else {
            true
        }
    }

    fn map_selection_multi(&mut self, selection: &Selection) -> TransformedMulti<Selection> {
        match selection {
            Selection::FragmentSpread(selection) => {
                self.transform_fragment_spread(selection).into()
            }
            Selection::InlineFragment(selection) => self.transform_inline_fragment_multi(selection),
            Selection::LinkedField(selection) => self.transform_linked_field(selection).into(),
            Selection::ScalarField(selection) => self.transform_scalar_field(selection).into(),
            Selection::Condition(condition) => self.transform_constant_conditions(condition),
        }
    }

    fn transform_inline_fragment_multi(
        &mut self,
        inline_fragment: &InlineFragment,
    ) -> TransformedMulti<Selection> {
        if let Some(directive) = inline_fragment
            .directives
            .named(DEFER_STREAM_CONSTANTS.defer_name)
        {
            assert!(inline_fragment.directives.len() == 1);
            if let Some(if_arg) = DeferDirective::from(directive).if_arg {
                if let Value::Constant(ConstantValue::Boolean(false)) = &if_arg.value.item {
                    return TransformedMulti::ReplaceMultiple(
                        self.transform_selections(&inline_fragment.selections)
                            .replace_or_else(|| inline_fragment.selections.clone()),
                    );
                }
            }
        };
        self.default_transform_inline_fragment(inline_fragment)
            .into()
    }

    fn transform_constant_conditions(
        &mut self,
        condition: &Condition,
    ) -> TransformedMulti<Selection> {
        if let ConditionValue::Constant(b) = condition.value {
            // passing_value == "what the value needs to be in order to include the contents"
            // in other words, @skip has passing_value of false
            // and @include has passing_value of true
            //
            // so, b == condition.passing_value implies that the condition is redundant,
            // and b != condition.passing_value implies that the whole node can be removed
            let keep_contents = b == condition.passing_value;
            if keep_contents {
                // @include(if: true) or @skip(if: false)
                let contents = match self.transform_selections(&condition.selections) {
                    TransformedValue::Keep => &condition.selections,
                    TransformedValue::Replace(ref t) => t,
                }
                .clone();
                TransformedMulti::ReplaceMultiple(contents)
            } else {
                // @include(if: false) or @skip(if: true)
                TransformedMulti::Delete
            }
        } else {
            self.transform_condition(condition).into()
        }
    }
}

#[derive(Clone, Debug, Error, Eq, PartialEq, Ord, PartialOrd, Hash)]
enum ValidationMessage {
    #[error(
        "After applying transforms to the {document} `{name}` selections of \
        the `{name}` that would be sent to the server are empty. \
        This is likely due to the use of `@skip`/`@include` directives with \
        constant values that remove all selections in the {document}. "
    )]
    EmptySelectionsInDocument {
        name: StringKey,
        document: &'static str,
    },
}
<commit_msg>Small refactoring in the `skip_unreachable_nodes` - store errors in the ValidationMode::Strict<commit_after>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::defer_stream::DEFER_STREAM_CONSTANTS;
use crate::{DeferDirective, NoInlineFragmentSpreadMetadata, StreamDirective};
use common::{Diagnostic, DiagnosticsResult, NamedItem};
use graphql_ir::{
    transform_list_multi, Condition, ConditionValue, ConstantValue, FragmentDefinition,
    FragmentSpread, InlineFragment, LinkedField, Program, Selection, Transformed, TransformedMulti,
    TransformedValue, Transformer, Value,
};
use intern::string_key::{StringKey, StringKeyMap};
use std::sync::Arc;
use thiserror::Error;

enum ValidationMode {
    Strict(Vec<Diagnostic>),
    Loose,
}

pub fn skip_unreachable_node_strict(program: &Program) -> DiagnosticsResult<Program> {
    let errors = vec![];
    let mut validation_mode = ValidationMode::Strict(errors);
    let next_program = skip_unreachable_node(program, &mut validation_mode);

    if let ValidationMode::Strict(errors) = validation_mode {
        if !errors.is_empty() {
            return Err(errors);
        }
    }
    Ok(next_program)
}

pub fn skip_unreachable_node_loose(program: &Program) -> Program {
    let mut validation_mode = ValidationMode::Loose;
    skip_unreachable_node(program, &mut validation_mode)
}

fn skip_unreachable_node(program: &Program, validation_mode: &mut ValidationMode) -> Program {
    let mut skip_unreachable_node_transform =
        SkipUnreachableNodeTransform::new(program, validation_mode);
    let transformed = skip_unreachable_node_transform.transform_program(program);

    transformed.replace_or_else(|| program.clone())
}

type VisitedFragments = StringKeyMap<(Arc<FragmentDefinition>, Transformed<FragmentDefinition>)>;

pub struct SkipUnreachableNodeTransform<'s> {
    visited_fragments: VisitedFragments,
    program: &'s Program,
    validation_mode: &'s mut ValidationMode,
}

impl<'s> Transformer for SkipUnreachableNodeTransform<'s> {
    const NAME: &'static str = "SkipUnreachableNodeTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_program(&mut self, program: &Program) -> TransformedValue<Program> {
        // Iterate over the operation ASTs. Whenever we encounter a FragmentSpread (by way of
        // fn transform_fragment_spread) look up that the associated FragmentDefinition.
        // Crawl that FragmentDefinition for other FragmentSpreads and repeat.
        // Delete that FragmentSpread if the associated FragmentDefinition was deleted.
        //
        // Throughout, when we encounter a condition node with a constant predicate, either remove
        // it or replace it with its contents.
        //
        // @include(if: false)  => remove
        // @include(if: true)   => replace with contents
        // @skip(if: true)      => remove
        // @skip(if: false)     => replace with contents
        // @defer(if: false)    => replace with contents
        //
        // Removal of a condition or spread can result in a FragmentDefinition being deleted.

        let mut next_program = Program::new(Arc::clone(&program.schema));
        let mut has_changes = false;

        for operation in program.operations() {
            match self.transform_operation(operation) {
                Transformed::Delete => {
                    if let ValidationMode::Strict(errors) = &mut self.validation_mode {
                        errors.push(Diagnostic::error(
                            ValidationMessage::EmptySelectionsInDocument {
                                document: "query",
                                name: operation.name.item,
                            },
                            operation.name.location,
                        ));
                    }
                    has_changes = true;
                }
                Transformed::Keep => next_program.insert_operation(Arc::clone(operation)),
                Transformed::Replace(replacement) => {
                    has_changes = true;
                    next_program.insert_operation(Arc::new(replacement))
                }
            }
        }
        for fragment in program.fragments() {
            if let Some(visited_fragment) = self.visited_fragments.get(&fragment.name.item) {
                match visited_fragment {
                    (_, Transformed::Delete) => {
                        has_changes = true;
                    }
                    (fragment, Transformed::Keep) => {
                        next_program.insert_fragment(Arc::clone(fragment));
                    }
                    (_, Transformed::Replace(replacement)) => {
                        next_program.insert_fragment(Arc::new(replacement.clone()));
                        has_changes = true;
                    }
                }
            } else {
                match self.transform_fragment(fragment) {
                    Transformed::Delete => {
                        if let ValidationMode::Strict(errors) = &mut self.validation_mode {
                            errors.push(Diagnostic::error(
                                ValidationMessage::EmptySelectionsInDocument {
                                    document: "fragment",
                                    name: fragment.name.item,
                                },
                                fragment.name.location,
                            ));
                        }
                        has_changes = true;
                    }
                    Transformed::Keep => {
                        next_program.insert_fragment(Arc::clone(fragment));
                    }
                    Transformed::Replace(replacement) => {
                        next_program.insert_fragment(Arc::new(replacement.clone()));
                        has_changes = true;
                    }
                }
            }
        }

        if has_changes {
            TransformedValue::Replace(next_program)
        } else {
            TransformedValue::Keep
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        // Remove the fragment with empty selections
        let selections = self.transform_selections(&fragment.selections);
        if let TransformedValue::Replace(selections) = &selections {
            if selections.is_empty() {
                return Transformed::Delete;
            }
        }
        let directives = self.transform_directives(&fragment.directives);
        if selections.should_keep() && directives.should_keep() {
            return Transformed::Keep;
        }
        Transformed::Replace(FragmentDefinition {
            directives: directives.replace_or_else(|| fragment.directives.clone()),
            selections: selections.replace_or_else(|| fragment.selections.clone()),
            ..fragment.clone()
        })
    }

    fn transform_selections(
        &mut self,
        selections: &[Selection],
    ) -> TransformedValue<Vec<Selection>> {
        transform_list_multi(selections, |selection| self.map_selection_multi(selection))
    }

    fn transform_fragment_spread(&mut self, spread: &FragmentSpread) -> Transformed<Selection> {
        if spread
            .directives
            .named(NoInlineFragmentSpreadMetadata::directive_name())
            .is_some()
        {
            return Transformed::Keep;
        }
        if self.should_delete_fragment_definition(spread.fragment.item) {
            Transformed::Delete
        } else {
            Transformed::Keep
        }
    }

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Selection> {
        let transformed_field = self.default_transform_linked_field(field);
        if let Some(directive) = field.directives.named(DEFER_STREAM_CONSTANTS.stream_name) {
            if let Some(if_arg) = StreamDirective::from(directive).if_arg {
                if let Value::Constant(ConstantValue::Boolean(false)) = &if_arg.value.item {
                    let mut next_field = match transformed_field {
                        Transformed::Delete => return Transformed::Delete,
                        Transformed::Keep => Arc::new(field.clone()),
                        Transformed::Replace(Selection::LinkedField(replacement)) => replacement,
                        Transformed::Replace(other) => {
                            panic!("unexpected replacement: {:?}", other)
                        }
                    };
                    let previous_directive_len = next_field.directives.len();
                    Arc::make_mut(&mut next_field)
                        .directives
                        .retain(|directive| {
                            directive.name.item != DEFER_STREAM_CONSTANTS.stream_name
                        });
                    assert_eq!(
                        previous_directive_len,
                        next_field.directives.len() + 1,
                        "should have removed exactly one directive"
                    );
                    return Transformed::Replace(Selection::LinkedField(next_field));
                }
            }
        }
        transformed_field
    }
}

impl<'s> SkipUnreachableNodeTransform<'s> {
    fn new(program: &'s Program, validation_mode: &'s mut ValidationMode) -> Self {
        Self {
            visited_fragments: Default::default(),
            program,
            validation_mode,
        }
    }

    // Visit the fragment definition and return true if it should be deleted
    fn should_delete_fragment_definition(&mut self, key: StringKey) -> bool {
        if let Some((_, transformed)) = self.visited_fragments.get(&key) {
            return matches!(transformed, Transformed::Delete);
        }
        if let Some(fragment) = self.program.fragment(key) {
            self.visited_fragments
                .insert(key, (Arc::clone(fragment), Transformed::Keep));
            let transformed = self.transform_fragment(fragment);
            let should_delete = matches!(transformed, Transformed::Delete);

            // N.B. we must call self.visited_fragments.get* twice, because we cannot have
            // a reference to visited_opt and call transform_segment, which requires an
            // exclusive reference to self.
            let (_, visited_opt) = self.visited_fragments.get_mut(&key).unwrap();
            *visited_opt = transformed;
            should_delete
        } else {
            true
        }
    }

    fn map_selection_multi(&mut self, selection: &Selection) -> TransformedMulti<Selection> {
        match selection {
            Selection::FragmentSpread(selection) => {
                self.transform_fragment_spread(selection).into()
            }
            Selection::InlineFragment(selection) => self.transform_inline_fragment_multi(selection),
            Selection::LinkedField(selection) => self.transform_linked_field(selection).into(),
            Selection::ScalarField(selection) => self.transform_scalar_field(selection).into(),
            Selection::Condition(condition) => self.transform_constant_conditions(condition),
        }
    }

    fn transform_inline_fragment_multi(
        &mut self,
        inline_fragment: &InlineFragment,
    ) -> TransformedMulti<Selection> {
        if let Some(directive) = inline_fragment
            .directives
            .named(DEFER_STREAM_CONSTANTS.defer_name)
        {
            assert!(inline_fragment.directives.len() == 1);
            if let Some(if_arg) = DeferDirective::from(directive).if_arg {
                if let Value::Constant(ConstantValue::Boolean(false)) = &if_arg.value.item {
                    return TransformedMulti::ReplaceMultiple(
                        self.transform_selections(&inline_fragment.selections)
                            .replace_or_else(|| inline_fragment.selections.clone()),
                    );
                }
            }
        };
        self.default_transform_inline_fragment(inline_fragment)
            .into()
    }

    fn transform_constant_conditions(
        &mut self,
        condition: &Condition,
    ) -> TransformedMulti<Selection> {
        if let ConditionValue::Constant(b) = condition.value {
            // passing_value == "what the value needs to be in order to include the contents"
            // in other words, @skip has passing_value of false
            // and @include has passing_value of true
            //
            // so, b == condition.passing_value implies that the condition is redundant,
            // and b != condition.passing_value implies that the whole node can be removed
            let keep_contents = b == condition.passing_value;
            if keep_contents {
                // @include(if: true) or @skip(if: false)
                let contents = match self.transform_selections(&condition.selections) {
                    TransformedValue::Keep => &condition.selections,
                    TransformedValue::Replace(ref t) => t,
                }
                .clone();
                TransformedMulti::ReplaceMultiple(contents)
            } else {
                // @include(if: false) or @skip(if: true)
                TransformedMulti::Delete
            }
        } else {
            self.transform_condition(condition).into()
        }
    }
}

#[derive(Clone, Debug, Error, Eq, PartialEq, Ord, PartialOrd, Hash)]
enum ValidationMessage {
    #[error(
        "After applying transforms to the {document} `{name}` selections of \
        the `{name}` that would be sent to the server are empty. \
        This is likely due to the use of `@skip`/`@include` directives with \
        constant values that remove all selections in the {document}. "
    )]
    EmptySelectionsInDocument {
        name: StringKey,
        document: &'static str,
    },
}
<|endoftext|>


<commit_before>use alloc::boxed::Box;

use core::clone::Clone;
use core::option::Option::{self, Some, None};

use super::super::interrupts::{on, off};

use super::super::data_structures::{LazyGlobal};

use super::Process;

// current process
static CURRENT_PROCESS: LazyGlobal<Option<Box<Process>>> = lazy_global!();

// Init the current process
pub fn init() {
    unsafe {
        CURRENT_PROCESS.init(None);
    }
}

// Get the current process
// Must use a Box because size of Process is unknown at compile time
pub fn current() -> Option<Box<Process>> {
    // disable interrupts
    //off();

    // TODO: fix circular reference... use cli/sti?

    let ret = unsafe {
        (*CURRENT_PROCESS.get()).clone()
    };

    // enable interrupts
    //on();

    ret
}

// use for updating the current process
pub fn current_mut<'a>() -> Option<&'a mut Box<Process>> {
    // disable interrupts
    //off();

    // TODO: fix circular reference... use cli/sti?

    let ret = unsafe {
        match *CURRENT_PROCESS.get_mut() {
            Some(ref mut cp) => Some(cp),
            None => None,
        }
    };

    // enable interrupts
    //on();

    ret
}

// set the current process from a Box
pub fn set_current(process: Option<Box<Process>>) {
    // disable interrupts
    off();

    unsafe {
        *CURRENT_PROCESS.get_mut() = process;
    }

    // enable interrupts
    on();
}
<commit_msg>Using cli/sti in process::current instead of on/off<commit_after>use alloc::boxed::Box;

use core::clone::Clone;
use core::option::Option::{self, Some, None};

use super::super::interrupts::{on, off};

use super::super::data_structures::{LazyGlobal};

use super::Process;

// current process
static CURRENT_PROCESS: LazyGlobal<Option<Box<Process>>> = lazy_global!();

// Init the current process
pub fn init() {
    unsafe {
        CURRENT_PROCESS.init(None);
    }
}

// Get the current process
// Must use a Box because size of Process is unknown at compile time
pub fn current() -> Option<Box<Process>> {
    // DO NOT USE on/off here b/c there will be circular reference.
    // We can simply use cli/sti here since the disable cound should
    // not change in the middle of the function.

    use super::super::machine::{cli, sti};
    // disable interrupts
    unsafe { cli(); }

    let ret = unsafe {
        (*CURRENT_PROCESS.get()).clone()
    };

    // enable interrupts
    unsafe { sti(); }

    ret
}

// use for updating the current process
pub fn current_mut<'a>() -> Option<&'a mut Box<Process>> {
    // DO NOT USE on/off here b/c there will be circular reference.
    // We can simply use cli/sti here since the disable cound should
    // not change in the middle of the function.

    use super::super::machine::{cli, sti};
    // disable interrupts
    unsafe { cli(); }

    let ret = unsafe {
        match *CURRENT_PROCESS.get_mut() {
            Some(ref mut cp) => Some(cp),
            None => None,
        }
    };

    // enable interrupts
    unsafe { sti(); }

    ret
}

// set the current process from a Box
pub fn set_current(process: Option<Box<Process>>) {
    // disable interrupts
    off();

    unsafe {
        *CURRENT_PROCESS.get_mut() = process;
    }

    // enable interrupts
    on();
}
<|endoftext|>


<commit_before>use bifrost::utils::async_locks::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use core::borrow::BorrowMut;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::Future;
use index::btree::*;
use itertools::Itertools;
use owning_ref::{OwningHandle, OwningRef, RcRef};
use ram::cell::Cell;
use ram::schema::{Field, Schema};
use ram::types::*;
use std::cell::Ref;
use std::cell::RefCell;
use std::cell::RefMut;
use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use utils::lru_cache::LRUCache;
use index::Slice;
use index::EntryKey;
use std::marker::PhantomData;

const PAGE_SCHEMA: &'static str = "NEB_BTREE_PAGE";
const KEYS_FIELD: &'static str = "keys";
const NEXT_FIELD: &'static str = "next";
const PREV_FIELD: &'static str = "prev";

lazy_static! {
    static ref KEYS_KEY_HASH: u64 = key_hash(KEYS_FIELD);
    static ref NEXT_PAGE_KEY_HASH: u64 = key_hash(NEXT_FIELD);
    static ref PREV_PAGE_KEY_HASH: u64 = key_hash(PREV_FIELD);
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct ExtNode<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub id: Id,
    pub keys: KS,
    pub next: NodeCellRef,
    pub prev: NodeCellRef,
    pub len: usize,
    pub dirty: bool,
    mark: PhantomData<PS>
}

pub struct ExtNodeSplit<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub node_2: ExtNode<KS, PS>,
    pub keys_1_len: usize,
}

impl <KS, PS>ExtNode<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub fn new(id: Id) -> ExtNode<KS, PS> {
        ExtNode {
            id,
            keys: KS::init(),
            next: Node::<KS, PS>::none_ref(),
            prev: Node::<KS, PS>::none_ref(),
            len: 0,
            dirty: false,
            mark: PhantomData
        }
    }

    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        let prev_id = read_node(&self.prev, |node: &NodeReadHandler<KS, PS>| node.extnode().id);
        let next_id = read_node(&self.next, |node: &NodeReadHandler<KS, PS>| node.extnode().id);

        value[*NEXT_PAGE_KEY_HASH] = Value::Id(prev_id);
        value[*PREV_PAGE_KEY_HASH] = Value::Id(next_id);
        value[*KEYS_KEY_HASH] = self.keys.as_slice_immute()[..self.len]
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        self.keys.as_slice_immute()[..self.len]
            .binary_search(key)
            .unwrap_or_else(|i| i)
    }
    pub fn remove_at(&mut self, pos: usize) {
        let mut cached_len = &mut self.len;
        debug!("Removing from external pos {}, len {}, key {:?}",
               pos, cached_len, self.keys.as_slice()[pos]);
        self.keys.remove_at(pos, cached_len);
    }

    pub fn split_insert(
        &mut self,
        key: EntryKey,
        pos: usize,
        self_ref: &NodeCellRef,
        self_next: &mut NodeWriteGuard<KS, PS>,
        tree: &BPlusTree<KS, PS>
    ) -> (NodeCellRef, EntryKey) {
        // cached.dump();
        let pivot = self.len / 2;
        let new_page_id = tree.new_page_id();
        let mut keys_1 = &mut self.keys;
        let mut keys_2 = keys_1.split_at_pivot(pivot, self.len);
        let mut keys_1_len = pivot;
        let mut keys_2_len = self.len - pivot;
        // modify next node point previous to new node
        insert_into_split(
            key,
            keys_1,
            &mut keys_2,
            &mut keys_1_len,
            &mut keys_2_len,
            pos,
        );
        let pivot_key = keys_2.as_slice()[0].clone();
        let extnode_2: ExtNode<KS, PS> = ExtNode {
            id: new_page_id,
            keys: keys_2,
            next: self.next.clone(),
            prev: self_ref.clone(),
            len: keys_2_len,
            dirty: true,
            mark: PhantomData
        };
        debug_assert!(pivot_key > smallvec!(0));
        debug_assert!(
            &pivot_key > &keys_1.as_slice()[keys_1_len - 1],
            "{:?} / {:?} @ {}",
            pivot_key,
            &keys_1.as_slice()[keys_1_len - 1],
            pos
        );
        debug_assert!(read_unchecked::<KS, PS>(&extnode_2.prev).is_ext());
        self.len = keys_1_len;
        debug!(
            "Split to left len {}, right len {}, right prev id: {:?}",
            self.len,
            extnode_2.len,
            read_unchecked::<KS, PS>(&extnode_2.prev).ext_id()
        );
        let node_2 = NodeCellRef::new(Node::external(extnode_2));
        {
            if !self_next.is_none() {
                self_next.extnode_mut().prev = node_2.clone();
            }
        }
        self.next = node_2.clone();

        (node_2, pivot_key)
    }

    pub fn insert(
        &mut self,
        key: &EntryKey,
        tree: &BPlusTree<KS, PS>,
        self_ref: &NodeCellRef,
        parent: &NodeCellRef,
    ) -> Option<NodeSplit<KS, PS>> {
        let key = key.clone();
        let pos = self.search(&key);
        debug_assert!(self.len <= KS::slice_len());
        debug_assert!(pos <= self.len);
        if self.len == KS::slice_len() {
            // need to split
            debug!("insert to external with split, key {:?}, pos {}", key, pos);
            let mut self_next: NodeWriteGuard<KS, PS> = write_node(&self.next);
            let mut parent_latch: NodeWriteGuard<KS, PS> = write_node(parent);
            let (node_2, pivot_key) = self.split_insert(key, pos, self_ref, &mut self_next, tree);
            Some(NodeSplit {
                new_right_node: node_2,
                pivot: pivot_key,
                parent_latch,
                left_node_latch: NodeWriteGuard::default(),
            })
        } else {
            debug!("insert to external without split at {}, key {:?}", pos, key);
            self.keys.insert_at(key, pos, &mut self.len);
            None
        }
    }

    pub fn merge_with(&mut self, right: &mut Self) {
        debug!(
            "Merge external node, left len {}:{:?}, right len {}:{:?}",
            self.len, self.keys, right.len, right.keys
        );
        let self_len = self.len;
        let new_len = self.len + right.len;
        debug_assert!(new_len <= KS::slice_len());
        for i in self.len..new_len {
            mem::swap(&mut self.keys.as_slice()[i], &mut right.keys.as_slice()[i - self_len]);
        }
        self.len = new_len;
    }

    pub fn merge_sort(&mut self, right: &[&EntryKey]) {
        let self_len_before_merge = self.len;
        debug_assert!(self_len_before_merge + right.len() <= KS::slice_len());
        let mut new_slice = Self::new(self.id);
        let mut pos = self.len;
        let mut left_pos = 0;
        let mut right_pos = 0;
        while left_pos < self.len && right_pos < right.len() {
            let left_key = &self.keys.as_slice_immute()[left_pos];
            let right_key = &self.keys.as_slice_immute()[right_pos];
            if left_key < right_key {
                new_slice.keys.as_slice()[pos] = left_key.clone();
                left_pos += 1;
            } else {
                new_slice.keys.as_slice()[pos] = right_key.clone();
                right_pos += 1;
            }
            pos += 1;
        }
        for key in &self.keys.as_slice()[left_pos..] {
            new_slice.keys.as_slice()[pos] = key.clone();
            pos += 1;
            left_pos += 1;
        }
        for key in &right[right_pos..] {
            new_slice.keys.as_slice()[pos] = (*key).clone();
            pos += 1;
            right_pos += 1;
        }
        self.len = pos;
        debug_assert_eq!(self_len_before_merge, left_pos);
        debug_assert_eq!(right.len(), right_pos);
        debug_assert_eq!(self.len, self_len_before_merge + right.len());
    }

    pub fn dump(&self) {
        debug!("Dumping {:?}, keys {}", self.id, self.len);
        for i in 0..KS::slice_len() {
            debug!("{}\t- {:?}", i, self.keys.as_slice_immute()[i]);
        }
    }
    pub fn is_dirty(&self) -> bool {
        unimplemented!()
    }
}

pub fn page_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new(
            "*",
            0,
            false,
            false,
            Some(vec![
                Field::new(NEXT_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(PREV_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(KEYS_FIELD, type_id_of(Type::SmallBytes), false, true, None),
            ]),
        ),
    }
}

pub struct NodeCache<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    nodes: RefCell<HashMap<Id, NodeCellRef>>,
    storage: Arc<AsyncClient>,
    marker: PhantomData<(KS, PS)>
}

impl <KS, PS>NodeCache<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        NodeCache {
            nodes: RefCell::new(HashMap::new()),
            storage: neb_client.clone(),
            marker: PhantomData
        }
    }

    pub fn get(&self, id: &Id) -> NodeCellRef {
        if id.is_unit_id() {
            return NodeCellRef::new(Node::<KS, PS>::none());
        }
        let mut nodes = self.nodes.borrow_mut();
        nodes
            .entry(*id)
            .or_insert_with(|| {
                let cell = self.storage.read_cell(*id).wait().unwrap().unwrap();
                NodeCellRef::new(Node::external(self.extnode_from_cell(cell)))
            })
            .clone()
    }

    fn extnode_from_cell(&self, cell: Cell) -> ExtNode<KS, PS> {
        let cell_id = cell.id();
        let cell_version = cell.header.version;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        let mut key_slice = KS::init();
        let mut key_count = 0;
        for (i, key_val) in keys_array.iter().enumerate() {
            key_slice.as_slice()[i] = EntryKey::from(key_val.as_slice());
            key_count += 1;
        }
        ExtNode {
            id: cell_id,
            keys: key_slice,
            next: self.get(next),
            prev: self.get(prev),
            len: key_count,
            dirty: false,
            mark: PhantomData
        }
    }
}
<commit_msg>node replacement<commit_after>use bifrost::utils::async_locks::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use core::borrow::BorrowMut;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::Future;
use index::btree::*;
use itertools::Itertools;
use owning_ref::{OwningHandle, OwningRef, RcRef};
use ram::cell::Cell;
use ram::schema::{Field, Schema};
use ram::types::*;
use std::cell::Ref;
use std::cell::RefCell;
use std::cell::RefMut;
use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use utils::lru_cache::LRUCache;
use index::Slice;
use index::EntryKey;
use std::marker::PhantomData;

const PAGE_SCHEMA: &'static str = "NEB_BTREE_PAGE";
const KEYS_FIELD: &'static str = "keys";
const NEXT_FIELD: &'static str = "next";
const PREV_FIELD: &'static str = "prev";

lazy_static! {
    static ref KEYS_KEY_HASH: u64 = key_hash(KEYS_FIELD);
    static ref NEXT_PAGE_KEY_HASH: u64 = key_hash(NEXT_FIELD);
    static ref PREV_PAGE_KEY_HASH: u64 = key_hash(PREV_FIELD);
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct ExtNode<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub id: Id,
    pub keys: KS,
    pub next: NodeCellRef,
    pub prev: NodeCellRef,
    pub len: usize,
    pub dirty: bool,
    mark: PhantomData<PS>
}

pub struct ExtNodeSplit<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub node_2: ExtNode<KS, PS>,
    pub keys_1_len: usize,
}

impl <KS, PS>ExtNode<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub fn new(id: Id) -> ExtNode<KS, PS> {
        ExtNode {
            id,
            keys: KS::init(),
            next: Node::<KS, PS>::none_ref(),
            prev: Node::<KS, PS>::none_ref(),
            len: 0,
            dirty: false,
            mark: PhantomData
        }
    }

    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        let prev_id = read_node(&self.prev, |node: &NodeReadHandler<KS, PS>| node.extnode().id);
        let next_id = read_node(&self.next, |node: &NodeReadHandler<KS, PS>| node.extnode().id);

        value[*NEXT_PAGE_KEY_HASH] = Value::Id(prev_id);
        value[*PREV_PAGE_KEY_HASH] = Value::Id(next_id);
        value[*KEYS_KEY_HASH] = self.keys.as_slice_immute()[..self.len]
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        self.keys.as_slice_immute()[..self.len]
            .binary_search(key)
            .unwrap_or_else(|i| i)
    }
    pub fn remove_at(&mut self, pos: usize) {
        let mut cached_len = &mut self.len;
        debug!("Removing from external pos {}, len {}, key {:?}",
               pos, cached_len, self.keys.as_slice()[pos]);
        self.keys.remove_at(pos, cached_len);
    }

    pub fn split_insert(
        &mut self,
        key: EntryKey,
        pos: usize,
        self_ref: &NodeCellRef,
        self_next: &mut NodeWriteGuard<KS, PS>,
        tree: &BPlusTree<KS, PS>
    ) -> (NodeCellRef, EntryKey) {
        // cached.dump();
        let pivot = self.len / 2;
        let new_page_id = tree.new_page_id();
        let mut keys_1 = &mut self.keys;
        let mut keys_2 = keys_1.split_at_pivot(pivot, self.len);
        let mut keys_1_len = pivot;
        let mut keys_2_len = self.len - pivot;
        // modify next node point previous to new node
        insert_into_split(
            key,
            keys_1,
            &mut keys_2,
            &mut keys_1_len,
            &mut keys_2_len,
            pos,
        );
        let pivot_key = keys_2.as_slice()[0].clone();
        let extnode_2: ExtNode<KS, PS> = ExtNode {
            id: new_page_id,
            keys: keys_2,
            next: self.next.clone(),
            prev: self_ref.clone(),
            len: keys_2_len,
            dirty: true,
            mark: PhantomData
        };
        debug_assert!(pivot_key > smallvec!(0));
        debug_assert!(
            &pivot_key > &keys_1.as_slice()[keys_1_len - 1],
            "{:?} / {:?} @ {}",
            pivot_key,
            &keys_1.as_slice()[keys_1_len - 1],
            pos
        );
        debug_assert!(read_unchecked::<KS, PS>(&extnode_2.prev).is_ext());
        self.len = keys_1_len;
        debug!(
            "Split to left len {}, right len {}, right prev id: {:?}",
            self.len,
            extnode_2.len,
            read_unchecked::<KS, PS>(&extnode_2.prev).ext_id()
        );
        let node_2 = NodeCellRef::new(Node::external(extnode_2));
        {
            if !self_next.is_none() {
                self_next.extnode_mut().prev = node_2.clone();
            }
        }
        self.next = node_2.clone();

        (node_2, pivot_key)
    }

    pub fn insert(
        &mut self,
        key: &EntryKey,
        tree: &BPlusTree<KS, PS>,
        self_ref: &NodeCellRef,
        parent: &NodeCellRef,
    ) -> Option<NodeSplit<KS, PS>> {
        let key = key.clone();
        let pos = self.search(&key);
        debug_assert!(self.len <= KS::slice_len());
        debug_assert!(pos <= self.len);
        if self.len == KS::slice_len() {
            // need to split
            debug!("insert to external with split, key {:?}, pos {}", key, pos);
            let mut self_next: NodeWriteGuard<KS, PS> = write_node(&self.next);
            let mut parent_latch: NodeWriteGuard<KS, PS> = write_node(parent);
            let (node_2, pivot_key) = self.split_insert(key, pos, self_ref, &mut self_next, tree);
            Some(NodeSplit {
                new_right_node: node_2,
                pivot: pivot_key,
                parent_latch,
                left_node_latch: NodeWriteGuard::default(),
            })
        } else {
            debug!("insert to external without split at {}, key {:?}", pos, key);
            self.keys.insert_at(key, pos, &mut self.len);
            None
        }
    }

    pub fn merge_with(&mut self, right: &mut Self) {
        debug!(
            "Merge external node, left len {}:{:?}, right len {}:{:?}",
            self.len, self.keys, right.len, right.keys
        );
        let self_len = self.len;
        let new_len = self.len + right.len;
        debug_assert!(new_len <= KS::slice_len());
        for i in self.len..new_len {
            mem::swap(&mut self.keys.as_slice()[i], &mut right.keys.as_slice()[i - self_len]);
        }
        self.len = new_len;
    }

    pub fn merge_sort(&mut self, right: &[&EntryKey]) {
        let self_len_before_merge = self.len;
        debug_assert!(self_len_before_merge + right.len() <= KS::slice_len());
        let mut new_node = Self::new(self.id);
        let mut pos = self.len;
        let mut left_pos = 0;
        let mut right_pos = 0;
        while left_pos < self.len && right_pos < right.len() {
            let left_key = &self.keys.as_slice_immute()[left_pos];
            let right_key = &self.keys.as_slice_immute()[right_pos];
            if left_key < right_key {
                new_node.keys.as_slice()[pos] = left_key.clone();
                left_pos += 1;
            } else {
                new_node.keys.as_slice()[pos] = right_key.clone();
                right_pos += 1;
            }
            pos += 1;
        }
        for key in &self.keys.as_slice()[left_pos..] {
            new_node.keys.as_slice()[pos] = key.clone();
            pos += 1;
            left_pos += 1;
        }
        for key in &right[right_pos..] {
            new_node.keys.as_slice()[pos] = (*key).clone();
            pos += 1;
            right_pos += 1;
        }
        new_node.len = pos;
        new_node.dirty = self.dirty;
        new_node.next = self.next.clone();
        new_node.prev = self.prev.clone();
        *self = new_node;
        debug_assert_eq!(self_len_before_merge, left_pos);
        debug_assert_eq!(right.len(), right_pos);
        debug_assert_eq!(self.len, self_len_before_merge + right.len());
    }

    pub fn dump(&self) {
        debug!("Dumping {:?}, keys {}", self.id, self.len);
        for i in 0..KS::slice_len() {
            debug!("{}\t- {:?}", i, self.keys.as_slice_immute()[i]);
        }
    }
    pub fn is_dirty(&self) -> bool {
        unimplemented!()
    }
}

pub fn page_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new(
            "*",
            0,
            false,
            false,
            Some(vec![
                Field::new(NEXT_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(PREV_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(KEYS_FIELD, type_id_of(Type::SmallBytes), false, true, None),
            ]),
        ),
    }
}

pub struct NodeCache<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    nodes: RefCell<HashMap<Id, NodeCellRef>>,
    storage: Arc<AsyncClient>,
    marker: PhantomData<(KS, PS)>
}

impl <KS, PS>NodeCache<KS, PS>
    where KS: Slice<EntryKey> + Debug + 'static,
          PS: Slice<NodeCellRef> + 'static
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        NodeCache {
            nodes: RefCell::new(HashMap::new()),
            storage: neb_client.clone(),
            marker: PhantomData
        }
    }

    pub fn get(&self, id: &Id) -> NodeCellRef {
        if id.is_unit_id() {
            return NodeCellRef::new(Node::<KS, PS>::none());
        }
        let mut nodes = self.nodes.borrow_mut();
        nodes
            .entry(*id)
            .or_insert_with(|| {
                let cell = self.storage.read_cell(*id).wait().unwrap().unwrap();
                NodeCellRef::new(Node::external(self.extnode_from_cell(cell)))
            })
            .clone()
    }

    fn extnode_from_cell(&self, cell: Cell) -> ExtNode<KS, PS> {
        let cell_id = cell.id();
        let cell_version = cell.header.version;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        let mut key_slice = KS::init();
        let mut key_count = 0;
        for (i, key_val) in keys_array.iter().enumerate() {
            key_slice.as_slice()[i] = EntryKey::from(key_val.as_slice());
            key_count += 1;
        }
        ExtNode {
            id: cell_id,
            keys: key_slice,
            next: self.get(next),
            prev: self.get(prev),
            len: key_count,
            dirty: false,
            mark: PhantomData
        }
    }
}
<|endoftext|>


<commit_before>// Implementation details for box writing

use std::io::Write;
use std::{io, cmp, mem, u16};
use sodiumoxide::crypto::secretbox;

use crypto::{CYPHER_HEADER_SIZE, MAX_PACKET_SIZE, MAX_PACKET_USIZE, encrypt_packet, final_header};

const WRITE_BUFFER_SIZE: usize = CYPHER_HEADER_SIZE + MAX_PACKET_USIZE;

// Buffers encrypted bytes, so that the stream can correctly resume even if the
// underlying stream is unable to write a whole header + packet combination at
// once.
pub struct WriterBuffer {
    // Stores the result of a call to `encrypt_packet`.
    buffer: [u8; WRITE_BUFFER_SIZE],
    // Indicates where to resume writing, or whether to encrypt the next packet (if greater than the size of `buffer`).
    offset: u16,
    // Length of the data that is actually relevant, everything from buffer[length] is useless data from a previous packet.
    length: u16,
}

impl WriterBuffer {
    pub fn new() -> WriterBuffer {
        WriterBuffer {
            buffer: unsafe { mem::uninitialized() },
            offset: CYPHER_HEADER_SIZE as u16 + MAX_PACKET_SIZE,
            length: 0,
        }
    }

    fn is_empty(&self) -> bool {
        self.offset >= self.length
    }

    // Adds data to this buffer and returns how many bytes of the data source were added.
    fn fill(&mut self, data: &[u8], key: &secretbox::Key, nonce: &mut secretbox::Nonce) -> u16 {
        let packet_length = cmp::min(data.len() as u16, MAX_PACKET_SIZE);
        self.length = packet_length + CYPHER_HEADER_SIZE as u16;
        self.offset = 0;

        unsafe {
            encrypt_packet(self.buffer.as_mut_ptr(),
                           data.as_ptr(),
                           packet_length,
                           &key.0,
                           &mut nonce.0);
        }

        packet_length
    }

    // Writes data from the buffer and updates the buffer's data offset.
    fn write_buffered_data<W: Write>(&mut self, writer: &mut W) -> io::Result<usize> {
        let written = writer
            .write(&self.buffer[self.offset as usize..self.length as usize])?;
        self.offset += written as u16;
        Ok(written)
    }

    // Flushes the buffer to a writer, propagating the first encountered error.
    // Does not call flush on the writer.
    fn flush_to<W: Write>(&mut self, writer: &mut W) -> io::Result<()> {
        while !self.is_empty() {
            self.write_buffered_data(writer)?;
        }
        Ok(())
    }
}

// Implements box writing. The different streams delegate to this in `write`.
pub fn do_write<W: Write>(data: &[u8],
                          writer: &mut W,
                          key: &secretbox::Key,
                          nonce: &mut secretbox::Nonce,
                          buffer: &mut WriterBuffer)
                          -> io::Result<usize> {

    let buffered: u16;

    if buffer.is_empty() {
        buffered = buffer.fill(data, key, nonce);
    } else {
        buffered = 0;
    }

    buffer.write_buffered_data(writer)?;
    Ok(buffered as usize)
}

// Implements box flushing. The different streams delegate to this in `flush`.
pub fn do_flush<W: Write>(writer: &mut W, buffer: &mut WriterBuffer) -> io::Result<()> {
    buffer.flush_to(writer)?;

    writer.flush()
}

// Implements box shutdown. The different streams delegate to this in `shutdown`.
// TODO make sure to only write one final header until this returns a ok
pub fn do_shutdown<W: Write>(writer: &mut W,
                             key: &secretbox::Key,
                             nonce: &secretbox::Nonce,
                             buffer: &mut WriterBuffer)
                             -> io::Result<()> {
    buffer.flush_to(writer)?;

    unsafe {
        final_header(&mut *(buffer.buffer.as_mut_ptr() as *mut [u8; CYPHER_HEADER_SIZE]),
                     &key.0,
                     &nonce.0);
    }
    buffer.offset = 0;
    buffer.length = CYPHER_HEADER_SIZE as u16;

    buffer.flush_to(writer)?;
    writer.flush()
}
<commit_msg>Prevent writing multiple final headers<commit_after>// Implementation details for box writing

use std::io::Write;
use std::{io, cmp, mem, u16};
use sodiumoxide::crypto::secretbox;

use crypto::{CYPHER_HEADER_SIZE, MAX_PACKET_SIZE, MAX_PACKET_USIZE, encrypt_packet, final_header};

const WRITE_BUFFER_SIZE: usize = CYPHER_HEADER_SIZE + MAX_PACKET_USIZE;

// Buffers encrypted bytes, so that the stream can correctly resume even if the
// underlying stream is unable to write a whole header + packet combination at
// once.
pub struct WriterBuffer {
    // Stores the result of a call to `encrypt_packet`.
    buffer: [u8; WRITE_BUFFER_SIZE],
    // Indicates where to resume writing, or whether to encrypt the next packet (if greater than the size of `buffer`).
    offset: u16,
    // Length of the data that is actually relevant, everything from buffer[length] is useless data from a previous packet.
    length: u16,
    // set to true to indicate that a final header has been written and should not be written again.
    done: bool,
}

impl WriterBuffer {
    pub fn new() -> WriterBuffer {
        WriterBuffer {
            buffer: unsafe { mem::uninitialized() },
            offset: CYPHER_HEADER_SIZE as u16 + MAX_PACKET_SIZE,
            length: 0,
            done: false,
        }
    }

    fn is_empty(&self) -> bool {
        self.offset >= self.length
    }

    // Adds data to this buffer and returns how many bytes of the data source were added.
    fn fill(&mut self, data: &[u8], key: &secretbox::Key, nonce: &mut secretbox::Nonce) -> u16 {
        let packet_length = cmp::min(data.len() as u16, MAX_PACKET_SIZE);
        self.length = packet_length + CYPHER_HEADER_SIZE as u16;
        self.offset = 0;

        unsafe {
            encrypt_packet(self.buffer.as_mut_ptr(),
                           data.as_ptr(),
                           packet_length,
                           &key.0,
                           &mut nonce.0);
        }

        packet_length
    }

    // Writes data from the buffer and updates the buffer's data offset.
    fn write_buffered_data<W: Write>(&mut self, writer: &mut W) -> io::Result<usize> {
        let written = writer
            .write(&self.buffer[self.offset as usize..self.length as usize])?;
        self.offset += written as u16;
        Ok(written)
    }

    // Flushes the buffer to a writer, propagating the first encountered error.
    // Does not call flush on the writer.
    fn flush_to<W: Write>(&mut self, writer: &mut W) -> io::Result<()> {
        while !self.is_empty() {
            self.write_buffered_data(writer)?;
        }
        Ok(())
    }
}

// Implements box writing. The different streams delegate to this in `write`.
pub fn do_write<W: Write>(data: &[u8],
                          writer: &mut W,
                          key: &secretbox::Key,
                          nonce: &mut secretbox::Nonce,
                          buffer: &mut WriterBuffer)
                          -> io::Result<usize> {

    let buffered: u16;

    if buffer.is_empty() {
        buffered = buffer.fill(data, key, nonce);
    } else {
        buffered = 0;
    }

    buffer.write_buffered_data(writer)?;
    Ok(buffered as usize)
}

// Implements box flushing. The different streams delegate to this in `flush`.
pub fn do_flush<W: Write>(writer: &mut W, buffer: &mut WriterBuffer) -> io::Result<()> {
    buffer.flush_to(writer)?;

    writer.flush()
}

// Implements box shutdown. The different streams delegate to this in `shutdown`.
pub fn do_shutdown<W: Write>(writer: &mut W,
                             key: &secretbox::Key,
                             nonce: &secretbox::Nonce,
                             buffer: &mut WriterBuffer)
                             -> io::Result<()> {
    if !buffer.done {
        buffer.flush_to(writer)?;

        unsafe {
            final_header(&mut *(buffer.buffer.as_mut_ptr() as *mut [u8; CYPHER_HEADER_SIZE]),
                         &key.0,
                         &nonce.0);
        }
        buffer.offset = 0;
        buffer.length = CYPHER_HEADER_SIZE as u16;
        buffer.done = true;
    }

    buffer.flush_to(writer)?;
    writer.flush()
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_id="cgmath#0.1"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![comment = "A mathematics library for computer graphics."]
#![license = "ASL2"]

#![feature(globs)]
#![feature(macro_rules)]

//! Computer graphics-centric math.
//!
//! This crate provides useful mathematical primitives and operations on them.
//! It is organized into one module per primitive. The core structures are
//! vectors and matrices. A strongly-typed interface is provided, to prevent
//! mixing units or violating mathematical invariants.
//!
//! Transformations are not usually done directly on matrices, but go through
//! transformation objects that can be converted to matrices. Rotations go
//! through the `Basis` types, which are guaranteed to be orthogonal matrices.
//! Despite this, one can directly create a limited rotation matrix using the
//! `look_at`, `from_angle`, `from_euler`, and `from_axis_angle` methods.
//! These are provided for convenience.

pub mod array;

pub mod matrix;
pub mod quaternion;
pub mod vector;

pub mod angle;
pub mod plane;
pub mod point;
pub mod line;
pub mod ray;
pub mod rotation;
pub mod transform;

pub mod projection;

pub mod aabb;
pub mod cylinder;
pub mod frustum;
pub mod intersect;
pub mod obb;
pub mod sphere;

pub mod approx;
pub mod num;
<commit_msg>Remove crate_id.<commit_after>// Copyright 2013-2014 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![comment = "A mathematics library for computer graphics."]
#![license = "ASL2"]

#![feature(globs)]
#![feature(macro_rules)]

//! Computer graphics-centric math.
//!
//! This crate provides useful mathematical primitives and operations on them.
//! It is organized into one module per primitive. The core structures are
//! vectors and matrices. A strongly-typed interface is provided, to prevent
//! mixing units or violating mathematical invariants.
//!
//! Transformations are not usually done directly on matrices, but go through
//! transformation objects that can be converted to matrices. Rotations go
//! through the `Basis` types, which are guaranteed to be orthogonal matrices.
//! Despite this, one can directly create a limited rotation matrix using the
//! `look_at`, `from_angle`, `from_euler`, and `from_axis_angle` methods.
//! These are provided for convenience.

pub mod array;

pub mod matrix;
pub mod quaternion;
pub mod vector;

pub mod angle;
pub mod plane;
pub mod point;
pub mod line;
pub mod ray;
pub mod rotation;
pub mod transform;

pub mod projection;

pub mod aabb;
pub mod cylinder;
pub mod frustum;
pub mod intersect;
pub mod obb;
pub mod sphere;

pub mod approx;
pub mod num;
<|endoftext|>


<commit_before>use devp2p::RLPxNode;
use ethereum_types::H512;
use hex_literal::hex;
use libsecp256k1::SecretKey;
use rand::rngs::OsRng;

const REMOTE_ID: H512 = H512(hex!("d860a01f9722d78051619d1e2351aba3f43f943f6f00718d1b9baa4101932a1f5011f16bb2b1bb35db20d6fe28fa0bf09636d26a87d31de9ec6203eeedb1f666"));

#[tokio::main]
async fn main() {
    let _ = env_logger::init();

    let addr = "18.138.108.67:30303".parse().unwrap();
    let client = RLPxNode::new(
        SecretKey::random(&mut OsRng),
        "rust-devp2p/0.1.0".to_string(),
        None,
    )
    .await
    .unwrap();

    client.add_peer(addr, REMOTE_ID).await.unwrap();

    let timeout = 5;
    tokio::time::delay_for(std::time::Duration::from_secs(timeout)).await;
}
<commit_msg>Improve local_connect example<commit_after>use devp2p::*;
use libsecp256k1::SecretKey;
use log::*;
use rand::{rngs::OsRng, seq::SliceRandom};
use std::time::Duration;
use tokio::time::delay_for;

#[tokio::main]
async fn main() {
    let _ = env_logger::init();

    // Bootnodes as used in OpenEthereum
    let bootnodes = vec![
        "enode://d860a01f9722d78051619d1e2351aba3f43f943f6f00718d1b9baa4101932a1f5011f16bb2b1bb35db20d6fe28fa0bf09636d26a87d31de9ec6203eeedb1f666@18.138.108.67:30303",
		"enode://22a8232c3abc76a16ae9d6c3b164f98775fe226f0917b0ca871128a74a8e9630b458460865bab457221f1d448dd9791d24c4e5d88786180ac185df813a68d4de@3.209.45.79:30303",
		"enode://ca6de62fce278f96aea6ec5a2daadb877e51651247cb96ee310a318def462913b653963c155a0ef6c7d50048bba6e6cea881130857413d9f50a621546b590758@34.255.23.113:30303",
		"enode://279944d8dcd428dffaa7436f25ca0ca43ae19e7bcf94a8fb7d1641651f92d121e972ac2e8f381414b80cc8e5555811c2ec6e1a99bb009b3f53c4c69923e11bd8@35.158.244.151:30303",
		"enode://8499da03c47d637b20eee24eec3c356c9a2e6148d6fe25ca195c7949ab8ec2c03e3556126b0d7ed644675e78c4318b08691b7b57de10e5f0d40d05b09238fa0a@52.187.207.27:30303",
		"enode://103858bdb88756c71f15e9b5e09b56dc1be52f0a5021d46301dbbfb7e130029cc9d0d6f73f693bc29b665770fff7da4d34f3c6379fe12721b5d7a0bcb5ca1fc1@191.234.162.198:30303",
		"enode://715171f50508aba88aecd1250af392a45a330af91d7b90701c436b618c86aaa1589c9184561907bebbb56439b8f8787bc01f49a7c77276c58c1b09822d75e8e8@52.231.165.108:30303",
		"enode://5d6d7cd20d6da4bb83a1d28cadb5d409b64edf314c0335df658c1a54e32c7c4a7ab7823d57c39b6a757556e68ff1df17c748b698544a55cb488b52479a92b60f@104.42.217.25:30303",
		"enode://68f46370191198b71a1595dd453c489bbfe28036a9951fc0397fabd1b77462930b3c5a5359b20e99677855939be47b39fc8edcf1e9ff2522a922b86d233bf2df@144.217.153.76:30303",
		"enode://ffed6382e05ee42854d862f08e4e39b8452c50a5a5d399072c40f9a0b2d4ad34b0eb5312455ad8bcf0dcb4ce969dc89a9a9fd00183eaf8abf46bbcc59dc6e9d5@51.195.3.238:30303",
		"enode://b47b197244c054d385f25d7740b33cc7e2a74d6f715befad2b789fd3e3594bb1c8dd2ca2faf1a3bf6b4c9ec03e53b52301f722a2316b78976be03ccbe703c581@54.37.94.238:30303",
		"enode://5f7d0794c464b2fcd514d41e16e4b535a98ac792a71ca9667c7cef35595dc34c9a1b793c0622554cf87f34006942abb526af7d2e37d715ac32ed02170556cce2@51.161.101.207:30303",
    ];

    let node = bootnodes
        .choose(&mut OsRng)
        .unwrap()
        .parse::<NodeRecord>()
        .unwrap();

    debug!("Connecting to {}", node.addr);

    let client = RLPxNode::new(
        SecretKey::random(&mut OsRng),
        "rust-devp2p/0.1.0".to_string(),
        None,
    )
    .await
    .unwrap();

    client.add_peer(node).await.unwrap();

    let timeout = 5;
    delay_for(Duration::from_secs(timeout)).await;
}
<|endoftext|>


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

//! Logging configuration types

/// Log level types
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum LogLevel {
	/// Error
	Error,
	/// Warning
	Warning,
	/// Info
	Info,
	/// Debug
	Debug,
	/// Trace
	Trace,
}

/// Logging config
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LoggingConfig {
	/// whether to log to stdout
	pub log_to_stdout: bool,
	/// logging level for stdout
	pub stdout_log_level: LogLevel,
	/// whether to log to file
	pub log_to_file: bool,
	/// log file level
	pub file_log_level: LogLevel,
	/// Log file path
	pub log_file_path: String,
	/// Whether to append to log or replace
	pub log_file_append: bool,
	/// Size of the log in bytes to rotate over (optional)
	pub log_max_size: Option<u64>,
	/// Whether the tui is running (optional)
	pub tui_running: Option<bool>,
}

impl Default for LoggingConfig {
	fn default() -> LoggingConfig {
		LoggingConfig {
			log_to_stdout: true,
			stdout_log_level: LogLevel::Warning,
			log_to_file: true,
			file_log_level: LogLevel::Debug,
			log_file_path: String::from("grin.log"),
			log_file_append: true,
			log_max_size: Some(1024 * 1024 * 16), // 16 megabytes default
			tui_running: None,
		}
	}
}

use std::ops::Deref;
use zeroize::Zeroize;
/// Zeroing string, mainly useful for password
#[derive(Clone, PartialEq, PartialOrd)]
pub struct ZeroingString(String);

impl Drop for ZeroingString {
	fn drop(&mut self) {
		self.0.zeroize();
	}
}

impl From<&str> for ZeroingString {
	fn from(s: &str) -> Self {
		ZeroingString(String::from(s))
	}
}

impl From<String> for ZeroingString {
	fn from(s: String) -> Self {
		ZeroingString(s)
	}
}

impl Deref for ZeroingString {
	type Target = str;

	fn deref(&self) -> &str {
		&self.0
	}
}
<commit_msg>default log level DEBUG -> INFO (#2420)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Logging configuration types

/// Log level types
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum LogLevel {
	/// Error
	Error,
	/// Warning
	Warning,
	/// Info
	Info,
	/// Debug
	Debug,
	/// Trace
	Trace,
}

/// Logging config
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LoggingConfig {
	/// whether to log to stdout
	pub log_to_stdout: bool,
	/// logging level for stdout
	pub stdout_log_level: LogLevel,
	/// whether to log to file
	pub log_to_file: bool,
	/// log file level
	pub file_log_level: LogLevel,
	/// Log file path
	pub log_file_path: String,
	/// Whether to append to log or replace
	pub log_file_append: bool,
	/// Size of the log in bytes to rotate over (optional)
	pub log_max_size: Option<u64>,
	/// Whether the tui is running (optional)
	pub tui_running: Option<bool>,
}

impl Default for LoggingConfig {
	fn default() -> LoggingConfig {
		LoggingConfig {
			log_to_stdout: true,
			stdout_log_level: LogLevel::Warning,
			log_to_file: true,
			file_log_level: LogLevel::Info,
			log_file_path: String::from("grin.log"),
			log_file_append: true,
			log_max_size: Some(1024 * 1024 * 16), // 16 megabytes default
			tui_running: None,
		}
	}
}

use std::ops::Deref;
use zeroize::Zeroize;
/// Zeroing string, mainly useful for password
#[derive(Clone, PartialEq, PartialOrd)]
pub struct ZeroingString(String);

impl Drop for ZeroingString {
	fn drop(&mut self) {
		self.0.zeroize();
	}
}

impl From<&str> for ZeroingString {
	fn from(s: &str) -> Self {
		ZeroingString(String::from(s))
	}
}

impl From<String> for ZeroingString {
	fn from(s: String) -> Self {
		ZeroingString(s)
	}
}

impl Deref for ZeroingString {
	type Target = str;

	fn deref(&self) -> &str {
		&self.0
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::io::{Result, Write};
use std::option::Option;
use std::sync::Arc;

use crate::rocks::RateLimiter;

const PRIORITY_HIGH: u8 = 1;
const REFILL_PERIOD: i64 = 100 * 1000;
const FARENESS: i32 = 10;
const SNAP_MAX_BYTES_PER_TIME: i64 = 4 * 1024 * 1024;
pub const DEFAULT_SNAP_MAX_BYTES_PER_SEC: u64 = 100 * 1024 * 1024;

/// The I/O rate limiter for RocksDB.
///
/// Throttles the maximum bytes per second written to disk.
pub struct IOLimiter {
    inner: RateLimiter,
}

impl IOLimiter {
    /// # Arguments
    ///
    /// - `bytes_per_sec`: controls the total write rate of compaction and flush in bytes per second.
    pub fn new(bytes_per_sec: u64) -> IOLimiter {
        IOLimiter {
            inner: RateLimiter::new(bytes_per_sec as i64, REFILL_PERIOD, FARENESS),
        }
    }

    /// Sets the rate limit in bytes per second
    pub fn set_bytes_per_second(&self, bytes_per_sec: i64) {
        self.inner.set_bytes_per_second(bytes_per_sec)
    }

    /// Requests an access token to read or write bytes. If this request can not be satisfied, the call is blocked.
    pub fn request(&self, bytes: i64) {
        self.inner.request(bytes, PRIORITY_HIGH)
    }

    /// Gets the max bytes that can be granted in a single burst.
    /// Note: it will be less than or equal to `SNAP_MAX_BYTES_PER_TIME`.
    pub fn get_max_bytes_per_time(&self) -> i64 {
        if self.inner.get_singleburst_bytes() > SNAP_MAX_BYTES_PER_TIME {
            SNAP_MAX_BYTES_PER_TIME
        } else {
            self.inner.get_singleburst_bytes()
        }
    }

    /// Gets the total bytes that have gone through the rate limiter.
    pub fn get_total_bytes_through(&self) -> i64 {
        self.inner.get_total_bytes_through(PRIORITY_HIGH)
    }

    /// Gets the rate limit in bytes per second.
    pub fn get_bytes_per_second(&self) -> i64 {
        self.inner.get_bytes_per_second()
    }

    /// Gets the total number of requests that have gone through rate limiter
    pub fn get_total_requests(&self) -> i64 {
        self.inner.get_total_requests(PRIORITY_HIGH)
    }
}

pub struct LimitWriter<'a, T: Write> {
    limiter: Option<Arc<IOLimiter>>,
    writer: &'a mut T,
}

impl<'a, T: Write + 'a> LimitWriter<'a, T> {
    pub fn new(limiter: Option<Arc<IOLimiter>>, writer: &'a mut T) -> LimitWriter<'a, T> {
        LimitWriter { limiter, writer }
    }
}

impl<'a, T: Write + 'a> Write for LimitWriter<'a, T> {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let total = buf.len();
        if let Some(ref limiter) = self.limiter {
            let single = limiter.get_max_bytes_per_time() as usize;
            let mut curr = 0;
            let mut end;
            while curr < total {
                if curr + single >= total {
                    end = total;
                } else {
                    end = curr + single;
                }
                limiter.request((end - curr) as i64);
                self.writer.write_all(&buf[curr..end])?;
                curr = end;
            }
        } else {
            self.writer.write_all(buf)?;
        }
        Ok(total)
    }

    fn flush(&mut self) -> Result<()> {
        self.writer.flush()?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::fs::{self, File};
    use std::io::Write;
    use std::sync::Arc;
    use tempfile::Builder;

    use super::{IOLimiter, LimitWriter, SNAP_MAX_BYTES_PER_TIME};

    #[test]
    fn test_io_limiter() {
        let limiter = IOLimiter::new(10 * 1024 * 1024);
        assert!(limiter.get_max_bytes_per_time() <= SNAP_MAX_BYTES_PER_TIME);

        limiter.set_bytes_per_second(20 * 1024 * 1024);
        assert_eq!(limiter.get_bytes_per_second(), 20 * 1024 * 1024);

        assert_eq!(limiter.get_total_bytes_through(), 0);

        limiter.request(1024 * 1024);
        assert_eq!(limiter.get_total_bytes_through(), 1024 * 1024);

        assert_eq!(limiter.get_total_requests(), 1);
    }

    #[test]
    fn test_limit_writer() {
        let dir = Builder::new()
            .prefix("_test_limit_writer")
            .tempdir()
            .unwrap();
        let path = dir.path().join("test-file");
        let mut file = File::create(&path).unwrap();
        let mut limit_writer = LimitWriter::new(Some(Arc::new(IOLimiter::new(1024))), &mut file);

        let mut s = String::new();
        for _ in 0..100 {
            s.push_str("Hello, World!");
        }
        limit_writer.write_all(s.as_bytes()).unwrap();
        limit_writer.flush().unwrap();

        let contents = fs::read_to_string(&path).unwrap();
        assert_eq!(contents, s);
    }
}
<commit_msg>engine: add LimitReader (#5434)<commit_after>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::io::{Read, Result, Write};
use std::option::Option;
use std::sync::Arc;

use crate::rocks::RateLimiter;

const PRIORITY_HIGH: u8 = 1;
const REFILL_PERIOD: i64 = 100 * 1000;
const FARENESS: i32 = 10;
const SNAP_MAX_BYTES_PER_TIME: i64 = 4 * 1024 * 1024;
pub const DEFAULT_SNAP_MAX_BYTES_PER_SEC: u64 = 100 * 1024 * 1024;

/// The I/O rate limiter for RocksDB.
///
/// Throttles the maximum bytes per second written or read.
pub struct IOLimiter {
    inner: RateLimiter,
}

impl IOLimiter {
    /// # Arguments
    ///
    /// - `bytes_per_sec`: controls the total write rate of compaction and flush in bytes per second.
    pub fn new(bytes_per_sec: u64) -> IOLimiter {
        IOLimiter {
            inner: RateLimiter::new(bytes_per_sec as i64, REFILL_PERIOD, FARENESS),
        }
    }

    /// Sets the rate limit in bytes per second
    pub fn set_bytes_per_second(&self, bytes_per_sec: i64) {
        self.inner.set_bytes_per_second(bytes_per_sec)
    }

    /// Requests an access token to read or write bytes. If this request can not be satisfied, the call is blocked.
    pub fn request(&self, bytes: i64) {
        self.inner.request(bytes, PRIORITY_HIGH)
    }

    /// Gets the max bytes that can be granted in a single burst.
    /// Note: it will be less than or equal to `SNAP_MAX_BYTES_PER_TIME`.
    pub fn get_max_bytes_per_time(&self) -> i64 {
        if self.inner.get_singleburst_bytes() > SNAP_MAX_BYTES_PER_TIME {
            SNAP_MAX_BYTES_PER_TIME
        } else {
            self.inner.get_singleburst_bytes()
        }
    }

    /// Gets the total bytes that have gone through the rate limiter.
    pub fn get_total_bytes_through(&self) -> i64 {
        self.inner.get_total_bytes_through(PRIORITY_HIGH)
    }

    /// Gets the rate limit in bytes per second.
    pub fn get_bytes_per_second(&self) -> i64 {
        self.inner.get_bytes_per_second()
    }

    /// Gets the total number of requests that have gone through rate limiter
    pub fn get_total_requests(&self) -> i64 {
        self.inner.get_total_requests(PRIORITY_HIGH)
    }
}

pub struct LimitWriter<'a, T: Write> {
    limiter: Option<Arc<IOLimiter>>,
    writer: &'a mut T,
}

impl<'a, T: Write + 'a> LimitWriter<'a, T> {
    pub fn new(limiter: Option<Arc<IOLimiter>>, writer: &'a mut T) -> LimitWriter<'a, T> {
        LimitWriter { limiter, writer }
    }
}

impl<'a, T: Write + 'a> Write for LimitWriter<'a, T> {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let total = buf.len();
        if let Some(ref limiter) = self.limiter {
            let single = limiter.get_max_bytes_per_time() as usize;
            let mut curr = 0;
            let mut end;
            while curr < total {
                if curr + single >= total {
                    end = total;
                } else {
                    end = curr + single;
                }
                limiter.request((end - curr) as i64);
                self.writer.write_all(&buf[curr..end])?;
                curr = end;
            }
        } else {
            self.writer.write_all(buf)?;
        }
        Ok(total)
    }

    fn flush(&mut self) -> Result<()> {
        self.writer.flush()?;
        Ok(())
    }
}

/// A limited reader.
///
/// The read limits the bytes per second read from an underlying reader.
pub struct LimitReader<'a, T: Read> {
    limiter: Option<Arc<IOLimiter>>,
    reader: &'a mut T,
}

impl<'a, T: Read + 'a> LimitReader<'a, T> {
    /// Create a new `LimitReader`.
    pub fn new(limiter: Option<Arc<IOLimiter>>, reader: &'a mut T) -> LimitReader<'a, T> {
        LimitReader { limiter, reader }
    }
}

impl<'a, T: Read + 'a> Read for LimitReader<'a, T> {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        if let Some(ref limiter) = self.limiter {
            let total = buf.len();
            let single = limiter.get_max_bytes_per_time() as usize;
            let mut count = 0;
            let mut curr = 0;
            let mut end;
            while curr < total {
                if curr + single >= total {
                    end = total;
                } else {
                    end = curr + single;
                }
                limiter.request((end - curr) as i64);
                count += self.reader.read(&mut buf[curr..end])?;
                curr = end;
            }
            Ok(count)
        } else {
            self.reader.read(buf)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::fs::{self, File};
    use std::io::Write;
    use std::sync::Arc;
    use tempfile::Builder;

    use super::*;

    #[test]
    fn test_io_limiter() {
        let limiter = IOLimiter::new(10 * 1024 * 1024);
        assert!(limiter.get_max_bytes_per_time() <= SNAP_MAX_BYTES_PER_TIME);

        limiter.set_bytes_per_second(20 * 1024 * 1024);
        assert_eq!(limiter.get_bytes_per_second(), 20 * 1024 * 1024);

        assert_eq!(limiter.get_total_bytes_through(), 0);

        limiter.request(1024 * 1024);
        assert_eq!(limiter.get_total_bytes_through(), 1024 * 1024);

        assert_eq!(limiter.get_total_requests(), 1);
    }

    #[test]
    fn test_limit_writer() {
        let dir = Builder::new()
            .prefix("_test_limit_writer")
            .tempdir()
            .unwrap();
        let path = dir.path().join("test-file");
        let mut file = File::create(&path).unwrap();
        let mut limit_writer = LimitWriter::new(Some(Arc::new(IOLimiter::new(1024))), &mut file);

        let mut s = String::new();
        for _ in 0..100 {
            s.push_str("Hello, World!");
        }
        limit_writer.write_all(s.as_bytes()).unwrap();
        limit_writer.flush().unwrap();

        let contents = fs::read_to_string(&path).unwrap();
        assert_eq!(contents, s);
    }

    #[test]
    fn test_limit_reader() {
        let mut buf = Vec::with_capacity(512);
        let bytes_per_sec = 10 * 1024 * 1024; // 10MB/s
        for c in 0..1024usize {
            let mut source = std::io::repeat(b'7').take(c as _);
            let mut limit_reader =
                LimitReader::new(Some(Arc::new(IOLimiter::new(bytes_per_sec))), &mut source);
            let count = limit_reader.read_to_end(&mut buf).unwrap();
            assert_eq!(count, c);
            assert_eq!(count, buf.len());
            buf.clear();
        }
    }
}
<|endoftext|>