<commit_before>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    _child: &Child,
    _signal: InternalSignal,
    _send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {task_id} has already finished by itself");
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {target_pid} snapShot");
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {target_pid} snapShot");
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!("Failed to suspend thread {tid} with error code {err_code}");
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!("Failed to resume thread {tid} with error code {err_code}");
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!("Failed to terminate process {pid} with error code {err_code}");
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    /// A test helper function, which ensures that a specific amount of subprocesses can be
    /// observed for a given PID in a given time window.
    /// If the correct amount can be observed, the process ids are then returned.
    ///
    /// The process count is checked every few milliseconds for the given duration.
    fn assert_process_ids(pid: u32, expected_processes: usize, millis: usize) -> Result<Vec<u32>> {
        // Check every 50 milliseconds.
        let interval = 50;
        let tries = millis / interval;
        let mut current_try = 0;

        while current_try <= tries {
            // Continue waiting if the count doesn't match.
            let process_ids = get_cur_task_processes(pid);
            if process_ids.len() != expected_processes {
                current_try += 1;
                sleep(Duration::from_millis(interval as u64));
                continue;
            }

            return Ok(process_ids);
        }

        let count = get_cur_task_processes(pid).len();
        bail!("{expected_processes} processes were expected. Last process count was {count}")
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[ignore]
    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    ///
    /// This test is ignored for now, as it is flaky from time to time.
    /// See https://github.com/Nukesor/pueue/issues/315
    fn test_shell_command_is_killed() -> Result<()> {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() -> Result<()> {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 2, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() -> Result<()> {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let _ = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));

        Ok(())
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() -> Result<()> {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }
}
<commit_msg>CI: Ignore flaky windows process tests<commit_after>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    _child: &Child,
    _signal: InternalSignal,
    _send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {task_id} has already finished by itself");
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {target_pid} snapShot");
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {target_pid} snapShot");
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!("Failed to suspend thread {tid} with error code {err_code}");
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!("Failed to resume thread {tid} with error code {err_code}");
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!("Failed to terminate process {pid} with error code {err_code}");
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    /// A test helper function, which ensures that a specific amount of subprocesses can be
    /// observed for a given PID in a given time window.
    /// If the correct amount can be observed, the process ids are then returned.
    ///
    /// The process count is checked every few milliseconds for the given duration.
    fn assert_process_ids(pid: u32, expected_processes: usize, millis: usize) -> Result<Vec<u32>> {
        // Check every 50 milliseconds.
        let interval = 50;
        let tries = millis / interval;
        let mut current_try = 0;

        while current_try <= tries {
            // Continue waiting if the count doesn't match.
            let process_ids = get_cur_task_processes(pid);
            if process_ids.len() != expected_processes {
                current_try += 1;
                sleep(Duration::from_millis(interval as u64));
                continue;
            }

            return Ok(process_ids);
        }

        let count = get_cur_task_processes(pid).len();
        bail!("{expected_processes} processes were expected. Last process count was {count}")
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[ignore]
    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    ///
    /// This test is ignored for now, as it is flaky from time to time.
    /// See https://github.com/Nukesor/pueue/issues/315
    fn test_shell_command_is_killed() -> Result<()> {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[ignore]
    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() -> Result<()> {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 2, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[ignore]
    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() -> Result<()> {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let _ = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));

        Ok(())
    }

    #[ignore]
    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() -> Result<()> {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate roboime_next;
#[macro_use] extern crate glium;
extern crate clock_ticks;
//extern crate image;
extern crate env_logger;
extern crate log;

use std::error::Error;
use std::process::exit;
//use std::io::Cursor;
use glium::Surface;
use glium::glutin;
//use image::image::GenericImage;
//use roboime_next::prelude::*;
use roboime_next::game;
use ::utils::*;
use ::models::TeamSide;

pub mod models;
pub mod colors;
mod utils;
mod draw;

fn main_loop() -> Result<(), Box<Error>> {
    use std::env;
    use std::thread;
    use std::time::Duration;
    use glium::DisplayBuild;
    use glium::glutin::GlRequest;
    use log::LogLevelFilter;
    use env_logger::LogBuilder;

    {
        let mut builder = LogBuilder::new();
        builder.format(|record| {
            format!("{} [{}] {}", record.level(), record.location().module_path(), record.args())
        });
        builder.filter(Some("roboime"), LogLevelFilter::Info);
        if env::var("RUST_LOG").is_ok() {
           builder.parse(&env::var("RUST_LOG").unwrap());
        }
        try!(builder.init());
    }

    // building the display, ie. the main object
    let display = {
        // 1 extra pixel to align the middle lines to the screen
        // TODO: refactor: deduplicate options
        if let Ok(display) = glutin::WindowBuilder::new()
            .with_title(format!("RoboIME Next"))
            .with_dimensions(1041, 741)
            .with_depth_buffer(24)
            .with_gl(GlRequest::Latest)
            .with_multisampling(4)
            .with_srgb(Some(true))
            .with_vsync()
            .build_glium() {
            display
        } else {
            try!(glutin::WindowBuilder::new()
                .with_title(format!("RoboIME Next"))
                .with_dimensions(1041, 741)
                .with_depth_buffer(24)
                .with_gl(GlRequest::Latest)
                .with_multisampling(4)
                .with_vsync()
                .build_glium())
        }
    };

    let team_side = TeamSide::BlueIsLeft;
    let mut draw_game = try!(draw::Game::new(&display));
    draw_game.team_side(&display, team_side);

    let mut accumulator = 0;
    let mut previous_clock = clock_ticks::precise_time_ns();
    let game_state = game::SharedState::new();

    // add some robots
    {
        let mut game_state = try!(game_state.write());
        add_initial_robots(game_state.get_robots_blue_mut(), 6, false);
        add_initial_robots(game_state.get_robots_yellow_mut(), 6, true);
    }

    // the main loop
    'main: loop {

        let mut target = display.draw();

        let view = view_matrix(&[0.0, 0.0, 10.0], &[0.0, 0.0, -1.0], &[0.0, 1.0, 0.0]);
        //let view = view_matrix(&[-4.0, -4.0, 4.0], &[1.0, 1.0, -1.0], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[4.0, -4.0, 4.0], &[-1.0, 1.0, -1.0], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[-0.2, -0.2, 0.6], &[0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[1.0, 2.0, 3.0], &[1.0, 1.0, -3.0], &[0.0, 1.0, 0.0]);
        //let view = view_matrix(&[-4.0, 0.7, 1.6], &[-0.2, -0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[-5.0, -0.7, 1.6], &[0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[4.0, 0.7, 1.6], &[0.2, -0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[5.0, -0.7, 1.6], &[-0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);

        {
            let state = try!(game_state.read());
            try!(draw_game.draw_to(&mut target, &state, view));
        }

        try!(target.finish());

        // polling and handling the events received by the window
        for event in display.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) |
                glutin::Event::Closed => break 'main,
                _ => ()
            }
        }

        let now = clock_ticks::precise_time_ns();
        accumulator += now - previous_clock;
        previous_clock = now;

        const FIXED_TIME_STAMP: u64 = 16666667;
        while accumulator >= FIXED_TIME_STAMP {
            accumulator -= FIXED_TIME_STAMP;

            // if you have a game, update the state here
        }

        thread::sleep(Duration::from_millis(((FIXED_TIME_STAMP - accumulator) / 1000000) as u64));
    }

    Ok(())
}

fn main() {
    if let Err(err) = main_loop() {
        use std::ops::Deref;

        println!("Error: {}.", err.description());

        let mut err: &Error = err.deref();
        while let Some(cause) = err.cause() {
            println!("- caused by {}", cause.description());
            err = cause;
        }

        exit(1);
    }
}
<commit_msg>No need to require srgb afer all.<commit_after>extern crate roboime_next;
#[macro_use] extern crate glium;
extern crate clock_ticks;
//extern crate image;
extern crate env_logger;
extern crate log;

use std::error::Error;
use std::process::exit;
//use std::io::Cursor;
use glium::Surface;
use glium::glutin;
//use image::image::GenericImage;
//use roboime_next::prelude::*;
use roboime_next::game;
use ::utils::*;
use ::models::TeamSide;

pub mod models;
pub mod colors;
mod utils;
mod draw;

fn main_loop() -> Result<(), Box<Error>> {
    use std::env;
    use std::thread;
    use std::time::Duration;
    use glium::DisplayBuild;
    use glium::glutin::GlRequest;
    use log::LogLevelFilter;
    use env_logger::LogBuilder;

    {
        let mut builder = LogBuilder::new();
        builder.format(|record| {
            format!("{} [{}] {}", record.level(), record.location().module_path(), record.args())
        });
        builder.filter(Some("roboime"), LogLevelFilter::Info);
        if env::var("RUST_LOG").is_ok() {
           builder.parse(&env::var("RUST_LOG").unwrap());
        }
        try!(builder.init());
    }

    // building the display, ie. the main object
    let display = try!(glutin::WindowBuilder::new()
        .with_title(format!("RoboIME Next"))
        // 1 extra pixel to align the middle lines to the screen
        .with_dimensions(1041, 741)
        .with_depth_buffer(24)
        .with_gl(GlRequest::Latest)
        .with_multisampling(4)
        .with_vsync()
        .build_glium());

    let team_side = TeamSide::BlueIsLeft;
    let mut draw_game = try!(draw::Game::new(&display));
    draw_game.team_side(&display, team_side);

    let mut accumulator = 0;
    let mut previous_clock = clock_ticks::precise_time_ns();
    let game_state = game::SharedState::new();

    // add some robots
    {
        let mut game_state = try!(game_state.write());
        add_initial_robots(game_state.get_robots_blue_mut(), 6, false);
        add_initial_robots(game_state.get_robots_yellow_mut(), 6, true);
    }

    // the main loop
    'main: loop {

        let mut target = display.draw();

        let view = view_matrix(&[0.0, 0.0, 10.0], &[0.0, 0.0, -1.0], &[0.0, 1.0, 0.0]);
        //let view = view_matrix(&[-4.0, -4.0, 4.0], &[1.0, 1.0, -1.0], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[4.0, -4.0, 4.0], &[-1.0, 1.0, -1.0], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[-0.2, -0.2, 0.6], &[0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[1.0, 2.0, 3.0], &[1.0, 1.0, -3.0], &[0.0, 1.0, 0.0]);
        //let view = view_matrix(&[-4.0, 0.7, 1.6], &[-0.2, -0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[-5.0, -0.7, 1.6], &[0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[4.0, 0.7, 1.6], &[0.2, -0.2, -0.6], &[0.0, 0.0, 1.0]);
        //let view = view_matrix(&[5.0, -0.7, 1.6], &[-0.2, 0.2, -0.6], &[0.0, 0.0, 1.0]);

        {
            let state = try!(game_state.read());
            try!(draw_game.draw_to(&mut target, &state, view));
        }

        try!(target.finish());

        // polling and handling the events received by the window
        for event in display.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) |
                glutin::Event::Closed => break 'main,
                _ => ()
            }
        }

        let now = clock_ticks::precise_time_ns();
        accumulator += now - previous_clock;
        previous_clock = now;

        const FIXED_TIME_STAMP: u64 = 16666667;
        while accumulator >= FIXED_TIME_STAMP {
            accumulator -= FIXED_TIME_STAMP;

            // if you have a game, update the state here
        }

        thread::sleep(Duration::from_millis(((FIXED_TIME_STAMP - accumulator) / 1000000) as u64));
    }

    Ok(())
}

fn main() {
    if let Err(err) = main_loop() {
        use std::ops::Deref;

        println!("Error: {}.", err.description());

        let mut err: &Error = err.deref();
        while let Some(cause) = err.cause() {
            println!("- caused by {}", cause.description());
            err = cause;
        }

        exit(1);
    }
}
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

use extra::json;
use std::hashmap::HashMap;
use serialize::Decodable;
use glfw;
use gl;
use gl::types::{
  GLfloat,
  GLint,
  GLuint,
};
use cgmath::vector::{
  Vec3,
  Vec2,
};
use glh = gl_helpers;
use camera::Camera;
use map::TileIterator;
use geom::Geom;
use tile_picker::TilePicker;
use obj;
use mesh::Mesh;
use misc::read_file;

fn build_hex_mesh(&geom: &Geom) -> ~[Vec3<GLfloat>] {
  let mut vertex_data = ~[];
  for tile_pos in TileIterator::new() {
    let pos = geom.map_pos_to_world_pos(tile_pos);
    for num in range(0, 6) {
      let vertex = geom.index_to_hex_vertex(num);
      let next_vertex = geom.index_to_hex_vertex(num + 1);
      vertex_data.push(pos + vertex);
      vertex_data.push(pos + next_vertex);
      vertex_data.push(pos + Vec3::zero());
    }
  }
  vertex_data
}

fn build_hex_tex_coord() -> ~[Vec2<GLfloat>] {
  let mut vertex_data = ~[];
  for _ in TileIterator::new() {
    for _ in range(0, 6) {
      vertex_data.push(Vec2{x: 0.0, y: 0.0});
      vertex_data.push(Vec2{x: 1.0, y: 0.0});
      vertex_data.push(Vec2{x: 0.5, y: 0.5});
    }
  }
  vertex_data
}

pub struct SceneNode {
  pos: Vec3<f32>,
}

#[deriving(Decodable)]
struct Size2<T> {
  x: T,
  y: T,
}

fn read_win_size(config_path: &str) -> Vec2<int> {
  let path = Path::new(config_path);
  let json = json::from_str(read_file(&path)).unwrap();
  let mut decoder = json::Decoder::new(json);
  let size: Size2<int> = Decodable::decode(&mut decoder);
  Vec2{x: size.x, y: size.y}
}

fn init_win(win_size: Vec2<int>) -> glfw::Window {
  glfw::set_error_callback(~glfw::LogErrorHandler);
  let init_status = glfw::init();
  if !init_status.is_ok() {
    fail!("Failed to initialize GLFW");
  }
  let win = glfw::Window::create(
    win_size.x as u32,
    win_size.y as u32,
    "Marauder",
    glfw::Windowed,
  ).unwrap();
  win.make_context_current();
  win.set_all_polling(true);
  win
}

pub struct Visualizer {
  program: GLuint,
  map_mesh: Mesh,
  unit_mesh: Mesh,
  mat_id: GLint,
  win: Option<glfw::Window>,
  mouse_pos: Vec2<f32>,
  camera: Camera,
  picker: TilePicker,
  selected_tile_pos: Option<Vec2<i32>>,
  geom: Geom,
  scene_nodes: HashMap<i32, SceneNode>,
}

impl Visualizer {
  pub fn new() -> ~Visualizer {
    let win_size = read_win_size("config.json");
    let win = init_win(win_size);
    let geom = Geom::new();
    let mut vis = ~Visualizer {
      program: 0,
      map_mesh: Mesh::new(),
      unit_mesh: Mesh::new(),
      mat_id: 0,
      win: Some(win),
      mouse_pos: Vec2::zero(),
      camera: Camera::new(),
      picker: TilePicker::new(),
      selected_tile_pos: None,
      geom: geom,
      scene_nodes: HashMap::new(),
    };
    vis.init_opengl();
    vis.picker.init(&geom);
    vis.init_models();
    vis.init_units();
    vis
  }

  fn win<'a>(&'a self) -> &'a glfw::Window {
    self.win.get_ref()
  }

  fn init_models(&mut self) {
    self.program = glh::compile_program(
      read_file(&Path::new("normal.vs.glsl")),
      read_file(&Path::new("normal.fs.glsl")),
    );
    gl::UseProgram(self.program);
    self.mat_id = glh::get_uniform(self.program, "mvp_mat");
    let position_attr = glh::get_attr(self.program, "in_vertex_coordinates");
    gl::EnableVertexAttribArray(position_attr);
    glh::vertex_attrib_pointer(position_attr, 3);
    let vt_attr = glh::get_attr(self.program, "in_texture_coordinates");
    gl::EnableVertexAttribArray(vt_attr);
    glh::vertex_attrib_pointer(vt_attr, 3);
    let map_vertex_data = build_hex_mesh(&self.geom);
    self.map_mesh.set_vertex_coords(map_vertex_data);
    self.map_mesh.set_texture_coords(build_hex_tex_coord());
    self.map_mesh.set_texture(glh::load_texture(~"data/floor.png"));
    let unit_obj = obj::Model::new("data/tank.obj");
    self.unit_mesh.set_vertex_coords(unit_obj.build());
    self.unit_mesh.set_texture_coords(unit_obj.build_tex_coord());
    self.unit_mesh.set_texture(glh::load_texture(~"data/tank.png"));
  }

  fn add_unit(&mut self, id: i32, pos: Vec2<i32>) {
    let world_pos = self.geom.map_pos_to_world_pos(pos);
    self.scene_nodes.insert(id, SceneNode{pos: world_pos});
  }

  fn init_units(&mut self) {
    self.add_unit(0, Vec2{x: 0, y: 0});
    self.add_unit(1, Vec2{x: 0, y: 1});
    self.add_unit(2, Vec2{x: 1, y: 0});
    self.add_unit(2, Vec2{x: 1, y: 1});
  }

  fn init_opengl(&mut self) {
    gl::load_with(glfw::get_proc_address);
    gl::Enable(gl::DEPTH_TEST);
  }

  fn cleanup_opengl(&self) {
    gl::DeleteProgram(self.program);
  }

  fn draw_units(&self) {
    gl::UseProgram(self.program);
    for (_, unit) in self.scene_nodes.iter() {
      let m = glh::tr(self.camera.mat(), unit.pos);
      glh::uniform_mat4f(self.mat_id, &m);
      self.unit_mesh.draw(self.program);
    }
  }

  fn draw_map(&self) {
    glh::uniform_mat4f(self.mat_id, &self.camera.mat());
    self.map_mesh.draw(self.program);
  }

  fn draw(&self) {
    gl::ClearColor(0.3, 0.3, 0.3, 1.0);
    gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
    gl::UseProgram(self.program);
    self.draw_units();
    self.draw_map();
    self.win().swap_buffers();
  }

  pub fn is_running(&self) -> bool {
    return !self.win().should_close()
  }

  fn handle_key_event(&mut self, key: glfw::Key) {
    match key {
      glfw::KeyEscape | glfw::KeyQ => {
        self.win().set_should_close(true);
      },
      glfw::KeySpace => println!("space"),
      glfw::KeyUp => self.camera.move(270.0),
      glfw::KeyDown => self.camera.move(90.0),
      glfw::KeyRight => self.camera.move(0.0),
      glfw::KeyLeft => self.camera.move(180.0),
      glfw::KeyMinus => self.camera.zoom += 1.0,
      glfw::KeyEqual => self.camera.zoom -= 1.0,
      _ => {},
    }
  }

  fn handle_cursor_pos_event(&mut self, pos: Vec2<f32>) {
    let button = self.win().get_mouse_button(glfw::MouseButtonRight);
    if button == glfw::Press {
      self.camera.z_angle += (self.mouse_pos.x - pos.x) / 2.0;
      self.camera.x_angle += (self.mouse_pos.y - pos.y) / 2.0;
    }
    self.mouse_pos = pos;
  }

  fn handle_mouse_button_event(&mut self) {
    if !self.selected_tile_pos.is_none() {
      let map_pos = self.selected_tile_pos.unwrap();
      let pos = self.geom.map_pos_to_world_pos(map_pos);
      let unit_id = 0;
      self.scene_nodes.get_mut(&unit_id).pos = pos;
    }
  }

  fn get_events(&mut self) -> ~[glfw::WindowEvent] {
    glfw::poll_events();
    let mut events = ~[];
    for (_, event) in self.win().flush_events() {
      events.push(event);
    }
    events
  }

  fn handle_event(&mut self, event: glfw::WindowEvent) {
    match event {
      glfw::KeyEvent(key, _, glfw::Press, _) => {
        self.handle_key_event(key);
      },
      glfw::CursorPosEvent(x, y) => {
        let p = Vec2{x: x as f32, y: y as f32};
        self.handle_cursor_pos_event(p);
      },
      glfw::MouseButtonEvent(glfw::MouseButtonLeft, glfw::Press, _) => {
        self.handle_mouse_button_event();
      },
      glfw::SizeEvent(w, h) => {
        gl::Viewport(0, 0, w, h);
      },
      _ => {},
    }
  }

  fn handle_events(&mut self) {
    for event in self.get_events().iter() {
      self.handle_event(*event);
    }
  }

  fn pick_tile(&mut self) {
    let mouse_pos = Vec2 {
      x: self.mouse_pos.x as i32,
      y: self.mouse_pos.y as i32,
    };
    let win_size = self.win().get_size();
    self.selected_tile_pos = self.picker.pick_tile(
      win_size, &self.camera, mouse_pos);
  }

  pub fn tick(&mut self) {
    self.handle_events();
    self.pick_tile();
    self.draw();
  }
}

impl Drop for Visualizer {
  fn drop(&mut self) {
    self.cleanup_opengl();
  }
}

// vim: set tabstop=2 shiftwidth=2 softtabstop=2 expandtab:
<commit_msg>Visualizer: init_models: Fixed *_attr var names<commit_after>// See LICENSE file for copyright and license details.

use extra::json;
use std::hashmap::HashMap;
use serialize::Decodable;
use glfw;
use gl;
use gl::types::{
  GLfloat,
  GLint,
  GLuint,
};
use cgmath::vector::{
  Vec3,
  Vec2,
};
use glh = gl_helpers;
use camera::Camera;
use map::TileIterator;
use geom::Geom;
use tile_picker::TilePicker;
use obj;
use mesh::Mesh;
use misc::read_file;

fn build_hex_mesh(&geom: &Geom) -> ~[Vec3<GLfloat>] {
  let mut vertex_data = ~[];
  for tile_pos in TileIterator::new() {
    let pos = geom.map_pos_to_world_pos(tile_pos);
    for num in range(0, 6) {
      let vertex = geom.index_to_hex_vertex(num);
      let next_vertex = geom.index_to_hex_vertex(num + 1);
      vertex_data.push(pos + vertex);
      vertex_data.push(pos + next_vertex);
      vertex_data.push(pos + Vec3::zero());
    }
  }
  vertex_data
}

fn build_hex_tex_coord() -> ~[Vec2<GLfloat>] {
  let mut vertex_data = ~[];
  for _ in TileIterator::new() {
    for _ in range(0, 6) {
      vertex_data.push(Vec2{x: 0.0, y: 0.0});
      vertex_data.push(Vec2{x: 1.0, y: 0.0});
      vertex_data.push(Vec2{x: 0.5, y: 0.5});
    }
  }
  vertex_data
}

pub struct SceneNode {
  pos: Vec3<f32>,
}

#[deriving(Decodable)]
struct Size2<T> {
  x: T,
  y: T,
}

fn read_win_size(config_path: &str) -> Vec2<int> {
  let path = Path::new(config_path);
  let json = json::from_str(read_file(&path)).unwrap();
  let mut decoder = json::Decoder::new(json);
  let size: Size2<int> = Decodable::decode(&mut decoder);
  Vec2{x: size.x, y: size.y}
}

fn init_win(win_size: Vec2<int>) -> glfw::Window {
  glfw::set_error_callback(~glfw::LogErrorHandler);
  let init_status = glfw::init();
  if !init_status.is_ok() {
    fail!("Failed to initialize GLFW");
  }
  let win = glfw::Window::create(
    win_size.x as u32,
    win_size.y as u32,
    "Marauder",
    glfw::Windowed,
  ).unwrap();
  win.make_context_current();
  win.set_all_polling(true);
  win
}

pub struct Visualizer {
  program: GLuint,
  map_mesh: Mesh,
  unit_mesh: Mesh,
  mat_id: GLint,
  win: Option<glfw::Window>,
  mouse_pos: Vec2<f32>,
  camera: Camera,
  picker: TilePicker,
  selected_tile_pos: Option<Vec2<i32>>,
  geom: Geom,
  scene_nodes: HashMap<i32, SceneNode>,
}

impl Visualizer {
  pub fn new() -> ~Visualizer {
    let win_size = read_win_size("config.json");
    let win = init_win(win_size);
    let geom = Geom::new();
    let mut vis = ~Visualizer {
      program: 0,
      map_mesh: Mesh::new(),
      unit_mesh: Mesh::new(),
      mat_id: 0,
      win: Some(win),
      mouse_pos: Vec2::zero(),
      camera: Camera::new(),
      picker: TilePicker::new(),
      selected_tile_pos: None,
      geom: geom,
      scene_nodes: HashMap::new(),
    };
    vis.init_opengl();
    vis.picker.init(&geom);
    vis.init_models();
    vis.init_units();
    vis
  }

  fn win<'a>(&'a self) -> &'a glfw::Window {
    self.win.get_ref()
  }

  fn init_models(&mut self) {
    self.program = glh::compile_program(
      read_file(&Path::new("normal.vs.glsl")),
      read_file(&Path::new("normal.fs.glsl")),
    );
    gl::UseProgram(self.program);
    self.mat_id = glh::get_uniform(self.program, "mvp_mat");
    let vertex_coordinates_attr = glh::get_attr(
      self.program, "in_vertex_coordinates");
    gl::EnableVertexAttribArray(vertex_coordinates_attr);
    glh::vertex_attrib_pointer(vertex_coordinates_attr, 3);
    let texture_coords_attr = glh::get_attr(
      self.program, "in_texture_coordinates");
    gl::EnableVertexAttribArray(texture_coords_attr);
    glh::vertex_attrib_pointer(texture_coords_attr, 3);
    let map_vertex_data = build_hex_mesh(&self.geom);
    self.map_mesh.set_vertex_coords(map_vertex_data);
    self.map_mesh.set_texture_coords(build_hex_tex_coord());
    self.map_mesh.set_texture(glh::load_texture(~"data/floor.png"));
    let unit_obj = obj::Model::new("data/tank.obj");
    self.unit_mesh.set_vertex_coords(unit_obj.build());
    self.unit_mesh.set_texture_coords(unit_obj.build_tex_coord());
    self.unit_mesh.set_texture(glh::load_texture(~"data/tank.png"));
  }

  fn add_unit(&mut self, id: i32, pos: Vec2<i32>) {
    let world_pos = self.geom.map_pos_to_world_pos(pos);
    self.scene_nodes.insert(id, SceneNode{pos: world_pos});
  }

  fn init_units(&mut self) {
    self.add_unit(0, Vec2{x: 0, y: 0});
    self.add_unit(1, Vec2{x: 0, y: 1});
    self.add_unit(2, Vec2{x: 1, y: 0});
    self.add_unit(2, Vec2{x: 1, y: 1});
  }

  fn init_opengl(&mut self) {
    gl::load_with(glfw::get_proc_address);
    gl::Enable(gl::DEPTH_TEST);
  }

  fn cleanup_opengl(&self) {
    gl::DeleteProgram(self.program);
  }

  fn draw_units(&self) {
    gl::UseProgram(self.program);
    for (_, unit) in self.scene_nodes.iter() {
      let m = glh::tr(self.camera.mat(), unit.pos);
      glh::uniform_mat4f(self.mat_id, &m);
      self.unit_mesh.draw(self.program);
    }
  }

  fn draw_map(&self) {
    glh::uniform_mat4f(self.mat_id, &self.camera.mat());
    self.map_mesh.draw(self.program);
  }

  fn draw(&self) {
    gl::ClearColor(0.3, 0.3, 0.3, 1.0);
    gl::Clear(gl::COLOR_BUFFER_BIT | gl::DEPTH_BUFFER_BIT);
    gl::UseProgram(self.program);
    self.draw_units();
    self.draw_map();
    self.win().swap_buffers();
  }

  pub fn is_running(&self) -> bool {
    return !self.win().should_close()
  }

  fn handle_key_event(&mut self, key: glfw::Key) {
    match key {
      glfw::KeyEscape | glfw::KeyQ => {
        self.win().set_should_close(true);
      },
      glfw::KeySpace => println!("space"),
      glfw::KeyUp => self.camera.move(270.0),
      glfw::KeyDown => self.camera.move(90.0),
      glfw::KeyRight => self.camera.move(0.0),
      glfw::KeyLeft => self.camera.move(180.0),
      glfw::KeyMinus => self.camera.zoom += 1.0,
      glfw::KeyEqual => self.camera.zoom -= 1.0,
      _ => {},
    }
  }

  fn handle_cursor_pos_event(&mut self, pos: Vec2<f32>) {
    let button = self.win().get_mouse_button(glfw::MouseButtonRight);
    if button == glfw::Press {
      self.camera.z_angle += (self.mouse_pos.x - pos.x) / 2.0;
      self.camera.x_angle += (self.mouse_pos.y - pos.y) / 2.0;
    }
    self.mouse_pos = pos;
  }

  fn handle_mouse_button_event(&mut self) {
    if !self.selected_tile_pos.is_none() {
      let map_pos = self.selected_tile_pos.unwrap();
      let pos = self.geom.map_pos_to_world_pos(map_pos);
      let unit_id = 0;
      self.scene_nodes.get_mut(&unit_id).pos = pos;
    }
  }

  fn get_events(&mut self) -> ~[glfw::WindowEvent] {
    glfw::poll_events();
    let mut events = ~[];
    for (_, event) in self.win().flush_events() {
      events.push(event);
    }
    events
  }

  fn handle_event(&mut self, event: glfw::WindowEvent) {
    match event {
      glfw::KeyEvent(key, _, glfw::Press, _) => {
        self.handle_key_event(key);
      },
      glfw::CursorPosEvent(x, y) => {
        let p = Vec2{x: x as f32, y: y as f32};
        self.handle_cursor_pos_event(p);
      },
      glfw::MouseButtonEvent(glfw::MouseButtonLeft, glfw::Press, _) => {
        self.handle_mouse_button_event();
      },
      glfw::SizeEvent(w, h) => {
        gl::Viewport(0, 0, w, h);
      },
      _ => {},
    }
  }

  fn handle_events(&mut self) {
    for event in self.get_events().iter() {
      self.handle_event(*event);
    }
  }

  fn pick_tile(&mut self) {
    let mouse_pos = Vec2 {
      x: self.mouse_pos.x as i32,
      y: self.mouse_pos.y as i32,
    };
    let win_size = self.win().get_size();
    self.selected_tile_pos = self.picker.pick_tile(
      win_size, &self.camera, mouse_pos);
  }

  pub fn tick(&mut self) {
    self.handle_events();
    self.pick_tile();
    self.draw();
  }
}

impl Drop for Visualizer {
  fn drop(&mut self) {
    self.cleanup_opengl();
  }
}

// vim: set tabstop=2 shiftwidth=2 softtabstop=2 expandtab:
<|endoftext|>


<commit_before>extern crate lichen;

use lichen::parse::Parser;
use lichen::var::Var;
use lichen::eval::Evaluator;
use lichen::fun::MutFn;

// Test for mutable state
#[derive(Debug)]
struct Player {
    coins: f32,
    name: String
}

#[test]
fn state_mut() {
    let src = "\n
def global\n
;\n
\n
root\n
    @global.coins 1\n
    emit \"step\"\n
    @global.coins + 5\n
    emit global.coins\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();
    assert_eq!(vars[0], 6.0 .into());
}

#[test]
fn parse_cust_fn() {
    let src = "root\n
    @root.five (inc) 1 2 3\n
    emit root.five\n
;\n
def root\n
    five 5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();
    let inc = MutFn::new(move |args, def| {
        let mut r = 0.;
        for n in args.iter() {
            if let Ok(v) = n.get_num(def) {
                r += v;
            }
        }

        return Some(Var::Num(r))
    });
    env.fun.insert("inc".to_owned(), inc);

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[0], 6.0 .into());
}

#[test]
fn parse_def_block() {
    let src = "root\n
    emit global.name global.size\n
;\n
\n
def global\n
    name \"my-game\"\n
    size 1.5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let mut ev = Evaluator::new(&mut env);
    let (vars,_) = ev.next().unwrap();

    assert_eq!(vars[0], Var::String("my-game".to_owned()));
    assert_eq!(vars[1], Var::Num(1.5 .to_owned()));
}

#[test]
fn validate_def_block() {
    let src = "root\n
    @global.size + 0.5\n
    @global.name \"other-game\"\n
    emit global.name global.size\n\n
;\n
\n
def global\n
    name \"my-game\"\n
    size 1.5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[0], Var::String("other-game".to_owned()));
    assert_eq!(vars[1], Var::Num(2.0 .to_owned()));
}

#[test]
fn validate_when_block() {
    let src = "root\n
    needs_coins global.coins < 1\n
    when {needs_coins @global.coins + 2, \n
         !global.name @global.name \"new-name\"}\n
    emit global.name global.coins\n
;\n
def global\n
    coins 0\n
;\n";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[1], 2.0 .into());
    assert_eq!(vars[0], "new-name".into());
}

#[test]
fn save_state() {
    let src = "root\n
    next:now other\n
;\n
\n
other\n
    @coins 1\n
    next:await end\n
    emit \"bye\"\n
;\n
\n
end\n
    emit true\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();
    let data = Player { coins: 0.0, name: "Pan".to_owned() };

    let state = {
        let mut ev = Evaluator::new(&mut env);
        let (_,_) = ev.next().unwrap();
        ev.save()
    };

    assert_eq!(data.coins, 0.0);

    let state = {
        let mut ev = state.to_eval(&mut env);
        let (_,_) = ev.next().unwrap();
        ev.save()
    };

    assert_eq!(data.coins, 1.0);

    //check if we held our place within node after await
    let ev = state.to_eval(&mut env);
    let (vars,_) = ev.last().unwrap();
    assert_eq!(vars[0],"bye".into());
}

#[test]
fn validate_or_block() {
    let src = "root\n
    if !global.drunk \"not drunk\"\n
    or \"is drunk\"\n
\n
    @global.drunk true\n
\n
    if !global.drunk \"not drunk\"\n
    or \"is drunk\"\n
;\n
\n
def global\n
    drunk false
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let mut ev = Evaluator::new(&mut env);
    let (vars,_) = ev.next().unwrap();

    assert_eq!(vars[0], "not drunk".into());

    let (vars,_) = ev.next().unwrap();
    assert_eq!(vars[0], "is drunk".into());
}
<commit_msg>wip: cust save<commit_after>extern crate lichen;

use lichen::parse::Parser;
use lichen::var::Var;
use lichen::eval::Evaluator;
use lichen::fun::MutFn;

// Test for mutable state
#[derive(Debug)]
struct Player {
    coins: f32,
    name: String
}

#[test]
fn state_mut() {
    let src = "\n
def global\n
;\n
\n
root\n
    @global.coins 1\n
    emit \"step\"\n
    @global.coins + 5\n
    emit global.coins\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();
    assert_eq!(vars[0], 6.0 .into());
}

#[test]
fn parse_cust_fn() {
    let src = "root\n
    @root.five (inc) 1 2 3\n
    emit root.five\n
;\n
def root\n
    five 5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();
    let inc = MutFn::new(move |args, def| {
        let mut r = 0.;
        for n in args.iter() {
            if let Ok(v) = n.get_num(def) {
                r += v;
            }
        }

        return Some(Var::Num(r))
    });
    env.fun.insert("inc".to_owned(), inc);

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[0], 6.0 .into());
}

#[test]
fn parse_def_block() {
    let src = "root\n
    emit global.name global.size\n
;\n
\n
def global\n
    name \"my-game\"\n
    size 1.5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let mut ev = Evaluator::new(&mut env);
    let (vars,_) = ev.next().unwrap();

    assert_eq!(vars[0], Var::String("my-game".to_owned()));
    assert_eq!(vars[1], Var::Num(1.5 .to_owned()));
}

#[test]
fn validate_def_block() {
    let src = "root\n
    @global.size + 0.5\n
    @global.name \"other-game\"\n
    emit global.name global.size\n\n
;\n
\n
def global\n
    name \"my-game\"\n
    size 1.5\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[0], Var::String("other-game".to_owned()));
    assert_eq!(vars[1], Var::Num(2.0 .to_owned()));
}

#[test]
fn validate_when_block() {
    let src = "root\n
    needs_coins global.coins < 1\n
    when {needs_coins @global.coins + 2, \n
         !global.name @global.name \"new-name\"}\n
    emit global.name global.coins\n
;\n
def global\n
    coins 0\n
;\n";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let ev = Evaluator::new(&mut env);
    let (vars,_) = ev.last().unwrap();

    assert_eq!(vars[1], 2.0 .into());
    assert_eq!(vars[0], "new-name".into());
}

#[test]
fn save_state() {
    let src = "root\n
    @ (inc) \"coins\" 1 2 3\n
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();
    let data = Player { coins: 0.0, name: "Pan".to_owned() };
    let data = Arc::new(Mutex::new(data));

    let dc = data.clone();
    let inc = MutFn::new(move |args, def| {
        let mut r = dc.lock().unwrap().coins;
        for n in args.iter() {
            if let Ok(v) = n.get_num(def) {
                *r += v;
            }
        }

        None
    });
    env.fun.insert("inc".to_owned(), inc);

    let state = {
        let mut ev = Evaluator::new(&mut env);
        let (_,_) = ev.next().unwrap();
        ev.save()
    };

    assert_eq!(data.coins, 0.0);

    let state = {
        let mut ev = state.to_eval(&mut env);
        let (_,_) = ev.next().unwrap();
        ev.save()
    };

    assert_eq!(data.coins, 1.0);

    //check if we held our place within node after await
    let ev = state.to_eval(&mut env);
    let (vars,_) = ev.last().unwrap();
    assert_eq!(vars[0],"bye".into());
}

#[test]
fn validate_or_block() {
    let src = "root\n
    if !global.drunk \"not drunk\"\n
    or \"is drunk\"\n
\n
    @global.drunk true\n
\n
    if !global.drunk \"not drunk\"\n
    or \"is drunk\"\n
;\n
\n
def global\n
    drunk false
;";

    let mut env = Parser::parse_blocks(src).expect("ERROR: Unable to parse source").into_env();

    let mut ev = Evaluator::new(&mut env);
    let (vars,_) = ev.next().unwrap();

    assert_eq!(vars[0], "not drunk".into());

    let (vars,_) = ev.next().unwrap();
    assert_eq!(vars[0], "is drunk".into());
}
<|endoftext|>


<commit_before>use std::env;
use std::process::Command;
use std::path::PathBuf;

fn get_bin_dir() -> PathBuf {
    env::current_exe()
        .expect("test bin's directory")
        .parent()
        .expect("test bin's parent directory")
        .parent()
        .expect("executable's directory")
        .to_path_buf()
}

fn cmd_aux2bib() -> Command {
    let path;
    if cfg!(not(windows)) {
        path = get_bin_dir().join("aux2bib");
    } else {
        path = get_bin_dir().join("aux2bib.exe");
    }
    if !path.is_file() {
        panic!("aux2bib binary {:?} was not found", path);
    }
    println!("{:?}", path);
    let mut cmd = Command::new(path);
    cmd.env_clear();

    cmd
}

fn cmd_blg2bib() -> Command {
    let path;
    if cfg!(not(windows)) {
        path = get_bin_dir().join("blg2bib");
    } else {
        path = get_bin_dir().join("blg2bib.exe");
    }
    if !path.is_file() {
        panic!("blg2bib binary {:?} was not found", path);
    }
    let mut cmd = Command::new(path);
    cmd.env_clear();

    cmd
}

#[cfg(not(windows))]
#[test]
fn aux2bib_runs() {
    let mut cmd = cmd_aux2bib().arg("--help").spawn().expect(
        "Failed to execute aux2bib",
    );

    let error_code = cmd.wait().expect("Failed to wait on aux2bib");

    assert!(error_code.success());

}

#[cfg(not(windows))]
#[test]
fn blg2bib_runs() {
    let mut cmd = cmd_blg2bib().arg("--help").spawn().expect(
        "Failed to execute blg2bib",
    );

    let error_code = cmd.wait().expect("Failed to wait on blg2bib");

    assert!(error_code.success());

}
<commit_msg>test: empty stdin input to binaries<commit_after>use std::env;
use std::process::{Command, Stdio};
use std::path::PathBuf;
use std::io::Write;

fn get_bin_dir() -> PathBuf {
    env::current_exe()
        .expect("test bin's directory")
        .parent()
        .expect("test bin's parent directory")
        .parent()
        .expect("executable's directory")
        .to_path_buf()
}

fn cmd_aux2bib() -> Command {
    let path;
    if cfg!(not(windows)) {
        path = get_bin_dir().join("aux2bib");
    } else {
        path = get_bin_dir().join("aux2bib.exe");
    }
    if !path.is_file() {
        panic!("aux2bib binary {:?} was not found", path);
    }
    println!("{:?}", path);
    let mut cmd = Command::new(path);
    cmd.env_clear().stderr(Stdio::piped()).stdout(
        Stdio::piped(),
    );

    cmd
}

fn cmd_blg2bib() -> Command {
    let path;
    if cfg!(not(windows)) {
        path = get_bin_dir().join("blg2bib");
    } else {
        path = get_bin_dir().join("blg2bib.exe");
    }
    if !path.is_file() {
        panic!("blg2bib binary {:?} was not found", path);
    }
    let mut cmd = Command::new(path);
    cmd.env_clear().stderr(Stdio::piped()).stdout(
        Stdio::piped(),
    );

    cmd
}

#[cfg(not(windows))]
#[test]
fn aux2bib_runs() {
    let mut cmd = cmd_aux2bib().arg("--help").spawn().expect(
        "Failed to execute aux2bib",
    );

    let error_code = cmd.wait().expect("Failed to wait on aux2bib");

    assert!(error_code.success());
}

#[cfg(not(windows))]
#[test]
fn blg2bib_runs() {
    let mut cmd = cmd_blg2bib().arg("--help").spawn().expect(
        "Failed to execute blg2bib",
    );

    let error_code = cmd.wait().expect("Failed to wait on blg2bib");

    assert!(error_code.success());
}

#[cfg(not(windows))]
#[test]
fn aux2bib_stdin_empty() {
    let mut child = cmd_aux2bib().stdin(Stdio::piped()).spawn().expect(
        "Failed to execute aux2bib",
    );

    {
        let stdin = child.stdin.as_mut().expect("Failed to get stdin");
        stdin.write_all(b"").expect("Failed to write to stdin");
    }

    let output = child.wait_with_output().expect("Failed to wait on aux2bib");

    assert!(output.status.success());
    assert_eq!(output.stdout, []);
}

#[cfg(not(windows))]
#[test]
fn blg2bib_stdin_empty() {
    let mut child = cmd_blg2bib().stdin(Stdio::piped()).spawn().expect(
        "Failed to execute blg2bib",
    );

    {
        let stdin = child.stdin.as_mut().expect("Failed to get stdin");
        stdin.write_all(b"").expect("Failed to write to stdin");
    }

    let output = child.wait_with_output().expect("Failed to wait on aux2bib");

    assert!(output.status.success());
    assert_eq!(output.stdout, []);
}
<|endoftext|>


<commit_before>use itertools::Itertools;
use lightning::map::{Map, ObjectMap};
use rayon::prelude::*;
use std::{
    collections::{BTreeSet, HashMap, HashSet},
    sync::Arc,
};
use tokio_stream::StreamExt;

use dovahkiin::types::{OwnedValue, SharedValue};

use crate::ram::{
    cell::{header_from_chunk_raw, select_from_chunk_raw},
    chunk::Chunk,
};

mod histogram;
pub mod sm;

pub struct SchemaStatistics {
    histogram: HashMap<u64, [OwnedValue; 10]>,
    count: usize,
    segs: usize,
    bytes: usize,
    timestamp: u64,
}

pub struct ChunkStatistics {
    schemas: ObjectMap<Arc<SchemaStatistics>>,
}

const HISTOGRAM_PARTITATION_SIZE: usize = 1024;
const HISTOGRAM_PARTITATION_DEPTH: usize = 128;

type HistogramKey = [u8; 8];

impl ChunkStatistics {
    pub fn from_chunk(chunk: &Chunk) -> Self {
        let histogram_partitations = chunk
            .cell_index
            .entries()
            .chunks(HISTOGRAM_PARTITATION_SIZE)
            .map(|s| s.to_vec())
            .collect_vec();
        let partitations: Vec<_> = histogram_partitations
            .into_par_iter()
            .map(|partitation| {
                // Build exact histogram for each of the partitation and then approximate overall histogram
                let mut sizes = HashMap::new();
                let mut segs = HashMap::new();
                let mut exact_accumlators = HashMap::new();
                let partitation_size = partitation.len();
                for (hash, _) in partitation {
                    let loc = if let Ok(ptr) = chunk.location_for_read(hash as u64) {
                        ptr
                    } else {
                        trace!("Cannot obtain cell lock {} for statistics", hash);
                        continue;
                    };
                    match header_from_chunk_raw(*loc) {
                        Ok((header, _, entry_header)) => {
                            let cell_size = entry_header.content_length as usize;
                            let cell_seg = chunk.allocator.id_by_addr(*loc);
                            let schema_id = header.schema;
                            if let Some(schema) = chunk.meta.schemas.get(&schema_id) {
                                let fields = schema.index_fields.keys().cloned().collect_vec();
                                if let Ok(partial_cell) =
                                    select_from_chunk_raw(*loc, chunk, fields.as_slice())
                                {
                                    let field_array = if fields.len() == 1 {
                                        vec![partial_cell]
                                    } else if let SharedValue::Array(arr) = partial_cell {
                                        arr
                                    } else {
                                        error!(
                                            "Cannot decode partial cell for statistics {:?}",
                                            partial_cell
                                        );
                                        continue;
                                    };
                                    for (i, val) in field_array.into_iter().enumerate() {
                                        if val == SharedValue::Null || val == SharedValue::NA {
                                            continue;
                                        }
                                        let field_id = fields[i];
                                        exact_accumlators
                                            .entry(schema_id)
                                            .or_insert_with(|| HashMap::new())
                                            .entry(field_id)
                                            .or_insert_with(|| Vec::with_capacity(partitation_size))
                                            .push(val.feature());
                                    }
                                    *sizes.entry(schema_id).or_insert(0) += cell_size;
                                    segs.entry(schema_id)
                                        .or_insert_with(|| HashSet::new())
                                        .insert(cell_seg);
                                }
                            } else {
                                warn!("Cannot get schema {} for statistics", schema_id);
                            }
                        }
                        Err(e) => {
                            warn!("Failed to read {} for statistics, error: {:?}", hash, e);
                        }
                    }
                }
                let histograms: HashMap<_, _> = exact_accumlators
                    .into_iter()
                    .map(|(schema_id, schema_histograms)| {
                        let compiled_histograms = schema_histograms
                            .into_iter()
                            .map(|(field, mut items)| {
                                items.sort();
                                let depth = items.len() / HISTOGRAM_PARTITATION_DEPTH;
                                let mut histogram = (0..HISTOGRAM_PARTITATION_DEPTH)
                                    .map(|tile| items[tile * depth])
                                    .collect_vec();
                                let last_item = &items[items.len() - 1];
                                if histogram.last().unwrap() != last_item {
                                    histogram.push(*last_item);
                                }
                                (field, (histogram, items.len(), depth))
                            })
                            .collect::<HashMap<_, _>>();
                        (schema_id, compiled_histograms)
                    })
                    .collect::<HashMap<_, _>>();
                (sizes, segs, histograms)
            })
            .collect();
        let schema_ids: Vec<_> = partitations
            .iter()
            .map(|(sizes, _, _)| sizes.keys())
            .flatten()
            .dedup()
            .collect();
        let total_size = schema_ids
            .iter()
            .map(|sid| {
                (
                    *sid,
                    partitations
                        .iter()
                        .map(|(sizes, _, _)| sizes.get(sid).unwrap_or(&0))
                        .sum::<usize>(),
                )
            })
            .collect::<HashMap<_, _>>();
        let total_segs = schema_ids
            .iter()
            .map(|sid| {
                (
                    *sid,
                    partitations
                        .iter()
                        .map(|(_, segs, _)| segs.get(sid).map(|set| set.len()).unwrap_or(0))
                        .sum::<usize>(),
                )
            })
            .collect::<HashMap<_, _>>();
        let empty_histo = Default::default();
        let schema_histograms = schema_ids
            .iter()
            .map(|sid| {
                (*sid, {
                    let parted_histos = partitations
                        .iter()
                        .map(|(_, _, histo)| histo.get(sid).unwrap_or(&empty_histo))
                        .collect_vec();
                    let field_ids = parted_histos
                        .iter()
                        .map(|histo_map| histo_map.keys())
                        .flatten()
                        .dedup()
                        .collect::<Vec<_>>();
                    field_ids
                        .par_iter()
                        .map(|field_id| {
                            let schema_field_histograms = parted_histos
                                .iter()
                                .map(|histo_map| &histo_map[field_id])
                                .collect_vec();
                            (**field_id, build_histogram(schema_field_histograms))
                        })
                        .collect::<HashMap<u64, Vec<HistogramKey>>>()
                })
            })
            .collect::<HashMap<_, _>>();
        unimplemented!()
    }
}

fn build_histogram(partitations: Vec<&(Vec<HistogramKey>, usize, usize)>) -> Vec<HistogramKey> {
    unimplemented!()
}
<commit_msg>impl the algorithm<commit_after>use itertools::Itertools;
use lightning::map::{Map, ObjectMap};
use rayon::prelude::*;
use std::{
    collections::{BTreeSet, HashMap, HashSet},
    sync::Arc,
};
use tokio_stream::StreamExt;

use dovahkiin::types::{OwnedValue, SharedValue};

use crate::ram::{
    cell::{header_from_chunk_raw, select_from_chunk_raw},
    chunk::Chunk,
};

mod histogram;
pub mod sm;

pub struct SchemaStatistics {
    histogram: HashMap<u64, [OwnedValue; 10]>,
    count: usize,
    segs: usize,
    bytes: usize,
    timestamp: u64,
}

pub struct ChunkStatistics {
    schemas: ObjectMap<Arc<SchemaStatistics>>,
}

const HISTOGRAM_PARTITATION_SIZE: usize = 1024;
const HISTOGRAM_PARTITATION_BUCKETS: usize = 256;
const HISTOGRAM_TARGET_BUCKETS: usize = 100;

type HistogramKey = [u8; 8];

impl ChunkStatistics {
    pub fn from_chunk(chunk: &Chunk) -> Self {
        let histogram_partitations = chunk
            .cell_index
            .entries()
            .chunks(HISTOGRAM_PARTITATION_SIZE)
            .map(|s| s.to_vec())
            .collect_vec();
        let partitations: Vec<_> = histogram_partitations
            .into_par_iter()
            .map(|partitation| {
                // Build exact histogram for each of the partitation and then approximate overall histogram
                let mut sizes = HashMap::new();
                let mut segs = HashMap::new();
                let mut exact_accumlators = HashMap::new();
                let partitation_size = partitation.len();
                for (hash, _) in partitation {
                    let loc = if let Ok(ptr) = chunk.location_for_read(hash as u64) {
                        ptr
                    } else {
                        trace!("Cannot obtain cell lock {} for statistics", hash);
                        continue;
                    };
                    match header_from_chunk_raw(*loc) {
                        Ok((header, _, entry_header)) => {
                            let cell_size = entry_header.content_length as usize;
                            let cell_seg = chunk.allocator.id_by_addr(*loc);
                            let schema_id = header.schema;
                            if let Some(schema) = chunk.meta.schemas.get(&schema_id) {
                                let fields = schema.index_fields.keys().cloned().collect_vec();
                                if let Ok(partial_cell) =
                                    select_from_chunk_raw(*loc, chunk, fields.as_slice())
                                {
                                    let field_array = if fields.len() == 1 {
                                        vec![partial_cell]
                                    } else if let SharedValue::Array(arr) = partial_cell {
                                        arr
                                    } else {
                                        error!(
                                            "Cannot decode partial cell for statistics {:?}",
                                            partial_cell
                                        );
                                        continue;
                                    };
                                    for (i, val) in field_array.into_iter().enumerate() {
                                        if val == SharedValue::Null || val == SharedValue::NA {
                                            continue;
                                        }
                                        let field_id = fields[i];
                                        exact_accumlators
                                            .entry(schema_id)
                                            .or_insert_with(|| HashMap::new())
                                            .entry(field_id)
                                            .or_insert_with(|| Vec::with_capacity(partitation_size))
                                            .push(val.feature());
                                    }
                                    *sizes.entry(schema_id).or_insert(0) += cell_size;
                                    segs.entry(schema_id)
                                        .or_insert_with(|| HashSet::new())
                                        .insert(cell_seg);
                                }
                            } else {
                                warn!("Cannot get schema {} for statistics", schema_id);
                            }
                        }
                        Err(e) => {
                            warn!("Failed to read {} for statistics, error: {:?}", hash, e);
                        }
                    }
                }
                let histograms: HashMap<_, _> = exact_accumlators
                    .into_iter()
                    .map(|(schema_id, schema_histograms)| {
                        let compiled_histograms = schema_histograms
                            .into_iter()
                            .map(|(field, mut items)| {
                                items.sort();
                                let depth = items.len() / HISTOGRAM_PARTITATION_BUCKETS;
                                let mut histogram = (0..HISTOGRAM_PARTITATION_BUCKETS)
                                    .map(|tile| items[tile * depth])
                                    .collect_vec();
                                let last_item = &items[items.len() - 1];
                                if histogram.last().unwrap() != last_item {
                                    histogram.push(*last_item);
                                }
                                (field, (histogram, items.len(), depth))
                            })
                            .collect::<HashMap<_, _>>();
                        (schema_id, compiled_histograms)
                    })
                    .collect::<HashMap<_, _>>();
                (sizes, segs, histograms)
            })
            .collect();
        let schema_ids: Vec<_> = partitations
            .iter()
            .map(|(sizes, _, _)| sizes.keys())
            .flatten()
            .dedup()
            .collect();
        let total_size = schema_ids
            .iter()
            .map(|sid| {
                (
                    *sid,
                    partitations
                        .iter()
                        .map(|(sizes, _, _)| sizes.get(sid).unwrap_or(&0))
                        .sum::<usize>(),
                )
            })
            .collect::<HashMap<_, _>>();
        let total_segs = schema_ids
            .iter()
            .map(|sid| {
                (
                    *sid,
                    partitations
                        .iter()
                        .map(|(_, segs, _)| segs.get(sid).map(|set| set.len()).unwrap_or(0))
                        .sum::<usize>(),
                )
            })
            .collect::<HashMap<_, _>>();
        let empty_histo = Default::default();
        let schema_histograms = schema_ids
            .iter()
            .map(|sid| {
                (*sid, {
                    let parted_histos = partitations
                        .iter()
                        .map(|(_, _, histo)| histo.get(sid).unwrap_or(&empty_histo))
                        .collect_vec();
                    let field_ids = parted_histos
                        .iter()
                        .map(|histo_map| histo_map.keys())
                        .flatten()
                        .dedup()
                        .collect::<Vec<_>>();
                    field_ids
                        .par_iter()
                        .map(|field_id| {
                            let schema_field_histograms = parted_histos
                                .iter()
                                .map(|histo_map| &histo_map[field_id])
                                .collect_vec();
                            (**field_id, build_histogram(schema_field_histograms))
                        })
                        .collect::<HashMap<u64, _>>()
                })
            })
            .collect::<HashMap<_, _>>();
        unimplemented!()
    }
}

fn build_histogram(
    partitations: Vec<&(Vec<HistogramKey>, usize, usize)>,
) -> [HistogramKey; HISTOGRAM_TARGET_BUCKETS + 1] {
    // Build the approximated histogram from partitation histograms
    // https://arxiv.org/abs/1606.05633
    let mut part_idxs = vec![0; partitations.len()];
    let num_total = partitations.iter().map(|(_, num, _)| num).sum::<usize>();
    let part_depths = partitations.iter().map(|(_, _, depth)| *depth).collect_vec();
    let target_width = num_total / HISTOGRAM_TARGET_BUCKETS;
    let part_histos = partitations
        .iter()
        .map(|(histo, _, _)| histo)
        .filter(|histo| !histo.is_empty())
        .collect_vec();
    let mut target_histogram = [[0u8; 8]; HISTOGRAM_TARGET_BUCKETS + 1];
    // Perform a merge sort for sorted pre-histogram
    let mut filled = target_width;
    let mut last_key = Default::default();
    'HISTO_CONST:
    for i in 0..HISTOGRAM_PARTITATION_BUCKETS {
        loop {
            let (key, ended) = if let Some((part_idx, histo)) =
                part_histos.iter().enumerate().filter(|(i, h)| {
                    let idx = part_idxs[*i];
                    idx < h.len()
                }).min_by(|(i1, h1), (i2, h2)| {
                    let h1_idx = part_idxs[*i1];
                    let h2_idx = part_idxs[*i2];
                    h1[h1_idx].cmp(&h2[h2_idx])
                }) {
                part_idxs[part_idx] += 1;
                ((histo[part_idx], part_idx), false)
            } else {
                (last_key, true)
            };
            last_key = key;
            let idx = last_key.1;
            if filled >= target_width || ended {
                target_histogram[i] = last_key.0;
                filled = 0;
                continue 'HISTO_CONST;
            }
            filled += part_depths[idx];
        }
    }
    target_histogram
}
<|endoftext|>


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

use std::fs::{self, File};
/// Wrappers around the `zip-rs` library to compress and decompress zip archives.
use std::io;
use std::path::Path;
use walkdir::WalkDir;

use self::zip_rs::result::{ZipError, ZipResult};
use self::zip_rs::write::FileOptions;
use zip as zip_rs;

/// Compress a source directory recursively into a zip file.
/// Permissions are set to 644 by default to avoid any
/// unwanted execution bits.
pub fn compress(src_dir: &Path, dst_file: &File) -> ZipResult<()> {
	if !Path::new(src_dir).is_dir() {
		return Err(ZipError::Io(io::Error::new(
			io::ErrorKind::Other,
			"Source must be a directory.",
		)));
	}

	let options = FileOptions::default()
		.compression_method(zip_rs::CompressionMethod::Stored)
		.unix_permissions(0o644);

	let mut zip = zip_rs::ZipWriter::new(dst_file);
	let walkdir = WalkDir::new(src_dir.to_str().unwrap());
	let it = walkdir.into_iter();

	for dent in it.filter_map(|e| e.ok()) {
		let path = dent.path();
		let name = path
			.strip_prefix(Path::new(src_dir))
			.unwrap()
			.to_str()
			.unwrap();

		if path.is_file() {
			zip.start_file(name, options)?;
			let mut f = File::open(path)?;
			io::copy(&mut f, &mut zip)?;
		}
	}

	zip.finish()?;
	dst_file.sync_all()?;
	Ok(())
}

/// Decompress a source file into the provided destination path.
pub fn decompress<R, F>(src_file: R, dest: &Path, expected: F) -> ZipResult<usize>
where
	R: io::Read + io::Seek,
	F: Fn(&Path) -> bool,
{
	let mut decompressed = 0;
	let mut archive = zip_rs::ZipArchive::new(src_file)?;

	for i in 0..archive.len() {
		let mut file = archive.by_index(i)?;
		let san_name = file.sanitized_name();
		if san_name.to_str().unwrap_or("").replace("\\", "/") != file.name().replace("\\", "/")
			|| !expected(&san_name)
		{
			info!(
				"ignoring a suspicious file: {}, got {:?}",
				file.name(),
				san_name.to_str()
			);
			continue;
		}
		let file_path = dest.join(san_name);

		if (&*file.name()).ends_with('/') {
			fs::create_dir_all(&file_path)?;
		} else {
			if let Some(p) = file_path.parent() {
				if !p.exists() {
					fs::create_dir_all(&p)?;
				}
			}
			let res = fs::File::create(&file_path);
			let mut outfile = match res {
				Err(e) => {
					error!("{:?}", e);
					return Err(zip::result::ZipError::Io(e));
				}
				Ok(r) => r,
			};
			io::copy(&mut file, &mut outfile)?;
			decompressed += 1;
		}

		// Get and Set permissions
		#[cfg(unix)]
		{
			use std::os::unix::fs::PermissionsExt;
			if let Some(mode) = file.unix_mode() {
				fs::set_permissions(
					&file_path.to_str().unwrap(),
					PermissionsExt::from_mode(mode),
				)?;
			}
		}
	}
	Ok(decompressed)
}
<commit_msg>catch the panic to avoid peer thread quit early (#2686)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fs::{self, File};
/// Wrappers around the `zip-rs` library to compress and decompress zip archives.
use std::io;
use std::panic;
use std::path::Path;
use walkdir::WalkDir;

use self::zip_rs::result::{ZipError, ZipResult};
use self::zip_rs::write::FileOptions;
use zip as zip_rs;

/// Compress a source directory recursively into a zip file.
/// Permissions are set to 644 by default to avoid any
/// unwanted execution bits.
pub fn compress(src_dir: &Path, dst_file: &File) -> ZipResult<()> {
	if !Path::new(src_dir).is_dir() {
		return Err(ZipError::Io(io::Error::new(
			io::ErrorKind::Other,
			"Source must be a directory.",
		)));
	}

	let options = FileOptions::default()
		.compression_method(zip_rs::CompressionMethod::Stored)
		.unix_permissions(0o644);

	let mut zip = zip_rs::ZipWriter::new(dst_file);
	let walkdir = WalkDir::new(src_dir.to_str().unwrap());
	let it = walkdir.into_iter();

	for dent in it.filter_map(|e| e.ok()) {
		let path = dent.path();
		let name = path
			.strip_prefix(Path::new(src_dir))
			.unwrap()
			.to_str()
			.unwrap();

		if path.is_file() {
			zip.start_file(name, options)?;
			let mut f = File::open(path)?;
			io::copy(&mut f, &mut zip)?;
		}
	}

	zip.finish()?;
	dst_file.sync_all()?;
	Ok(())
}

/// Decompress a source file into the provided destination path.
pub fn decompress<R, F>(src_file: R, dest: &Path, expected: F) -> ZipResult<usize>
where
	R: io::Read + io::Seek + panic::UnwindSafe,
	F: Fn(&Path) -> bool + panic::UnwindSafe,
{
	let mut decompressed = 0;

	// catch the panic to avoid the thread quit
	panic::set_hook(Box::new(|panic_info| {
		error!(
			"panic occurred: {:?}",
			panic_info.payload().downcast_ref::<&str>().unwrap()
		);
	}));
	let result = panic::catch_unwind(move || {
		let mut archive = zip_rs::ZipArchive::new(src_file)?;

		for i in 0..archive.len() {
			let mut file = archive.by_index(i)?;
			let san_name = file.sanitized_name();
			if san_name.to_str().unwrap_or("").replace("\\", "/") != file.name().replace("\\", "/")
				|| !expected(&san_name)
			{
				info!(
					"ignoring a suspicious file: {}, got {:?}",
					file.name(),
					san_name.to_str()
				);
				continue;
			}
			let file_path = dest.join(san_name);

			if (&*file.name()).ends_with('/') {
				fs::create_dir_all(&file_path)?;
			} else {
				if let Some(p) = file_path.parent() {
					if !p.exists() {
						fs::create_dir_all(&p)?;
					}
				}
				let res = fs::File::create(&file_path);
				let mut outfile = match res {
					Err(e) => {
						error!("{:?}", e);
						return Err(zip::result::ZipError::Io(e));
					}
					Ok(r) => r,
				};
				io::copy(&mut file, &mut outfile)?;
				decompressed += 1;
			}

			// Get and Set permissions
			#[cfg(unix)]
			{
				use std::os::unix::fs::PermissionsExt;
				if let Some(mode) = file.unix_mode() {
					fs::set_permissions(
						&file_path.to_str().unwrap(),
						PermissionsExt::from_mode(mode),
					)?;
				}
			}
		}
		Ok(decompressed)
	});
	match result {
		Ok(res) => match res {
			Err(e) => Err(e.into()),
			Ok(_) => res,
		},
		Err(_) => {
			error!("panic occurred on zip::decompress!");
			Err(zip::result::ZipError::InvalidArchive(
				"panic occurred on zip::decompress",
			))
		}
	}
}
<|endoftext|>