<commit_before>use std::borrow::Borrow;
use std::env;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

fn run_command_or_fail<P, S>(dir: &str, cmd: P, args: &[S])
where
    P: AsRef<Path>,
    S: Borrow<str> + AsRef<OsStr>,
{
    let cmd = cmd.as_ref();
    let cmd = if cmd.components().count() > 1 && cmd.is_relative() {
        // If `cmd` is a relative path (and not a bare command that should be
        // looked up in PATH), absolutize it relative to `dir`, as otherwise the
        // behavior of std::process::Command is undefined.
        // https://github.com/rust-lang/rust/issues/37868
        PathBuf::from(dir)
            .join(cmd)
            .canonicalize()
            .expect("canonicalization failed")
    } else {
        PathBuf::from(cmd)
    };
    eprintln!(
        "Running command: \"{} {}\" in dir: {}",
        cmd.display(),
        args.join(" "),
        dir
    );
    let ret = Command::new(cmd).current_dir(dir).args(args).status();
    match ret.map(|status| (status.success(), status.code())) {
        Ok((true, _)) => return,
        Ok((false, Some(c))) => panic!("Command failed with error code {}", c),
        Ok((false, None)) => panic!("Command got killed"),
        Err(e) => panic!("Command failed with error: {}", e),
    }
}

fn main() {
    let librdkafka_version = env!("CARGO_PKG_VERSION")
        .split('-')
        .next()
        .expect("Crate version is not valid");

    if env::var("DEP_OPENSSL_VENDORED").is_ok() {
        let openssl_root = env::var("DEP_OPENSSL_ROOT").expect("DEP_OPENSSL_ROOT is not set");
        env::set_var("CFLAGS", format!("-I{}/include", openssl_root));
        env::set_var("LDFLAGS", format!("-L{}/lib", openssl_root));
    }

    if env::var("CARGO_FEATURE_DYNAMIC_LINKING").is_ok() {
        eprintln!("librdkafka will be linked dynamically");
        let pkg_probe = pkg_config::Config::new()
            .cargo_metadata(true)
            .atleast_version(librdkafka_version)
            .probe("rdkafka");

        match pkg_probe {
            Ok(library) => {
                eprintln!("librdkafka found on the system:");
                eprintln!("  Name: {:?}", library.libs);
                eprintln!("  Path: {:?}", library.link_paths);
                eprintln!("  Version: {}", library.version);
            }
            Err(_) => {
                eprintln!(
                    "librdkafka {} cannot be found on the system",
                    librdkafka_version
                );
                eprintln!("Dynamic linking failed. Exiting.");
                process::exit(1);
            }
        }
    } else {
        if !Path::new("librdkafka/LICENSE").exists() {
            eprintln!("Setting up submodules");
            run_command_or_fail("../", "git", &["submodule", "update", "--init"]);
        }
        eprintln!("Building and linking librdkafka statically");
        build_librdkafka();
    }
}

#[cfg(not(feature = "cmake-build"))]
fn build_librdkafka() {
    let mut configure_flags: Vec<String> = Vec::new();

    let mut cflags = Vec::new();
    if let Ok(var) = env::var("CFLAGS") {
        cflags.push(var);
    }

    let mut ldflags = Vec::new();
    if let Ok(var) = env::var("LDFLAGS") {
        ldflags.push(var);
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        configure_flags.push("--enable-ssl".into());
        if let Ok(openssl_root) = env::var("DEP_OPENSSL_ROOT") {
            cflags.push(format!("-I{}/include", openssl_root));
            ldflags.push(format!("-L{}/lib", openssl_root));
        }
    } else {
        configure_flags.push("--disable-ssl".into());
    }

    if env::var("CARGO_FEATURE_GSSAPI").is_ok() {
        configure_flags.push("--enable-gssapi".into());
        println!("cargo:rustc-link-lib=sasl2");
    } else {
        configure_flags.push("--disable-gssapi".into());
    }

    if env::var("CARGO_FEATURE_LIBZ").is_ok() {
        // There is no --enable-zlib option, but it is enabled by default.
        if let Ok(z_root) = env::var("DEP_Z_ROOT") {
            cflags.push(format!("-I{}/include", z_root));
            ldflags.push(format!("-L{}/build", z_root));
        }
    } else {
        configure_flags.push("--disable-zlib".into());
    }

    if env::var("CARGO_FEATURE_ZSTD").is_ok() {
        configure_flags.push("--enable-zstd".into());
        if let Ok(zstd_root) = env::var("DEP_ZSTD_ROOT") {
            cflags.push(format!("-I{}/include", zstd_root));
            ldflags.push(format!("-L{}", zstd_root));
        }
    } else {
        configure_flags.push("--disable-zstd".into());
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        configure_flags.push("--enable-lz4-ext".into());
        if let Ok(lz4_root) = env::var("DEP_LZ4_ROOT") {
            cflags.push(format!("-I{}/include", lz4_root));
            ldflags.push(format!("-L{}", lz4_root));
        }
    } else {
        configure_flags.push("--disable-lz4-ext".into());
    }

    env::set_var("CFLAGS", cflags.join(" "));
    env::set_var("LDFLAGS", ldflags.join(" "));

    println!("Configuring librdkafka");
    run_command_or_fail("librdkafka", "./configure", configure_flags.as_slice());

    println!("Compiling librdkafka");
    env::set_var("MAKEFLAGS", env::var_os("CARGO_MAKEFLAGS").expect("CARGO_MAKEFLAGS env var missing"));
    run_command_or_fail(
        "librdkafka",
        if cfg!(target_os = "freebsd") { "gmake" } else { "make" },
        &["libs"],
    );

    println!(
        "cargo:rustc-link-search=native={}/librdkafka/src",
        env::current_dir()
            .expect("Can't find current dir")
            .display()
    );
    println!("cargo:rustc-link-lib=static=rdkafka");
    println!("cargo:root={}", env::var("OUT_DIR").expect("OUT_DIR missing"));
}

#[cfg(feature = "cmake-build")]
fn build_librdkafka() {
    let mut config = cmake::Config::new("librdkafka");

    config
        .define("RDKAFKA_BUILD_STATIC", "1")
        .build_target("rdkafka");

    if env::var("CARGO_FEATURE_LIBZ").is_ok() {
        config.define("WITH_ZLIB", "1");
        config.register_dep("z");
        if let Ok(z_root) = env::var("DEP_Z_ROOT") {
            env::set_var("CMAKE_LIBRARY_PATH", format!("{}/build", z_root));
        }
    } else {
        config.define("WITH_ZLIB", "0");
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        config.define("WITH_SSL", "1");
        config.register_dep("openssl");
    } else {
        config.define("WITH_SSL", "0");
    }

    if env::var("CARGO_FEATURE_GSSAPI").is_ok() {
        config.define("WITH_SASL", "1");
        println!("cargo:rustc-link-lib=sasl2");
    } else {
        config.define("WITH_SASL", "0");
    }

    if env::var("CARGO_FEATURE_ZSTD").is_ok() {
        config.define("WITH_ZSTD", "1");
        config.register_dep("zstd");
    } else {
        config.define("WITH_ZSTD", "0");
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        config.define("ENABLE_LZ4_EXT", "1");
        config.register_dep("lz4");
    } else {
        config.define("ENABLE_LZ4_EXT", "0");
    }

    if let Ok(system_name) = env::var("CMAKE_SYSTEM_NAME") {
        config.define("CMAKE_SYSTEM_NAME", system_name);
    }

    // The CMake build will incorrectly use config.h from the non-CMake build,
    // if it exists, so remove it if it does.
    match std::fs::remove_file("librdkafka/config.h") {
        Ok(()) => (),
        Err(err) => match err.kind() {
            std::io::ErrorKind::NotFound => (),
            _ => panic!("Unable to remove config.h from non-CMake build: {}", err),
        }
    }

    println!("Configuring and compiling librdkafka");
    let dst = config.build();

    if cfg!(target_env = "msvc") {
        let profile = match &env::var("PROFILE").expect("Cannot determine build profile")[..] {
            "release" | "bench" => "Release",
            _ => "Debug"
        };
        println!("cargo:rustc-link-search=native={}/build/src/{}", dst.display(), profile);
    } else {
        println!("cargo:rustc-link-search=native={}/build/src", dst.display());
    }
    println!("cargo:rustc-link-lib=static=rdkafka");
}
<commit_msg>Use CMake's install target for proper MSVC handling<commit_after>use std::borrow::Borrow;
use std::env;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::process::{self, Command};

fn run_command_or_fail<P, S>(dir: &str, cmd: P, args: &[S])
where
    P: AsRef<Path>,
    S: Borrow<str> + AsRef<OsStr>,
{
    let cmd = cmd.as_ref();
    let cmd = if cmd.components().count() > 1 && cmd.is_relative() {
        // If `cmd` is a relative path (and not a bare command that should be
        // looked up in PATH), absolutize it relative to `dir`, as otherwise the
        // behavior of std::process::Command is undefined.
        // https://github.com/rust-lang/rust/issues/37868
        PathBuf::from(dir)
            .join(cmd)
            .canonicalize()
            .expect("canonicalization failed")
    } else {
        PathBuf::from(cmd)
    };
    eprintln!(
        "Running command: \"{} {}\" in dir: {}",
        cmd.display(),
        args.join(" "),
        dir
    );
    let ret = Command::new(cmd).current_dir(dir).args(args).status();
    match ret.map(|status| (status.success(), status.code())) {
        Ok((true, _)) => return,
        Ok((false, Some(c))) => panic!("Command failed with error code {}", c),
        Ok((false, None)) => panic!("Command got killed"),
        Err(e) => panic!("Command failed with error: {}", e),
    }
}

fn main() {
    let librdkafka_version = env!("CARGO_PKG_VERSION")
        .split('-')
        .next()
        .expect("Crate version is not valid");

    if env::var("DEP_OPENSSL_VENDORED").is_ok() {
        let openssl_root = env::var("DEP_OPENSSL_ROOT").expect("DEP_OPENSSL_ROOT is not set");
        env::set_var("CFLAGS", format!("-I{}/include", openssl_root));
        env::set_var("LDFLAGS", format!("-L{}/lib", openssl_root));
    }

    if env::var("CARGO_FEATURE_DYNAMIC_LINKING").is_ok() {
        eprintln!("librdkafka will be linked dynamically");
        let pkg_probe = pkg_config::Config::new()
            .cargo_metadata(true)
            .atleast_version(librdkafka_version)
            .probe("rdkafka");

        match pkg_probe {
            Ok(library) => {
                eprintln!("librdkafka found on the system:");
                eprintln!("  Name: {:?}", library.libs);
                eprintln!("  Path: {:?}", library.link_paths);
                eprintln!("  Version: {}", library.version);
            }
            Err(_) => {
                eprintln!(
                    "librdkafka {} cannot be found on the system",
                    librdkafka_version
                );
                eprintln!("Dynamic linking failed. Exiting.");
                process::exit(1);
            }
        }
    } else {
        if !Path::new("librdkafka/LICENSE").exists() {
            eprintln!("Setting up submodules");
            run_command_or_fail("../", "git", &["submodule", "update", "--init"]);
        }
        eprintln!("Building and linking librdkafka statically");
        build_librdkafka();
    }
}

#[cfg(not(feature = "cmake-build"))]
fn build_librdkafka() {
    let mut configure_flags: Vec<String> = Vec::new();

    let mut cflags = Vec::new();
    if let Ok(var) = env::var("CFLAGS") {
        cflags.push(var);
    }

    let mut ldflags = Vec::new();
    if let Ok(var) = env::var("LDFLAGS") {
        ldflags.push(var);
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        configure_flags.push("--enable-ssl".into());
        if let Ok(openssl_root) = env::var("DEP_OPENSSL_ROOT") {
            cflags.push(format!("-I{}/include", openssl_root));
            ldflags.push(format!("-L{}/lib", openssl_root));
        }
    } else {
        configure_flags.push("--disable-ssl".into());
    }

    if env::var("CARGO_FEATURE_GSSAPI").is_ok() {
        configure_flags.push("--enable-gssapi".into());
        println!("cargo:rustc-link-lib=sasl2");
    } else {
        configure_flags.push("--disable-gssapi".into());
    }

    if env::var("CARGO_FEATURE_LIBZ").is_ok() {
        // There is no --enable-zlib option, but it is enabled by default.
        if let Ok(z_root) = env::var("DEP_Z_ROOT") {
            cflags.push(format!("-I{}/include", z_root));
            ldflags.push(format!("-L{}/build", z_root));
        }
    } else {
        configure_flags.push("--disable-zlib".into());
    }

    if env::var("CARGO_FEATURE_ZSTD").is_ok() {
        configure_flags.push("--enable-zstd".into());
        if let Ok(zstd_root) = env::var("DEP_ZSTD_ROOT") {
            cflags.push(format!("-I{}/include", zstd_root));
            ldflags.push(format!("-L{}", zstd_root));
        }
    } else {
        configure_flags.push("--disable-zstd".into());
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        configure_flags.push("--enable-lz4-ext".into());
        if let Ok(lz4_root) = env::var("DEP_LZ4_ROOT") {
            cflags.push(format!("-I{}/include", lz4_root));
            ldflags.push(format!("-L{}", lz4_root));
        }
    } else {
        configure_flags.push("--disable-lz4-ext".into());
    }

    env::set_var("CFLAGS", cflags.join(" "));
    env::set_var("LDFLAGS", ldflags.join(" "));

    println!("Configuring librdkafka");
    run_command_or_fail("librdkafka", "./configure", configure_flags.as_slice());

    println!("Compiling librdkafka");
    env::set_var("MAKEFLAGS", env::var_os("CARGO_MAKEFLAGS").expect("CARGO_MAKEFLAGS env var missing"));
    run_command_or_fail(
        "librdkafka",
        if cfg!(target_os = "freebsd") { "gmake" } else { "make" },
        &["libs"],
    );

    println!(
        "cargo:rustc-link-search=native={}/librdkafka/src",
        env::current_dir()
            .expect("Can't find current dir")
            .display()
    );
    println!("cargo:rustc-link-lib=static=rdkafka");
    println!("cargo:root={}", env::var("OUT_DIR").expect("OUT_DIR missing"));
}

#[cfg(feature = "cmake-build")]
fn build_librdkafka() {
    let mut config = cmake::Config::new("librdkafka");

    config.define("RDKAFKA_BUILD_STATIC", "1");

    if env::var("CARGO_FEATURE_LIBZ").is_ok() {
        config.define("WITH_ZLIB", "1");
        config.register_dep("z");
        if let Ok(z_root) = env::var("DEP_Z_ROOT") {
            env::set_var("CMAKE_LIBRARY_PATH", format!("{}/build", z_root));
        }
    } else {
        config.define("WITH_ZLIB", "0");
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        config.define("WITH_SSL", "1");
        config.register_dep("openssl");
    } else {
        config.define("WITH_SSL", "0");
    }

    if env::var("CARGO_FEATURE_GSSAPI").is_ok() {
        config.define("WITH_SASL", "1");
        println!("cargo:rustc-link-lib=sasl2");
    } else {
        config.define("WITH_SASL", "0");
    }

    if env::var("CARGO_FEATURE_ZSTD").is_ok() {
        config.define("WITH_ZSTD", "1");
        config.register_dep("zstd");
    } else {
        config.define("WITH_ZSTD", "0");
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        config.define("ENABLE_LZ4_EXT", "1");
        config.register_dep("lz4");
    } else {
        config.define("ENABLE_LZ4_EXT", "0");
    }

    if let Ok(system_name) = env::var("CMAKE_SYSTEM_NAME") {
        config.define("CMAKE_SYSTEM_NAME", system_name);
    }

    // The CMake build will incorrectly use config.h from the non-CMake build,
    // if it exists, so remove it if it does.
    match std::fs::remove_file("librdkafka/config.h") {
        Ok(()) => (),
        Err(err) => match err.kind() {
            std::io::ErrorKind::NotFound => (),
            _ => panic!("Unable to remove config.h from non-CMake build: {}", err),
        }
    }

    println!("Configuring and compiling librdkafka");
    let dst = config.build();

    println!("cargo:rustc-link-search=native={}/lib", dst.display());
    println!("cargo:rustc-link-lib=static=rdkafka");
}
<|endoftext|>


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

use anyhow::{bail, Context, Error};
use fehler::throws;
use libtest_mimic::{run_tests, Arguments, Outcome, Test};
use std::fs::{read_dir, read_to_string};
use std::path::Path;
use wasmbin::{
    io::DecodeError,
    visit::{Visit, VisitError},
    Module,
};
use wast::parser::{parse, ParseBuffer};
use wast::Wast;

const IGNORED_ERRORS: &[&str] = &[
    // We don't perform cross-section analysis.
    "function and code section have inconsistent lengths",
    "data count section required",
    "data count and data section have inconsistent lengths",
    // We allow non-zero table and memory IDs already.
    "zero flag expected",
    // We don't perform full function analysis either.
    "too many locals",
];

struct WasmTest {
    module: Vec<u8>,
    expect_result: Result<(), String>,
}

#[throws]
fn read_tests_from_file(path: &Path, dest: &mut Vec<Test<WasmTest>>, ignore_malformed: bool) {
    let src = read_to_string(path)?;
    let set_err_path_text = |mut err: wast::Error| {
        err.set_path(path);
        err.set_text(&src);
        err
    };
    let buf = ParseBuffer::new(&src).map_err(set_err_path_text)?;
    let wast = parse::<Wast>(&buf).map_err(set_err_path_text)?;
    for directive in wast.directives {
        let is_ignored = match directive {
            wast::WastDirective::AssertMalformed { .. } => ignore_malformed,
            _ => false,
        };
        let (span, mut module, expect_result) = match directive {
            // Expect errors for assert_malformed on binary or AST modules.
            wast::WastDirective::AssertMalformed {
                span,
                module: wast::QuoteModule::Module(module),
                message,
            }
            // Unlike other AssertInvalid, this is actually something we
            // check at the parsing time, because it's part of our
            // typesystem and doesn't require cross-function or
            // cross-section checks.
            | wast::WastDirective::AssertInvalid {
                span,
                module,
                message: message @ "invalid lane index",
            } => (span, module, Err(message)),
            // Expect successful parsing for regular AST modules.
            wast::WastDirective::Module(module) => (module.span, module, Ok(())),
            // Counter-intuitively, expect successful parsing for modules that are supposed
            // to error out at runtime or linking stage, too.
            wast::WastDirective::AssertInvalid { span, module, .. }
            | wast::WastDirective::AssertUnlinkable { span, module, .. } => (span, module, Ok(())),
            _ => {
                // Skipping interpreted
                continue;
            }
        };
        let (line, col) = span.linecol_in(&src);
        let module = module.encode()?;
        dest.push(Test {
            name: format!("{}:{}:{}", path.display(), line + 1, col + 1),
            kind: String::default(),
            is_ignored: is_ignored || match expect_result {
                // see https://github.com/WebAssembly/bulk-memory-operations/issues/153
                Ok(()) => match module[..] {
                    | [0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00, 0x05, 0x03, 0x01, 0x00, 0x00, 0x0B, 0x07, 0x01, 0x80, 0x00, 0x41, 0x00, 0x0B, 0x00]
                    | [0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00, 0x04, 0x04, 0x01, 0x70, 0x00, 0x00, 0x09, 0x07, 0x01, 0x80, 0x00, 0x41, 0x00, 0x0B, 0x00]
                    => true,
                    _ => false,
                },
                Err(err) => IGNORED_ERRORS.contains(&err),
            },
            is_bench: false,
            data: WasmTest {
                module,
                expect_result: expect_result.map_err(|err| err.to_owned()),
            },
        });
    }
}

#[throws]
fn read_tests_from_dir(path: &Path, dest: &mut Vec<Test<WasmTest>>, ignore_malformed: bool) {
    for file in read_dir(path)? {
        let path = file?.path();
        if path.extension().map_or(false, |ext| ext == "wast") {
            read_tests_from_file(&path, dest, ignore_malformed)?;
        }
    }
}

#[throws]
fn read_all_tests(path: &Path) -> (Vec<Test<WasmTest>>, bool) {
    let mut tests = Vec::new();
    let proposals_dir = path.join("proposals");

    macro_rules! read_proposal_tests {
        ($name:literal) => {
            if cfg!(feature = $name) {
                read_tests_from_dir(&proposals_dir.join($name), &mut tests, false).context($name)?
            }
        };
    }

    read_proposal_tests!("bulk-memory-operations");
    read_proposal_tests!("reference-types");
    read_proposal_tests!("simd");
    read_proposal_tests!("tail-call");

    let has_proposals = !tests.is_empty();

    read_tests_from_dir(path, &mut tests, has_proposals)?;

    if tests.is_empty() {
        bail!("Couldn't find any tests. Did you run `git submodule update --init`?");
    }

    (tests, has_proposals)
}

fn unlazify<T: Visit>(mut wasm: T) -> Result<T, DecodeError> {
    match wasm.visit_mut(|()| {}) {
        Ok(()) => Ok(wasm),
        Err(err) => match err {
            VisitError::LazyDecode(err) => Err(err),
            VisitError::Custom(err) => match err {},
        },
    }
}

#[throws]
fn run_test(test: &WasmTest) {
    let mut slice = test.module.as_slice();
    let module = match (Module::decode_from(&mut slice).and_then(unlazify), &test.expect_result) {
        (Ok(_), Err(err)) => bail!("Expected an invalid module definition with an error: {}", err),
        (Err(err), Ok(())) => bail!(
            "Expected a valid module definition, but got an error\nParsed part: {:02X?}\nUnparsed part: {:02X?}\nError: {:#}",
            &test.module[..test.module.len() - slice.len()],
            slice,
            err
        ),
        (Ok(module), Ok(())) => module,
        (Err(_), Err(_)) => return,
    };
    let out = module.encode_into(Vec::new())?;
    if out != test.module {
        // In the rare case that binary representation doesn't match, it
        // might be because the test uses longer LEB128 form than
        // required. Verify that at least decoding it back produces the
        // same module.
        let module2 = Module::decode_from(out.as_slice())?;
        if module != module2 {
            bail!(
                "Roundtrip mismatch. Old: {:#?}\nNew: {:#?}",
                module,
                module2
            );
        }
    }
}

#[throws]
fn main() {
    let (tests, has_proposals) = read_all_tests(&Path::new("tests").join("testsuite"))?;

    run_tests(&Arguments::from_args(), tests, |test| {
        match run_test(&test.data) {
            Ok(()) => Outcome::Passed,
            Err(err) => Outcome::Failed {
                msg: Some(err.to_string()),
            },
        }
    })
    .exit_if_failed();

    if has_proposals {
        eprintln!(
            "Proposals were enabled, so upstream `assert_malformed` tests were ignored.\n\
             Re-run without any proposals enabled for the full upstream test coverage."
        );
    }
}
<commit_msg>Simplify test ignore code<commit_after>// Copyright 2020 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use anyhow::{bail, Context, Error};
use fehler::throws;
use libtest_mimic::{run_tests, Arguments, Outcome, Test};
use std::fs::{read_dir, read_to_string};
use std::path::Path;
use wasmbin::{
    io::DecodeError,
    visit::{Visit, VisitError},
    Module,
};
use wast::parser::{parse, ParseBuffer};
use wast::Wast;

const IGNORED_ERRORS: &[&str] = &[
    // We don't perform cross-section analysis.
    "function and code section have inconsistent lengths",
    "data count section required",
    "data count and data section have inconsistent lengths",
    // We allow non-zero table and memory IDs already.
    "zero flag expected",
    // We don't perform full function analysis either.
    "too many locals",
];

struct WasmTest {
    module: Vec<u8>,
    expect_result: Result<(), String>,
}

#[throws]
fn read_tests_from_file(path: &Path, dest: &mut Vec<Test<WasmTest>>, ignore_malformed: bool) {
    let src = read_to_string(path)?;
    let set_err_path_text = |mut err: wast::Error| {
        err.set_path(path);
        err.set_text(&src);
        err
    };
    let buf = ParseBuffer::new(&src).map_err(set_err_path_text)?;
    let wast = parse::<Wast>(&buf).map_err(set_err_path_text)?;
    for directive in wast.directives {
        let (span, mut module, expect_result) = match directive {
            // Expect errors for assert_malformed on binary or AST modules.
            wast::WastDirective::AssertMalformed {
                span,
                module: wast::QuoteModule::Module(module),
                message,
            }
            // Unlike other AssertInvalid, this is actually something we
            // check at the parsing time, because it's part of our
            // typesystem and doesn't require cross-function or
            // cross-section checks.
            | wast::WastDirective::AssertInvalid {
                span,
                module,
                message: message @ "invalid lane index",
            } => (span, module, Err(message.to_owned())),
            // Expect successful parsing for regular AST modules.
            wast::WastDirective::Module(module) => (module.span, module, Ok(())),
            // Counter-intuitively, expect successful parsing for modules that are supposed
            // to error out at runtime or linking stage, too.
            wast::WastDirective::AssertInvalid { span, module, .. }
            | wast::WastDirective::AssertUnlinkable { span, module, .. } => (span, module, Ok(())),
            _ => {
                // Skipping interpreted
                continue;
            }
        };
        let (line, col) = span.linecol_in(&src);
        let module = module.encode()?;
        dest.push(Test {
            name: format!("{}:{}:{}", path.display(), line + 1, col + 1),
            kind: String::default(),
            is_ignored: match &expect_result {
                // see https://github.com/WebAssembly/bulk-memory-operations/issues/153
                Ok(()) => match module[..] {
                    | [0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00, 0x05, 0x03, 0x01, 0x00, 0x00, 0x0B, 0x07, 0x01, 0x80, 0x00, 0x41, 0x00, 0x0B, 0x00]
                    | [0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00, 0x04, 0x04, 0x01, 0x70, 0x00, 0x00, 0x09, 0x07, 0x01, 0x80, 0x00, 0x41, 0x00, 0x0B, 0x00]
                    => true,
                    _ => false,
                },
                Err(err) => ignore_malformed || IGNORED_ERRORS.contains(&err.as_str()),
            },
            is_bench: false,
            data: WasmTest {
                module,
                expect_result,
            },
        });
    }
}

#[throws]
fn read_tests_from_dir(path: &Path, dest: &mut Vec<Test<WasmTest>>, ignore_malformed: bool) {
    for file in read_dir(path)? {
        let path = file?.path();
        if path.extension().map_or(false, |ext| ext == "wast") {
            read_tests_from_file(&path, dest, ignore_malformed)?;
        }
    }
}

#[throws]
fn read_all_tests(path: &Path) -> (Vec<Test<WasmTest>>, bool) {
    let mut tests = Vec::new();
    let proposals_dir = path.join("proposals");

    macro_rules! read_proposal_tests {
        ($name:literal) => {
            if cfg!(feature = $name) {
                read_tests_from_dir(&proposals_dir.join($name), &mut tests, false).context($name)?
            }
        };
    }

    read_proposal_tests!("bulk-memory-operations");
    read_proposal_tests!("reference-types");
    read_proposal_tests!("simd");
    read_proposal_tests!("tail-call");

    let has_proposals = !tests.is_empty();

    read_tests_from_dir(path, &mut tests, has_proposals)?;

    if tests.is_empty() {
        bail!("Couldn't find any tests. Did you run `git submodule update --init`?");
    }

    (tests, has_proposals)
}

fn unlazify<T: Visit>(mut wasm: T) -> Result<T, DecodeError> {
    match wasm.visit_mut(|()| {}) {
        Ok(()) => Ok(wasm),
        Err(err) => match err {
            VisitError::LazyDecode(err) => Err(err),
            VisitError::Custom(err) => match err {},
        },
    }
}

#[throws]
fn run_test(test: &WasmTest) {
    let mut slice = test.module.as_slice();
    let module = match (Module::decode_from(&mut slice).and_then(unlazify), &test.expect_result) {
        (Ok(_), Err(err)) => bail!("Expected an invalid module definition with an error: {}", err),
        (Err(err), Ok(())) => bail!(
            "Expected a valid module definition, but got an error\nParsed part: {:02X?}\nUnparsed part: {:02X?}\nError: {:#}",
            &test.module[..test.module.len() - slice.len()],
            slice,
            err
        ),
        (Ok(module), Ok(())) => module,
        (Err(_), Err(_)) => return,
    };
    let out = module.encode_into(Vec::new())?;
    if out != test.module {
        // In the rare case that binary representation doesn't match, it
        // might be because the test uses longer LEB128 form than
        // required. Verify that at least decoding it back produces the
        // same module.
        let module2 = Module::decode_from(out.as_slice())?;
        if module != module2 {
            bail!(
                "Roundtrip mismatch. Old: {:#?}\nNew: {:#?}",
                module,
                module2
            );
        }
    }
}

#[throws]
fn main() {
    let (tests, has_proposals) = read_all_tests(&Path::new("tests").join("testsuite"))?;

    run_tests(&Arguments::from_args(), tests, |test| {
        match run_test(&test.data) {
            Ok(()) => Outcome::Passed,
            Err(err) => Outcome::Failed {
                msg: Some(err.to_string()),
            },
        }
    })
    .exit_if_failed();

    if has_proposals {
        eprintln!(
            "Proposals were enabled, so upstream `assert_malformed` tests were ignored.\n\
             Re-run without any proposals enabled for the full upstream test coverage."
        );
    }
}
<|endoftext|>


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

use kvproto::coprocessor as coppb;
use tipb::ColumnInfo;

use crate::codec::datum::Datum;

/// Convert the key to the smallest key which is larger than the key given.
pub fn convert_to_prefix_next(key: &mut Vec<u8>) {
    if key.is_empty() {
        key.push(0);
        return;
    }
    let mut i = key.len() - 1;

    // Add 1 to the last byte that is not 255, and set it's following bytes to 0.
    loop {
        if key[i] == 255 {
            key[i] = 0;
        } else {
            key[i] += 1;
            return;
        }
        if i == 0 {
            // All bytes are 255. Append a 0 to the key.
            for byte in key.iter_mut() {
                *byte = 255;
            }
            key.push(0);
            return;
        }
        i -= 1;
    }
}

/// Check if `key`'s prefix next equals to `next`
pub fn is_prefix_next(key: &[u8], next: &[u8]) -> bool {
    let len = key.len();
    let next_len = next.len();

    if len == next_len {
        // Find the last non-255 byte
        let mut carry_pos = len;
        loop {
            if carry_pos == 0 {
                // All bytes of `key` are 255. `next` couldn't be `key`'s prefix_next since their
                // lengths are equal.
                return false;
            }

            carry_pos -= 1;
            if key[carry_pos] != 255 {
                break;
            }
        }

        // Now `carry_pos` is the index of the last byte that is not 255. For example:
        //   key: [1, 2, 3, 255, 255, 255]
        //               ^ carry_pos == 2

        // So they are equal when:
        // * `key`'s value at `carry_pos` is that of `next` - 1 and
        // * `next`'s part after carry_pos is all 0
        // * `key` and `next`'s parts before `carry_pos` are equal.
        // For example:
        //   key:  [1, 2, 3, 255, 255, 255]
        //   next: [1, 2, 4,   0,   0,   0]
        //                ^ carry_pos == 2
        // The part before `carry_pos` is all [1, 2],
        // the bytes at `carry_pos` differs by 1 (4 == 3 + 1), and
        // the remaining bytes of next ([0, 0, 0]) is all 0.
        // so [1, 2, 4, 0, 0, 0] is prefix_next of [1, 2, 3, 255, 255, 255]
        key[carry_pos] + 1 == next[carry_pos]
            && next[carry_pos + 1..].iter().all(|byte| *byte == 0)
            && key[..carry_pos] == next[..carry_pos]
    } else if len + 1 == next_len {
        // `next` must has one more 0 than `key`, and the first `len` bytes must be all 255.
        // The case that `len == 0` is also covered here.
        *next.last().unwrap() == 0
            && key.iter().all(|byte| *byte == 255)
            && next.iter().take(len).all(|byte| *byte == 255)
    } else {
        // Length not match.
        false
    }
}

/// `is_point` checks if the key range represents a point.
#[inline]
pub fn is_point(range: &coppb::KeyRange) -> bool {
    is_prefix_next(range.get_start(), range.get_end())
}

#[inline]
pub fn get_pk(col: &ColumnInfo, h: i64) -> Datum {
    use tidb_query_datatype::{FieldTypeAccessor, FieldTypeFlag};

    if col.as_accessor().flag().contains(FieldTypeFlag::UNSIGNED) {
        // PK column is unsigned
        Datum::U64(h as u64)
    } else {
        Datum::I64(h)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_prefix_next_once(key: &[u8], expected: &[u8]) {
        let mut key = key.to_vec();
        convert_to_prefix_next(&mut key);
        assert_eq!(key.as_slice(), expected);
    }

    #[test]
    fn test_prefix_next() {
        test_prefix_next_once(&[], &[0]);
        test_prefix_next_once(&[0], &[1]);
        test_prefix_next_once(&[1], &[2]);
        test_prefix_next_once(&[255], &[255, 0]);
        test_prefix_next_once(&[255, 255, 255], &[255, 255, 255, 0]);
        test_prefix_next_once(&[1, 255], &[2, 0]);
        test_prefix_next_once(&[0, 1, 255], &[0, 2, 0]);
        test_prefix_next_once(&[0, 1, 255, 5], &[0, 1, 255, 6]);
        test_prefix_next_once(&[0, 1, 5, 255], &[0, 1, 6, 0]);
        test_prefix_next_once(&[0, 1, 255, 255], &[0, 2, 0, 0]);
        test_prefix_next_once(&[0, 255, 255, 255], &[1, 0, 0, 0]);
    }

    fn test_is_prefix_next_case(lhs: &[u8], expected: &[u8], unexpected: &[&[u8]]) {
        assert!(is_prefix_next(lhs, expected));
        for rhs in unexpected {
            assert!(!is_prefix_next(lhs, rhs));
        }
    }

    #[test]
    fn test_is_prefix_next() {
        test_is_prefix_next_case(&[], &[0], &[&[], &[1], &[2]]);
        test_is_prefix_next_case(&[0], &[1], &[&[], &[0], &[0, 0], &[2], &[255]]);
        test_is_prefix_next_case(&[1], &[2], &[&[], &[1], &[3], &[1, 0]]);
        test_is_prefix_next_case(&[255], &[255, 0], &[&[0], &[255, 255, 0]]);
        test_is_prefix_next_case(
            &[255, 255, 255],
            &[255, 255, 255, 0],
            &[
                &[],
                &[0],
                &[0, 0, 0],
                &[255, 255, 0],
                &[255, 255, 255, 255, 0],
            ],
        );
        test_is_prefix_next_case(
            &[1, 255],
            &[2, 0],
            &[&[], &[1, 255, 0], &[2, 255], &[1, 255], &[2, 0, 0]],
        );
        test_is_prefix_next_case(
            &[0, 255],
            &[1, 0],
            &[&[], &[0, 255, 0], &[1, 255], &[0, 255], &[1, 0, 0]],
        );
        test_is_prefix_next_case(
            &[1, 2, 3, 4, 255, 255],
            &[1, 2, 3, 5, 0, 0],
            &[
                &[],
                &[1, 2, 3, 4, 255, 255],
                &[1, 2, 3, 4, 0, 0],
                &[1, 2, 3, 5, 255, 255],
                &[1, 2, 3, 5, 0, 1],
                &[1, 2, 3, 5, 1, 0],
                &[1, 2, 4, 0, 0, 0],
            ],
        );
    }
}
<commit_msg>copr: modify is_prefix_next fn comment (#6301)<commit_after>// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use kvproto::coprocessor as coppb;
use tipb::ColumnInfo;

use crate::codec::datum::Datum;

/// Convert the key to the smallest key which is larger than the key given.
pub fn convert_to_prefix_next(key: &mut Vec<u8>) {
    if key.is_empty() {
        key.push(0);
        return;
    }
    let mut i = key.len() - 1;

    // Add 1 to the last byte that is not 255, and set it's following bytes to 0.
    loop {
        if key[i] == 255 {
            key[i] = 0;
        } else {
            key[i] += 1;
            return;
        }
        if i == 0 {
            // All bytes are 255. Append a 0 to the key.
            for byte in key.iter_mut() {
                *byte = 255;
            }
            key.push(0);
            return;
        }
        i -= 1;
    }
}

/// Check if `key`'s prefix next equals to `next`.
pub fn is_prefix_next(key: &[u8], next: &[u8]) -> bool {
    let len = key.len();
    let next_len = next.len();

    if len == next_len {
        // Find the last non-255 byte
        let mut carry_pos = len;
        loop {
            if carry_pos == 0 {
                // All bytes of `key` are 255. `next` couldn't be `key`'s prefix_next since their
                // lengths are equal.
                return false;
            }

            carry_pos -= 1;
            if key[carry_pos] != 255 {
                break;
            }
        }

        // Now `carry_pos` is the index of the last byte that is not 255. For example:
        //   key: [1, 2, 3, 255, 255, 255]
        //               ^ carry_pos == 2

        // So they are equal when:
        // * `key`'s value at `carry_pos` is that of `next` - 1 and
        // * `next`'s part after carry_pos is all 0
        // * `key` and `next`'s parts before `carry_pos` are equal.
        // For example:
        //   key:  [1, 2, 3, 255, 255, 255]
        //   next: [1, 2, 4,   0,   0,   0]
        //                ^ carry_pos == 2
        // The part before `carry_pos` is all [1, 2],
        // the bytes at `carry_pos` differs by 1 (4 == 3 + 1), and
        // the remaining bytes of next ([0, 0, 0]) is all 0.
        // so [1, 2, 4, 0, 0, 0] is prefix_next of [1, 2, 3, 255, 255, 255]
        key[carry_pos] + 1 == next[carry_pos]
            && next[carry_pos + 1..].iter().all(|byte| *byte == 0)
            && key[..carry_pos] == next[..carry_pos]
    } else if len + 1 == next_len {
        // `next` must has one more 0 than `key`, and the first `len` bytes must be all 255.
        // The case that `len == 0` is also covered here.
        *next.last().unwrap() == 0
            && key.iter().all(|byte| *byte == 255)
            && next.iter().take(len).all(|byte| *byte == 255)
    } else {
        // Length not match.
        false
    }
}

/// `is_point` checks if the key range represents a point.
#[inline]
pub fn is_point(range: &coppb::KeyRange) -> bool {
    is_prefix_next(range.get_start(), range.get_end())
}

#[inline]
pub fn get_pk(col: &ColumnInfo, h: i64) -> Datum {
    use tidb_query_datatype::{FieldTypeAccessor, FieldTypeFlag};

    if col.as_accessor().flag().contains(FieldTypeFlag::UNSIGNED) {
        // PK column is unsigned
        Datum::U64(h as u64)
    } else {
        Datum::I64(h)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_prefix_next_once(key: &[u8], expected: &[u8]) {
        let mut key = key.to_vec();
        convert_to_prefix_next(&mut key);
        assert_eq!(key.as_slice(), expected);
    }

    #[test]
    fn test_prefix_next() {
        test_prefix_next_once(&[], &[0]);
        test_prefix_next_once(&[0], &[1]);
        test_prefix_next_once(&[1], &[2]);
        test_prefix_next_once(&[255], &[255, 0]);
        test_prefix_next_once(&[255, 255, 255], &[255, 255, 255, 0]);
        test_prefix_next_once(&[1, 255], &[2, 0]);
        test_prefix_next_once(&[0, 1, 255], &[0, 2, 0]);
        test_prefix_next_once(&[0, 1, 255, 5], &[0, 1, 255, 6]);
        test_prefix_next_once(&[0, 1, 5, 255], &[0, 1, 6, 0]);
        test_prefix_next_once(&[0, 1, 255, 255], &[0, 2, 0, 0]);
        test_prefix_next_once(&[0, 255, 255, 255], &[1, 0, 0, 0]);
    }

    fn test_is_prefix_next_case(lhs: &[u8], expected: &[u8], unexpected: &[&[u8]]) {
        assert!(is_prefix_next(lhs, expected));
        for rhs in unexpected {
            assert!(!is_prefix_next(lhs, rhs));
        }
    }

    #[test]
    fn test_is_prefix_next() {
        test_is_prefix_next_case(&[], &[0], &[&[], &[1], &[2]]);
        test_is_prefix_next_case(&[0], &[1], &[&[], &[0], &[0, 0], &[2], &[255]]);
        test_is_prefix_next_case(&[1], &[2], &[&[], &[1], &[3], &[1, 0]]);
        test_is_prefix_next_case(&[255], &[255, 0], &[&[0], &[255, 255, 0]]);
        test_is_prefix_next_case(
            &[255, 255, 255],
            &[255, 255, 255, 0],
            &[
                &[],
                &[0],
                &[0, 0, 0],
                &[255, 255, 0],
                &[255, 255, 255, 255, 0],
            ],
        );
        test_is_prefix_next_case(
            &[1, 255],
            &[2, 0],
            &[&[], &[1, 255, 0], &[2, 255], &[1, 255], &[2, 0, 0]],
        );
        test_is_prefix_next_case(
            &[0, 255],
            &[1, 0],
            &[&[], &[0, 255, 0], &[1, 255], &[0, 255], &[1, 0, 0]],
        );
        test_is_prefix_next_case(
            &[1, 2, 3, 4, 255, 255],
            &[1, 2, 3, 5, 0, 0],
            &[
                &[],
                &[1, 2, 3, 4, 255, 255],
                &[1, 2, 3, 4, 0, 0],
                &[1, 2, 3, 5, 255, 255],
                &[1, 2, 3, 5, 0, 1],
                &[1, 2, 3, 5, 1, 0],
                &[1, 2, 4, 0, 0, 0],
            ],
        );
    }
}
<|endoftext|>


<commit_before>//! Log all requests in a format similar to Heroku's router, but with additional
//! information that we care about like User-Agent

use super::prelude::*;
use crate::util::request_header;

use conduit::{header, RequestExt, StatusCode};
use conduit_cookie::RequestSession;

use crate::middleware::response_timing::ResponseTime;
use std::fmt::{self, Display, Formatter};

const SLOW_REQUEST_THRESHOLD_MS: u64 = 1000;

#[derive(Default)]
pub(super) struct LogRequests();

struct OriginalPath(String);

impl Middleware for LogRequests {
    fn before(&self, req: &mut dyn RequestExt) -> BeforeResult {
        let path = OriginalPath(req.path().to_string());
        req.mut_extensions().insert(path);

        if let Some(request_id) = req
            .headers()
            .get("x-request-id")
            .and_then(|value| value.to_str().ok())
        {
            sentry::configure_scope(|scope| scope.set_tag("request.id", request_id));
        }

        Ok(())
    }

    fn after(&self, req: &mut dyn RequestExt, res: AfterResult) -> AfterResult {
        println!("{}", RequestLine { req, res: &res });

        report_to_sentry(req, &res);

        res
    }
}

struct CustomMetadata {
    entries: Vec<(&'static str, String)>,
}

pub fn add_custom_metadata<V: Display>(req: &mut dyn RequestExt, key: &'static str, value: V) {
    if let Some(metadata) = req.mut_extensions().find_mut::<CustomMetadata>() {
        metadata.entries.push((key, value.to_string()));
    } else {
        let mut metadata = CustomMetadata {
            entries: Vec::new(),
        };
        metadata.entries.push((key, value.to_string()));
        req.mut_extensions().insert(metadata);
    }

    sentry::configure_scope(|scope| scope.set_extra(key, value.to_string().into()));
}

fn report_to_sentry(req: &dyn RequestExt, res: &AfterResult) {
    sentry::configure_scope(|scope| {
        {
            let id = req.session().get("user_id").map(|str| str.to_string());

            let user = sentry::User {
                id,
                ..Default::default()
            };

            scope.set_user(Some(user));
        }

        {
            let status = res
                .as_ref()
                .map(|resp| resp.status())
                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

            scope.set_tag("response.status", status.as_str());
        }

        if let Some(response_time) = req.extensions().find::<ResponseTime>() {
            scope.set_extra("Response time [ms]", response_time.as_millis().into());
        }
    });
}

#[cfg(test)]
pub(crate) fn get_log_message(req: &dyn RequestExt, key: &'static str) -> String {
    // Unwrap shouldn't panic as no other code has access to the private struct to remove it
    for (k, v) in &req.extensions().find::<CustomMetadata>().unwrap().entries {
        if key == *k {
            return v.clone();
        }
    }
    panic!("expected log message for {} not found", key);
}

struct RequestLine<'r> {
    req: &'r dyn RequestExt,
    res: &'r AfterResult,
}

impl Display for RequestLine<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut line = LogLine::new(f);

        let status = self.res.as_ref().map(|res| res.status());
        let status = status.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

        let at = if status.is_server_error() {
            "error"
        } else {
            "info"
        };

        line.add_field("at", at)?;
        line.add_field("method", self.req.method())?;
        line.add_quoted_field("path", FullPath(self.req))?;

        // The request_id is not logged for successful download requests
        if !(self.req.path().ends_with("/download")
            && self
                .res
                .as_ref()
                .ok()
                .map(|ok| ok.status().is_redirection())
                == Some(true))
        {
            line.add_field("request_id", request_header(self.req, "x-request-id"))?;
        }

        line.add_quoted_field("fwd", request_header(self.req, "x-real-ip"))?;

        let response_time = self.req.extensions().find::<ResponseTime>();
        if let Some(response_time) = response_time {
            line.add_field("service", response_time)?;
        }
        line.add_field("status", status.as_str())?;
        line.add_quoted_field("user_agent", request_header(self.req, header::USER_AGENT))?;

        if let Some(metadata) = self.req.extensions().find::<CustomMetadata>() {
            for (key, value) in &metadata.entries {
                line.add_quoted_field(key, value)?;
            }
        }

        if let Err(err) = self.res {
            line.add_quoted_field("error", err)?;
        }

        if let Some(response_time) = response_time {
            if response_time.as_millis() > SLOW_REQUEST_THRESHOLD_MS {
                line.add_marker("SLOW REQUEST")?;
            }
        }

        Ok(())
    }
}

struct FullPath<'a>(&'a dyn RequestExt);

impl<'a> Display for FullPath<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let request = self.0;

        let original_path = request.extensions().find::<OriginalPath>();
        // Unwrap shouldn't panic as no other code has access to the private struct to remove it
        let path = original_path.map(|p| p.0.as_str()).unwrap();

        write!(f, "{}", path)?;

        if let Some(q_string) = request.query_string() {
            write!(f, "?{}", q_string)?;
        }
        Ok(())
    }
}

struct LogLine<'f, 'g> {
    f: &'f mut Formatter<'g>,
    first: bool,
}

impl<'f, 'g> LogLine<'f, 'g> {
    fn new(f: &'f mut Formatter<'g>) -> Self {
        Self { f, first: true }
    }

    fn add_field<K: Display, V: Display>(&mut self, key: K, value: V) -> fmt::Result {
        self.start_item()?;

        key.fmt(self.f)?;
        self.f.write_str("=")?;
        value.fmt(self.f)?;

        Ok(())
    }

    fn add_quoted_field<K: Display, V: Display>(&mut self, key: K, value: V) -> fmt::Result {
        self.start_item()?;

        key.fmt(self.f)?;
        self.f.write_str("=\"")?;
        value.fmt(self.f)?;
        self.f.write_str("\"")?;

        Ok(())
    }

    fn add_marker<M: Display>(&mut self, marker: M) -> fmt::Result {
        self.start_item()?;

        marker.fmt(self.f)?;

        Ok(())
    }

    fn start_item(&mut self) -> fmt::Result {
        if !self.first {
            self.f.write_str(" ")?;
        }
        self.first = false;
        Ok(())
    }
}
<commit_msg>middleware::log_request: Replace `unwrap()` call with `request.path()` fallback<commit_after>//! Log all requests in a format similar to Heroku's router, but with additional
//! information that we care about like User-Agent

use super::prelude::*;
use crate::util::request_header;

use conduit::{header, RequestExt, StatusCode};
use conduit_cookie::RequestSession;

use crate::middleware::response_timing::ResponseTime;
use std::fmt::{self, Display, Formatter};

const SLOW_REQUEST_THRESHOLD_MS: u64 = 1000;

#[derive(Default)]
pub(super) struct LogRequests();

struct OriginalPath(String);

impl Middleware for LogRequests {
    fn before(&self, req: &mut dyn RequestExt) -> BeforeResult {
        let path = OriginalPath(req.path().to_string());
        req.mut_extensions().insert(path);

        if let Some(request_id) = req
            .headers()
            .get("x-request-id")
            .and_then(|value| value.to_str().ok())
        {
            sentry::configure_scope(|scope| scope.set_tag("request.id", request_id));
        }

        Ok(())
    }

    fn after(&self, req: &mut dyn RequestExt, res: AfterResult) -> AfterResult {
        println!("{}", RequestLine { req, res: &res });

        report_to_sentry(req, &res);

        res
    }
}

struct CustomMetadata {
    entries: Vec<(&'static str, String)>,
}

pub fn add_custom_metadata<V: Display>(req: &mut dyn RequestExt, key: &'static str, value: V) {
    if let Some(metadata) = req.mut_extensions().find_mut::<CustomMetadata>() {
        metadata.entries.push((key, value.to_string()));
    } else {
        let mut metadata = CustomMetadata {
            entries: Vec::new(),
        };
        metadata.entries.push((key, value.to_string()));
        req.mut_extensions().insert(metadata);
    }

    sentry::configure_scope(|scope| scope.set_extra(key, value.to_string().into()));
}

fn report_to_sentry(req: &dyn RequestExt, res: &AfterResult) {
    sentry::configure_scope(|scope| {
        {
            let id = req.session().get("user_id").map(|str| str.to_string());

            let user = sentry::User {
                id,
                ..Default::default()
            };

            scope.set_user(Some(user));
        }

        {
            let status = res
                .as_ref()
                .map(|resp| resp.status())
                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

            scope.set_tag("response.status", status.as_str());
        }

        if let Some(response_time) = req.extensions().find::<ResponseTime>() {
            scope.set_extra("Response time [ms]", response_time.as_millis().into());
        }
    });
}

#[cfg(test)]
pub(crate) fn get_log_message(req: &dyn RequestExt, key: &'static str) -> String {
    // Unwrap shouldn't panic as no other code has access to the private struct to remove it
    for (k, v) in &req.extensions().find::<CustomMetadata>().unwrap().entries {
        if key == *k {
            return v.clone();
        }
    }
    panic!("expected log message for {} not found", key);
}

struct RequestLine<'r> {
    req: &'r dyn RequestExt,
    res: &'r AfterResult,
}

impl Display for RequestLine<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut line = LogLine::new(f);

        let status = self.res.as_ref().map(|res| res.status());
        let status = status.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

        let at = if status.is_server_error() {
            "error"
        } else {
            "info"
        };

        line.add_field("at", at)?;
        line.add_field("method", self.req.method())?;
        line.add_quoted_field("path", FullPath(self.req))?;

        // The request_id is not logged for successful download requests
        if !(self.req.path().ends_with("/download")
            && self
                .res
                .as_ref()
                .ok()
                .map(|ok| ok.status().is_redirection())
                == Some(true))
        {
            line.add_field("request_id", request_header(self.req, "x-request-id"))?;
        }

        line.add_quoted_field("fwd", request_header(self.req, "x-real-ip"))?;

        let response_time = self.req.extensions().find::<ResponseTime>();
        if let Some(response_time) = response_time {
            line.add_field("service", response_time)?;
        }
        line.add_field("status", status.as_str())?;
        line.add_quoted_field("user_agent", request_header(self.req, header::USER_AGENT))?;

        if let Some(metadata) = self.req.extensions().find::<CustomMetadata>() {
            for (key, value) in &metadata.entries {
                line.add_quoted_field(key, value)?;
            }
        }

        if let Err(err) = self.res {
            line.add_quoted_field("error", err)?;
        }

        if let Some(response_time) = response_time {
            if response_time.as_millis() > SLOW_REQUEST_THRESHOLD_MS {
                line.add_marker("SLOW REQUEST")?;
            }
        }

        Ok(())
    }
}

struct FullPath<'a>(&'a dyn RequestExt);

impl<'a> Display for FullPath<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let request = self.0;

        let original_path = request.extensions().find::<OriginalPath>();
        let path = original_path
            .map(|p| p.0.as_str())
            .unwrap_or_else(|| request.path());

        write!(f, "{}", path)?;

        if let Some(q_string) = request.query_string() {
            write!(f, "?{}", q_string)?;
        }
        Ok(())
    }
}

struct LogLine<'f, 'g> {
    f: &'f mut Formatter<'g>,
    first: bool,
}

impl<'f, 'g> LogLine<'f, 'g> {
    fn new(f: &'f mut Formatter<'g>) -> Self {
        Self { f, first: true }
    }

    fn add_field<K: Display, V: Display>(&mut self, key: K, value: V) -> fmt::Result {
        self.start_item()?;

        key.fmt(self.f)?;
        self.f.write_str("=")?;
        value.fmt(self.f)?;

        Ok(())
    }

    fn add_quoted_field<K: Display, V: Display>(&mut self, key: K, value: V) -> fmt::Result {
        self.start_item()?;

        key.fmt(self.f)?;
        self.f.write_str("=\"")?;
        value.fmt(self.f)?;
        self.f.write_str("\"")?;

        Ok(())
    }

    fn add_marker<M: Display>(&mut self, marker: M) -> fmt::Result {
        self.start_item()?;

        marker.fmt(self.f)?;

        Ok(())
    }

    fn start_item(&mut self) -> fmt::Result {
        if !self.first {
            self.f.write_str(" ")?;
        }
        self.first = false;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::net::{SocketAddr,IpAddr};
use std::io::Write;
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;
use pool::{Pool,Checkout,Reset};
use time::{Duration, precise_time_s, precise_time_ns};
use uuid::Uuid;
use network::{ClientResult,Protocol};
use network::buffer_queue::BufferQueue;
use network::session::{Readiness,SessionMetrics};
use network::socket::{SocketHandler,SocketResult};
use nom::HexDisplay;

type BackendToken = Token;

#[derive(PartialEq)]
pub enum ClientStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            TcpStream,
  token:              Option<Token>,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<BufferQueue>,
  back_buf:           Checkout<BufferQueue>,
  front_buf_position: usize,
  back_buf_position:  usize,
  pub app_id:         Option<String>,
  pub request_id:     String,
  pub readiness:      Readiness,
  pub log_ctx:        String,
  public_address:     Option<IpAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(frontend: Front, backend: TcpStream, front_buf: Checkout<BufferQueue>, back_buf: Checkout<BufferQueue>, public_address: Option<IpAddr>) -> Option<Pipe<Front>> {
    let request_id = Uuid::new_v4().hyphenated().to_string();
    let log_ctx    = format!("{}\tunknown\t", &request_id);
    let client = Pipe {
      frontend:           frontend,
      backend:            backend,
      token:              None,
      backend_token:      None,
      front_buf:          front_buf,
      back_buf:           back_buf,
      front_buf_position: 0,
      back_buf_position:  0,
      app_id:             None,
      request_id:         request_id,
      readiness:          Readiness {
                            front_interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            back_interest:   UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            front_readiness: UnixReady::from(Ready::empty()),
                            back_readiness:  UnixReady::from(Ready::empty()),
      },
      log_ctx:            log_ctx,
      public_address:     public_address,
    };

    trace!("created pipe");
    Some(client)
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(front) = self.token {
      if let Some(back) = self.backend_token {
        return Some((front, back))
      }
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    Some(&self.backend)
  }

  pub fn front_token(&self)  -> Option<Token> {
    self.token
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn log_context(&self) -> String {
    if let Some(ref app_id) = self.app_id {
      format!("{}\t{}\t", self.request_id, app_id)
    } else {
      format!("{}\tunknown\t", self.request_id)
    }
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.token         = Some(token);
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  pub fn front_hup(&mut self) -> ClientResult {
    if self.backend_token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::CloseBoth
    }
  }

  pub fn back_hup(&mut self) -> ClientResult {
    ClientResult::CloseBoth
  }

  // Read content from the client
  pub fn readable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe readable");
    if self.front_buf.buffer.available_space() == 0 {
      self.readiness.front_interest.remove(Ready::readable());
      self.readiness.back_interest.insert(Ready::writable());
      return ClientResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.buffer.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.token, sz);

    if sz > 0 {
      //FIXME: replace with copy()
      self.front_buf.buffer.fill(sz);
      self.front_buf.sliced_input(sz);
      self.front_buf.consume_parsed_data(sz);
      self.front_buf.slice_output(sz);

      count!("bytes_in", sz as i64);
      metrics.bin += sz;

      if self.front_buf.buffer.available_space() == 0 {
        self.readiness.front_interest.remove(Ready::readable());
      }
      self.readiness.back_interest.insert(Ready::writable());
    } else {
      self.readiness.front_readiness.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] front socket error, closing the connection", self.log_ctx, self.token);
        metrics.service_stop();
        incr_ereq!();
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.readiness.back_interest.insert(Ready::writable());
    ClientResult::Continue
  }

  // Forward content to client
  pub fn writable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe writable");
    if self.back_buf.output_data_size() == 0 || self.back_buf.next_output_data().len() == 0 {
      self.readiness.back_interest.insert(Ready::readable());
      self.readiness.front_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue && self.back_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.back_buf.next_output_data().len() == 0 {
        self.readiness.back_interest.insert(Ready::readable());
        self.readiness.front_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.next_output_data());
      res = current_res;
      self.back_buf.consume_output_data(current_sz);
      self.back_buf_position += current_sz;
      sz += current_sz;
    }

    if sz > 0 {
      self.readiness.back_interest.insert(Ready::readable());
      metrics.bout += sz;
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}, buffer position {} restart position {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.output_data_size(),
        self.back_buf.buffer_position, self.back_buf.start_parsing_position);
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.token);
        incr_ereq!();
        metrics.service_stop();
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    ClientResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe back_writable");
    if self.front_buf.output_data_size() == 0 || self.front_buf.next_output_data().len() == 0 {
      self.readiness.front_interest.insert(Ready::readable());
      self.readiness.back_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let tokens = self.tokens().clone();
    let output_size = self.front_buf.output_data_size();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    while socket_res == SocketResult::Continue && self.front_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.front_buf.next_output_data().len() == 0 {
        self.readiness.front_interest.insert(Ready::readable());
        self.readiness.back_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.backend.socket_write(self.front_buf.next_output_data());
      socket_res = current_res;
      self.front_buf.consume_output_data(current_sz);
      self.front_buf_position += current_sz;
      sz += current_sz;
    }

    metrics.backend_bout += sz;

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        metrics.service_stop();
        incr_ereq!();
        self.readiness.reset();
        return ClientResult::CloseBoth;
      },
      SocketResult::WouldBlock => {
        self.readiness.back_readiness.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    ClientResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe back_readable");
    if self.back_buf.buffer.available_space() == 0 {
      self.readiness.back_interest.remove(Ready::readable());
      return ClientResult::Continue;
    }

    let tokens     = self.tokens().clone();
    let (sz, r) = self.backend.socket_read(&mut self.back_buf.buffer.space());
    self.back_buf.buffer.fill(sz);
    self.back_buf.sliced_input(sz);
    self.back_buf.consume_parsed_data(sz);
    self.back_buf.slice_output(sz);

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
    }

    if r != SocketResult::Continue || sz == 0 {
      self.readiness.back_readiness.remove(Ready::readable());
    }
    if sz > 0 {
      self.readiness.front_interest.insert(Ready::writable());
      metrics.backend_bin += sz;
    }

    if r == SocketResult::Error {
      error!("{}\tback socket read error, closing connection", self.log_ctx);
      metrics.service_stop();
      incr_ereq!();
      self.readiness.reset();
      return ClientResult::CloseBoth;
    }

    ClientResult::Continue
  }
}

<commit_msg>Change backend socket in Pipe to be optional and now, creating a Pipe just return a Client instead of Option<Client><commit_after>use std::cmp::min;
use std::net::{SocketAddr,IpAddr};
use std::io::Write;
use mio::*;
use mio::tcp::TcpStream;
use mio::unix::UnixReady;
use pool::{Pool,Checkout,Reset};
use time::{Duration, precise_time_s, precise_time_ns};
use uuid::Uuid;
use network::{ClientResult,Protocol};
use network::buffer_queue::BufferQueue;
use network::session::{Readiness,SessionMetrics};
use network::socket::{SocketHandler,SocketResult};
use nom::HexDisplay;

type BackendToken = Token;

#[derive(PartialEq)]
pub enum ClientStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            Option<TcpStream>,
  token:              Option<Token>,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<BufferQueue>,
  back_buf:           Checkout<BufferQueue>,
  front_buf_position: usize,
  back_buf_position:  usize,
  pub app_id:         Option<String>,
  pub request_id:     String,
  pub readiness:      Readiness,
  pub log_ctx:        String,
  public_address:     Option<IpAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(frontend: Front, backend: Option<TcpStream>, front_buf: Checkout<BufferQueue>, back_buf: Checkout<BufferQueue>, public_address: Option<IpAddr>) -> Pipe<Front> {
    let request_id = Uuid::new_v4().hyphenated().to_string();
    let log_ctx    = format!("{}\tunknown\t", &request_id);
    let client = Pipe {
      frontend:           frontend,
      backend:            backend,
      token:              None,
      backend_token:      None,
      front_buf:          front_buf,
      back_buf:           back_buf,
      front_buf_position: 0,
      back_buf_position:  0,
      app_id:             None,
      request_id:         request_id,
      readiness:          Readiness {
                            front_interest:  UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            back_interest:   UnixReady::from(Ready::readable() | Ready::writable()) | UnixReady::hup() | UnixReady::error(),
                            front_readiness: UnixReady::from(Ready::empty()),
                            back_readiness:  UnixReady::from(Ready::empty()),
      },
      log_ctx:            log_ctx,
      public_address:     public_address,
    };

    trace!("created pipe");
    client
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(front) = self.token {
      if let Some(back) = self.backend_token {
        return Some((front, back))
      }
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    self.backend.as_ref()
  }

  pub fn set_back_socket(&mut self, socket: TcpStream) {
    self.backend = Some(socket);
  }

  pub fn front_token(&self)  -> Option<Token> {
    self.token
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn log_context(&self) -> String {
    if let Some(ref app_id) = self.app_id {
      format!("{}\t{}\t", self.request_id, app_id)
    } else {
      format!("{}\tunknown\t", self.request_id)
    }
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.token         = Some(token);
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  pub fn front_hup(&mut self) -> ClientResult {
    if self.backend_token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::CloseBoth
    }
  }

  pub fn back_hup(&mut self) -> ClientResult {
    ClientResult::CloseBoth
  }

  // Read content from the client
  pub fn readable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe readable");
    if self.front_buf.buffer.available_space() == 0 {
      self.readiness.front_interest.remove(Ready::readable());
      self.readiness.back_interest.insert(Ready::writable());
      return ClientResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.buffer.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.token, sz);

    if sz > 0 {
      //FIXME: replace with copy()
      self.front_buf.buffer.fill(sz);
      self.front_buf.sliced_input(sz);
      self.front_buf.consume_parsed_data(sz);
      self.front_buf.slice_output(sz);

      count!("bytes_in", sz as i64);
      metrics.bin += sz;

      if self.front_buf.buffer.available_space() == 0 {
        self.readiness.front_interest.remove(Ready::readable());
      }
      self.readiness.back_interest.insert(Ready::writable());
    } else {
      self.readiness.front_readiness.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] front socket error, closing the connection", self.log_ctx, self.token);
        metrics.service_stop();
        incr_ereq!();
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.readiness.back_interest.insert(Ready::writable());
    ClientResult::Continue
  }

  // Forward content to client
  pub fn writable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe writable");
    if self.back_buf.output_data_size() == 0 || self.back_buf.next_output_data().len() == 0 {
      self.readiness.back_interest.insert(Ready::readable());
      self.readiness.front_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue && self.back_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.back_buf.next_output_data().len() == 0 {
        self.readiness.back_interest.insert(Ready::readable());
        self.readiness.front_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.next_output_data());
      res = current_res;
      self.back_buf.consume_output_data(current_sz);
      self.back_buf_position += current_sz;
      sz += current_sz;
    }

    if sz > 0 {
      self.readiness.back_interest.insert(Ready::readable());
      metrics.bout += sz;
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}, buffer position {} restart position {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.output_data_size(),
        self.back_buf.buffer_position, self.back_buf.start_parsing_position);
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.token);
        incr_ereq!();
        metrics.service_stop();
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    ClientResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe back_writable");
    if self.front_buf.output_data_size() == 0 || self.front_buf.next_output_data().len() == 0 {
      self.readiness.front_interest.insert(Ready::readable());
      self.readiness.back_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let tokens = self.tokens().clone();
    let output_size = self.front_buf.output_data_size();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    while socket_res == SocketResult::Continue && self.front_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.front_buf.next_output_data().len() == 0 {
        self.readiness.front_interest.insert(Ready::readable());
        self.readiness.back_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }

      if let Some(ref mut backend) = self.backend {
        let (current_sz, current_res) = backend.socket_write(self.front_buf.next_output_data());
        socket_res = current_res;
        self.front_buf.consume_output_data(current_sz);
        self.front_buf_position += current_sz;
        sz += current_sz;
      }
    }

    metrics.backend_bout += sz;

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        metrics.service_stop();
        incr_ereq!();
        self.readiness.reset();
        return ClientResult::CloseBoth;
      },
      SocketResult::WouldBlock => {
        self.readiness.back_readiness.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    ClientResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self, metrics: &mut SessionMetrics) -> ClientResult {
    trace!("pipe back_readable");
    if self.back_buf.buffer.available_space() == 0 {
      self.readiness.back_interest.remove(Ready::readable());
      return ClientResult::Continue;
    }

    let tokens     = self.tokens().clone();

    if let Some(ref mut backend) = self.backend {
      let (sz, r) = backend.socket_read(&mut self.back_buf.buffer.space());
      self.back_buf.buffer.fill(sz);
      self.back_buf.sliced_input(sz);
      self.back_buf.consume_parsed_data(sz);
      self.back_buf.slice_output(sz);

      if let Some((front,back)) = tokens {
        debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
      }

      if r != SocketResult::Continue || sz == 0 {
        self.readiness.back_readiness.remove(Ready::readable());
      }
      if sz > 0 {
        self.readiness.front_interest.insert(Ready::writable());
        metrics.backend_bin += sz;
      }

      if r == SocketResult::Error {
        error!("{}\tback socket read error, closing connection", self.log_ctx);
        metrics.service_stop();
        incr_ereq!();
        self.readiness.reset();
        return ClientResult::CloseBoth;
      }
    }

    ClientResult::Continue
  }
}

<|endoftext|>


<commit_before>extern crate r2d2;
extern crate time;

use std::default::Default;
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::thread;
use time::Duration;

mod config;

#[derive(Debug, PartialEq)]
struct FakeConnection;

struct OkManager;

impl r2d2::ConnectionManager for OkManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        Ok(FakeConnection)
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl r2d2::ConnectionManager for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection)
        } else {
            Err(())
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = r2d2::Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).unwrap();
    let mut conns = vec![];
    for _ in 0..config.pool_size() {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = r2d2::Config::builder().pool_size(2).build();
    let pool = r2d2::Pool::new(config, OkManager, Box::new(r2d2::NoopErrorHandler)).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send+Sync>() {}
    is_send_sync::<r2d2::Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl r2d2::ConnectionManager for BlockingChecker {
        type Connection = FakeConnection;
        type Error = ();

        fn connect(&self) -> Result<FakeConnection, ()> {
            Ok(FakeConnection)
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

    let (s1, r1) = mpsc::sync_channel(0);
    let (s2, r2) = mpsc::sync_channel(0);

    let config = r2d2::Config::builder()
        .test_on_check_out(true)
        .pool_size(2)
        .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).unwrap());

    let p2 = pool.clone();
    let _t = thread::scoped(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl r2d2::ConnectionManager for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = r2d2::Pool::new(Default::default(), Handler, Box::new(r2d2::NoopErrorHandler)).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = r2d2::Config::builder()
        .connection_timeout(Duration::seconds(1))
        .build();
    let manager = NthConnectFailManager {
        n: Mutex::new(0),
    };
    r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).err().unwrap();
}

#[test]
fn test_get_timeout() {
    let config = r2d2::Config::builder()
        .pool_size(1)
        .connection_timeout(Duration::seconds(1))
        .build();
    let pool = r2d2::Pool::new(config, OkManager, Box::new(r2d2::NoopErrorHandler)).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}
<commit_msg>use spawn instead of scoped in test<commit_after>extern crate r2d2;
extern crate time;

use std::default::Default;
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::thread;
use time::Duration;

mod config;

#[derive(Debug, PartialEq)]
struct FakeConnection;

struct OkManager;

impl r2d2::ConnectionManager for OkManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        Ok(FakeConnection)
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl r2d2::ConnectionManager for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection)
        } else {
            Err(())
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = r2d2::Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).unwrap();
    let mut conns = vec![];
    for _ in 0..config.pool_size() {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = r2d2::Config::builder().pool_size(2).build();
    let pool = r2d2::Pool::new(config, OkManager, Box::new(r2d2::NoopErrorHandler)).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send+Sync>() {}
    is_send_sync::<r2d2::Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl r2d2::ConnectionManager for BlockingChecker {
        type Connection = FakeConnection;
        type Error = ();

        fn connect(&self) -> Result<FakeConnection, ()> {
            Ok(FakeConnection)
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

    let (s1, r1) = mpsc::sync_channel(0);
    let (s2, r2) = mpsc::sync_channel(0);

    let config = r2d2::Config::builder()
        .test_on_check_out(true)
        .pool_size(2)
        .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).unwrap());

    let p2 = pool.clone();
    let t = thread::spawn(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();

    t.join().ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl r2d2::ConnectionManager for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = r2d2::Pool::new(Default::default(), Handler, Box::new(r2d2::NoopErrorHandler)).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = r2d2::Config::builder()
        .connection_timeout(Duration::seconds(1))
        .build();
    let manager = NthConnectFailManager {
        n: Mutex::new(0),
    };
    r2d2::Pool::new(config, manager, Box::new(r2d2::NoopErrorHandler)).err().unwrap();
}

#[test]
fn test_get_timeout() {
    let config = r2d2::Config::builder()
        .pool_size(1)
        .connection_timeout(Duration::seconds(1))
        .build();
    let pool = r2d2::Pool::new(config, OkManager, Box::new(r2d2::NoopErrorHandler)).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}
<|endoftext|>


<commit_before>use game::{Position, Move, ZobristHash, NumPlies};

use negamax::Bound;

#[derive(Clone, Debug)]
pub struct Data {
    pub hash: ZobristHash,
    pub depth: NumPlies,
    pub bound: Bound,
    pub best_move_opt: Option<Move>,
}

pub struct TranspositionTable(Vec<Option<Data>>);
impl TranspositionTable {
    pub fn with_capacity(capacity: usize) -> Self {
        TranspositionTable(vec![None; capacity])
    }
    pub fn get<'a>(&'a self, pos: &Position) -> Option<&'a Data> {
        let hash = pos.hash();
        let idx = (hash.0 % (self.0.len() as u64)) as usize;
        self.0[idx].as_ref().and_then(|x| if x.hash == hash { Some(x) } else { None })
    }
    // TODO implement a better replacement scheme
    pub fn set(&mut self, pos: &Position, depth: NumPlies, best_move_opt: Option<Move>, bound: Bound) {
        let hash = pos.hash();
        let idx = (hash.0 % (self.0.len() as u64)) as usize;
        self.0[idx] = Some(Data { hash: hash, depth: depth, best_move_opt: best_move_opt, bound: bound });
    }
}
<commit_msg>Check if best move in transposition table is valid<commit_after>use game::{Position, Move, ZobristHash, NumPlies};

use negamax::Bound;

#[derive(Clone, Debug)]
pub struct Data {
    pub hash: ZobristHash,
    pub depth: NumPlies,
    pub bound: Bound,
    pub best_move_opt: Option<Move>,
}

pub struct TranspositionTable(Vec<Option<Data>>);
impl TranspositionTable {
    pub fn with_capacity(capacity: usize) -> Self {
        TranspositionTable(vec![None; capacity])
    }
    pub fn get<'a>(&'a self, pos: &Position) -> Option<&'a Data> {
        let hash = pos.hash();
        let idx = (hash.0 % (self.0.len() as u64)) as usize;
        self.0[idx].as_ref().and_then(|x| {
            let is_correct_pos = x.hash == hash &&
                                 x.best_move_opt.as_ref()
                                  .map_or(true, |y| pos.legal_iter().any(|z| *y == z));
            if is_correct_pos { Some(x) } else { None }
        })
    }
    // TODO implement a better replacement scheme
    pub fn set(&mut self, pos: &Position, depth: NumPlies, best_move_opt: Option<Move>, bound: Bound) {
        let hash = pos.hash();
        let idx = (hash.0 % (self.0.len() as u64)) as usize;
        self.0[idx] = Some(Data { hash: hash, depth: depth, best_move_opt: best_move_opt, bound: bound });
    }
}
<|endoftext|>