<commit_before>extern crate syntex;
extern crate syntex_syntax;
extern crate syntex_pos;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;

mod parse_utils;
mod jni;
mod core;
mod rust_generator;
mod java_generator;

use std::collections::HashMap;
use std::env;
use std::sync::atomic::AtomicBool;
use std::sync::Mutex;

use syntex::Registry;
use syntex_syntax::ext::base::{ExtCtxt, MacResult, DummyResult};
use syntex_syntax::parse::{token, parser};
use syntex_syntax::tokenstream::TokenTree;
use syntex_syntax::codemap::Span;
use syntex_syntax::{parse, ast};
use syntex_syntax::ptr::P;
use syntex_syntax::print::pprust;

use parse_utils::*;
use core::*;
use rust_generator::{generate_rust_code, RUST_OBJECT_TO_JOBJECT};
use java_generator::generate_java_code;

lazy_static! {
    pub static ref COMMON_CODE_GENERATED: AtomicBool = AtomicBool::new(false);
    static ref TYPE_HANDLERS: Mutex<Vec<TypeHandler>> = {
        Mutex::new(vec![
            TypeHandler{rust_type_name: "bool".into(), jni_type_name: "jboolean", java_type_name: "boolean".into(),
                        from_jni_converter: Some(FromForeignArgConverter("let {arg_name} = {arg_name} != 0;".into())),
                        to_jni_converter: Some(ToForeignRetConverter("let ret = if ret { 1 as jboolean } else { 0 as jboolean };".into()))},
            TypeHandler{rust_type_name: "i32".into(), jni_type_name: "jint", java_type_name: "int".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "f32".into(), jni_type_name: "jfloat", java_type_name: "float".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "f64".into(), jni_type_name: "jdouble", java_type_name: "double".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "&str".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: Some(FromForeignArgConverter(r#"
  let {arg_name} = JavaString::new(env, {arg_name});
  let {arg_name} = {arg_name}.to_str();
"#.into())),
                        to_jni_converter: Some(ToForeignRetConverter(r#"
  let ret = ::std::ffi::CString::new(ret).unwrap();
  let ret = unsafe { (**env).NewStringUTF.unwrap()(env, ret.as_ptr()) };
"#.into()))},
            TypeHandler{rust_type_name: "&Path".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: Some(FromForeignArgConverter(
                            r#"
  let {arg_name} = JavaString::new(env, {arg_name});
  let {arg_name} = Path::new({arg_name}.to_str());
"#.into()
                        )),
                        to_jni_converter: None},
            TypeHandler{rust_type_name: "String".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: None,
                        to_jni_converter: Some(ToForeignRetConverter(
                            r#"
  let ret = ret.into_bytes();
  let ret = unsafe { ::std::ffi::CString::from_vec_unchecked(ret) };
  let ret = unsafe { (**env).NewStringUTF.unwrap()(env, ret.as_ptr()) };
"#.into()
                        ))},
            ])
    };
}

pub fn register(registry: &mut Registry) {
    registry.add_macro("foreigner_class", expand_foreigner_class);
}

fn expand_foreigner_class<'cx>(cx: &'cx mut ExtCtxt,
                               _: Span,
                               tokens: &[TokenTree])
                               -> Box<MacResult + 'cx> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess, cx.cfg.clone(), tokens.to_vec());
    let class_ident = parser.parse_ident().unwrap();
    if class_ident.name.as_str() != "class" {
        debug!("class_indent {:?}", class_ident);
        cx.span_err(parser.span, "expect class here");
        return DummyResult::any(parser.span);
    }
    let class_name_indent = parser.parse_ident().unwrap();
    debug!("CLASS NAME {:?}", class_name_indent);
    parser.expect(&token::Token::OpenDelim(token::DelimToken::Brace)).unwrap();
    let mut methods = Vec::new();
    let mut rust_self_type = ast::Path{span: parser.span, global: false, segments: Vec::new()};
    let alias = ast::Ident::with_empty_ctxt(token::intern("alias"));
    let mut constructor_ret_type: Option<ast::Ty> = None;
    let mut this_type_for_method: Option<ast::Ty> = None;

    loop {
        if parser.eat(&token::Token::CloseDelim(token::DelimToken::Brace)) {
            break;
        }
        let func_type_name = parser.parse_ident().unwrap();
        debug!("func_type {:?}", func_type_name);
        if func_type_name.name.as_str() == "self_type" {
            rust_self_type = parser.parse_path(parser::PathStyle::Type).expect("Can not parse self_type");
            debug!("self_type: {:?}", rust_self_type);
            parser.expect(&token::Token::Semi).unwrap();
            continue;
        }
        let func_type = FuncVariant::from_str(&func_type_name.name.as_str());
        if func_type.is_none() {
            println!("unknown func type: {:?}", func_type_name);
            cx.span_err(parser.span, "expect 'constructor' or 'method' or 'static_method' here");
            return DummyResult::any(parser.span);
        }
        let func_type = func_type.unwrap();
        let func_name = parser.parse_path(parser::PathStyle::Mod).unwrap();
        debug!("func_name {:?}", func_name);

        let func_decl = match func_type {
            FuncVariant::Constructor | FuncVariant::StaticMethod => parser.parse_fn_decl(false).unwrap(),
            FuncVariant::Method => parse_fn_decl_with_self(&mut parser, |p| p.parse_arg()).unwrap(),
        };
        debug!("func_decl {:?}", func_decl);
        parser.expect(&token::Token::Semi).unwrap();
        let mut func_name_alias = None;
        if parser.eat_contextual_keyword(alias) {
            if func_type == FuncVariant::Constructor {
                cx.span_err(parser.span, "alias not supported for 'constructor'");
                return DummyResult::any(parser.span);
            }
            func_name_alias = Some(parser.parse_ident().unwrap());
            debug!("we have ALIAS `{:?}`", func_name_alias.unwrap());
            parser.expect(&token::Token::Semi).expect("no ; at the end of alias");
        }
        let (may_return_error, ret_type) = match &func_decl.output {
            &ast::FunctionRetTy::Default(_) => (false, None),
            &ast::FunctionRetTy::Ty(ref ret_type) => (is_type_std_result(ret_type), Some(ret_type.clone())),
        };
        if let FuncVariant::Constructor = func_type {
            let ret_type = ret_type.expect(&format!("{}: constructor should return value", class_name_indent));
            if let Some(ref constructor_ret_type) = constructor_ret_type {
                if pprust::ty_to_string(constructor_ret_type) != pprust::ty_to_string(&*ret_type) {
                    cx.span_err(parser.span, &format!("mismatched types of construtors: {:?} {:?}", constructor_ret_type, ret_type));
                    return DummyResult::any(parser.span);
                }
            } else {
                constructor_ret_type = Some(ret_type.unwrap());
                this_type_for_method = Some(unpack_if_type_is_result(constructor_ret_type.as_ref().unwrap()));
            }
        }
        methods.push(ForeignerMethod{
            func_type: func_type, path: func_name, in_out_type: func_decl,
            name_alias: func_name_alias.map(|v| v.name.as_str()),
            may_return_error: may_return_error,
        });
    }
    let mut type_handlers = TYPE_HANDLERS.lock().unwrap();

    let package_name = env::var("RUST_SWIG_JNI_JAVA_PACKAGE").unwrap();

    let class_info = ForeignerClassInfo {
        package_name: package_name,
        class_name: &class_name_indent.name.as_str(),
        methods: methods,
        self_rust_type: rust_self_type,
        constructor_ret_type: constructor_ret_type,
        this_type_for_method: this_type_for_method,
    };

    if class_info.this_type_for_method.is_some() {
        let mut class_th: TypeHandler = (&class_info).into();
        class_th.to_jni_converter = Some(ToForeignRetConverter(
            RUST_OBJECT_TO_JOBJECT.replace(
                "{full_class_name}", &class_info.full_java_class_name())));
        type_handlers.push(class_th);
    }

    let mut rust_java_types_map = HashMap::new();
    for it in type_handlers.iter() {
        rust_java_types_map.insert(it.rust_type_name.as_str(), &*it);
    }

    let java_output_dir = env::var("RUST_SWIG_JNI_JAVA_OUTPUT_DIR").unwrap();

    generate_java_code(&rust_java_types_map, &class_info, &java_output_dir);
    generate_rust_code(cx, &rust_java_types_map, &class_info)
}

fn is_type_std_result(ty: &ast::Ty) -> bool {
    match ty.node {
        ast::TyKind::Path(_/*self info*/, ref path) => {
            debug!("is_type_std_result_result: path: {:?}, ident {:?}", path.segments, path.segments[0].identifier.name.as_str());
            path.segments.first().map(|v| v.identifier.name.as_str() == "Result").unwrap_or(false)
        }
        _ => false,
    }
}

fn unpack_if_type_is_result(ty: &ast::Ty) -> ast::Ty {
    match ty.node {
        ast::TyKind::Path(_/*self info*/, ref path) => {
            debug!("unpack_if_type_is_result: path: {:?}, ident {:?}", path.segments, path.segments[0].identifier.name.as_str());
            path.segments.first().map(|v|
                                      if v.identifier.name.as_str() == "Result" {
                                          match v.parameters {
                                              ast::PathParameters::AngleBracketed(ref params) => {
                                                  params.types.first().map(|v: &P<ast::Ty>| {
                                                      debug!("unpack_if_type_is_result: result param {:?}", *v);
                                                      (**v).clone()
                                                  }).unwrap_or(ty.clone())
                                              }
                                              _ => ty.clone(),
                                          }
                                      } else {
                                          ty.clone()
                                      }).unwrap_or(ty.clone())
        }
        _ => ty.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syntex_syntax::{parse, ast};

    #[test]
    fn test_is_type_std_result() {
        let session = parse::ParseSess::new();
        let mut parser = parse::new_parser_from_source_str(
            &session, ast::CrateConfig::new(), "test".into(),
            "Result<Foo, String>".into());
        let ty = parser.parse_ty().unwrap();
        assert!(is_type_std_result(&*ty));
    }

    #[test]
    fn test_unpack_if_type_is_result() {
        let session = parse::ParseSess::new();
        let mut parser = parse::new_parser_from_source_str(
            &session, ast::CrateConfig::new(), "test".into(),
            "Result<Foo, String>".into());
        let ty = parser.parse_ty().unwrap();
        assert!(is_type_std_result(&*ty));

        let ok_ty = unpack_if_type_is_result(&*ty);
        assert!(!is_type_std_result(&ok_ty));

        assert_eq!(pprust::ty_to_string(&ok_ty), "Foo");
    }
}
<commit_msg>get u64 handling<commit_after>extern crate syntex;
extern crate syntex_syntax;
extern crate syntex_pos;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;

mod parse_utils;
mod jni;
mod core;
mod rust_generator;
mod java_generator;

use std::collections::HashMap;
use std::env;
use std::sync::atomic::AtomicBool;
use std::sync::Mutex;

use syntex::Registry;
use syntex_syntax::ext::base::{ExtCtxt, MacResult, DummyResult};
use syntex_syntax::parse::{token, parser};
use syntex_syntax::tokenstream::TokenTree;
use syntex_syntax::codemap::Span;
use syntex_syntax::{parse, ast};
use syntex_syntax::ptr::P;
use syntex_syntax::print::pprust;

use parse_utils::*;
use core::*;
use rust_generator::{generate_rust_code, RUST_OBJECT_TO_JOBJECT};
use java_generator::generate_java_code;

lazy_static! {
    pub static ref COMMON_CODE_GENERATED: AtomicBool = AtomicBool::new(false);
    static ref TYPE_HANDLERS: Mutex<Vec<TypeHandler>> = {
        Mutex::new(vec![
            TypeHandler{rust_type_name: "bool".into(), jni_type_name: "jboolean", java_type_name: "boolean".into(),
                        from_jni_converter: Some(FromForeignArgConverter("let {arg_name} = {arg_name} != 0;".into())),
                        to_jni_converter: Some(ToForeignRetConverter("let ret = if ret { 1 as jboolean } else { 0 as jboolean };".into()))},
            TypeHandler{rust_type_name: "i32".into(), jni_type_name: "jint", java_type_name: "int".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "u64".into(), jni_type_name: "jlong", java_type_name: "long".into(),
                        from_jni_converter: None,
                        to_jni_converter: Some(ToForeignRetConverter(r#"
  let ret: i64 = if (::std::i64::MAX as u64) < ret {
                    error!("u64->jlong type overflow: {}", ret);
                    ::std::i64::MAX
                 } else { ret as i64 };
"#.into())),
            },
            TypeHandler{rust_type_name: "f32".into(), jni_type_name: "jfloat", java_type_name: "float".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "f64".into(), jni_type_name: "jdouble", java_type_name: "double".into(),
                        from_jni_converter: None, to_jni_converter: None},
            TypeHandler{rust_type_name: "&str".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: Some(FromForeignArgConverter(r#"
  let {arg_name} = JavaString::new(env, {arg_name});
  let {arg_name} = {arg_name}.to_str();
"#.into())),
                        to_jni_converter: Some(ToForeignRetConverter(r#"
  let ret = ::std::ffi::CString::new(ret).unwrap();
  let ret = unsafe { (**env).NewStringUTF.unwrap()(env, ret.as_ptr()) };
"#.into()))},
            TypeHandler{rust_type_name: "&Path".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: Some(FromForeignArgConverter(
                            r#"
  let {arg_name} = JavaString::new(env, {arg_name});
  let {arg_name} = Path::new({arg_name}.to_str());
"#.into()
                        )),
                        to_jni_converter: None},
            TypeHandler{rust_type_name: "String".into(), jni_type_name: "jstring", java_type_name: "String".into(),
                        from_jni_converter: None,
                        to_jni_converter: Some(ToForeignRetConverter(
                            r#"
  let ret = ret.into_bytes();
  let ret = unsafe { ::std::ffi::CString::from_vec_unchecked(ret) };
  let ret = unsafe { (**env).NewStringUTF.unwrap()(env, ret.as_ptr()) };
"#.into()
                        ))},
            ])
    };
}

pub fn register(registry: &mut Registry) {
    registry.add_macro("foreigner_class", expand_foreigner_class);
}

fn expand_foreigner_class<'cx>(cx: &'cx mut ExtCtxt,
                               _: Span,
                               tokens: &[TokenTree])
                               -> Box<MacResult + 'cx> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess, cx.cfg.clone(), tokens.to_vec());
    let class_ident = parser.parse_ident().unwrap();
    if class_ident.name.as_str() != "class" {
        debug!("class_indent {:?}", class_ident);
        cx.span_err(parser.span, "expect class here");
        return DummyResult::any(parser.span);
    }
    let class_name_indent = parser.parse_ident().unwrap();
    debug!("CLASS NAME {:?}", class_name_indent);
    parser.expect(&token::Token::OpenDelim(token::DelimToken::Brace)).unwrap();
    let mut methods = Vec::new();
    let mut rust_self_type = ast::Path{span: parser.span, global: false, segments: Vec::new()};
    let alias = ast::Ident::with_empty_ctxt(token::intern("alias"));
    let mut constructor_ret_type: Option<ast::Ty> = None;
    let mut this_type_for_method: Option<ast::Ty> = None;

    loop {
        if parser.eat(&token::Token::CloseDelim(token::DelimToken::Brace)) {
            break;
        }
        let func_type_name = parser.parse_ident().unwrap();
        debug!("func_type {:?}", func_type_name);
        if func_type_name.name.as_str() == "self_type" {
            rust_self_type = parser.parse_path(parser::PathStyle::Type).expect("Can not parse self_type");
            debug!("self_type: {:?}", rust_self_type);
            parser.expect(&token::Token::Semi).unwrap();
            continue;
        }
        let func_type = FuncVariant::from_str(&func_type_name.name.as_str());
        if func_type.is_none() {
            println!("unknown func type: {:?}", func_type_name);
            cx.span_err(parser.span, "expect 'constructor' or 'method' or 'static_method' here");
            return DummyResult::any(parser.span);
        }
        let func_type = func_type.unwrap();
        let func_name = parser.parse_path(parser::PathStyle::Mod).unwrap();
        debug!("func_name {:?}", func_name);

        let func_decl = match func_type {
            FuncVariant::Constructor | FuncVariant::StaticMethod => parser.parse_fn_decl(false).unwrap(),
            FuncVariant::Method => parse_fn_decl_with_self(&mut parser, |p| p.parse_arg()).unwrap(),
        };
        debug!("func_decl {:?}", func_decl);
        parser.expect(&token::Token::Semi).unwrap();
        let mut func_name_alias = None;
        if parser.eat_contextual_keyword(alias) {
            if func_type == FuncVariant::Constructor {
                cx.span_err(parser.span, "alias not supported for 'constructor'");
                return DummyResult::any(parser.span);
            }
            func_name_alias = Some(parser.parse_ident().unwrap());
            debug!("we have ALIAS `{:?}`", func_name_alias.unwrap());
            parser.expect(&token::Token::Semi).expect("no ; at the end of alias");
        }
        let (may_return_error, ret_type) = match &func_decl.output {
            &ast::FunctionRetTy::Default(_) => (false, None),
            &ast::FunctionRetTy::Ty(ref ret_type) => (is_type_std_result(ret_type), Some(ret_type.clone())),
        };
        if let FuncVariant::Constructor = func_type {
            let ret_type = ret_type.expect(&format!("{}: constructor should return value", class_name_indent));
            if let Some(ref constructor_ret_type) = constructor_ret_type {
                if pprust::ty_to_string(constructor_ret_type) != pprust::ty_to_string(&*ret_type) {
                    cx.span_err(parser.span, &format!("mismatched types of construtors: {:?} {:?}", constructor_ret_type, ret_type));
                    return DummyResult::any(parser.span);
                }
            } else {
                constructor_ret_type = Some(ret_type.unwrap());
                this_type_for_method = Some(unpack_if_type_is_result(constructor_ret_type.as_ref().unwrap()));
            }
        }
        methods.push(ForeignerMethod{
            func_type: func_type, path: func_name, in_out_type: func_decl,
            name_alias: func_name_alias.map(|v| v.name.as_str()),
            may_return_error: may_return_error,
        });
    }
    let mut type_handlers = TYPE_HANDLERS.lock().unwrap();

    let package_name = env::var("RUST_SWIG_JNI_JAVA_PACKAGE").unwrap();

    let class_info = ForeignerClassInfo {
        package_name: package_name,
        class_name: &class_name_indent.name.as_str(),
        methods: methods,
        self_rust_type: rust_self_type,
        constructor_ret_type: constructor_ret_type,
        this_type_for_method: this_type_for_method,
    };

    if class_info.this_type_for_method.is_some() {
        let mut class_th: TypeHandler = (&class_info).into();
        class_th.to_jni_converter = Some(ToForeignRetConverter(
            RUST_OBJECT_TO_JOBJECT.replace(
                "{full_class_name}", &class_info.full_java_class_name())));
        type_handlers.push(class_th);
    }

    let mut rust_java_types_map = HashMap::new();
    for it in type_handlers.iter() {
        rust_java_types_map.insert(it.rust_type_name.as_str(), &*it);
    }

    let java_output_dir = env::var("RUST_SWIG_JNI_JAVA_OUTPUT_DIR").unwrap();

    generate_java_code(&rust_java_types_map, &class_info, &java_output_dir);
    generate_rust_code(cx, &rust_java_types_map, &class_info)
}

fn is_type_std_result(ty: &ast::Ty) -> bool {
    match ty.node {
        ast::TyKind::Path(_/*self info*/, ref path) => {
            debug!("is_type_std_result_result: path: {:?}, ident {:?}", path.segments, path.segments[0].identifier.name.as_str());
            path.segments.first().map(|v| v.identifier.name.as_str() == "Result").unwrap_or(false)
        }
        _ => false,
    }
}

fn unpack_if_type_is_result(ty: &ast::Ty) -> ast::Ty {
    match ty.node {
        ast::TyKind::Path(_/*self info*/, ref path) => {
            debug!("unpack_if_type_is_result: path: {:?}, ident {:?}", path.segments, path.segments[0].identifier.name.as_str());
            path.segments.first().map(|v|
                                      if v.identifier.name.as_str() == "Result" {
                                          match v.parameters {
                                              ast::PathParameters::AngleBracketed(ref params) => {
                                                  params.types.first().map(|v: &P<ast::Ty>| {
                                                      debug!("unpack_if_type_is_result: result param {:?}", *v);
                                                      (**v).clone()
                                                  }).unwrap_or(ty.clone())
                                              }
                                              _ => ty.clone(),
                                          }
                                      } else {
                                          ty.clone()
                                      }).unwrap_or(ty.clone())
        }
        _ => ty.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syntex_syntax::{parse, ast};

    #[test]
    fn test_is_type_std_result() {
        let session = parse::ParseSess::new();
        let mut parser = parse::new_parser_from_source_str(
            &session, ast::CrateConfig::new(), "test".into(),
            "Result<Foo, String>".into());
        let ty = parser.parse_ty().unwrap();
        assert!(is_type_std_result(&*ty));
    }

    #[test]
    fn test_unpack_if_type_is_result() {
        let session = parse::ParseSess::new();
        let mut parser = parse::new_parser_from_source_str(
            &session, ast::CrateConfig::new(), "test".into(),
            "Result<Foo, String>".into());
        let ty = parser.parse_ty().unwrap();
        assert!(is_type_std_result(&*ty));

        let ok_ty = unpack_if_type_is_result(&*ty);
        assert!(!is_type_std_result(&ok_ty));

        assert_eq!(pprust::ty_to_string(&ok_ty), "Foo");
    }
}
<|endoftext|>


<commit_before>use raft::{SyncClient, ClientClusterInfo, RaftMsg, RaftStateMachine, LogEntry, ClientQryResponse};
use raft::state_machine::OpType;
use std::collections::{HashMap, BTreeMap, HashSet};
use std::iter::FromIterator;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use bifrost_plugins::hash_str;
use rand;

const ORDERING: Ordering = Ordering::Relaxed;

struct QryMeta {
    last_log_id: u64,
    last_log_term: u64,
    pos: AtomicU64
}

struct Members {
    clients: BTreeMap<u64, Mutex<SyncClient>>,
    id_map: HashMap<u64, String>,
}

pub struct RaftClient {
    qry_meta: QryMeta,
    members: Members,
    term: u64, // last member list update term
}

impl RaftClient {
    pub fn new(servers: Vec<String>) -> Option<RaftClient> {
        let mut client = RaftClient {
            qry_meta: QryMeta {
                last_log_id: 0,
                last_log_term: 0,
                pos: AtomicU64::new(rand::random::<u64>())
            },
            members: Members {
                clients: BTreeMap::new(),
                id_map: HashMap::new()
            },
            term: 0,
        };
        match client.update_info(&HashSet::from_iter(servers)) {
            Ok(_) => Some(client),
            Err(_) => None
        }
    }

    fn update_info(&mut self, addrs: &HashSet<String>) -> Result<(), ()> {
        let info: ClientClusterInfo;
        let mut servers = None;
        for server_addr in addrs {
            let id = hash_str(server_addr.clone());
            let mut client = self.members.clients.entry(id).or_insert_with(|| {
                Mutex::new(SyncClient::new(server_addr))
            }).lock().unwrap();
            if let Some(Ok(info)) = client.c_server_cluster_info() {
                servers = Some(info);
                break;
            }
        }
        match servers {
            Some(servers) => {
                let remote_members = servers.members;
                let mut remote_ids = HashSet::with_capacity(remote_members.len());
                self.members.id_map.clear();
                for (id, addr) in remote_members {
                    self.members.id_map.insert(id, addr);
                    remote_ids.insert(id);
                }
                let mut connected_ids = HashSet::with_capacity(self.members.clients.len());
                for id in self.members.clients.keys() {connected_ids.insert(*id);}
                let ids_to_remove = connected_ids.difference(&remote_ids);
                for id in ids_to_remove {self.members.clients.remove(id);}
                for id in remote_ids.difference(&connected_ids) {
                    let addr = self.members.id_map.get(id).unwrap();
                    self.members.clients.entry(id.clone()).or_insert_with(|| {
                        Mutex::new(SyncClient::new(addr))
                    }).lock().unwrap();
                }
                Ok(())
            },
            None => Err(()),
        }
    }

    pub fn execute<R>(&mut self, sm: RaftStateMachine, msg: &RaftMsg<R>) -> Option<R> {
        let (fn_id, op, req_data) = msg.encode();
        let sm_id = sm.id;
        let response = match op {
            OpType::QUERY => {

            },
            OpType::COMMAND => {

            },
        };
        None
    }

    fn query(&mut self, sm_id: u64, fn_id: u64, data: &Vec<u8>) -> Option<Vec<u8>> {
        let pos = self.qry_meta.pos.fetch_add(1, ORDERING);
        let clients_count = self.members.clients.len();
        let res = {
            let mut client = self.members.clients.values()
                .nth(pos as usize % clients_count)
                .unwrap()
                .lock().unwrap();
            client.c_query(LogEntry {
                id: self.qry_meta.last_log_id,
                term: self.qry_meta.last_log_term,
                sm_id: sm_id,
                fn_id: fn_id,
                data: data.clone()
            })
        };
        match res {
            Some(Ok(res)) => {
                match res {
                    ClientQryResponse::LeftBehind => {
                        self.query(sm_id, fn_id, data)
                    },
                    ClientQryResponse::Success{
                        data: data,
                        last_log_term: last_log_term,
                        last_log_id: last_log_id
                    } => {
                        self.qry_meta.last_log_id = last_log_id;
                        self.qry_meta.last_log_term = last_log_term;
                        Some(data)
                    },
                }
            },
            _ => None
        }
    }
}<commit_msg>lock free for parent client<commit_after>use raft::{SyncClient, ClientClusterInfo, RaftMsg, RaftStateMachine, LogEntry, ClientQryResponse};
use raft::state_machine::OpType;
use std::collections::{HashMap, BTreeMap, HashSet};
use std::iter::FromIterator;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::cell::RefCell;
use bifrost_plugins::hash_str;
use rand;

const ORDERING: Ordering = Ordering::Relaxed;

struct QryMeta {
    last_log_id: AtomicU64,
    last_log_term: AtomicU64,
    pos: AtomicU64
}

struct Members {
    clients: BTreeMap<u64, Mutex<SyncClient>>,
    id_map: HashMap<u64, String>,
}

pub struct RaftClient {
    qry_meta: QryMeta,
    members: Members,
    term: u64, // last member list update term
}

impl RaftClient {
    pub fn new(servers: Vec<String>) -> Option<RaftClient> {
        let mut client = RaftClient {
            qry_meta: QryMeta {
                last_log_id: AtomicU64::new(0),
                last_log_term: AtomicU64::new(0),
                pos: AtomicU64::new(rand::random::<u64>())
            },
            members: Members {
                clients: BTreeMap::new(),
                id_map: HashMap::new()
            },
            term: 0,
        };
        match client.update_info(&HashSet::from_iter(servers)) {
            Ok(_) => Some(client),
            Err(_) => None
        }
    }

    fn update_info(&mut self, addrs: &HashSet<String>) -> Result<(), ()> {
        let info: ClientClusterInfo;
        let mut servers = None;
        for server_addr in addrs {
            let id = hash_str(server_addr.clone());
            let mut client = self.members.clients.entry(id).or_insert_with(|| {
                Mutex::new(SyncClient::new(server_addr))
            });
            if let Some(Ok(info)) = client.lock().unwrap().c_server_cluster_info() {
                servers = Some(info);
                break;
            }
        }
        match servers {
            Some(servers) => {
                let remote_members = servers.members;
                let mut remote_ids = HashSet::with_capacity(remote_members.len());
                self.members.id_map.clear();
                for (id, addr) in remote_members {
                    self.members.id_map.insert(id, addr);
                    remote_ids.insert(id);
                }
                let mut connected_ids = HashSet::with_capacity(self.members.clients.len());
                for id in self.members.clients.keys() {connected_ids.insert(*id);}
                let ids_to_remove = connected_ids.difference(&remote_ids);
                for id in ids_to_remove {self.members.clients.remove(id);}
                for id in remote_ids.difference(&connected_ids) {
                    let addr = self.members.id_map.get(id).unwrap();
                    self.members.clients.entry(id.clone()).or_insert_with(|| {
                        Mutex::new(SyncClient::new(addr))
                    });
                }
                Ok(())
            },
            None => Err(()),
        }
    }

    pub fn execute<R>(&self, sm: RaftStateMachine, msg: &RaftMsg<R>) -> Option<R> {
        let (fn_id, op, req_data) = msg.encode();
        let sm_id = sm.id;
        let response = match op {
            OpType::QUERY => {

            },
            OpType::COMMAND => {

            },
        };
        None
    }

    fn query(&self, sm_id: u64, fn_id: u64, data: &Vec<u8>) -> Option<Vec<u8>> {
        let pos = self.qry_meta.pos.fetch_add(1, ORDERING);
        let clients_count = self.members.clients.len();
        let res = {
            let mut client = self.members.clients.values()
                .nth(pos as usize % clients_count)
                .unwrap().lock().unwrap();
            client.c_query(LogEntry {
                id: self.qry_meta.last_log_id.load(ORDERING),
                term: self.qry_meta.last_log_term.load(ORDERING),
                sm_id: sm_id,
                fn_id: fn_id,
                data: data.clone()
            })
        };
        match res {
            Some(Ok(res)) => {
                match res {
                    ClientQryResponse::LeftBehind => {
                        self.query(sm_id, fn_id, data)
                    },
                    ClientQryResponse::Success{
                        data: data,
                        last_log_term: last_log_term,
                        last_log_id: last_log_id
                    } => {
                        self.qry_meta.last_log_id.store(last_log_id, ORDERING);
                        self.qry_meta.last_log_term.store(last_log_term, ORDERING);
                        Some(data)
                    },
                }
            },
            _ => None
        }
    }
}<|endoftext|>


<commit_before>use clap::arg_enum;
use config::config::NodeConfig;
use logger::prelude::*;
use regex::Regex;
use std::{fs, net::IpAddr, path::PathBuf, str::FromStr};
use structopt::StructOpt;

arg_enum! {
    #[derive(Debug)]
    pub enum Executable {
        TestLiveness,
        MeasureThroughput,
    }
}

/// CLI options for RuBen.
#[derive(Debug, StructOpt)]
#[structopt(
    name = "RuBen",
    author = "Libra",
    about = "RuBen (Ru)ns The Libra (Ben)chmarker For You."
)]
pub struct Opt {
    /// Validator address list seperated by whitespace: `ip_address:port ip_address:port ...`.
    /// It is requried unless (and hence conflict with) swarm_config_dir is present.
    #[structopt(
        short = "a",
        long = "validator_addresses",
        conflicts_with = "swarm_config_dir",
        requires = "debug_address",
        required_unless = "swarm_config_dir"
    )]
    pub validator_addresses: Vec<String>,
    /// Debug interface address in the form of ip_address:port.
    /// It is requried unless (and hence conflict with) swarm_config_dir is present.
    #[structopt(
        short = "d",
        long = "debug_address",
        conflicts_with = "swarm_config_dir",
        requires = "validator_addresses",
        required_unless = "swarm_config_dir"
    )]
    pub debug_address: Option<String>,
    /// libra_swarm's config file directory, which holds libra_node's config .toml file(s).
    /// It is requried unless (and hence conflict with)
    /// validator_addresses and debug_address are both present.
    #[structopt(
        short = "s",
        long = "swarm_config_dir",
        raw(conflicts_with_all = r#"&["validator_addresses", "debug_address"]"#),
        raw(required_unless_all = r#"&["validator_addresses", "debug_address"]"#)
    )]
    pub swarm_config_dir: Option<String>,
    /// Valid faucet key file path.
    #[structopt(short = "f", long = "faucet_key_file_path", required = true)]
    pub faucet_key_file_path: String,
    /// Number of accounts to create in Libra.
    #[structopt(short = "n", long = "num_accounts", default_value = "32")]
    pub num_accounts: u64,
    /// Free lunch amount to accounts.
    #[structopt(short = "l", long = "free_lunch", default_value = "1000000")]
    pub free_lunch: u64,
    /// Number of AC clients.
    /// If not specified or equals 0, it will be set to validator_addresses.len().
    #[structopt(short = "c", long = "num_clients", default_value = "0")]
    pub num_clients: usize,
    /// Number of repetition to attempt, in one epoch, to increase overal number of sent TXNs.
    #[structopt(short = "r", long = "num_rounds", default_value = "1")]
    pub num_rounds: u64,
    /// Number of epochs to measure the TXN throughput, each time with newly created Benchmarker.
    #[structopt(short = "e", long = "num_epochs", default_value = "10")]
    pub num_epochs: u64,
    /// Supported application of Benchmarker: `TestLiveness` or `MeasureThroughput`.
    #[structopt(
        short = "x",
        long = "executable",
        raw(possible_values = "&Executable::variants()"),
        case_insensitive = true,
        default_value = "MeasureThroughput"
    )]
    pub executable: Executable,
}

/// Helper that checks if address is valid, and converts unspecified address to localhost.
fn parse_socket_address(address: &str, port: u16) -> String {
    let ip_address = IpAddr::from_str(address).expect("invalid network address");
    if ip_address.is_unspecified() {
        match ip_address.is_ipv4() {
            true => format!("127.0.0.1:{}", port),
            false => format!("::1:{}", port),
        }
    } else {
        format!("{}:{}", address, port)
    }
}

/// Scan *.node.config.toml files under config_dir_name, parse them as node config
/// and return libra_swarm's configuration info as a tuple:
/// (addresses for all nodes, debug_address)
fn parse_swarm_config_from_dir(config_dir_name: &str) -> (Vec<String>, String) {
    let mut validator_addresses: Vec<String> = Vec::new();
    let mut debug_address = String::new();
    let re = Regex::new(r"[[:alnum:]]{64}\.node\.config\.toml").expect("failed to build regex");
    let config_dir = PathBuf::from(config_dir_name);
    if config_dir.is_dir() {
        for entry in fs::read_dir(config_dir).expect("invalid config directory") {
            let path = entry.expect("invalid path under config directory").path();
            if path.is_file() {
                let filename = path
                    .file_name()
                    .expect("failed to convert filename to string")
                    .to_str()
                    .expect("failed to convert filename to string");
                if re.is_match(filename) {
                    debug!("Parsing node config file {:?}.", filename);
                    let config_string = fs::read_to_string(&path)
                        .unwrap_or_else(|_| panic!("failed to load config file {:?}", filename));
                    let config = NodeConfig::parse(&config_string).unwrap_or_else(|_| {
                        panic!("failed to parse NodeConfig from {:?}", filename)
                    });
                    if debug_address.is_empty() {
                        debug_address = parse_socket_address(
                            &config.debug_interface.address,
                            config.debug_interface.admission_control_node_debug_port,
                        );
                    }
                    let address = parse_socket_address(
                        &config.admission_control.address,
                        config.admission_control.admission_control_service_port,
                    );
                    validator_addresses.push(address);
                }
            }
        }
    }
    info!(
        "Parsed arguments from {:?}: {:?}, {:?}",
        config_dir_name, validator_addresses, debug_address,
    );
    (validator_addresses, debug_address)
}

impl Opt {
    pub fn new_from_args() -> Self {
        let mut args = Opt::from_args();
        // Override validator_addresses and debug_address if swarm_config_dir is provided.
        if let Some(swarm_config_dir) = args.swarm_config_dir.as_ref() {
            let (validator_addresses, debug_address) =
                parse_swarm_config_from_dir(swarm_config_dir);
            args.validator_addresses = validator_addresses;
            args.debug_address = Some(debug_address);
        }
        if args.num_clients == 0 {
            args.num_clients = args.validator_addresses.len();
        }
        args
    }
}
<commit_msg>[benchmark] Fix two bugs in RuBen that are releated to arguments<commit_after>use clap::arg_enum;
use config::config::NodeConfig;
use failure::prelude::*;
use logger::prelude::*;
use regex::Regex;
use std::{fs, net::IpAddr, path::PathBuf, str::FromStr};
use structopt::StructOpt;

arg_enum! {
    #[derive(Debug)]
    pub enum Executable {
        TestLiveness,
        MeasureThroughput,
    }
}

/// CLI options for RuBen.
#[derive(Debug, StructOpt)]
#[structopt(
    name = "RuBen",
    author = "Libra",
    about = "RuBen (Ru)ns The Libra (Ben)chmarker For You."
)]
pub struct Opt {
    /// Validator address list seperated by whitespace: `ip_address:port ip_address:port ...`.
    /// It is requried unless (and hence conflict with) swarm_config_dir is present.
    #[structopt(
        short = "a",
        long = "validator_addresses",
        conflicts_with = "swarm_config_dir",
        requires = "debug_address",
        required_unless = "swarm_config_dir"
    )]
    pub validator_addresses: Vec<String>,
    /// Debug interface address in the form of ip_address:port.
    /// It is requried unless (and hence conflict with) swarm_config_dir is present.
    #[structopt(
        short = "d",
        long = "debug_address",
        conflicts_with = "swarm_config_dir",
        requires = "validator_addresses",
        required_unless = "swarm_config_dir"
    )]
    pub debug_address: Option<String>,
    /// libra_swarm's config file directory, which holds libra_node's config .toml file(s).
    /// It is requried unless (and hence conflict with)
    /// validator_addresses and debug_address are both present.
    #[structopt(
        short = "s",
        long = "swarm_config_dir",
        raw(conflicts_with_all = r#"&["validator_addresses", "debug_address"]"#),
        raw(required_unless_all = r#"&["validator_addresses", "debug_address"]"#)
    )]
    pub swarm_config_dir: Option<String>,
    /// Valid faucet key file path.
    #[structopt(short = "f", long = "faucet_key_file_path", required = true)]
    pub faucet_key_file_path: String,
    /// Number of accounts to create in Libra.
    #[structopt(short = "n", long = "num_accounts", default_value = "32")]
    pub num_accounts: u64,
    /// Free lunch amount to accounts.
    #[structopt(short = "l", long = "free_lunch", default_value = "1000000")]
    pub free_lunch: u64,
    /// Number of AC clients.
    /// If not specified or equals 0, it will be set to validator_addresses.len().
    #[structopt(short = "c", long = "num_clients", default_value = "0")]
    pub num_clients: usize,
    /// Number of repetition to attempt, in one epoch, to increase overal number of sent TXNs.
    #[structopt(short = "r", long = "num_rounds", default_value = "1")]
    pub num_rounds: u64,
    /// Number of epochs to measure the TXN throughput, each time with newly created Benchmarker.
    #[structopt(short = "e", long = "num_epochs", default_value = "10")]
    pub num_epochs: u64,
    /// Supported application of Benchmarker: `TestLiveness` or `MeasureThroughput`.
    #[structopt(
        short = "x",
        long = "executable",
        raw(possible_values = "&Executable::variants()"),
        case_insensitive = true,
        default_value = "MeasureThroughput"
    )]
    pub executable: Executable,
}

/// Helper that checks if address is valid, and converts unspecified address to localhost.
/// If parsing as numeric network address fails, treat as valid domain name.
fn parse_socket_address(address: &str, port: u16) -> String {
    match IpAddr::from_str(address) {
        Ok(ip_address) => {
            if ip_address.is_unspecified() {
                format!("localhost:{}", port)
            } else {
                format!("{}:{}", address, port)
            }
        }
        Err(_) => format!("{}:{}", address, port),
    }
}

/// Scan *.node.config.toml files under config_dir_name, parse them as node config
/// and return libra_swarm's configuration info as a tuple:
/// (addresses for all nodes, debug_address)
fn parse_swarm_config_from_dir(config_dir_name: &str) -> Result<(Vec<String>, String)> {
    let mut validator_addresses: Vec<String> = Vec::new();
    let mut debug_address = None;
    let re = Regex::new(r"[[:alnum:]]{64}\.node\.config\.toml").expect("failed to build regex");
    let config_dir = PathBuf::from(config_dir_name);
    if config_dir.is_dir() {
        for entry in fs::read_dir(config_dir).expect("invalid config directory") {
            let path = entry.expect("invalid path under config directory").path();
            if path.is_file() {
                let filename = path
                    .file_name()
                    .expect("failed to convert filename to string")
                    .to_str()
                    .expect("failed to convert filename to string");
                if re.is_match(filename) {
                    debug!("Parsing node config file {:?}.", filename);
                    let config_string = fs::read_to_string(&path)
                        .unwrap_or_else(|_| panic!("failed to load config file {:?}", filename));
                    let config = NodeConfig::parse(&config_string).unwrap_or_else(|_| {
                        panic!("failed to parse NodeConfig from {:?}", filename)
                    });
                    if debug_address.is_none() {
                        debug_address = Some(parse_socket_address(
                            &config.debug_interface.address,
                            config.debug_interface.admission_control_node_debug_port,
                        ));
                    }
                    let address = parse_socket_address(
                        &config.admission_control.address,
                        config.admission_control.admission_control_service_port,
                    );
                    validator_addresses.push(address);
                }
            }
        }
    }
    if validator_addresses.is_empty() {
        bail!(
            "unable to parse validator_addresses from {}",
            config_dir_name
        )
    }
    Ok((
        validator_addresses,
        debug_address
            .ok_or_else(|| format_err!("unable to parse debug_address from {}", config_dir_name))?,
    ))
}

impl Opt {
    pub fn new_from_args() -> Self {
        let mut args = Opt::from_args();
        // Override validator_addresses and debug_address if swarm_config_dir is provided.
        if let Some(swarm_config_dir) = args.swarm_config_dir.as_ref() {
            let (validator_addresses, debug_address) =
                parse_swarm_config_from_dir(swarm_config_dir).expect("invalid arguments");
            args.validator_addresses = validator_addresses;
            args.debug_address = Some(debug_address);
        }
        if args.num_clients == 0 {
            args.num_clients = args.validator_addresses.len();
        }
        args
    }
}

#[cfg(test)]
mod tests {
    use crate::ruben_opt::{parse_socket_address, parse_swarm_config_from_dir};

    #[test]
    fn test_parse_socket_address() {
        assert_eq!(
            parse_socket_address("216.10.234.56", 12345),
            "216.10.234.56:12345"
        );
        assert_eq!(parse_socket_address("0.0.0.0", 12345), "localhost:12345");
        assert_eq!(parse_socket_address("::", 12345), "localhost:12345");
        assert_eq!(
            parse_socket_address("face:booc::0", 12345),
            "face:booc::0:12345"
        );
        assert_eq!(
            parse_socket_address("2401:dbff:121f:a2f1:face:d:6c:0", 12345),
            "2401:dbff:121f:a2f1:face:d:6c:0:12345"
        );
        assert_eq!(parse_socket_address("localhost", 12345), "localhost:12345");
        assert_eq!(
            parse_socket_address("www.facebook.com", 12345),
            "www.facebook.com:12345"
        );
    }

    #[test]
    fn test_parse_swarm_config_from_invalid_dir() {
        // Directory doesn't exist at all.
        let non_exist_dir_name = String::from("NonExistDir");
        let mut result = parse_swarm_config_from_dir(&non_exist_dir_name);
        assert_eq!(result.is_err(), true);

        // Directory exists but config file does not.
        let path = std::env::current_dir().expect("unable to get current dir");
        if let Some(dir_without_config) = path.to_str() {
            result = parse_swarm_config_from_dir(&dir_without_config);
            assert_eq!(result.is_err(), true);
        }
    }
}
<|endoftext|>


<commit_before>//! Implementation of the `#[simd_test]` macro
//!
//! This macro expands to a `#[test]` function which tests the local machine
//! for the appropriate cfg before calling the inner test function.

extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate quote;

use proc_macro2::{Ident, Literal, Span, TokenStream, TokenTree};
use quote::ToTokens;
use std::env;

fn string(s: &str) -> TokenTree {
    Literal::string(s).into()
}

#[proc_macro_attribute]
pub fn simd_test(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let tokens = TokenStream::from(attr).into_iter().collect::<Vec<_>>();
    if tokens.len() != 3 {
        panic!("expected #[simd_test(enable = \"feature\")]");
    }
    match &tokens[0] {
        TokenTree::Ident(tt) if *tt == "enable" => {}
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    }
    match &tokens[1] {
        TokenTree::Punct(tt) if tt.as_char() == '=' => {}
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    }
    let enable_feature = match &tokens[2] {
        TokenTree::Literal(tt) => tt.to_string(),
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    };
    let enable_feature = enable_feature.trim_start_matches('"').trim_end_matches('"');
    let target_features: Vec<String> = enable_feature
        .replace('+', "")
        .split(',')
        .map(String::from)
        .collect();

    let mmx = target_features.iter().any(|s| s.starts_with("mmx"));

    let enable_feature = string(enable_feature);
    let item = TokenStream::from(item);
    let name = find_name(item.clone());

    let name: TokenStream = name
        .to_string()
        .parse()
        .unwrap_or_else(|_| panic!("failed to parse name: {}", name.to_string()));

    let target = env::var("TARGET").expect("TARGET environment variable should be set for rustc");
    let mut force_test = false;
    let macro_test = match target
        .split('-')
        .next()
        .unwrap_or_else(|| panic!("target triple contained no \"-\": {}", target))
    {
        "i686" | "x86_64" | "i586" => "is_x86_feature_detected",
        "arm" | "armv7" => "is_arm_feature_detected",
        "aarch64" => "is_aarch64_feature_detected",
        "powerpc" | "powerpcle" => "is_powerpc_feature_detected",
        "powerpc64" | "powerpc64le" => "is_powerpc64_feature_detected",
        "mips" | "mipsel" | "mipsisa32r6" | "mipsisa32r6el" => {
            // FIXME:
            // On MIPS CI run-time feature detection always returns false due
            // to this qemu bug: https://bugs.launchpad.net/qemu/+bug/1754372
            //
            // This is a workaround to force the MIPS tests to always run on
            // CI.
            force_test = true;
            "is_mips_feature_detected"
        }
        "mips64" | "mips64el" | "mipsisa64r6" | "mipsisa64r6el" => {
            // FIXME: see above
            force_test = true;
            "is_mips64_feature_detected"
        }
        t => panic!("unknown target: {}", t),
    };
    let macro_test = Ident::new(macro_test, Span::call_site());

    let mut cfg_target_features = TokenStream::new();
    for feature in target_features {
        let q = quote_spanned! {
            proc_macro2::Span::call_site() =>
            #macro_test!(#feature) &&
        };
        q.to_tokens(&mut cfg_target_features);
    }
    let q = quote! { true };
    q.to_tokens(&mut cfg_target_features);

    let test_norun = std::env::var("STDSIMD_TEST_NORUN").is_ok();
    let maybe_ignore = if test_norun {
        quote! { #[ignore] }
    } else {
        TokenStream::new()
    };

    let emms = if mmx {
        // note: if the test requires MMX we need to clear the FPU
        // registers once the test finishes before interfacing with
        // other x87 code:
        quote! { unsafe { super::_mm_empty() }; }
    } else {
        TokenStream::new()
    };

    let ret: TokenStream = quote_spanned! {
        proc_macro2::Span::call_site() =>
        #[allow(non_snake_case)]
        #[test]
        #maybe_ignore
        fn #name() {
            if #force_test | (#cfg_target_features) {
                let v = unsafe { #name() };
                #emms
                return v;
            } else {
                ::stdsimd_test::assert_skip_test_ok(stringify!(#name));
            }

            #[target_feature(enable = #enable_feature)]
            #item
        }
    };
    ret.into()
}

fn find_name(item: TokenStream) -> Ident {
    let mut tokens = item.into_iter();
    while let Some(tok) = tokens.next() {
        if let TokenTree::Ident(word) = tok {
            if word == "fn" {
                break;
            }
        }
    }

    match tokens.next() {
        Some(TokenTree::Ident(word)) => word,
        _ => panic!("failed to find function name"),
    }
}
<commit_msg>Better error message when TARGET is missing<commit_after>//! Implementation of the `#[simd_test]` macro
//!
//! This macro expands to a `#[test]` function which tests the local machine
//! for the appropriate cfg before calling the inner test function.

extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate quote;

use proc_macro2::{Ident, Literal, Span, TokenStream, TokenTree};
use quote::ToTokens;
use std::env;

fn string(s: &str) -> TokenTree {
    Literal::string(s).into()
}

#[proc_macro_attribute]
pub fn simd_test(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let tokens = TokenStream::from(attr).into_iter().collect::<Vec<_>>();
    if tokens.len() != 3 {
        panic!("expected #[simd_test(enable = \"feature\")]");
    }
    match &tokens[0] {
        TokenTree::Ident(tt) if *tt == "enable" => {}
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    }
    match &tokens[1] {
        TokenTree::Punct(tt) if tt.as_char() == '=' => {}
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    }
    let enable_feature = match &tokens[2] {
        TokenTree::Literal(tt) => tt.to_string(),
        _ => panic!("expected #[simd_test(enable = \"feature\")]"),
    };
    let enable_feature = enable_feature.trim_start_matches('"').trim_end_matches('"');
    let target_features: Vec<String> = enable_feature
        .replace('+', "")
        .split(',')
        .map(String::from)
        .collect();

    let mmx = target_features.iter().any(|s| s.starts_with("mmx"));

    let enable_feature = string(enable_feature);
    let item = TokenStream::from(item);
    let name = find_name(item.clone());

    let name: TokenStream = name
        .to_string()
        .parse()
        .unwrap_or_else(|_| panic!("failed to parse name: {}", name.to_string()));

    let target = env::var("TARGET").expect(
        "TARGET environment variable should be set for rustc (e.g. TARGET=x86_64-apple-darwin cargo test)"
    );
    let mut force_test = false;
    let macro_test = match target
        .split('-')
        .next()
        .unwrap_or_else(|| panic!("target triple contained no \"-\": {}", target))
    {
        "i686" | "x86_64" | "i586" => "is_x86_feature_detected",
        "arm" | "armv7" => "is_arm_feature_detected",
        "aarch64" => "is_aarch64_feature_detected",
        "powerpc" | "powerpcle" => "is_powerpc_feature_detected",
        "powerpc64" | "powerpc64le" => "is_powerpc64_feature_detected",
        "mips" | "mipsel" | "mipsisa32r6" | "mipsisa32r6el" => {
            // FIXME:
            // On MIPS CI run-time feature detection always returns false due
            // to this qemu bug: https://bugs.launchpad.net/qemu/+bug/1754372
            //
            // This is a workaround to force the MIPS tests to always run on
            // CI.
            force_test = true;
            "is_mips_feature_detected"
        }
        "mips64" | "mips64el" | "mipsisa64r6" | "mipsisa64r6el" => {
            // FIXME: see above
            force_test = true;
            "is_mips64_feature_detected"
        }
        t => panic!("unknown target: {}", t),
    };
    let macro_test = Ident::new(macro_test, Span::call_site());

    let mut cfg_target_features = TokenStream::new();
    for feature in target_features {
        let q = quote_spanned! {
            proc_macro2::Span::call_site() =>
            #macro_test!(#feature) &&
        };
        q.to_tokens(&mut cfg_target_features);
    }
    let q = quote! { true };
    q.to_tokens(&mut cfg_target_features);

    let test_norun = std::env::var("STDSIMD_TEST_NORUN").is_ok();
    let maybe_ignore = if test_norun {
        quote! { #[ignore] }
    } else {
        TokenStream::new()
    };

    let emms = if mmx {
        // note: if the test requires MMX we need to clear the FPU
        // registers once the test finishes before interfacing with
        // other x87 code:
        quote! { unsafe { super::_mm_empty() }; }
    } else {
        TokenStream::new()
    };

    let ret: TokenStream = quote_spanned! {
        proc_macro2::Span::call_site() =>
        #[allow(non_snake_case)]
        #[test]
        #maybe_ignore
        fn #name() {
            if #force_test | (#cfg_target_features) {
                let v = unsafe { #name() };
                #emms
                return v;
            } else {
                ::stdsimd_test::assert_skip_test_ok(stringify!(#name));
            }

            #[target_feature(enable = #enable_feature)]
            #item
        }
    };
    ret.into()
}

fn find_name(item: TokenStream) -> Ident {
    let mut tokens = item.into_iter();
    while let Some(tok) = tokens.next() {
        if let TokenTree::Ident(word) = tok {
            if word == "fn" {
                break;
            }
        }
    }

    match tokens.next() {
        Some(TokenTree::Ident(word)) => word,
        _ => panic!("failed to find function name"),
    }
}
<|endoftext|>


<commit_before>#![allow(unused_parens)] // FIXME: Remove this attribute once false positive is resolved.

use super::{PgConnection, PgTypeMetadata};
use crate::prelude::*;

use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::HashMap;

/// Determines the OID of types at runtime
#[allow(missing_debug_implementations)]
#[repr(transparent)]
pub struct PgMetadataLookup {
    conn: PgConnection,
}

impl PgMetadataLookup {
    #[allow(clippy::new_ret_no_self)]
    pub(crate) fn new(conn: &PgConnection) -> &Self {
        unsafe { &*(conn as *const PgConnection as *const PgMetadataLookup) }
    }

    /// Determine the type metadata for the given `type_name`
    ///
    /// This function should only be used for user defined types, or types which
    /// come from an extension. This function may perform a SQL query to look
    /// up the type. For built-in types, a static OID should be preferred.
    pub fn lookup_type(&self, type_name: &str, schema: Option<&str>) -> PgTypeMetadata {
        let metadata_cache = self.conn.get_metadata_cache();
        let cache_key = PgMetadataCacheKey {
            schema: schema.map(Cow::Borrowed),
            type_name: Cow::Borrowed(type_name),
        };

        metadata_cache.lookup_type(&cache_key).unwrap_or_else(|| {
            let r = lookup_type(&cache_key, &self.conn);

            match r {
                Ok(type_metadata) => {
                    metadata_cache.store_type(cache_key, type_metadata);
                    type_metadata
                }
                Err(_e) => {
                    eprintln!("Failed to find oid for type `{}`", type_name);
                    // We return 0 as type oid here
                    // because most of the time postgres
                    // can figure out the type on its own
                    PgTypeMetadata {
                        oid: 0,
                        array_oid: 0,
                    }
                }
            }
        })
    }
}

fn lookup_type(
    cache_key: &PgMetadataCacheKey<'_>,
    conn: &PgConnection,
) -> QueryResult<PgTypeMetadata> {
    let search_path: String;

    let search_schema = if let Some(schema) = cache_key.schema.as_ref() {
        vec![schema.as_ref()]
    } else {
        search_path = crate::dsl::sql("SHOW search_path").get_result::<String>(conn)?;

        search_path
            .split(',')
            // skip the `$user` entry for now
            .filter(|f| !f.starts_with("\"$"))
            .map(|s| s.trim())
            .collect()
    };

    let metadata = pg_type::table
        .inner_join(pg_namespace::table)
        .filter(pg_namespace::nspname.eq(crate::dsl::any(search_schema)))
        .filter(pg_type::typname.eq(&cache_key.type_name))
        .select((pg_type::oid, pg_type::typarray))
        .first(conn)?;

    Ok(metadata)
}

#[derive(Hash, PartialEq, Eq)]
struct PgMetadataCacheKey<'a> {
    schema: Option<Cow<'a, str>>,
    type_name: Cow<'a, str>,
}

impl<'a> PgMetadataCacheKey<'a> {
    fn as_owned(&self) -> PgMetadataCacheKey<'static> {
        PgMetadataCacheKey {
            schema: self
                .schema
                .as_ref()
                .map(|schema| Cow::Owned(schema.as_ref().to_owned())),
            type_name: Cow::Owned(self.type_name.as_ref().to_owned()),
        }
    }
}

/// Stores a cache for the OID of custom types
#[allow(missing_debug_implementations)]
pub struct PgMetadataCache {
    cache: RefCell<HashMap<PgMetadataCacheKey<'static>, PgTypeMetadata>>,
}

impl PgMetadataCache {
    pub(crate) fn new() -> Self {
        PgMetadataCache {
            cache: RefCell::new(HashMap::new()),
        }
    }

    fn lookup_type(&self, type_name: &PgMetadataCacheKey) -> Option<PgTypeMetadata> {
        Some(*self.cache.borrow().get(type_name)?)
    }

    fn store_type(&self, type_name: PgMetadataCacheKey, type_metadata: PgTypeMetadata) {
        self.cache
            .borrow_mut()
            .insert(type_name.as_owned(), type_metadata);
    }
}

table! {
    pg_type (oid) {
        oid -> Oid,
        typname -> Text,
        typarray -> Oid,
        typnamespace -> Oid,
    }
}

table! {
    pg_namespace (oid) {
        oid -> Oid,
        nspname -> Text,
    }
}

joinable!(pg_type -> pg_namespace(typnamespace));
allow_tables_to_appear_in_same_query!(pg_type, pg_namespace);
<commit_msg>Rename `PgMetadataCacheKey` to `into_owned` + change implementation slightly<commit_after>#![allow(unused_parens)] // FIXME: Remove this attribute once false positive is resolved.

use super::{PgConnection, PgTypeMetadata};
use crate::prelude::*;

use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::HashMap;

/// Determines the OID of types at runtime
#[allow(missing_debug_implementations)]
#[repr(transparent)]
pub struct PgMetadataLookup {
    conn: PgConnection,
}

impl PgMetadataLookup {
    #[allow(clippy::new_ret_no_self)]
    pub(crate) fn new(conn: &PgConnection) -> &Self {
        unsafe { &*(conn as *const PgConnection as *const PgMetadataLookup) }
    }

    /// Determine the type metadata for the given `type_name`
    ///
    /// This function should only be used for user defined types, or types which
    /// come from an extension. This function may perform a SQL query to look
    /// up the type. For built-in types, a static OID should be preferred.
    pub fn lookup_type(&self, type_name: &str, schema: Option<&str>) -> PgTypeMetadata {
        let metadata_cache = self.conn.get_metadata_cache();
        let cache_key = PgMetadataCacheKey {
            schema: schema.map(Cow::Borrowed),
            type_name: Cow::Borrowed(type_name),
        };

        metadata_cache.lookup_type(&cache_key).unwrap_or_else(|| {
            let r = lookup_type(&cache_key, &self.conn);

            match r {
                Ok(type_metadata) => {
                    metadata_cache.store_type(cache_key, type_metadata);
                    type_metadata
                }
                Err(_e) => {
                    eprintln!("Failed to find oid for type `{}`", type_name);
                    // We return 0 as type oid here
                    // because most of the time postgres
                    // can figure out the type on its own
                    PgTypeMetadata {
                        oid: 0,
                        array_oid: 0,
                    }
                }
            }
        })
    }
}

fn lookup_type(
    cache_key: &PgMetadataCacheKey<'_>,
    conn: &PgConnection,
) -> QueryResult<PgTypeMetadata> {
    let search_path: String;

    let search_schema = if let Some(schema) = cache_key.schema.as_ref() {
        vec![schema.as_ref()]
    } else {
        search_path = crate::dsl::sql("SHOW search_path").get_result::<String>(conn)?;

        search_path
            .split(',')
            // skip the `$user` entry for now
            .filter(|f| !f.starts_with("\"$"))
            .map(|s| s.trim())
            .collect()
    };

    let metadata = pg_type::table
        .inner_join(pg_namespace::table)
        .filter(pg_namespace::nspname.eq(crate::dsl::any(search_schema)))
        .filter(pg_type::typname.eq(&cache_key.type_name))
        .select((pg_type::oid, pg_type::typarray))
        .first(conn)?;

    Ok(metadata)
}

#[derive(Hash, PartialEq, Eq)]
struct PgMetadataCacheKey<'a> {
    schema: Option<Cow<'a, str>>,
    type_name: Cow<'a, str>,
}

impl<'a> PgMetadataCacheKey<'a> {
    fn into_owned(self) -> PgMetadataCacheKey<'static> {
        let PgMetadataCacheKey { schema, type_name } = self;
        PgMetadataCacheKey {
            schema: schema.map(|s| Cow::Owned(s.into_owned())),
            type_name: Cow::Owned(type_name.into_owned()),
        }
    }
}

/// Stores a cache for the OID of custom types
#[allow(missing_debug_implementations)]
pub struct PgMetadataCache {
    cache: RefCell<HashMap<PgMetadataCacheKey<'static>, PgTypeMetadata>>,
}

impl PgMetadataCache {
    pub(crate) fn new() -> Self {
        PgMetadataCache {
            cache: RefCell::new(HashMap::new()),
        }
    }

    fn lookup_type(&self, type_name: &PgMetadataCacheKey) -> Option<PgTypeMetadata> {
        Some(*self.cache.borrow().get(type_name)?)
    }

    fn store_type(&self, type_name: PgMetadataCacheKey, type_metadata: PgTypeMetadata) {
        self.cache
            .borrow_mut()
            .insert(type_name.into_owned(), type_metadata);
    }
}

table! {
    pg_type (oid) {
        oid -> Oid,
        typname -> Text,
        typarray -> Oid,
        typnamespace -> Oid,
    }
}

table! {
    pg_namespace (oid) {
        oid -> Oid,
        nspname -> Text,
    }
}

joinable!(pg_type -> pg_namespace(typnamespace));
allow_tables_to_appear_in_same_query!(pg_type, pg_namespace);
<|endoftext|>


<commit_before>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ptr;
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that both `hHeap` is allowed to be any value, and `lpMem` is allowed to be null,
    // both of which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocates a block of optionally zeroed memory for a given `layout`.
// Returns a pointer satisfying the guarantees of `System` about allocated pointers.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = unsafe { GetProcessHeap() };
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = false;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = true;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(GetProcessHeap(), 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            let heap = unsafe { GetProcessHeap() };
            if heap.is_null() {
                // Reallocation has failed, could not get the current process heap.
                return ptr::null_mut();
            }

            // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`,
            // `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<commit_msg>Cache `GetProcessHeap`<commit_after>#![deny(unsafe_op_in_unsafe_fn)]

use crate::alloc::{GlobalAlloc, Layout, System};
use crate::ffi::c_void;
use crate::ptr;
use crate::sync::atomic::{AtomicPtr, Ordering};
use crate::sys::c;
use crate::sys_common::alloc::{realloc_fallback, MIN_ALIGN};

#[cfg(test)]
mod tests;

// Heap memory management on Windows is done by using the system Heap API (heapapi.h)
// See https://docs.microsoft.com/windows/win32/api/heapapi/

// Flag to indicate that the memory returned by `HeapAlloc` should be zeroed.
const HEAP_ZERO_MEMORY: c::DWORD = 0x00000008;

extern "system" {
    // Get a handle to the default heap of the current process, or null if the operation fails.
    //
    // SAFETY: Successful calls to this function within the same process are assumed to
    // always return the same handle, which remains valid for the entire lifetime of the process.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-getprocessheap
    fn GetProcessHeap() -> c::HANDLE;

    // Allocate a block of `dwBytes` bytes of memory from a given heap `hHeap`.
    // The allocated memory may be uninitialized, or zeroed if `dwFlags` is
    // set to `HEAP_ZERO_MEMORY`.
    //
    // Returns a pointer to the newly-allocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to either zero or `HEAP_ZERO_MEMORY`.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapalloc
    fn HeapAlloc(hHeap: c::HANDLE, dwFlags: c::DWORD, dwBytes: c::SIZE_T) -> c::LPVOID;

    // Reallocate a block of memory behind a given pointer `lpMem` from a given heap `hHeap`,
    // to a block of at least `dwBytes` bytes, either shrinking the block in place,
    // or allocating at a new location, copying memory, and freeing the original location.
    //
    // Returns a pointer to the reallocated memory or null if the operation fails.
    // The returned pointer will be aligned to at least `MIN_ALIGN`.
    // If the operation fails the given block will never have been freed.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a non-null pointer to an allocated block returned by `HeapAlloc` or
    //     `HeapReAlloc`, that has not already been freed.
    // If the block was successfully reallocated at a new location, pointers pointing to
    // the freed memory, such as `lpMem`, must not be dereferenced ever again.
    //
    // Note that `dwBytes` is allowed to be zero, contrary to some other allocators.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heaprealloc
    fn HeapReAlloc(
        hHeap: c::HANDLE,
        dwFlags: c::DWORD,
        lpMem: c::LPVOID,
        dwBytes: c::SIZE_T,
    ) -> c::LPVOID;

    // Free a block of memory behind a given pointer `lpMem` from a given heap `hHeap`.
    // Returns a nonzero value if the operation is successful, and zero if the operation fails.
    //
    // SAFETY:
    //  - `hHeap` must be a non-null handle returned by `GetProcessHeap`.
    //  - `dwFlags` must be set to zero.
    //  - `lpMem` must be a pointer to an allocated block returned by `HeapAlloc` or `HeapReAlloc`,
    //     that has not already been freed.
    // If the block was successfully freed, pointers pointing to the freed memory, such as `lpMem`,
    // must not be dereferenced ever again.
    //
    // Note that `lpMem` is allowed to be null, which will not cause the operation to fail.
    //
    // See https://docs.microsoft.com/windows/win32/api/heapapi/nf-heapapi-heapfree
    fn HeapFree(hHeap: c::HANDLE, dwFlags: c::DWORD, lpMem: c::LPVOID) -> c::BOOL;
}

// Cached handle to the default heap of the current process.
// Either a non-null handle returned by `GetProcessHeap`, or null when not yet initialized or `GetProcessHeap` failed.
static HEAP: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());

// Get a handle to the default heap of the current process, or null if the operation fails.
// SAFETY: If this operation is successful, `HEAP` will be successfully initialized and contain
// a non-null handle returned by `GetProcessHeap`.
#[inline]
unsafe fn init_or_get_process_heap() -> c::HANDLE {
    let heap = HEAP.load(Ordering::Relaxed);
    if heap.is_null() {
        // `HEAP` has not yet been successfully initialized
        let heap = unsafe { GetProcessHeap() };
        if !heap.is_null() {
            // SAFETY: No locking is needed because within the same process,
            // successful calls to `GetProcessHeap` will always return the same value, even on different threads.
            HEAP.store(heap, Ordering::Relaxed);

            // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
            heap
        } else {
            // Could not get the current process heap.
            ptr::null_mut()
        }
    } else {
        // SAFETY: `HEAP` contains a non-null handle returned by `GetProcessHeap`
        heap
    }
}

// Header containing a pointer to the start of an allocated block.
// SAFETY: Size and alignment must be <= `MIN_ALIGN`.
#[repr(C)]
struct Header(*mut u8);

// Allocate a block of optionally zeroed memory for a given `layout`.
// SAFETY: Returns a pointer satisfying the guarantees of `System` about allocated pointers.
#[inline]
unsafe fn allocate(layout: Layout, zeroed: bool) -> *mut u8 {
    let heap = unsafe { init_or_get_process_heap() };
    if heap.is_null() {
        // Allocation has failed, could not get the current process heap.
        return ptr::null_mut();
    }

    // Allocated memory will be either zeroed or uninitialized.
    let flags = if zeroed { HEAP_ZERO_MEMORY } else { 0 };

    if layout.align() <= MIN_ALIGN {
        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        // The returned pointer points to the start of an allocated block.
        unsafe { HeapAlloc(heap, flags, layout.size()) as *mut u8 }
    } else {
        // Allocate extra padding in order to be able to satisfy the alignment.
        let total = layout.align() + layout.size();

        // SAFETY: `heap` is a non-null handle returned by `GetProcessHeap`.
        let ptr = unsafe { HeapAlloc(heap, flags, total) as *mut u8 };
        if ptr.is_null() {
            // Allocation has failed.
            return ptr::null_mut();
        }

        // Create a correctly aligned pointer offset from the start of the allocated block,
        // and write a header before it.

        let offset = layout.align() - (ptr as usize & (layout.align() - 1));
        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
        // least `MIN_ALIGN` bytes of padding before it.
        let aligned = unsafe { ptr.add(offset) };
        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
        // it, it is safe to write a header directly before it.
        unsafe { ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };

        // SAFETY: The returned pointer does not point to the to the start of an allocated block,
        // but there is a header readable directly before it containing the location of the start
        // of the block.
        aligned
    }
}

// All pointers returned by this allocator have, in addition to the guarantees of `GlobalAlloc`, the
// following properties:
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment <= `MIN_ALIGN`
// the pointer will be aligned to at least `MIN_ALIGN` and point to the start of the allocated block.
//
// If the pointer was allocated or reallocated with a `layout` specifying an alignment > `MIN_ALIGN`
// the pointer will be aligned to the specified alignment and not point to the start of the allocated block.
// Instead there will be a header readable directly before the returned pointer, containing the actual
// location of the start of the block.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = false;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: Pointers returned by `allocate` satisfy the guarantees of `System`
        let zeroed = true;
        unsafe { allocate(layout, zeroed) }
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let block = {
            if layout.align() <= MIN_ALIGN {
                ptr
            } else {
                // The location of the start of the block is stored in the padding before `ptr`.

                // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
                // and have a header readable directly before it.
                unsafe { ptr::read((ptr as *mut Header).offset(-1)).0 }
            }
        };

        // SAFETY: because `ptr` has been successfully allocated with this allocator,
        // `HEAP` must have been successfully initialized and contain a non-null handle
        // returned by `GetProcessHeap`.
        let heap = HEAP.load(Ordering::Relaxed);

        // SAFETY: `block` is a pointer to the start of an allocated block.
        unsafe {
            let err = HeapFree(heap, 0, block as c::LPVOID);
            debug_assert!(err != 0, "Failed to free heap memory: {}", c::GetLastError());
        }
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            // SAFETY: because `ptr` has been successfully allocated with this allocator,
            // `HEAP` must have been successfully initialized and contain a non-null handle
            // returned by `GetProcessHeap`.
            let heap = HEAP.load(Ordering::Relaxed);

            // SAFETY: `ptr` is a pointer to the start of an allocated block.
            // The returned pointer points to the start of an allocated block.
            unsafe { HeapReAlloc(heap, 0, ptr as c::LPVOID, new_size) as *mut u8 }
        } else {
            // SAFETY: `realloc_fallback` is implemented using `dealloc` and `alloc`, which will
            // correctly handle `ptr` and return a pointer satisfying the guarantees of `System`
            unsafe { realloc_fallback(self, ptr, layout, new_size) }
        }
    }
}
<|endoftext|>


<commit_before>use async_io::{Async, Timer};
use lapin::{
    executor::Executor,
    heartbeat::Heartbeat,
    reactor::{Reactor, ReactorBuilder, ReactorHandle, Slot},
    socket_state::{SocketEvent, SocketStateHandle},
    tcp::{TcpStream, TcpStreamWrapper},
    ConnectionProperties, Result,
};
use parking_lot::Mutex;
use std::{collections::HashMap, fmt, sync::Arc};

// ConnectionProperties extension

pub trait LapinAsyncIoExt {
    fn with_async_io(self, executor: impl Executor + 'static) -> Self
    where
        Self: Sized,
    {
        self.with_async_io_reactor(executor)
    }

    fn with_async_io_reactor(self, executor: impl Executor + 'static) -> Self
    where
        Self: Sized;
}

impl LapinAsyncIoExt for ConnectionProperties {
    fn with_async_io_reactor(self, executor: impl Executor + 'static) -> Self {
        self.with_reactor(AsyncIoReactorBuilder(Arc::new(executor)))
    }
}

// Reactor

struct AsyncIoReactorBuilder(Arc<dyn Executor>);

impl fmt::Debug for AsyncIoReactorBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AsyncIoReactorBuilder").finish()
    }
}

#[derive(Debug)]
struct AsyncIoReactor(AsyncIoReactorHandle);

#[derive(Clone)]
struct AsyncIoReactorHandle {
    heartbeat: Heartbeat,
    executor: Arc<dyn Executor>,
    inner: Arc<Mutex<Inner>>,
}

#[derive(Default)]
struct Inner {
    slot: Slot,
    slots: HashMap<usize, (Arc<Async<TcpStreamWrapper>>, SocketStateHandle)>,
}

impl fmt::Debug for AsyncIoReactorHandle {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AsyncIoReactorHandle").finish()
    }
}

impl Inner {
    fn register(
        &mut self,
        socket: Arc<Async<TcpStreamWrapper>>,
        socket_state: SocketStateHandle,
    ) -> Result<usize> {
        let slot = self.slot;
        self.slot += 1;
        self.slots.insert(slot, (socket, socket_state));
        Ok(slot)
    }
}

impl ReactorBuilder for AsyncIoReactorBuilder {
    fn build(&self, heartbeat: Heartbeat) -> Result<Box<dyn Reactor + Send>> {
        Ok(Box::new(AsyncIoReactor(AsyncIoReactorHandle {
            heartbeat,
            executor: self.0.clone(),
            inner: Arc::new(Mutex::new(Default::default())),
        })))
    }
}

impl Reactor for AsyncIoReactor {
    fn register(
        &mut self,
        socket: &mut TcpStream,
        socket_state: SocketStateHandle,
    ) -> Result<usize> {
        let socket = Arc::new(Async::new(unsafe { TcpStreamWrapper::new(socket) })?);
        let slot = self.0.inner.lock().register(socket, socket_state)?;
        self.0.poll_read(slot);
        self.0.poll_write(slot);
        Ok(slot)
    }

    fn handle(&self) -> Box<dyn ReactorHandle + Send> {
        Box::new(self.0.clone())
    }
}

impl ReactorHandle for AsyncIoReactorHandle {
    fn start_heartbeat(&self) {
        self.executor
            .spawn(Box::pin(heartbeat(self.heartbeat.clone())))
            .expect("start_heartbeat");
    }

    fn poll_read(&self, slot: usize) {
        if let Some((socket, socket_state)) = self.inner.lock().slots.get(&slot) {
            self.executor
                .spawn(Box::pin(poll_read(socket.clone(), socket_state.clone())))
                .expect("poll_read");
        }
    }

    fn poll_write(&self, slot: usize) {
        if let Some((socket, socket_state)) = self.inner.lock().slots.get(&slot) {
            self.executor
                .spawn(Box::pin(poll_write(socket.clone(), socket_state.clone())))
                .expect("poll_write");
        }
    }
}

async fn heartbeat(heartbeat: Heartbeat) {
    while let Ok(Some(timeout)) = heartbeat.poll_timeout() {
        Timer::after(timeout).await;
    }
}

async fn poll_read(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
    socket.readable().await.unwrap();
    socket_state.send(SocketEvent::Readable);
}

async fn poll_write(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
    socket.writable().await.unwrap();
    socket_state.send(SocketEvent::Writable);
}
<commit_msg>async-lapin: be consistent with type aliases<commit_after>use async_io::{Async, Timer};
use lapin::{
    executor::Executor,
    heartbeat::Heartbeat,
    reactor::{Reactor, ReactorBuilder, ReactorHandle, Slot},
    socket_state::{SocketEvent, SocketStateHandle},
    tcp::{TcpStream, TcpStreamWrapper},
    ConnectionProperties, Result,
};
use parking_lot::Mutex;
use std::{collections::HashMap, fmt, sync::Arc};

// ConnectionProperties extension

pub trait LapinAsyncIoExt {
    fn with_async_io(self, executor: impl Executor + 'static) -> Self
    where
        Self: Sized,
    {
        self.with_async_io_reactor(executor)
    }

    fn with_async_io_reactor(self, executor: impl Executor + 'static) -> Self
    where
        Self: Sized;
}

impl LapinAsyncIoExt for ConnectionProperties {
    fn with_async_io_reactor(self, executor: impl Executor + 'static) -> Self {
        self.with_reactor(AsyncIoReactorBuilder(Arc::new(executor)))
    }
}

// Reactor

struct AsyncIoReactorBuilder(Arc<dyn Executor>);

impl fmt::Debug for AsyncIoReactorBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AsyncIoReactorBuilder").finish()
    }
}

#[derive(Debug)]
struct AsyncIoReactor(AsyncIoReactorHandle);

#[derive(Clone)]
struct AsyncIoReactorHandle {
    heartbeat: Heartbeat,
    executor: Arc<dyn Executor>,
    inner: Arc<Mutex<Inner>>,
}

#[derive(Default)]
struct Inner {
    slot: Slot,
    slots: HashMap<Slot, (Arc<Async<TcpStreamWrapper>>, SocketStateHandle)>,
}

impl fmt::Debug for AsyncIoReactorHandle {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AsyncIoReactorHandle").finish()
    }
}

impl Inner {
    fn register(
        &mut self,
        socket: Arc<Async<TcpStreamWrapper>>,
        socket_state: SocketStateHandle,
    ) -> Result<usize> {
        let slot = self.slot;
        self.slot += 1;
        self.slots.insert(slot, (socket, socket_state));
        Ok(slot)
    }
}

impl ReactorBuilder for AsyncIoReactorBuilder {
    fn build(&self, heartbeat: Heartbeat) -> Result<Box<dyn Reactor + Send>> {
        Ok(Box::new(AsyncIoReactor(AsyncIoReactorHandle {
            heartbeat,
            executor: self.0.clone(),
            inner: Arc::new(Mutex::new(Default::default())),
        })))
    }
}

impl Reactor for AsyncIoReactor {
    fn register(
        &mut self,
        socket: &mut TcpStream,
        socket_state: SocketStateHandle,
    ) -> Result<usize> {
        let socket = Arc::new(Async::new(unsafe { TcpStreamWrapper::new(socket) })?);
        let slot = self.0.inner.lock().register(socket, socket_state)?;
        self.0.poll_read(slot);
        self.0.poll_write(slot);
        Ok(slot)
    }

    fn handle(&self) -> Box<dyn ReactorHandle + Send> {
        Box::new(self.0.clone())
    }
}

impl ReactorHandle for AsyncIoReactorHandle {
    fn start_heartbeat(&self) {
        self.executor
            .spawn(Box::pin(heartbeat(self.heartbeat.clone())))
            .expect("start_heartbeat");
    }

    fn poll_read(&self, slot: usize) {
        if let Some((socket, socket_state)) = self.inner.lock().slots.get(&slot) {
            self.executor
                .spawn(Box::pin(poll_read(socket.clone(), socket_state.clone())))
                .expect("poll_read");
        }
    }

    fn poll_write(&self, slot: usize) {
        if let Some((socket, socket_state)) = self.inner.lock().slots.get(&slot) {
            self.executor
                .spawn(Box::pin(poll_write(socket.clone(), socket_state.clone())))
                .expect("poll_write");
        }
    }
}

async fn heartbeat(heartbeat: Heartbeat) {
    while let Ok(Some(timeout)) = heartbeat.poll_timeout() {
        Timer::after(timeout).await;
    }
}

async fn poll_read(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
    socket.readable().await.unwrap();
    socket_state.send(SocketEvent::Readable);
}

async fn poll_write(socket: Arc<Async<TcpStreamWrapper>>, socket_state: SocketStateHandle) {
    socket.writable().await.unwrap();
    socket_state.send(SocketEvent::Writable);
}
<|endoftext|>