<commit_before>#[macro_use]
mod macros;

use std::str::FromStr;

use proc_macro2::{Delimiter, Group, Punct, Spacing, TokenStream, TokenTree};
use quote::quote;
use std::iter::FromIterator;
use syn::{Expr, ExprRange};

#[test]
fn test_expr_parse() {
    let code = "..100u32";
    let tt = TokenStream::from_str(code).unwrap();
    let expr: Expr = syn::parse2(tt.clone()).unwrap();
    let expr_range: ExprRange = syn::parse2(tt).unwrap();
    assert_eq!(expr, Expr::Range(expr_range));
}

#[test]
fn test_await() {
    // Must not parse as Expr::Field.
    let expr = syn::parse_str::<Expr>("fut.await").unwrap();

    snapshot!(expr, @r###"
    Expr::Await {
        base: Expr::Path {
            path: Path {
                segments: [
                    PathSegment {
                        ident: "fut",
                        arguments: None,
                    },
                ],
            },
        },
    }
    "###);
}

#[test]
fn test_macro_variable_func() {
    // mimics the token stream corresponding to `$fn()`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { f })),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Call {
        func: Expr::Group {
            expr: Expr::Path {
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "f",
                            arguments: None,
                        },
                    ],
                },
            },
        },
    }
    "###);

    let tokens = TokenStream::from_iter(vec![
        TokenTree::Punct(Punct::new('#', Spacing::Alone)),
        TokenTree::Group(Group::new(Delimiter::Bracket, quote! { outside })),
        TokenTree::Group(Group::new(Delimiter::None, quote! { #[inside] f })),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Call {
        attrs: [
            Attribute {
                style: Outer,
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "outside",
                            arguments: None,
                        },
                    ],
                },
                tokens: TokenStream(``),
            },
        ],
        func: Expr::Group {
            expr: Expr::Path {
                attrs: [
                    Attribute {
                        style: Outer,
                        path: Path {
                            segments: [
                                PathSegment {
                                    ident: "inside",
                                    arguments: None,
                                },
                            ],
                        },
                        tokens: TokenStream(``),
                    },
                ],
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "f",
                            arguments: None,
                        },
                    ],
                },
            },
        },
    }
    "###);
}

#[test]
fn test_macro_variable_macro() {
    // mimics the token stream corresponding to `$macro!()`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { m })),
        TokenTree::Punct(Punct::new('!', Spacing::Alone)),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Macro {
        mac: Macro {
            path: Path {
                segments: [
                    PathSegment {
                        ident: "m",
                        arguments: None,
                    },
                ],
            },
            delimiter: Paren,
            tokens: TokenStream(``),
        },
    }
    "###);
}
<commit_msg>Add test for struct init expr containing macro metavariable<commit_after>#[macro_use]
mod macros;

use std::str::FromStr;

use proc_macro2::{Delimiter, Group, Punct, Spacing, TokenStream, TokenTree};
use quote::quote;
use std::iter::FromIterator;
use syn::{Expr, ExprRange};

#[test]
fn test_expr_parse() {
    let code = "..100u32";
    let tt = TokenStream::from_str(code).unwrap();
    let expr: Expr = syn::parse2(tt.clone()).unwrap();
    let expr_range: ExprRange = syn::parse2(tt).unwrap();
    assert_eq!(expr, Expr::Range(expr_range));
}

#[test]
fn test_await() {
    // Must not parse as Expr::Field.
    let expr = syn::parse_str::<Expr>("fut.await").unwrap();

    snapshot!(expr, @r###"
    Expr::Await {
        base: Expr::Path {
            path: Path {
                segments: [
                    PathSegment {
                        ident: "fut",
                        arguments: None,
                    },
                ],
            },
        },
    }
    "###);
}

#[test]
fn test_macro_variable_func() {
    // mimics the token stream corresponding to `$fn()`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { f })),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Call {
        func: Expr::Group {
            expr: Expr::Path {
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "f",
                            arguments: None,
                        },
                    ],
                },
            },
        },
    }
    "###);

    let tokens = TokenStream::from_iter(vec![
        TokenTree::Punct(Punct::new('#', Spacing::Alone)),
        TokenTree::Group(Group::new(Delimiter::Bracket, quote! { outside })),
        TokenTree::Group(Group::new(Delimiter::None, quote! { #[inside] f })),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Call {
        attrs: [
            Attribute {
                style: Outer,
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "outside",
                            arguments: None,
                        },
                    ],
                },
                tokens: TokenStream(``),
            },
        ],
        func: Expr::Group {
            expr: Expr::Path {
                attrs: [
                    Attribute {
                        style: Outer,
                        path: Path {
                            segments: [
                                PathSegment {
                                    ident: "inside",
                                    arguments: None,
                                },
                            ],
                        },
                        tokens: TokenStream(``),
                    },
                ],
                path: Path {
                    segments: [
                        PathSegment {
                            ident: "f",
                            arguments: None,
                        },
                    ],
                },
            },
        },
    }
    "###);
}

#[test]
fn test_macro_variable_macro() {
    // mimics the token stream corresponding to `$macro!()`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { m })),
        TokenTree::Punct(Punct::new('!', Spacing::Alone)),
        TokenTree::Group(Group::new(Delimiter::Parenthesis, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Macro {
        mac: Macro {
            path: Path {
                segments: [
                    PathSegment {
                        ident: "m",
                        arguments: None,
                    },
                ],
            },
            delimiter: Paren,
            tokens: TokenStream(``),
        },
    }
    "###);
}

#[test]
fn test_macro_variable_struct() {
    // mimics the token stream corresponding to `$struct {}`
    let tokens = TokenStream::from_iter(vec![
        TokenTree::Group(Group::new(Delimiter::None, quote! { S })),
        TokenTree::Group(Group::new(Delimiter::Brace, TokenStream::new())),
    ]);

    snapshot!(tokens as Expr, @r###"
    Expr::Struct {
        path: Path {
            segments: [
                PathSegment {
                    ident: "S",
                    arguments: None,
                },
            ],
        },
    }
    "###);
}
<|endoftext|>


<commit_before>use super::super::trees::*;
use super::btree::level::LEVEL_M as BLOCK_SIZE;
use super::btree::storage;
use super::tree::*;
use super::super::sm::client::SMClient;
use crate::client::AsyncClient;
use crate::ram::types::Id;
use crate::ram::types::RandValue;
use bifrost::conshash::ConsistentHashing;
use bifrost::utils::time::get_time;
use bifrost_plugins::hash_ident;
use futures::future::BoxFuture;
use futures::prelude::*;
use lightning::map::Map;
use lightning::map::{HashMap, ObjectMap};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Duration;

pub type EntryKeyBlock = [EntryKey; BLOCK_SIZE];
pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Clone, Serialize, Deserialize)]
pub struct Boundary {
    upper: EntryKey,
    lower: EntryKey,
}

#[derive(Clone, Serialize, Deserialize)]
pub enum OpResult<T> {
    Successful(T),
    Failed,
    NotFound,
    OutOfBound,
    Migrating,
}

#[derive(Clone, Copy, Serialize, Deserialize)]
pub struct ServCursor {
    cursor_id: u64,
}

pub struct DistLSMTree {
    id: Id,
    tree: LSMTree,
    prop: RwLock<DistProp>,
}

struct DistProp {
    boundary: Boundary,
    migration: Option<Migration>,
}

struct Migration {
    pivot: EntryKey,
}

pub struct CursorMemo {
    tree_cursor: LSMTreeCursor,
    expires: i64,
}

service! {
    rpc crate_tree(id: Id, boundary: Boundary);
    rpc load_tree(id: Id, boundary: Boundary);
    rpc insert(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc delete(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc seek(id: Id, entry: EntryKey, ordering: Ordering, cursor_lifetime: u16)
        -> OpResult<(ServCursor, Option<EntryKey>)>;
    rpc renew_cursor(cursor: ServCursor, time: u16) -> bool;
    rpc dispose_cursor(cursor: ServCursor) -> bool;
    rpc cursor_next(cursor: ServCursor) -> Option<EntryKeyBlock>;
}

pub struct LSMTreeService {
    client: Arc<AsyncClient>,
    cursor_counter: AtomicUsize,
    cursors: ObjectMap<Arc<RefCell<CursorMemo>>>,
    trees: Arc<HashMap<Id, Arc<DistLSMTree>>>,
}

impl Service for LSMTreeService {
    fn crate_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::create(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn load_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::recover(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn insert(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.insert(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn delete(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.delete(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn seek(
        &self,
        id: Id,
        entry: EntryKey,
        ordering: Ordering,
        cursor_lifetime: u16,
    ) -> BoxFuture<OpResult<(ServCursor, Option<EntryKey>)>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            let mut tree_cursor = tree.seek(&entry, ordering);
            let cursor_id = self.cursor_counter.fetch_add(1, Relaxed);
            let expires = get_time() + cursor_lifetime as i64;
            let first_entry = tree_cursor.next();
            let cursor_memo = CursorMemo {
                tree_cursor,
                expires,
            };
            self.cursors
                .insert(&(cursor_id), Arc::new(RefCell::new(cursor_memo)));
            let cursor = ServCursor {
                cursor_id: cursor_id as u64,
            };
            OpResult::Successful((cursor, first_entry))
        })
    }

    fn renew_cursor(&self, cursor: ServCursor, time: u16) -> BoxFuture<bool> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                cursor.borrow_mut().expires = get_time() + time as i64;
                true
            } else {
                false
            },
        )
        .boxed()
    }

    fn dispose_cursor(&self, cursor: ServCursor) -> BoxFuture<bool> {
        future::ready(self.cursors.remove(&(cursor.cursor_id as usize)).is_some()).boxed()
    }

    fn cursor_next(&self, cursor: ServCursor) -> BoxFuture<Option<EntryKeyBlock>> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                let mut cursor_ref = cursor.borrow_mut();
                Some(entry_block_from_cursor_memo(&mut *cursor_ref))
            } else {
                None
            },
        )
        .boxed()
    }
}

impl LSMTreeService {
    pub fn new(client: &Arc<AsyncClient>, sm_client: &Arc<SMClient>) -> Self {
        let trees_map = Arc::new(HashMap::with_capacity(32));
        super::btree::storage::start_external_nodes_write_back(client);
        Self::start_tree_balancer(&trees_map, client);
        Self {
            client: client.clone(),
            cursor_counter: AtomicUsize::new(0),
            cursors: ObjectMap::with_capacity(64),
            trees: trees_map,
        }
    }

    pub fn start_tree_balancer(
        trees_map: &Arc<HashMap<Id, Arc<DistLSMTree>>>,
        client: &Arc<AsyncClient>,
        sm_client: &Arc<SMClient>,
    ) {
        let trees_map = trees_map.clone();
        let client = client.clone();
        let sm_client = sm_client.clone();
        tokio::spawn(async move {
            loop {
                for (_, dist_tree) in trees_map.entries() {
                    let tree = &dist_tree.tree;
                    tree.merge_levels();
                    if tree.oversized() {
                        // Tree oversized, need to migrate
                        let mid_key = tree.mid_key().unwrap();
                        let migration_target_id = Id::rand();
                        let migration_tree = LSMTree::create(&client, &migration_target_id).await;
                        {
                            let mut dist_tree_prop = dist_tree.prop.write();
                            dist_tree_prop.migration = Some(Migration {
                                pivot: mid_key.clone(),
                            });
                        }
                        tree.mark_migration(&dist_tree.id, Some(migration_target_id), &client)
                            .await;
                        let buffer_size = BLOCK_SIZE << 2;
                        let mut cursor = tree.seek(&mid_key, Ordering::Forward);
                        let mut entry_buffer = Vec::with_capacity(buffer_size);
                        // Moving keys
                        while cursor.current().is_some() {
                            if let Some(entry) = cursor.next() {
                                entry_buffer.push(entry);
                                if entry_buffer.len() >= buffer_size {
                                    migration_tree.merge_keys(Box::new(entry_buffer));
                                    entry_buffer = Vec::with_capacity(buffer_size);
                                }
                            }
                        }
                        storage::wait_until_updated().await;
                        sm_client.split(&migration_target_id, &mid_key).await.unwrap();
                        // Reset state on current tree
                        tree.mark_migration(&dist_tree.id, None, &client).await;
                        {
                            let mut dist_prop = dist_tree.prop.write();
                            dist_prop.boundary.upper = mid_key;
                            dist_prop.migration = None;
                        }
                    }
                }
                // Sleep for a while to check for trees to be merge in levels
                tokio::time::delay_for(Duration::from_secs(5)).await;
            }
        });
    }

    fn apply_in_ranged_tree<F, R>(&self, id: Id, entry: EntryKey, func: F) -> BoxFuture<OpResult<R>>
    where
        F: Fn(&EntryKey, &LSMTree) -> OpResult<R>,
        R: Send + 'static,
    {
        future::ready(if let Some(tree) = self.trees.get(&id) {
            let tree_prop = tree.prop.read();
            if tree_prop.boundary.in_boundary(&entry) {
                if let &Some(ref migration) = &tree_prop.migration {
                    if entry < migration.pivot {
                        // Entries lower than pivot should be safe to work on
                        func(&entry, &tree.tree)
                    } else {
                        OpResult::Migrating
                    }
                } else {
                    func(&entry, &tree.tree)
                }
            } else {
                OpResult::OutOfBound
            }
        } else {
            OpResult::NotFound
        })
        .boxed()
    }
}

fn entry_block_from_cursor_memo(cursor_memo: &mut CursorMemo) -> EntryKeyBlock {
    let mut res = EntryKeyBlock::default();
    for entry in res.iter_mut() {
        if let Some(tree_entry) = cursor_memo.tree_cursor.next() {
            *entry = tree_entry.clone();
        } else {
            break;
        }
    }
    res
}

impl DistLSMTree {
    fn new(id: Id, tree: LSMTree, boundary: Boundary, migration: Option<Migration>) -> Self {
        let prop = RwLock::new(DistProp {
            boundary,
            migration,
        });
        Self { id, tree, prop }
    }
}

impl Boundary {
    pub fn new(lower: EntryKey, upper: EntryKey) -> Self {
        Boundary { lower, upper }
    }
    fn in_boundary(&self, entry: &EntryKey) -> bool {
        return entry >= &self.lower && entry < &self.upper;
    }
}

dispatch_rpc_service_functions!(LSMTreeService);

unsafe impl Send for DistLSMTree {}
unsafe impl Sync for DistLSMTree {}

pub fn client_by_rpc_client(rpc: &Arc<RPCClient>) -> Arc<AsyncServiceClient> {
    AsyncServiceClient::new(DEFAULT_SERVICE_ID, rpc)
}

pub async fn locate_tree_server_from_conshash(
    id: &Id,
    conshash: &Arc<ConsistentHashing>,
) -> Result<Arc<AsyncServiceClient>, RPCError> {
    let server_id = conshash.get_server_id_by(id).unwrap();
    DEFAULT_CLIENT_POOL
        .get_by_id(server_id, move |sid| conshash.to_server_name(sid))
        .await
        .map_err(|e| RPCError::IOError(e))
        .map(|c| client_by_rpc_client(&c))
}
<commit_msg>record on state machine for split<commit_after>use super::super::trees::*;
use super::btree::level::LEVEL_M as BLOCK_SIZE;
use super::btree::storage;
use super::tree::*;
use super::super::sm::client::SMClient;
use crate::client::AsyncClient;
use crate::ram::types::Id;
use crate::ram::types::RandValue;
use bifrost::conshash::ConsistentHashing;
use bifrost::utils::time::get_time;
use bifrost_plugins::hash_ident;
use futures::future::BoxFuture;
use futures::prelude::*;
use lightning::map::Map;
use lightning::map::{HashMap, ObjectMap};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Duration;

pub type EntryKeyBlock = [EntryKey; BLOCK_SIZE];
pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(LSM_TREE_RPC_SERVICE) as u64;

#[derive(Clone, Serialize, Deserialize)]
pub struct Boundary {
    upper: EntryKey,
    lower: EntryKey,
}

#[derive(Clone, Serialize, Deserialize)]
pub enum OpResult<T> {
    Successful(T),
    Failed,
    NotFound,
    OutOfBound,
    Migrating,
}

#[derive(Clone, Copy, Serialize, Deserialize)]
pub struct ServCursor {
    cursor_id: u64,
}

pub struct DistLSMTree {
    id: Id,
    tree: LSMTree,
    prop: RwLock<DistProp>,
}

struct DistProp {
    boundary: Boundary,
    migration: Option<Migration>,
}

struct Migration {
    pivot: EntryKey,
}

pub struct CursorMemo {
    tree_cursor: LSMTreeCursor,
    expires: i64,
}

service! {
    rpc crate_tree(id: Id, boundary: Boundary);
    rpc load_tree(id: Id, boundary: Boundary);
    rpc insert(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc delete(id: Id, entry: EntryKey) -> OpResult<()>;
    rpc seek(id: Id, entry: EntryKey, ordering: Ordering, cursor_lifetime: u16)
        -> OpResult<(ServCursor, Option<EntryKey>)>;
    rpc renew_cursor(cursor: ServCursor, time: u16) -> bool;
    rpc dispose_cursor(cursor: ServCursor) -> bool;
    rpc cursor_next(cursor: ServCursor) -> Option<EntryKeyBlock>;
}

pub struct LSMTreeService {
    client: Arc<AsyncClient>,
    cursor_counter: AtomicUsize,
    cursors: ObjectMap<Arc<RefCell<CursorMemo>>>,
    trees: Arc<HashMap<Id, Arc<DistLSMTree>>>,
}

impl Service for LSMTreeService {
    fn crate_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::create(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn load_tree(&self, id: Id, boundary: Boundary) -> BoxFuture<()> {
        async move {
            if self.trees.contains(&id) {
                return;
            }
            let tree = LSMTree::recover(&self.client, &id).await;
            self.trees
                .insert(&id, Arc::new(DistLSMTree::new(id, tree, boundary, None)));
        }
        .boxed()
    }

    fn insert(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.insert(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn delete(&self, id: Id, entry: EntryKey) -> BoxFuture<OpResult<()>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            if tree.delete(&entry) {
                OpResult::Successful(())
            } else {
                OpResult::Failed
            }
        })
    }

    fn seek(
        &self,
        id: Id,
        entry: EntryKey,
        ordering: Ordering,
        cursor_lifetime: u16,
    ) -> BoxFuture<OpResult<(ServCursor, Option<EntryKey>)>> {
        self.apply_in_ranged_tree(id, entry, |entry, tree| {
            let mut tree_cursor = tree.seek(&entry, ordering);
            let cursor_id = self.cursor_counter.fetch_add(1, Relaxed);
            let expires = get_time() + cursor_lifetime as i64;
            let first_entry = tree_cursor.next();
            let cursor_memo = CursorMemo {
                tree_cursor,
                expires,
            };
            self.cursors
                .insert(&(cursor_id), Arc::new(RefCell::new(cursor_memo)));
            let cursor = ServCursor {
                cursor_id: cursor_id as u64,
            };
            OpResult::Successful((cursor, first_entry))
        })
    }

    fn renew_cursor(&self, cursor: ServCursor, time: u16) -> BoxFuture<bool> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                cursor.borrow_mut().expires = get_time() + time as i64;
                true
            } else {
                false
            },
        )
        .boxed()
    }

    fn dispose_cursor(&self, cursor: ServCursor) -> BoxFuture<bool> {
        future::ready(self.cursors.remove(&(cursor.cursor_id as usize)).is_some()).boxed()
    }

    fn cursor_next(&self, cursor: ServCursor) -> BoxFuture<Option<EntryKeyBlock>> {
        future::ready(
            if let Some(cursor) = self.cursors.write(cursor.cursor_id as usize) {
                let mut cursor_ref = cursor.borrow_mut();
                Some(entry_block_from_cursor_memo(&mut *cursor_ref))
            } else {
                None
            },
        )
        .boxed()
    }
}

impl LSMTreeService {
    pub fn new(client: &Arc<AsyncClient>, sm_client: &Arc<SMClient>) -> Self {
        let trees_map = Arc::new(HashMap::with_capacity(32));
        super::btree::storage::start_external_nodes_write_back(client);
        Self::start_tree_balancer(&trees_map, client, sm_client);
        Self {
            client: client.clone(),
            cursor_counter: AtomicUsize::new(0),
            cursors: ObjectMap::with_capacity(64),
            trees: trees_map,
        }
    }

    pub fn start_tree_balancer(
        trees_map: &Arc<HashMap<Id, Arc<DistLSMTree>>>,
        client: &Arc<AsyncClient>,
        sm_client: &Arc<SMClient>,
    ) {
        let trees_map = trees_map.clone();
        let client = client.clone();
        let sm_client = sm_client.clone();
        tokio::spawn(async move {
            loop {
                for (_, dist_tree) in trees_map.entries() {
                    let tree = &dist_tree.tree;
                    tree.merge_levels();
                    if tree.oversized() {
                        // Tree oversized, need to migrate
                        let mid_key = tree.mid_key().unwrap();
                        let migration_target_id = Id::rand();
                        let migration_tree = LSMTree::create(&client, &migration_target_id).await;
                        {
                            let mut dist_tree_prop = dist_tree.prop.write();
                            dist_tree_prop.migration = Some(Migration {
                                pivot: mid_key.clone(),
                            });
                        }
                        tree.mark_migration(&dist_tree.id, Some(migration_target_id), &client)
                            .await;
                        let buffer_size = BLOCK_SIZE << 2;
                        let mut cursor = tree.seek(&mid_key, Ordering::Forward);
                        let mut entry_buffer = Vec::with_capacity(buffer_size);
                        // Moving keys
                        while cursor.current().is_some() {
                            if let Some(entry) = cursor.next() {
                                entry_buffer.push(entry);
                                if entry_buffer.len() >= buffer_size {
                                    migration_tree.merge_keys(Box::new(entry_buffer));
                                    entry_buffer = Vec::with_capacity(buffer_size);
                                }
                            }
                        }
                        storage::wait_until_updated().await;
                        sm_client.split(&migration_target_id, &mid_key).await.unwrap();
                        // Reset state on current tree
                        tree.mark_migration(&dist_tree.id, None, &client).await;
                        {
                            let mut dist_prop = dist_tree.prop.write();
                            dist_prop.boundary.upper = mid_key;
                            dist_prop.migration = None;
                        }
                    }
                }
                // Sleep for a while to check for trees to be merge in levels
                tokio::time::delay_for(Duration::from_secs(5)).await;
            }
        });
    }

    fn apply_in_ranged_tree<F, R>(&self, id: Id, entry: EntryKey, func: F) -> BoxFuture<OpResult<R>>
    where
        F: Fn(&EntryKey, &LSMTree) -> OpResult<R>,
        R: Send + 'static,
    {
        future::ready(if let Some(tree) = self.trees.get(&id) {
            let tree_prop = tree.prop.read();
            if tree_prop.boundary.in_boundary(&entry) {
                if let &Some(ref migration) = &tree_prop.migration {
                    if entry < migration.pivot {
                        // Entries lower than pivot should be safe to work on
                        func(&entry, &tree.tree)
                    } else {
                        OpResult::Migrating
                    }
                } else {
                    func(&entry, &tree.tree)
                }
            } else {
                OpResult::OutOfBound
            }
        } else {
            OpResult::NotFound
        })
        .boxed()
    }
}

fn entry_block_from_cursor_memo(cursor_memo: &mut CursorMemo) -> EntryKeyBlock {
    let mut res = EntryKeyBlock::default();
    for entry in res.iter_mut() {
        if let Some(tree_entry) = cursor_memo.tree_cursor.next() {
            *entry = tree_entry.clone();
        } else {
            break;
        }
    }
    res
}

impl DistLSMTree {
    fn new(id: Id, tree: LSMTree, boundary: Boundary, migration: Option<Migration>) -> Self {
        let prop = RwLock::new(DistProp {
            boundary,
            migration,
        });
        Self { id, tree, prop }
    }
}

impl Boundary {
    pub fn new(lower: EntryKey, upper: EntryKey) -> Self {
        Boundary { lower, upper }
    }
    fn in_boundary(&self, entry: &EntryKey) -> bool {
        return entry >= &self.lower && entry < &self.upper;
    }
}

dispatch_rpc_service_functions!(LSMTreeService);

unsafe impl Send for DistLSMTree {}
unsafe impl Sync for DistLSMTree {}

pub fn client_by_rpc_client(rpc: &Arc<RPCClient>) -> Arc<AsyncServiceClient> {
    AsyncServiceClient::new(DEFAULT_SERVICE_ID, rpc)
}

pub async fn locate_tree_server_from_conshash(
    id: &Id,
    conshash: &Arc<ConsistentHashing>,
) -> Result<Arc<AsyncServiceClient>, RPCError> {
    let server_id = conshash.get_server_id_by(id).unwrap();
    DEFAULT_CLIENT_POOL
        .get_by_id(server_id, move |sid| conshash.to_server_name(sid))
        .await
        .map_err(|e| RPCError::IOError(e))
        .map(|c| client_by_rpc_client(&c))
}
<|endoftext|>


<commit_before>use lazy_static::lazy_static;
use petgraph::Direction;
use proc_macro2::{Span, TokenStream};
use quote::quote;
use rustc_hash::FxHashMap;
use std::io::Write;
use syn::{spanned::Spanned, Ident};

use super::{
    java_code, map_type::map_type, map_write_err, rust_code, JavaContext, JavaForeignTypeInfo,
    JniForeignMethodSignature,
};
use crate::{
    error::{panic_on_syn_error, DiagnosticError, Result},
    file_cache::FileWriteCache,
    source_registry::SourceId,
    typemap::{
        ast::{DisplayToTokens, TypeName},
        utils::rust_to_foreign_convert_method_inputs,
        ForeignTypeInfo,
    },
    types::ForeignInterface,
    WRITE_TO_MEM_FAILED_MSG,
};

pub(in crate::java_jni) fn generate_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
) -> Result<()> {
    let f_methods = find_suitable_ftypes_for_interace_methods(ctx, interface)?;
    generate_java_code_for_interface(
        ctx,
        interface,
        &f_methods,
        ctx.cfg.null_annotation_package.as_ref().map(String::as_str),
    )
    .map_err(|err| DiagnosticError::new(interface.src_id, interface.span(), err))?;
    generate_rust_code_for_interface(ctx, interface, &f_methods)?;

    let my_jobj_ti = ctx.conv_map.find_or_alloc_rust_type_with_suffix(
        &parse_type! { jobject },
        &interface.name.to_string(),
        SourceId::none(),
    );
    ctx.conv_map.add_foreign(
        my_jobj_ti,
        TypeName::from_ident(&interface.name, interface.src_id),
    )?;
    Ok(())
}

fn find_suitable_ftypes_for_interace_methods(
    ctx: &mut JavaContext,
    interace: &ForeignInterface,
) -> Result<Vec<JniForeignMethodSignature>> {
    let void_sym = "void";
    let dummy_ty = parse_type! { () };
    let dummy_rust_ty = ctx.conv_map.find_or_alloc_rust_type_no_src_id(&dummy_ty);
    let mut f_methods = Vec::with_capacity(interace.items.len());
    let jobject_ty = ctx
        .conv_map
        .find_or_alloc_rust_type_no_src_id(&parse_type! { jobject });

    for method in &interace.items {
        let mut input = Vec::<JavaForeignTypeInfo>::with_capacity(method.fn_decl.inputs.len() - 1);
        for arg in method.fn_decl.inputs.iter().skip(1) {
            let named_arg = arg
                .as_named_arg()
                .map_err(|err| DiagnosticError::from_syn_err(interace.src_id, err))?;
            let arg_rust_ty = ctx
                .conv_map
                .find_or_alloc_rust_type(&named_arg.ty, interace.src_id);
            let arg_span = (interace.src_id, named_arg.ty.span());
            let mut f_arg_type = map_type(ctx, &arg_rust_ty, Direction::Outgoing, arg_span)?;
            if f_arg_type.java_converter.is_some() {
                // it is hard to use Java code during callback, so may be
                // there is way to convert it to jobject ?
                ctx.conv_map.convert_rust_types(
                    arg_rust_ty.to_idx(),
                    jobject_ty.to_idx(),
                    "x",
                    "y",
                    "ret",
                    arg_span,
                )?;
                f_arg_type.java_converter = None;
                f_arg_type.base.correspoding_rust_type = jobject_ty.clone();
            }

            input.push(f_arg_type);
        }
        let output = match method.fn_decl.output {
            syn::ReturnType::Default => ForeignTypeInfo {
                name: void_sym.into(),
                correspoding_rust_type: dummy_rust_ty.clone(),
            }
            .into(),
            _ => unimplemented!(),
        };
        f_methods.push(JniForeignMethodSignature { output, input });
    }
    Ok(f_methods)
}

fn generate_java_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
    use_null_annotation: Option<&str>,
) -> std::result::Result<(), String> {
    let path = ctx.cfg.output_dir.join(format!("{}.java", interface.name));
    let mut file = FileWriteCache::new(&path, ctx.generated_foreign_files);
    let imports = java_code::get_null_annotation_imports(use_null_annotation, methods_sign);
    let interface_comments =
        java_code::doc_comments_to_java_comments(&interface.doc_comments, true);
    writeln!(
        file,
        r#"// Automatically generated by flapigen
package {package_name};
{imports}
{doc_comments}
public interface {interface_name} {{"#,
        package_name = ctx.cfg.package_name,
        interface_name = interface.name,
        doc_comments = interface_comments,
        imports = imports,
    )
    .expect(WRITE_TO_MEM_FAILED_MSG);

    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            file,
            r#"
{doc_comments}
    void {method_name}({single_args_with_types});"#,
            method_name = method.name,
            doc_comments = java_code::doc_comments_to_java_comments(&method.doc_comments, false),
            single_args_with_types = java_code::args_with_java_types(
                f_method,
                method.arg_names_without_self(),
                java_code::ArgsFormatFlags::EXTERNAL,
                use_null_annotation.is_some()
            ),
        )
        .expect(WRITE_TO_MEM_FAILED_MSG);
    }

    file.write_all(b"\n}\n").expect(WRITE_TO_MEM_FAILED_MSG);
    file.update_file_if_necessary().map_err(&map_write_err)?;
    Ok(())
}

fn generate_rust_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
) -> Result<()> {
    use std::fmt::Write;

    let mut new_conv_code = format!(
        r#"
#[swig_from_foreigner_hint = "{interface_name}"]
impl SwigFrom<jobject> for Box<dyn {trait_name}> {{
    fn swig_from(this: jobject, env: *mut JNIEnv) -> Self {{
        let mut cb = JavaCallback::new(this, env);
        cb.methods.reserve({methods_len});
        let class = unsafe {{ (**env).GetObjectClass.unwrap()(env, cb.this) }};
        assert!(!class.is_null(), "GetObjectClass return null class for {interface_name}");
"#,
        interface_name = interface.name,
        trait_name = DisplayToTokens(&interface.self_type),
        methods_len = interface.items.len(),
    );
    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            &mut new_conv_code,
            r#"
        let method_id: jmethodID = unsafe {{
            (**env).GetMethodID.unwrap()(env, class, swig_c_str!("{method_name}"),
                                         swig_c_str!("{method_sig}"))
        }};
        assert!(!method_id.is_null(), "Can not find {method_name} id");
        cb.methods.push(method_id);"#,
            method_name = method.name,
            method_sig = rust_code::jni_method_signature(ctx, f_method),
        )
        .unwrap();
    }
    new_conv_code.push_str(
        r#"
        Box::new(cb)
    }
}
"#,
    );
    ctx.conv_map
        .merge(SourceId::none(), &new_conv_code, ctx.pointer_target_width)?;

    let mut trait_impl_funcs = Vec::<TokenStream>::with_capacity(interface.items.len());

    for (method_idx, (method, f_method)) in interface.items.iter().zip(methods_sign).enumerate() {
        let func_name = &method
            .rust_name
            .segments
            .last()
            .ok_or_else(|| {
                DiagnosticError::new(
                    interface.src_id,
                    method.rust_name.span(),
                    "Empty trait function name",
                )
            })?
            .ident;

        let self_arg: TokenStream = method.fn_decl.inputs[0]
            .as_self_arg(interface.src_id)?
            .into();
        let mut args_with_types = Vec::with_capacity(method.fn_decl.inputs.len());
        args_with_types.push(self_arg);
        args_with_types.extend(
            method
                .fn_decl
                .inputs
                .iter()
                .skip(1)
                .enumerate()
                .map(|(i, v)| {
                    let arg_ty = &v.as_named_arg().unwrap().ty;
                    let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
                    quote!(#arg_name: #arg_ty)
                }),
        );
        assert!(!method.fn_decl.inputs.is_empty());
        let n_args = method.fn_decl.inputs.len() - 1;
        let (args, type_size_asserts) = convert_args_for_variadic_function_call(f_method);

        let (mut conv_deps, convert_args_code) = rust_to_foreign_convert_method_inputs(
            ctx.conv_map,
            interface.src_id,
            method,
            f_method,
            (0..n_args).map(|v| format!("a{}", v)),
            "()",
        )?;
        ctx.rust_code.append(&mut conv_deps);
        let convert_args: TokenStream = syn::parse_str(&convert_args_code).unwrap_or_else(|err| {
            panic_on_syn_error(
                "java/jni internal parse failed for convert arguments code",
                convert_args_code,
                err,
            )
        });

        trait_impl_funcs.push(quote! {
            #[allow(unused_mut)]
            fn #func_name(#(#args_with_types),*) {
                #type_size_asserts
                let env = self.get_jni_env();
                if let Some(env) = env.env {
                    #convert_args
                    unsafe {
                        (**env).CallVoidMethod.unwrap()(env, self.this,
                                                        self.methods[#method_idx],
                                                        #(#args),*);
                        if (**env).ExceptionCheck.unwrap()(env) != 0 {
                            log::error!(concat!(stringify!(#func_name), ": java throw exception"));
                            (**env).ExceptionDescribe.unwrap()(env);
                            (**env).ExceptionClear.unwrap()(env);
                        }
                    };
                }
            }
        });
    }

    let self_type_name = &interface.self_type.bounds[0];
    let tt: TokenStream = quote! {
        impl #self_type_name for JavaCallback {
            #(#trait_impl_funcs)*
        }
    };
    ctx.rust_code.push(tt);
    Ok(())
}

lazy_static! {
    static ref JNI_FOR_VARIADIC_C_FUNC_CALL: FxHashMap<&'static str, &'static str> = {
        let mut m = FxHashMap::default();
        m.insert("jboolean", "::std::os::raw::c_uint");
        m.insert("jbyte", "::std::os::raw::c_int");
        m.insert("jshort", "::std::os::raw::c_int");
        m.insert("jfloat", "f64");
        m
    };
}

// To use `C` function with variable number of arguments,
// we need automatic type conversation, see
// http://en.cppreference.com/w/c/language/conversion#Default_argument_promotions
// for more details.
// return arg with conversation plus asserts
fn convert_args_for_variadic_function_call(
    f_method: &JniForeignMethodSignature,
) -> (Vec<TokenStream>, TokenStream) {
    let mut ret = Vec::with_capacity(f_method.input.len());
    for (i, arg) in f_method.input.iter().enumerate() {
        let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
        if let Some(conv_type_str) = JNI_FOR_VARIADIC_C_FUNC_CALL
            .get(&*arg.as_ref().correspoding_rust_type.normalized_name.as_str())
        {
            let conv_type: TokenStream = syn::parse_str(*conv_type_str).unwrap_or_else(|err| {
                panic_on_syn_error(
                    "java/jni internal error: can not parse type for variable conversation",
                    conv_type_str.to_string(),
                    err,
                )
            });
            ret.push(quote!(#arg_name as #conv_type));
        } else {
            ret.push(quote!(#arg_name));
        }
    }
    let check_sizes = quote! {
        swig_assert_eq_size!(::std::os::raw::c_uint, u32);
        swig_assert_eq_size!(::std::os::raw::c_int, i32);
    };
    (ret, check_sizes)
}
<commit_msg>java/jni/feature: more details if type mappnig for callback impossible<commit_after>use lazy_static::lazy_static;
use petgraph::Direction;
use proc_macro2::{Span, TokenStream};
use quote::quote;
use rustc_hash::FxHashMap;
use std::io::Write;
use syn::{spanned::Spanned, Ident};

use super::{
    java_code, map_type::map_type, map_write_err, rust_code, JavaContext, JavaForeignTypeInfo,
    JniForeignMethodSignature,
};
use crate::{
    error::{panic_on_syn_error, DiagnosticError, Result},
    file_cache::FileWriteCache,
    source_registry::SourceId,
    typemap::{
        ast::{DisplayToTokens, TypeName},
        utils::rust_to_foreign_convert_method_inputs,
        ForeignTypeInfo,
    },
    types::ForeignInterface,
    WRITE_TO_MEM_FAILED_MSG,
};

pub(in crate::java_jni) fn generate_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
) -> Result<()> {
    let f_methods = find_suitable_ftypes_for_interace_methods(ctx, interface)?;
    generate_java_code_for_interface(
        ctx,
        interface,
        &f_methods,
        ctx.cfg.null_annotation_package.as_ref().map(String::as_str),
    )
    .map_err(|err| DiagnosticError::new(interface.src_id, interface.span(), err))?;
    generate_rust_code_for_interface(ctx, interface, &f_methods)?;

    let my_jobj_ti = ctx.conv_map.find_or_alloc_rust_type_with_suffix(
        &parse_type! { jobject },
        &interface.name.to_string(),
        SourceId::none(),
    );
    ctx.conv_map.add_foreign(
        my_jobj_ti,
        TypeName::from_ident(&interface.name, interface.src_id),
    )?;
    Ok(())
}

fn find_suitable_ftypes_for_interace_methods(
    ctx: &mut JavaContext,
    interace: &ForeignInterface,
) -> Result<Vec<JniForeignMethodSignature>> {
    let void_sym = "void";
    let dummy_ty = parse_type! { () };
    let dummy_rust_ty = ctx.conv_map.find_or_alloc_rust_type_no_src_id(&dummy_ty);
    let mut f_methods = Vec::with_capacity(interace.items.len());
    let jobject_ty = ctx
        .conv_map
        .find_or_alloc_rust_type_no_src_id(&parse_type! { jobject });

    for method in &interace.items {
        let mut input = Vec::<JavaForeignTypeInfo>::with_capacity(method.fn_decl.inputs.len() - 1);
        for arg in method.fn_decl.inputs.iter().skip(1) {
            let named_arg = arg
                .as_named_arg()
                .map_err(|err| DiagnosticError::from_syn_err(interace.src_id, err))?;
            let arg_rust_ty = ctx
                .conv_map
                .find_or_alloc_rust_type(&named_arg.ty, interace.src_id);
            let arg_span = (interace.src_id, named_arg.ty.span());
            let mut f_arg_type = map_type(ctx, &arg_rust_ty, Direction::Outgoing, arg_span)?;
            if let Some(java_conv) = f_arg_type.java_converter.as_ref() {
                // it is hard to use Java code during callback, so may be
                // there is way to convert it to jobject ?
                ctx.conv_map
                    .convert_rust_types(
                        arg_rust_ty.to_idx(),
                        jobject_ty.to_idx(),
                        "x",
                        "y",
                        "ret",
                        arg_span,
                    )
                    .map_err(|err| {
                        err.add_span_note(
                            arg_span,
                            format!(
                                "Java code required to convert type to jobject
It is impossible to use this Java code:{}\nfor callback types conversation",
                                java_conv.converter
                            ),
                        )
                    })?;
                f_arg_type.java_converter = None;
                f_arg_type.base.correspoding_rust_type = jobject_ty.clone();
            }

            input.push(f_arg_type);
        }
        let output = match method.fn_decl.output {
            syn::ReturnType::Default => ForeignTypeInfo {
                name: void_sym.into(),
                correspoding_rust_type: dummy_rust_ty.clone(),
            }
            .into(),
            _ => unimplemented!(),
        };
        f_methods.push(JniForeignMethodSignature { output, input });
    }
    Ok(f_methods)
}

fn generate_java_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
    use_null_annotation: Option<&str>,
) -> std::result::Result<(), String> {
    let path = ctx.cfg.output_dir.join(format!("{}.java", interface.name));
    let mut file = FileWriteCache::new(&path, ctx.generated_foreign_files);
    let imports = java_code::get_null_annotation_imports(use_null_annotation, methods_sign);
    let interface_comments =
        java_code::doc_comments_to_java_comments(&interface.doc_comments, true);
    writeln!(
        file,
        r#"// Automatically generated by flapigen
package {package_name};
{imports}
{doc_comments}
public interface {interface_name} {{"#,
        package_name = ctx.cfg.package_name,
        interface_name = interface.name,
        doc_comments = interface_comments,
        imports = imports,
    )
    .expect(WRITE_TO_MEM_FAILED_MSG);

    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            file,
            r#"
{doc_comments}
    void {method_name}({single_args_with_types});"#,
            method_name = method.name,
            doc_comments = java_code::doc_comments_to_java_comments(&method.doc_comments, false),
            single_args_with_types = java_code::args_with_java_types(
                f_method,
                method.arg_names_without_self(),
                java_code::ArgsFormatFlags::EXTERNAL,
                use_null_annotation.is_some()
            ),
        )
        .expect(WRITE_TO_MEM_FAILED_MSG);
    }

    file.write_all(b"\n}\n").expect(WRITE_TO_MEM_FAILED_MSG);
    file.update_file_if_necessary().map_err(&map_write_err)?;
    Ok(())
}

fn generate_rust_code_for_interface(
    ctx: &mut JavaContext,
    interface: &ForeignInterface,
    methods_sign: &[JniForeignMethodSignature],
) -> Result<()> {
    use std::fmt::Write;

    let mut new_conv_code = format!(
        r#"
#[swig_from_foreigner_hint = "{interface_name}"]
impl SwigFrom<jobject> for Box<dyn {trait_name}> {{
    fn swig_from(this: jobject, env: *mut JNIEnv) -> Self {{
        let mut cb = JavaCallback::new(this, env);
        cb.methods.reserve({methods_len});
        let class = unsafe {{ (**env).GetObjectClass.unwrap()(env, cb.this) }};
        assert!(!class.is_null(), "GetObjectClass return null class for {interface_name}");
"#,
        interface_name = interface.name,
        trait_name = DisplayToTokens(&interface.self_type),
        methods_len = interface.items.len(),
    );
    for (method, f_method) in interface.items.iter().zip(methods_sign) {
        writeln!(
            &mut new_conv_code,
            r#"
        let method_id: jmethodID = unsafe {{
            (**env).GetMethodID.unwrap()(env, class, swig_c_str!("{method_name}"),
                                         swig_c_str!("{method_sig}"))
        }};
        assert!(!method_id.is_null(), "Can not find {method_name} id");
        cb.methods.push(method_id);"#,
            method_name = method.name,
            method_sig = rust_code::jni_method_signature(ctx, f_method),
        )
        .unwrap();
    }
    new_conv_code.push_str(
        r#"
        Box::new(cb)
    }
}
"#,
    );
    ctx.conv_map
        .merge(SourceId::none(), &new_conv_code, ctx.pointer_target_width)?;

    let mut trait_impl_funcs = Vec::<TokenStream>::with_capacity(interface.items.len());

    for (method_idx, (method, f_method)) in interface.items.iter().zip(methods_sign).enumerate() {
        let func_name = &method
            .rust_name
            .segments
            .last()
            .ok_or_else(|| {
                DiagnosticError::new(
                    interface.src_id,
                    method.rust_name.span(),
                    "Empty trait function name",
                )
            })?
            .ident;

        let self_arg: TokenStream = method.fn_decl.inputs[0]
            .as_self_arg(interface.src_id)?
            .into();
        let mut args_with_types = Vec::with_capacity(method.fn_decl.inputs.len());
        args_with_types.push(self_arg);
        args_with_types.extend(
            method
                .fn_decl
                .inputs
                .iter()
                .skip(1)
                .enumerate()
                .map(|(i, v)| {
                    let arg_ty = &v.as_named_arg().unwrap().ty;
                    let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
                    quote!(#arg_name: #arg_ty)
                }),
        );
        assert!(!method.fn_decl.inputs.is_empty());
        let n_args = method.fn_decl.inputs.len() - 1;
        let (args, type_size_asserts) = convert_args_for_variadic_function_call(f_method);

        let (mut conv_deps, convert_args_code) = rust_to_foreign_convert_method_inputs(
            ctx.conv_map,
            interface.src_id,
            method,
            f_method,
            (0..n_args).map(|v| format!("a{}", v)),
            "()",
        )?;
        ctx.rust_code.append(&mut conv_deps);
        let convert_args: TokenStream = syn::parse_str(&convert_args_code).unwrap_or_else(|err| {
            panic_on_syn_error(
                "java/jni internal parse failed for convert arguments code",
                convert_args_code,
                err,
            )
        });

        trait_impl_funcs.push(quote! {
            #[allow(unused_mut)]
            fn #func_name(#(#args_with_types),*) {
                #type_size_asserts
                let env = self.get_jni_env();
                if let Some(env) = env.env {
                    #convert_args
                    unsafe {
                        (**env).CallVoidMethod.unwrap()(env, self.this,
                                                        self.methods[#method_idx],
                                                        #(#args),*);
                        if (**env).ExceptionCheck.unwrap()(env) != 0 {
                            log::error!(concat!(stringify!(#func_name), ": java throw exception"));
                            (**env).ExceptionDescribe.unwrap()(env);
                            (**env).ExceptionClear.unwrap()(env);
                        }
                    };
                }
            }
        });
    }

    let self_type_name = &interface.self_type.bounds[0];
    let tt: TokenStream = quote! {
        impl #self_type_name for JavaCallback {
            #(#trait_impl_funcs)*
        }
    };
    ctx.rust_code.push(tt);
    Ok(())
}

lazy_static! {
    static ref JNI_FOR_VARIADIC_C_FUNC_CALL: FxHashMap<&'static str, &'static str> = {
        let mut m = FxHashMap::default();
        m.insert("jboolean", "::std::os::raw::c_uint");
        m.insert("jbyte", "::std::os::raw::c_int");
        m.insert("jshort", "::std::os::raw::c_int");
        m.insert("jfloat", "f64");
        m
    };
}

// To use `C` function with variable number of arguments,
// we need automatic type conversation, see
// http://en.cppreference.com/w/c/language/conversion#Default_argument_promotions
// for more details.
// return arg with conversation plus asserts
fn convert_args_for_variadic_function_call(
    f_method: &JniForeignMethodSignature,
) -> (Vec<TokenStream>, TokenStream) {
    let mut ret = Vec::with_capacity(f_method.input.len());
    for (i, arg) in f_method.input.iter().enumerate() {
        let arg_name = Ident::new(&format!("a{}", i), Span::call_site());
        if let Some(conv_type_str) = JNI_FOR_VARIADIC_C_FUNC_CALL
            .get(&*arg.as_ref().correspoding_rust_type.normalized_name.as_str())
        {
            let conv_type: TokenStream = syn::parse_str(*conv_type_str).unwrap_or_else(|err| {
                panic_on_syn_error(
                    "java/jni internal error: can not parse type for variable conversation",
                    conv_type_str.to_string(),
                    err,
                )
            });
            ret.push(quote!(#arg_name as #conv_type));
        } else {
            ret.push(quote!(#arg_name));
        }
    }
    let check_sizes = quote! {
        swig_assert_eq_size!(::std::os::raw::c_uint, u32);
        swig_assert_eq_size!(::std::os::raw::c_int, i32);
    };
    (ret, check_sizes)
}
<|endoftext|>


<commit_before>// Copyright 2014-2017 The Rooster 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.

use password::v2::{Password, PasswordStore};
use std::io::Write;
use std::io::stdin;

/// Used to indicate lists should have a number, ie: 23 Google my.account@gmail.com
pub const WITH_NUMBERS: bool = true;

/// Used to indicate lists should not have a number, ie: Google my.account@gmail.com
pub const WITHOUT_NUMBERS: bool = false;

pub enum OutputStream {
    Stdout,
    Stderr,
}

pub fn print_list_of_passwords(
    passwords: &Vec<&Password>,
    with_numbers: bool,
    output_stream: OutputStream,
) {
    // Find the app name column length
    let longest_app_name = passwords.iter().fold(0, |acc, p| if p.name.len() > acc {
        p.name.len()
    } else {
        acc
    });

    // Find the number column length
    let i_width = ((passwords.len() as f64).log10() + 1 as f64).floor() as usize;

    for (i, p) in passwords.iter().enumerate() {
        let s = match with_numbers {
            WITH_NUMBERS => {
                format!(
                    "{:i_width$} {:app_name_width$} {:30}",
                    i + 1,
                    p.name,
                    p.username,
                    i_width = i_width,
                    app_name_width = longest_app_name
                )
            }
            WITHOUT_NUMBERS => {
                format!(
                    "{:app_name_width$} {:30}",
                    p.name,
                    p.username,
                    app_name_width = longest_app_name
                )
            }
        };

        match output_stream {
            OutputStream::Stdout => println!("{}", s),
            OutputStream::Stderr => println_stderr!("{}", s),
        }
    }
}

fn request_password_index_from_stdin(passwords: &Vec<&Password>, prompt: &str) -> usize {
    assert!(!passwords.is_empty());

    // Read the index from the command line and convert to a number
    let mut line = String::new();
    loop {

        if passwords.len() > 1 {
            println_stderr!("{}", prompt);
            println_stderr!("Type in number from 1 to {}", passwords.len());
        } else if passwords.len() == 1 {
            println_stderr!("If this is the password you mean, type number 1.");
        }

        line.clear();
        match stdin().read_line(&mut line) {
            Ok(_) => {
                match line.trim().parse::<usize>() {
                    Ok(index) => {
                        if index == 0 || index > passwords.len() {
                            println_err!(
                                "I need a number between 1 and {}. Let's try again:",
                                passwords.len()
                            );
                            continue;
                        }

                        return index - 1;
                    }
                    Err(err) => {
                        println_err!(
                            "This isn't a valid number (reason: {}). Let's try again (1 to {}): ",
                            err,
                            passwords.len()
                        );
                        continue;
                    }
                };
            }
            Err(err) => {
                println_err!(
                    "I couldn't read that (reason: {}). Let's try again (1 to {}): ",
                    err,
                    passwords.len()
                );
            }
        }
    }
}

pub fn choose_password_in_list(
    passwords: &Vec<&Password>,
    with_numbers: bool,
    prompt: &str,
) -> usize {
    print_list_of_passwords(passwords, with_numbers, OutputStream::Stderr);
    println_stderr!("");
    request_password_index_from_stdin(passwords, prompt)
}

pub fn search_and_choose_password<'a>(
    store: &'a PasswordStore,
    query: &str,
    with_numbers: bool,
    prompt: &str,
) -> Option<&'a Password> {
    let passwords = store.search_passwords(query);
    if passwords.len() == 0 {
        println_stderr!("I can't find any passwords for \"{}\"", query);
        return None;
    }

    if let Some(&password) = passwords.iter().find(|p| {
        p.name.to_lowercase() == query.to_lowercase()
    })
    {
        return Some(&password);
    }

    let index = choose_password_in_list(&passwords, with_numbers, prompt);
    Some(passwords[index])
}
<commit_msg>30DaysOfTests: password list display test<commit_after>// Copyright 2014-2017 The Rooster 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.

use password::v2::{Password, PasswordStore};
use std::io::Write;
use std::io::stdin;

/// Used to indicate lists should have a number, ie: 23 Google my.account@gmail.com
pub const WITH_NUMBERS: bool = true;

/// Used to indicate lists should not have a number, ie: Google my.account@gmail.com
pub const WITHOUT_NUMBERS: bool = false;

pub enum OutputStream {
    Stdout,
    Stderr,
}

fn get_list_of_passwords(passwords: &Vec<&Password>, with_numbers: bool) -> Vec<String> {
    // Find the app name column length
    let longest_app_name = passwords.iter().fold(0, |acc, p| if p.name.len() > acc {
        p.name.len()
    } else {
        acc
    });

    // Find the username column length
    let longest_username = passwords.iter().fold(
        0,
        |acc, p| if p.username.len() > acc {
            p.username.len()
        } else {
            acc
        },
    );

    // Find the number column length
    let i_width = ((passwords.len() as f64).log10() + 1 as f64).floor() as usize;

    let mut list = Vec::new();

    for (i, p) in passwords.iter().enumerate() {
        let s = match with_numbers {
            WITH_NUMBERS => {
                format!(
                    "{:i_width$} {:app_name_width$} {:username_width$}",
                    i + 1,
                    p.name,
                    p.username,
                    i_width = i_width,
                    app_name_width = longest_app_name,
                    username_width = longest_username,
                )
            }
            WITHOUT_NUMBERS => {
                format!(
                    "{:app_name_width$} {:username_width$}",
                    p.name,
                    p.username,
                    app_name_width = longest_app_name,
                    username_width = longest_username,
                )
            }
        };

        list.push(s);
    }

    list
}

pub fn print_list_of_passwords(
    passwords: &Vec<&Password>,
    with_numbers: bool,
    output_stream: OutputStream,
) {
    let list = get_list_of_passwords(passwords, with_numbers);

    for s in list {
        match output_stream {
            OutputStream::Stdout => println!("{}", s),
            OutputStream::Stderr => println_stderr!("{}", s),
        }
    }
}

fn request_password_index_from_stdin(passwords: &Vec<&Password>, prompt: &str) -> usize {
    assert!(!passwords.is_empty());

    // Read the index from the command line and convert to a number
    let mut line = String::new();
    loop {

        if passwords.len() > 1 {
            println_stderr!("{}", prompt);
            println_stderr!("Type in number from 1 to {}", passwords.len());
        } else if passwords.len() == 1 {
            println_stderr!("If this is the password you mean, type number 1.");
        }

        line.clear();
        match stdin().read_line(&mut line) {
            Ok(_) => {
                match line.trim().parse::<usize>() {
                    Ok(index) => {
                        if index == 0 || index > passwords.len() {
                            println_err!(
                                "I need a number between 1 and {}. Let's try again:",
                                passwords.len()
                            );
                            continue;
                        }

                        return index - 1;
                    }
                    Err(err) => {
                        println_err!(
                            "This isn't a valid number (reason: {}). Let's try again (1 to {}): ",
                            err,
                            passwords.len()
                        );
                        continue;
                    }
                };
            }
            Err(err) => {
                println_err!(
                    "I couldn't read that (reason: {}). Let's try again (1 to {}): ",
                    err,
                    passwords.len()
                );
            }
        }
    }
}

pub fn choose_password_in_list(
    passwords: &Vec<&Password>,
    with_numbers: bool,
    prompt: &str,
) -> usize {
    print_list_of_passwords(passwords, with_numbers, OutputStream::Stderr);
    println_stderr!("");
    request_password_index_from_stdin(passwords, prompt)
}

pub fn search_and_choose_password<'a>(
    store: &'a PasswordStore,
    query: &str,
    with_numbers: bool,
    prompt: &str,
) -> Option<&'a Password> {
    let passwords = store.search_passwords(query);
    if passwords.len() == 0 {
        println_stderr!("I can't find any passwords for \"{}\"", query);
        return None;
    }

    if let Some(&password) = passwords.iter().find(|p| {
        p.name.to_lowercase() == query.to_lowercase()
    })
    {
        return Some(&password);
    }

    let index = choose_password_in_list(&passwords, with_numbers, prompt);
    Some(passwords[index])
}

#[cfg(test)]
mod test {
    use super::get_list_of_passwords;
    use list::{WITH_NUMBERS, WITHOUT_NUMBERS};
    use password::v2::Password;
    use safe_string::SafeString;

    // Creates a list of at least two passwords, and more if specified
    fn get_passwords(mut additional: i32) -> Vec<Password> {
        let google = Password::new(
            format!("google"),
            format!("short un"),
            SafeString::new(format!("xxxx")),
        );

        let mut list = vec![
            Password::new(
                format!("youtube.com"),
                format!("that long username"),
                SafeString::new(format!("xxxx"))
            ),
            google.clone(),
        ];

        while additional > 0 {
            list.push(google.clone());
            additional -= 1;
        }

        list
    }

    #[test]
    fn password_list_has_right_format_with_numbers() {
        // With 2 passwords (number width 1)
        let passwords = get_passwords(0);
        let list = get_list_of_passwords(&passwords.iter().collect(), WITH_NUMBERS);

        assert_eq!(
            list,
            &[
                "1 youtube.com that long username",
                "2 google      short un          ",
            ]
        );

        // Now with 10 passwords (number width 2)
        let passwords = get_passwords(8);
        let list = get_list_of_passwords(&passwords.iter().collect(), WITH_NUMBERS);

        assert_eq!(
            list,
            &[
                " 1 youtube.com that long username",
                " 2 google      short un          ",
                " 3 google      short un          ",
                " 4 google      short un          ",
                " 5 google      short un          ",
                " 6 google      short un          ",
                " 7 google      short un          ",
                " 8 google      short un          ",
                " 9 google      short un          ",
                "10 google      short un          ",
            ]
        );
    }

    #[test]
    fn password_list_has_right_format_without_numbers() {
        let passwords = get_passwords(0);
        let list = get_list_of_passwords(&passwords.iter().collect(), WITHOUT_NUMBERS);

        assert_eq!(
            list,
            &[
                "youtube.com that long username",
                "google      short un          ",
            ]
        );
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use sys::component::{self, Component};
use sys::disk::Disk;
use sys::processor::*;

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem::{size_of, zeroed};

use DiskExt;
use Pid;
use ProcessExt;
use RefreshKind;
use SystemExt;

use windows::network::{self, NetworkData};
use windows::process::{
    compute_cpu_usage, get_handle, get_parent_process_id, update_proc_info, Process,
};
use windows::processor::CounterValue;
use windows::tools::*;

use winapi::shared::minwindef::{DWORD, FALSE};
use winapi::shared::winerror::ERROR_SUCCESS;
use winapi::um::minwinbase::STILL_ACTIVE;
use winapi::um::pdh::PdhEnumObjectItemsW;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::psapi::K32EnumProcesses;
use winapi::um::sysinfoapi::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
use winapi::um::winnt::HANDLE;

use rayon::prelude::*;

const PROCESS_LEN: usize = 10192;

/// Struct containing system's information.
pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_free: u64,
    swap_total: u64,
    swap_free: u64,
    processors: Vec<Processor>,
    temperatures: Vec<Component>,
    disks: Vec<Disk>,
    query: Option<Query>,
    network: NetworkData,
    uptime: u64,
}

impl System {
    fn clear_procs(&mut self) {
        if self.processors.len() > 0 {
            let mut to_delete = Vec::new();

            for (pid, proc_) in self.process_list.iter_mut() {
                if !is_proc_running(get_handle(proc_)) {
                    to_delete.push(*pid);
                } else {
                    compute_cpu_usage(proc_, self.processors.len() as u64 - 1);
                }
            }
            for pid in to_delete {
                self.process_list.remove(&pid);
            }
        }
    }
}

struct Wrap(Process);

unsafe impl Send for Wrap {}

struct WrapSystem<'a>(UnsafeCell<&'a mut System>);

impl<'a> WrapSystem<'a> {
    fn get(&self) -> &'a mut System {
        unsafe { *(self.0.get()) }
    }
}

unsafe impl<'a> Send for WrapSystem<'a> {}
unsafe impl<'a> Sync for WrapSystem<'a> {}

impl SystemExt for System {
    #[allow(non_snake_case)]
    fn new_with_specifics(refreshes: RefreshKind) -> System {
        let mut s = System {
            process_list: HashMap::new(),
            mem_total: 0,
            mem_free: 0,
            swap_total: 0,
            swap_free: 0,
            processors: init_processors(),
            temperatures: component::get_components(),
            disks: Vec::new(),
            query: Query::new(),
            network: network::new(),
            uptime: get_uptime(),
        };
        // TODO: in case a translation fails, it might be nice to log it somewhere...
        if let Some(ref mut query) = s.query {
            let x = unsafe { load_symbols() };
            if let Some(processor_trans) = get_translation(&"Processor".to_owned(), &x) {
                let idle_time_trans = get_translation(&"% Idle Time".to_owned(), &x);
                let proc_time_trans = get_translation(&"% Processor Time".to_owned(), &x);
                if let Some(ref proc_time_trans) = proc_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, proc_time_trans),
                        query,
                        get_key_used(&mut s.processors[0]),
                        "tot_0".to_owned(),
                        CounterValue::Float(0.),
                    );
                }
                if let Some(ref idle_time_trans) = idle_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, idle_time_trans),
                        query,
                        get_key_idle(&mut s.processors[0]),
                        "tot_1".to_owned(),
                        CounterValue::Float(0.),
                    );
                }
                for (pos, proc_) in s.processors.iter_mut().skip(1).enumerate() {
                    if let Some(ref proc_time_trans) = proc_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, proc_time_trans),
                            query,
                            get_key_used(proc_),
                            format!("{}_0", pos),
                            CounterValue::Float(0.),
                        );
                    }
                    if let Some(ref idle_time_trans) = idle_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, idle_time_trans),
                            query,
                            get_key_idle(proc_),
                            format!("{}_1", pos),
                            CounterValue::Float(0.),
                        );
                    }
                }
            }

            if let Some(network_trans) = get_translation(&"Network Interface".to_owned(), &x) {
                let network_in_trans = get_translation(&"Bytes Received/Sec".to_owned(), &x);
                let network_out_trans = get_translation(&"Bytes Sent/sec".to_owned(), &x);

                const PERF_DETAIL_WIZARD: DWORD = 400;
                const PDH_MORE_DATA: DWORD = 0x800007D2;

                let mut network_trans_utf16: Vec<u16> = network_trans.encode_utf16().collect();
                network_trans_utf16.push(0);
                let mut dwCounterListSize: DWORD = 0;
                let mut dwInstanceListSize: DWORD = 0;
                let status = unsafe {
                    PdhEnumObjectItemsW(
                        ::std::ptr::null(),
                        ::std::ptr::null(),
                        network_trans_utf16.as_ptr(),
                        ::std::ptr::null_mut(),
                        &mut dwCounterListSize,
                        ::std::ptr::null_mut(),
                        &mut dwInstanceListSize,
                        PERF_DETAIL_WIZARD,
                        0,
                    )
                };
                if status != PDH_MORE_DATA as i32 {
                    eprintln!("PdhEnumObjectItems invalid status: {:x}", status);
                }
                let mut pwsCounterListBuffer: Vec<u16> =
                    Vec::with_capacity(dwCounterListSize as usize);
                let mut pwsInstanceListBuffer: Vec<u16> =
                    Vec::with_capacity(dwInstanceListSize as usize);
                unsafe {
                    pwsCounterListBuffer.set_len(dwCounterListSize as usize);
                    pwsInstanceListBuffer.set_len(dwInstanceListSize as usize);
                }
                let status = unsafe {
                    PdhEnumObjectItemsW(
                        ::std::ptr::null(),
                        ::std::ptr::null(),
                        network_trans_utf16.as_ptr(),
                        pwsCounterListBuffer.as_mut_ptr(),
                        &mut dwCounterListSize,
                        pwsInstanceListBuffer.as_mut_ptr(),
                        &mut dwInstanceListSize,
                        PERF_DETAIL_WIZARD,
                        0,
                    )
                };
                if status != ERROR_SUCCESS as i32 {
                    eprintln!("PdhEnumObjectItems invalid status: {:x}", status);
                }

                for (pos, x) in pwsInstanceListBuffer
                    .split(|x| *x == 0)
                    .filter(|x| x.len() > 0)
                    .enumerate()
                {
                    let net_interface = String::from_utf16(x).expect("invalid utf16");
                    if let Some(ref network_in_trans) = network_in_trans {
                        let mut key_in = None;
                        add_counter(
                            format!(
                                "\\{}({})\\{}",
                                network_trans, net_interface, network_in_trans
                            ),
                            query,
                            &mut key_in,
                            format!("net{}_in", pos),
                            CounterValue::Integer(0),
                        );
                        if key_in.is_some() {
                            network::get_keys_in(&mut s.network).push(key_in.unwrap());
                        }
                    }
                    if let Some(ref network_out_trans) = network_out_trans {
                        let mut key_out = None;
                        add_counter(
                            format!(
                                "\\{}({})\\{}",
                                network_trans, net_interface, network_out_trans
                            ),
                            query,
                            &mut key_out,
                            format!("net{}_out", pos),
                            CounterValue::Integer(0),
                        );
                        if key_out.is_some() {
                            network::get_keys_out(&mut s.network).push(key_out.unwrap());
                        }
                    }
                }
            }
            query.start();
        }
        s.refresh_specifics(refreshes);
        s
    }

    fn refresh_system(&mut self) {
        unsafe {
            let mut mem_info: MEMORYSTATUSEX = zeroed();
            mem_info.dwLength = size_of::<MEMORYSTATUSEX>() as u32;
            GlobalMemoryStatusEx(&mut mem_info);
            self.mem_total = auto_cast!(mem_info.ullTotalPhys, u64);
            self.mem_free = auto_cast!(mem_info.ullAvailPhys, u64);
            //self.swap_total = auto_cast!(mem_info.ullTotalPageFile - mem_info.ullTotalPhys, u64);
            //self.swap_free = auto_cast!(mem_info.ullAvailPageFile, u64);
        }
        self.uptime = get_uptime();
        if let Some(ref mut query) = self.query {
            for p in self.processors.iter_mut() {
                let mut idle_time = None;
                if let &mut Some(ref key_idle) = get_key_idle(p) {
                    idle_time = Some(query.get(&key_idle.unique_id).expect("key disappeared"));
                }
                if let Some(idle_time) = idle_time {
                    set_cpu_usage(p, 1. - idle_time);
                }
            }
        }
    }

    fn refresh_network(&mut self) {
        network::refresh(&mut self.network, &self.query);
    }

    fn refresh_process(&mut self, pid: Pid) -> bool {
        if refresh_existing_process(self, pid, true) == false {
            self.process_list.remove(&pid);
            false
        } else {
            true
        }
    }

    fn refresh_processes(&mut self) {
        // I think that 10192 as length will be enough to get all processes at once...
        let mut process_ids: Vec<DWORD> = Vec::with_capacity(PROCESS_LEN);
        let mut cb_needed = 0;

        unsafe {
            process_ids.set_len(PROCESS_LEN);
        }
        let size = ::std::mem::size_of::<DWORD>() * process_ids.len();
        unsafe {
            if K32EnumProcesses(process_ids.as_mut_ptr(), size as DWORD, &mut cb_needed) == 0 {
                return;
            }
        }
        let nb_processes = cb_needed / ::std::mem::size_of::<DWORD>() as DWORD;
        unsafe {
            process_ids.set_len(nb_processes as usize);
        }

        {
            let this = WrapSystem(UnsafeCell::new(self));

            process_ids
                .par_iter()
                .filter_map(|pid| {
                    let pid = *pid as usize;
                    if !refresh_existing_process(this.get(), pid, false) {
                        let ppid = unsafe { get_parent_process_id(pid) };
                        let mut p = Process::new(pid, ppid, 0);
                        update_proc_info(&mut p);
                        Some(Wrap(p))
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>()
        }
        .into_iter()
        .for_each(|p| {
            self.process_list.insert(p.0.pid(), p.0);
        });
        self.clear_procs();
    }

    fn refresh_disks(&mut self) {
        for disk in &mut self.disks {
            disk.update();
        }
    }

    fn refresh_disk_list(&mut self) {
        self.disks = unsafe { get_disks() };
    }

    fn get_process_list(&self) -> &HashMap<Pid, Process> {
        &self.process_list
    }

    fn get_process(&self, pid: Pid) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    fn get_processor_list(&self) -> &[Processor] {
        &self.processors[..]
    }

    fn get_total_memory(&self) -> u64 {
        self.mem_total >> 10
    }

    fn get_free_memory(&self) -> u64 {
        self.mem_free >> 10
    }

    fn get_used_memory(&self) -> u64 {
        (self.mem_total - self.mem_free) >> 10
    }

    fn get_total_swap(&self) -> u64 {
        self.swap_total >> 10
    }

    fn get_free_swap(&self) -> u64 {
        self.swap_free >> 10
    }

    fn get_used_swap(&self) -> u64 {
        (self.swap_total - self.swap_free) >> 10
    }

    fn get_components_list(&self) -> &[Component] {
        &self.temperatures[..]
    }

    fn get_disks(&self) -> &[Disk] {
        &self.disks[..]
    }

    fn get_network(&self) -> &NetworkData {
        &self.network
    }

    fn get_uptime(&self) -> u64 {
        self.uptime
    }
}

fn is_proc_running(handle: HANDLE) -> bool {
    let mut exit_code = 0;
    let ret = unsafe { GetExitCodeProcess(handle, &mut exit_code) };
    !(ret == FALSE || exit_code != STILL_ACTIVE)
}

fn refresh_existing_process(s: &mut System, pid: Pid, compute_cpu: bool) -> bool {
    if let Some(ref mut entry) = s.process_list.get_mut(&(pid as usize)) {
        if !is_proc_running(get_handle(entry)) {
            return false;
        }
        update_proc_info(entry);
        if compute_cpu {
            compute_cpu_usage(entry, s.processors.len() as u64 - 1);
        }
        true
    } else {
        false
    }
}
<commit_msg>Don't blindly progress past errors...<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use sys::component::{self, Component};
use sys::disk::Disk;
use sys::processor::*;

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem::{size_of, zeroed};

use DiskExt;
use Pid;
use ProcessExt;
use RefreshKind;
use SystemExt;

use windows::network::{self, NetworkData};
use windows::process::{
    compute_cpu_usage, get_handle, get_parent_process_id, update_proc_info, Process,
};
use windows::processor::CounterValue;
use windows::tools::*;

use winapi::shared::minwindef::{DWORD, FALSE};
use winapi::shared::winerror::ERROR_SUCCESS;
use winapi::um::minwinbase::STILL_ACTIVE;
use winapi::um::pdh::PdhEnumObjectItemsW;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::psapi::K32EnumProcesses;
use winapi::um::sysinfoapi::{GlobalMemoryStatusEx, MEMORYSTATUSEX};
use winapi::um::winnt::HANDLE;

use rayon::prelude::*;

const PROCESS_LEN: usize = 10192;

/// Struct containing system's information.
pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_free: u64,
    swap_total: u64,
    swap_free: u64,
    processors: Vec<Processor>,
    temperatures: Vec<Component>,
    disks: Vec<Disk>,
    query: Option<Query>,
    network: NetworkData,
    uptime: u64,
}

impl System {
    fn clear_procs(&mut self) {
        if self.processors.len() > 0 {
            let mut to_delete = Vec::new();

            for (pid, proc_) in self.process_list.iter_mut() {
                if !is_proc_running(get_handle(proc_)) {
                    to_delete.push(*pid);
                } else {
                    compute_cpu_usage(proc_, self.processors.len() as u64 - 1);
                }
            }
            for pid in to_delete {
                self.process_list.remove(&pid);
            }
        }
    }
}

struct Wrap(Process);

unsafe impl Send for Wrap {}

struct WrapSystem<'a>(UnsafeCell<&'a mut System>);

impl<'a> WrapSystem<'a> {
    fn get(&self) -> &'a mut System {
        unsafe { *(self.0.get()) }
    }
}

unsafe impl<'a> Send for WrapSystem<'a> {}
unsafe impl<'a> Sync for WrapSystem<'a> {}

impl SystemExt for System {
    #[allow(non_snake_case)]
    fn new_with_specifics(refreshes: RefreshKind) -> System {
        let mut s = System {
            process_list: HashMap::new(),
            mem_total: 0,
            mem_free: 0,
            swap_total: 0,
            swap_free: 0,
            processors: init_processors(),
            temperatures: component::get_components(),
            disks: Vec::new(),
            query: Query::new(),
            network: network::new(),
            uptime: get_uptime(),
        };
        // TODO: in case a translation fails, it might be nice to log it somewhere...
        if let Some(ref mut query) = s.query {
            let x = unsafe { load_symbols() };
            if let Some(processor_trans) = get_translation(&"Processor".to_owned(), &x) {
                let idle_time_trans = get_translation(&"% Idle Time".to_owned(), &x);
                let proc_time_trans = get_translation(&"% Processor Time".to_owned(), &x);
                if let Some(ref proc_time_trans) = proc_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, proc_time_trans),
                        query,
                        get_key_used(&mut s.processors[0]),
                        "tot_0".to_owned(),
                        CounterValue::Float(0.),
                    );
                }
                if let Some(ref idle_time_trans) = idle_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, idle_time_trans),
                        query,
                        get_key_idle(&mut s.processors[0]),
                        "tot_1".to_owned(),
                        CounterValue::Float(0.),
                    );
                }
                for (pos, proc_) in s.processors.iter_mut().skip(1).enumerate() {
                    if let Some(ref proc_time_trans) = proc_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, proc_time_trans),
                            query,
                            get_key_used(proc_),
                            format!("{}_0", pos),
                            CounterValue::Float(0.),
                        );
                    }
                    if let Some(ref idle_time_trans) = idle_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, idle_time_trans),
                            query,
                            get_key_idle(proc_),
                            format!("{}_1", pos),
                            CounterValue::Float(0.),
                        );
                    }
                }
            }

            if let Some(network_trans) = get_translation(&"Network Interface".to_owned(), &x) {
                let network_in_trans = get_translation(&"Bytes Received/Sec".to_owned(), &x);
                let network_out_trans = get_translation(&"Bytes Sent/sec".to_owned(), &x);

                const PERF_DETAIL_WIZARD: DWORD = 400;
                const PDH_MORE_DATA: DWORD = 0x800007D2;

                let mut network_trans_utf16: Vec<u16> = network_trans.encode_utf16().collect();
                network_trans_utf16.push(0);
                let mut dwCounterListSize: DWORD = 0;
                let mut dwInstanceListSize: DWORD = 0;
                let status = unsafe {
                    PdhEnumObjectItemsW(
                        ::std::ptr::null(),
                        ::std::ptr::null(),
                        network_trans_utf16.as_ptr(),
                        ::std::ptr::null_mut(),
                        &mut dwCounterListSize,
                        ::std::ptr::null_mut(),
                        &mut dwInstanceListSize,
                        PERF_DETAIL_WIZARD,
                        0,
                    )
                };
                if status != PDH_MORE_DATA as i32 {
                    eprintln!("PdhEnumObjectItems invalid status: {:x}", status);
                } else {
                    let mut pwsCounterListBuffer: Vec<u16> =
                        Vec::with_capacity(dwCounterListSize as usize);
                    let mut pwsInstanceListBuffer: Vec<u16> =
                        Vec::with_capacity(dwInstanceListSize as usize);
                    unsafe {
                        pwsCounterListBuffer.set_len(dwCounterListSize as usize);
                        pwsInstanceListBuffer.set_len(dwInstanceListSize as usize);
                    }
                    let status = unsafe {
                        PdhEnumObjectItemsW(
                            ::std::ptr::null(),
                            ::std::ptr::null(),
                            network_trans_utf16.as_ptr(),
                            pwsCounterListBuffer.as_mut_ptr(),
                            &mut dwCounterListSize,
                            pwsInstanceListBuffer.as_mut_ptr(),
                            &mut dwInstanceListSize,
                            PERF_DETAIL_WIZARD,
                            0,
                        )
                    };
                    if status != ERROR_SUCCESS as i32 {
                        eprintln!("PdhEnumObjectItems invalid status: {:x}", status);
                    } else {
                        for (pos, x) in pwsInstanceListBuffer
                            .split(|x| *x == 0)
                            .filter(|x| x.len() > 0)
                            .enumerate()
                        {
                            let net_interface = String::from_utf16(x).expect("invalid utf16");
                            if let Some(ref network_in_trans) = network_in_trans {
                                let mut key_in = None;
                                add_counter(
                                    format!(
                                        "\\{}({})\\{}",
                                        network_trans, net_interface, network_in_trans
                                    ),
                                    query,
                                    &mut key_in,
                                    format!("net{}_in", pos),
                                    CounterValue::Integer(0),
                                );
                                if key_in.is_some() {
                                    network::get_keys_in(&mut s.network).push(key_in.unwrap());
                                }
                            }
                            if let Some(ref network_out_trans) = network_out_trans {
                                let mut key_out = None;
                                add_counter(
                                    format!(
                                        "\\{}({})\\{}",
                                        network_trans, net_interface, network_out_trans
                                    ),
                                    query,
                                    &mut key_out,
                                    format!("net{}_out", pos),
                                    CounterValue::Integer(0),
                                );
                                if key_out.is_some() {
                                    network::get_keys_out(&mut s.network).push(key_out.unwrap());
                                }
                            }
                        }
                    }
                }
            }
            query.start();
        }
        s.refresh_specifics(refreshes);
        s
    }

    fn refresh_system(&mut self) {
        unsafe {
            let mut mem_info: MEMORYSTATUSEX = zeroed();
            mem_info.dwLength = size_of::<MEMORYSTATUSEX>() as u32;
            GlobalMemoryStatusEx(&mut mem_info);
            self.mem_total = auto_cast!(mem_info.ullTotalPhys, u64);
            self.mem_free = auto_cast!(mem_info.ullAvailPhys, u64);
            //self.swap_total = auto_cast!(mem_info.ullTotalPageFile - mem_info.ullTotalPhys, u64);
            //self.swap_free = auto_cast!(mem_info.ullAvailPageFile, u64);
        }
        self.uptime = get_uptime();
        if let Some(ref mut query) = self.query {
            for p in self.processors.iter_mut() {
                let mut idle_time = None;
                if let &mut Some(ref key_idle) = get_key_idle(p) {
                    idle_time = Some(query.get(&key_idle.unique_id).expect("key disappeared"));
                }
                if let Some(idle_time) = idle_time {
                    set_cpu_usage(p, 1. - idle_time);
                }
            }
        }
    }

    fn refresh_network(&mut self) {
        network::refresh(&mut self.network, &self.query);
    }

    fn refresh_process(&mut self, pid: Pid) -> bool {
        if refresh_existing_process(self, pid, true) == false {
            self.process_list.remove(&pid);
            false
        } else {
            true
        }
    }

    fn refresh_processes(&mut self) {
        // I think that 10192 as length will be enough to get all processes at once...
        let mut process_ids: Vec<DWORD> = Vec::with_capacity(PROCESS_LEN);
        let mut cb_needed = 0;

        unsafe {
            process_ids.set_len(PROCESS_LEN);
        }
        let size = ::std::mem::size_of::<DWORD>() * process_ids.len();
        unsafe {
            if K32EnumProcesses(process_ids.as_mut_ptr(), size as DWORD, &mut cb_needed) == 0 {
                return;
            }
        }
        let nb_processes = cb_needed / ::std::mem::size_of::<DWORD>() as DWORD;
        unsafe {
            process_ids.set_len(nb_processes as usize);
        }

        {
            let this = WrapSystem(UnsafeCell::new(self));

            process_ids
                .par_iter()
                .filter_map(|pid| {
                    let pid = *pid as usize;
                    if !refresh_existing_process(this.get(), pid, false) {
                        let ppid = unsafe { get_parent_process_id(pid) };
                        let mut p = Process::new(pid, ppid, 0);
                        update_proc_info(&mut p);
                        Some(Wrap(p))
                    } else {
                        None
                    }
                })
                .collect::<Vec<_>>()
        }
        .into_iter()
        .for_each(|p| {
            self.process_list.insert(p.0.pid(), p.0);
        });
        self.clear_procs();
    }

    fn refresh_disks(&mut self) {
        for disk in &mut self.disks {
            disk.update();
        }
    }

    fn refresh_disk_list(&mut self) {
        self.disks = unsafe { get_disks() };
    }

    fn get_process_list(&self) -> &HashMap<Pid, Process> {
        &self.process_list
    }

    fn get_process(&self, pid: Pid) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    fn get_processor_list(&self) -> &[Processor] {
        &self.processors[..]
    }

    fn get_total_memory(&self) -> u64 {
        self.mem_total >> 10
    }

    fn get_free_memory(&self) -> u64 {
        self.mem_free >> 10
    }

    fn get_used_memory(&self) -> u64 {
        (self.mem_total - self.mem_free) >> 10
    }

    fn get_total_swap(&self) -> u64 {
        self.swap_total >> 10
    }

    fn get_free_swap(&self) -> u64 {
        self.swap_free >> 10
    }

    fn get_used_swap(&self) -> u64 {
        (self.swap_total - self.swap_free) >> 10
    }

    fn get_components_list(&self) -> &[Component] {
        &self.temperatures[..]
    }

    fn get_disks(&self) -> &[Disk] {
        &self.disks[..]
    }

    fn get_network(&self) -> &NetworkData {
        &self.network
    }

    fn get_uptime(&self) -> u64 {
        self.uptime
    }
}

fn is_proc_running(handle: HANDLE) -> bool {
    let mut exit_code = 0;
    let ret = unsafe { GetExitCodeProcess(handle, &mut exit_code) };
    !(ret == FALSE || exit_code != STILL_ACTIVE)
}

fn refresh_existing_process(s: &mut System, pid: Pid, compute_cpu: bool) -> bool {
    if let Some(ref mut entry) = s.process_list.get_mut(&(pid as usize)) {
        if !is_proc_running(get_handle(entry)) {
            return false;
        }
        update_proc_info(entry);
        if compute_cpu {
            compute_cpu_usage(entry, s.processors.len() as u64 - 1);
        }
        true
    } else {
        false
    }
}
<|endoftext|>