<commit_before>#![feature(linkage, naked_functions, asm, const_fn, test, proc_macro)]
// Execute using: RUSTFLAGS="-C soft-float -C relocation-model=static -C code-model=kernel" RUST_BACKTRACE=1 cargo test --verbose --test kvm -- --nocapture

extern crate kvm;
extern crate memmap;
extern crate x86;
extern crate core;

#[macro_use]
extern crate klogger;

extern crate test_macros;
use test_macros::panics_note;
use test_macros::kvmattrs;

use kvm::{Capability, Exit, IoDirection, Segment, System, Vcpu, VirtualMachine};
use memmap::{Mmap, Protection};
use std::fs::File;
use std::io::{BufRead, BufReader};

use x86::shared::control_regs::*;
use x86::shared::paging::*;
use x86::bits64::paging::*;

#[kvmattrs(identity_map)]
fn use_the_port() {
    use_the_port_setup();
    log!("1");
    unsafe {
        asm!("inb $0, %al" :: "i"(0x01) :: "volatile");
    }
}

#[test]
fn io_example() {
    let vaddr = VAddr::from_usize(use_the_port as *const () as usize);
    static PAGE_TABLE_P: PAddr = PAddr::from_u64(0x1000);

    // Set up a page table that identity maps the lower half of the address space
    let mut anon_mmap = Mmap::anonymous(3 * (1 << 20), Protection::ReadWrite).unwrap(); // Map 1 MiB
    let page_table_memory = unsafe { anon_mmap.as_mut_slice() };

    // "physical" layout of PageTable is: PML4Table at 0x1000, 1st PDPT at 0x2000, 2nd PDPT at 0x3000 ... 512th PDPT at 0x1000+512*0x1000
    type PageTable = (PML4, [PDPT; 512]);

    let page_table: &mut PageTable =
        unsafe { ::std::mem::transmute(&mut page_table_memory[PAGE_TABLE_P.as_u64() as usize]) };
    let (ref mut pml4, ref mut pdpts) = *page_table;
    // Identity map everything in PML4 slots 0..256
    for i in 0..512 {
        let offset = 0x2000 + 0x1000 * i;
        pml4[i] = PML4Entry::new(PAddr::from_u64(offset as _), PML4_P | PML4_RW);
        let pdpt = &mut pdpts[i];
        for j in 0..512 {
            pdpt[j] = PDPTEntry::new(PAddr::from_u64(((512 * i + j) as u64) << 30),
                                     PDPT_P | PDPT_RW | PDPT_PS); // Set-up 1 GiB page mappings
            if i == pml4_index(vaddr) && j == pdpt_index(vaddr) {
                println!("pml4_index(fn) {:x}", pml4[i].get_address());
                println!("pdpt_index(fn) {:x}", pdpt[j].get_address());
            }
        }
    }

    let mut stack_mmap = Mmap::anonymous(65536, Protection::ReadWrite).unwrap();
    let stack_size = stack_mmap.len();
    let stack_memory = unsafe { stack_mmap.as_mut_slice() };
    let stack_base = VAddr::from_usize(stack_memory.as_mut_ptr() as usize);
    static STACK_BASE_T: PAddr = PAddr::from_u64(0x2000000);
    println!("Stack base {:x} with size {:x}", stack_base, stack_size);

    // Initialize the KVM system
    let sys = System::initialize().unwrap();

    // Create a Virtual Machine
    let mut vm = VirtualMachine::create(&sys).unwrap();

    // Ensure that the VM supports memory backing with user memory
    assert!(vm.check_capability(Capability::UserMemory) > 0);

    // Once the memory is set we can't even call length.
    let page_table_memory_limit = page_table_memory.len() - 1;

    // Map the page table memory
    vm.set_user_memory_region(0, page_table_memory, 0).unwrap();
    // Map stack space
    vm.set_user_memory_region(STACK_BASE_T.as_u64(), stack_memory, 0).unwrap();

    // Map the process
    let f = File::open("/proc/self/maps").unwrap();
    let reader = BufReader::new(f);

    for line in reader.lines() {
        let line = line.unwrap();
        println!("{}", line);
        let mut s = line.split(' ');
        let mut s2 = s.next().unwrap().split('-');
        let begin = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        let end = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        if end < 0x800000000000 {
            let perm = s.next().unwrap();
            //println!("{:#X}-{:#X} {}", begin, end, perm);
            let slice = {
                let begin_ptr: *mut u8 = begin as *const u8 as _;
                unsafe { ::std::slice::from_raw_parts_mut(begin_ptr, end - begin) }
            };
            // Make sure process doesn't overlap with page table
            assert!(begin > page_table_memory_limit);
            vm.set_user_memory_region(begin as _, slice, 0).unwrap();
        }
    }

    // Create a new VCPU
    let mut vcpu = Vcpu::create(&mut vm).unwrap();

    // Set supported CPUID (KVM fails without doing this)
    let mut cpuid = sys.get_supported_cpuid().unwrap();
    /*for mut entry in cpuid.entries_mut() {
        if entry.function == 0x80000008 {
            entry.eax = entry.eax | 0xff;
        }
        println!("{:?}", entry);
    }*/
    vcpu.set_cpuid2(&mut cpuid).unwrap();

    // Setup the special registers
    let mut sregs = vcpu.get_sregs().unwrap();

    // Set the code segment to have base 0, limit 4GB (flat segmentation)
    let segment_template = Segment {
        base: 0x0,
        limit: 0xffffffff,
        selector: 0,
        _type: 0,
        present: 0,
        dpl: 0,
        db: 1,
        s: 0,
        l: 0,
        g: 1,
        avl: 0,
        ..Default::default()
    };

    sregs.cs = Segment {
        selector: 0x8,
        _type: 0xb,
        present: 1,
        db: 0,
        s: 1,
        l: 1,
        ..segment_template
    };
    sregs.ss = Segment { ..segment_template };
    sregs.ds = Segment { ..segment_template };
    sregs.es = Segment { ..segment_template };
    sregs.fs = Segment { ..segment_template };
    sregs.gs = Segment { ..segment_template };

    // We don't need to populate the GDT if we have our segments setup
    // cr0 - protected mode on, paging enabled
    sregs.cr0 = (CR0_PROTECTED_MODE | CR0_MONITOR_COPROCESSOR | CR0_EXTENSION_TYPE |
                 CR0_ENABLE_PAGING | CR0_NUMERIC_ERROR | CR0_WRITE_PROTECT |
                 CR0_ALIGNMENT_MASK | CR0_ENABLE_PAGING)
        .bits() as u64;
    sregs.cr3 = PAGE_TABLE_P.as_u64();
    sregs.cr4 = (CR4_ENABLE_PSE | CR4_ENABLE_PAE | CR4_ENABLE_GLOBAL_PAGES | CR4_ENABLE_SSE |
                 CR4_UNMASKED_SSE |
                 CR4_ENABLE_OS_XSAVE | CR4_ENABLE_SMEP | CR4_ENABLE_VME)
        .bits() as u64;
    sregs.efer = 0xd01;

    // Set the special registers
    vcpu.set_sregs(&sregs).unwrap();

    let mut regs = vcpu.get_regs().unwrap();
    // set the instruction pointer to 1 MB
    regs.rip = vaddr.as_usize() as u64;
    println!("regs.rip = 0x{:x}", regs.rip); // but is at: 0x40cd60
    //println!("regs.rip = 0x{:x}", unsafe { *(regs.rip as *const u64) }); // but is at: 0x40cd60
    //regs.rip = 0x40cd60;
    regs.rflags = 0x246;
    regs.rsp = STACK_BASE_T.as_u64() + stack_size as u64;
    regs.rbp = regs.rsp;
    vcpu.set_regs(&regs).unwrap();

    // Actually run the VCPU

    println!("size of run {:?}", std::mem::size_of::<kvm::Run>());
    let mut vm_is_done = false;
    let mut new_regs = kvm::Regs::default();
    while !vm_is_done {
        {
            let (run, mut regs) = unsafe { vcpu.run_regs() }.unwrap();
            match run.exit_reason {
                Exit::Io => {
                    let io = unsafe { *run.io() };
                    match io.direction {
                        IoDirection::In => {
                            if io.port == 0x3fd {
                                regs.rax = 0x20; // Mark serial line as ready to write
                            } else {
                                println!("IO on unknown port: {}", io.port);
                            }
                        }
                        IoDirection::Out => {
                            if io.port == 0x3f8 {
                                println!("got char {:#?}", regs.rax as u8 as char);
                            }
                        }
                    }
                }
                Exit::Shutdown => {
                    println!("Shutting down");
                    vm_is_done = true;
                }
                _ => {
                    println!("Unknown exit reason: {:?}", run.exit_reason);
                }
            }

            new_regs = regs;
        }
        vcpu.set_regs(&new_regs).unwrap();
    }

    // Ensure that the exit reason we get back indicates that the I/O
    // instruction was executed
    /*let regs = vcpu.get_regs().unwrap();
    println!("vcpu.rip: {:x}", regs.rip);

    assert!(run.exit_reason == Exit::Io);
    let io = unsafe { *run.io() };
    assert!(io.direction == IoDirection::In);
    assert!(io.size == 1);
    assert!(io.port == 0x1);
    unsafe {
        println!("{:#?}", *run.io());
    }*/
}
<commit_msg>Update cmd line execution #20.<commit_after>#![feature(linkage, naked_functions, asm, const_fn, test, proc_macro)]
// Execute using: RUSTFLAGS="-C relocation-model=dynamic-no-pic -C code-model=kernel" RUST_BACKTRACE=1 cargo test --verbose --test kvm -- --nocapture

extern crate kvm;
extern crate memmap;
extern crate x86;
extern crate core;

#[macro_use]
extern crate klogger;

extern crate test_macros;
use test_macros::panics_note;
use test_macros::kvmattrs;

use kvm::{Capability, Exit, IoDirection, Segment, System, Vcpu, VirtualMachine};
use memmap::{Mmap, Protection};
use std::fs::File;
use std::io::{BufRead, BufReader};

use x86::shared::control_regs::*;
use x86::shared::paging::*;
use x86::bits64::paging::*;

#[kvmattrs(identity_map)]
fn use_the_port() {
    use_the_port_setup();
    log!("1");
    unsafe {
        asm!("inb $0, %al" :: "i"(0x01) :: "volatile");
    }
}

#[test]
fn io_example() {
    let vaddr = VAddr::from_usize(use_the_port as *const () as usize);
    static PAGE_TABLE_P: PAddr = PAddr::from_u64(0x1000);

    // Set up a page table that identity maps the lower half of the address space
    let mut anon_mmap = Mmap::anonymous(3 * (1 << 20), Protection::ReadWrite).unwrap(); // Map 1 MiB
    let page_table_memory = unsafe { anon_mmap.as_mut_slice() };

    // "physical" layout of PageTable is: PML4Table at 0x1000, 1st PDPT at 0x2000, 2nd PDPT at 0x3000 ... 512th PDPT at 0x1000+512*0x1000
    type PageTable = (PML4, [PDPT; 512]);

    let page_table: &mut PageTable =
        unsafe { ::std::mem::transmute(&mut page_table_memory[PAGE_TABLE_P.as_u64() as usize]) };
    let (ref mut pml4, ref mut pdpts) = *page_table;
    // Identity map everything in PML4 slots 0..256
    for i in 0..512 {
        let offset = 0x2000 + 0x1000 * i;
        pml4[i] = PML4Entry::new(PAddr::from_u64(offset as _), PML4_P | PML4_RW);
        let pdpt = &mut pdpts[i];
        for j in 0..512 {
            pdpt[j] = PDPTEntry::new(PAddr::from_u64(((512 * i + j) as u64) << 30),
                                     PDPT_P | PDPT_RW | PDPT_PS); // Set-up 1 GiB page mappings
            if i == pml4_index(vaddr) && j == pdpt_index(vaddr) {
                println!("pml4_index(fn) {:x}", pml4[i].get_address());
                println!("pdpt_index(fn) {:x}", pdpt[j].get_address());
            }
        }
    }

    let mut stack_mmap = Mmap::anonymous(65536, Protection::ReadWrite).unwrap();
    let stack_size = stack_mmap.len();
    let stack_memory = unsafe { stack_mmap.as_mut_slice() };
    let stack_base = VAddr::from_usize(stack_memory.as_mut_ptr() as usize);
    static STACK_BASE_T: PAddr = PAddr::from_u64(0x2000000);
    println!("Stack base {:x} with size {:x}", stack_base, stack_size);

    // Initialize the KVM system
    let sys = System::initialize().unwrap();

    // Create a Virtual Machine
    let mut vm = VirtualMachine::create(&sys).unwrap();

    // Ensure that the VM supports memory backing with user memory
    assert!(vm.check_capability(Capability::UserMemory) > 0);

    // Once the memory is set we can't even call length.
    let page_table_memory_limit = page_table_memory.len() - 1;

    // Map the page table memory
    vm.set_user_memory_region(0, page_table_memory, 0).unwrap();
    // Map stack space
    vm.set_user_memory_region(STACK_BASE_T.as_u64(), stack_memory, 0).unwrap();

    // Map the process
    let f = File::open("/proc/self/maps").unwrap();
    let reader = BufReader::new(f);

    for line in reader.lines() {
        let line = line.unwrap();
        println!("{}", line);
        let mut s = line.split(' ');
        let mut s2 = s.next().unwrap().split('-');
        let begin = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        let end = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        if end < 0x800000000000 {
            let perm = s.next().unwrap();
            //println!("{:#X}-{:#X} {}", begin, end, perm);
            let slice = {
                let begin_ptr: *mut u8 = begin as *const u8 as _;
                unsafe { ::std::slice::from_raw_parts_mut(begin_ptr, end - begin) }
            };
            // Make sure process doesn't overlap with page table
            assert!(begin > page_table_memory_limit);
            vm.set_user_memory_region(begin as _, slice, 0).unwrap();
        }
    }

    // Create a new VCPU
    let mut vcpu = Vcpu::create(&mut vm).unwrap();

    // Set supported CPUID (KVM fails without doing this)
    let mut cpuid = sys.get_supported_cpuid().unwrap();
    /*for mut entry in cpuid.entries_mut() {
        if entry.function == 0x80000008 {
            entry.eax = entry.eax | 0xff;
        }
        println!("{:?}", entry);
    }*/
    vcpu.set_cpuid2(&mut cpuid).unwrap();

    // Setup the special registers
    let mut sregs = vcpu.get_sregs().unwrap();

    // Set the code segment to have base 0, limit 4GB (flat segmentation)
    let segment_template = Segment {
        base: 0x0,
        limit: 0xffffffff,
        selector: 0,
        _type: 0,
        present: 0,
        dpl: 0,
        db: 1,
        s: 0,
        l: 0,
        g: 1,
        avl: 0,
        ..Default::default()
    };

    sregs.cs = Segment {
        selector: 0x8,
        _type: 0xb,
        present: 1,
        db: 0,
        s: 1,
        l: 1,
        ..segment_template
    };
    sregs.ss = Segment { ..segment_template };
    sregs.ds = Segment { ..segment_template };
    sregs.es = Segment { ..segment_template };
    sregs.fs = Segment { ..segment_template };
    sregs.gs = Segment { ..segment_template };

    // We don't need to populate the GDT if we have our segments setup
    // cr0 - protected mode on, paging enabled
    sregs.cr0 = (CR0_PROTECTED_MODE | CR0_MONITOR_COPROCESSOR | CR0_EXTENSION_TYPE |
                 CR0_ENABLE_PAGING | CR0_NUMERIC_ERROR | CR0_WRITE_PROTECT |
                 CR0_ALIGNMENT_MASK | CR0_ENABLE_PAGING)
        .bits() as u64;
    sregs.cr3 = PAGE_TABLE_P.as_u64();
    sregs.cr4 = (CR4_ENABLE_PSE | CR4_ENABLE_PAE | CR4_ENABLE_GLOBAL_PAGES | CR4_ENABLE_SSE |
                 CR4_UNMASKED_SSE |
                 CR4_ENABLE_OS_XSAVE | CR4_ENABLE_SMEP | CR4_ENABLE_VME)
        .bits() as u64;
    sregs.efer = 0xd01;

    // Set the special registers
    vcpu.set_sregs(&sregs).unwrap();

    let mut regs = vcpu.get_regs().unwrap();
    // set the instruction pointer to 1 MB
    regs.rip = vaddr.as_usize() as u64;
    println!("regs.rip = 0x{:x}", regs.rip); // but is at: 0x40cd60
    //println!("regs.rip = 0x{:x}", unsafe { *(regs.rip as *const u64) }); // but is at: 0x40cd60
    //regs.rip = 0x40cd60;
    regs.rflags = 0x246;
    regs.rsp = STACK_BASE_T.as_u64() + stack_size as u64;
    regs.rbp = regs.rsp;
    vcpu.set_regs(&regs).unwrap();

    // Actually run the VCPU

    println!("size of run {:?}", std::mem::size_of::<kvm::Run>());
    let mut vm_is_done = false;
    let mut new_regs = kvm::Regs::default();
    while !vm_is_done {
        {
            let (run, mut regs) = unsafe { vcpu.run_regs() }.unwrap();
            match run.exit_reason {
                Exit::Io => {
                    let io = unsafe { *run.io() };
                    match io.direction {
                        IoDirection::In => {
                            if io.port == 0x3fd {
                                regs.rax = 0x20; // Mark serial line as ready to write
                            } else {
                                println!("IO on unknown port: {}", io.port);
                            }
                        }
                        IoDirection::Out => {
                            if io.port == 0x3f8 {
                                println!("got char {:#?}", regs.rax as u8 as char);
                            }
                        }
                    }
                }
                Exit::Shutdown => {
                    println!("Shutting down");
                    vm_is_done = true;
                }
                _ => {
                    println!("Unknown exit reason: {:?}", run.exit_reason);
                }
            }

            new_regs = regs;
        }
        vcpu.set_regs(&new_regs).unwrap();
    }

    // Ensure that the exit reason we get back indicates that the I/O
    // instruction was executed
    /*let regs = vcpu.get_regs().unwrap();
    println!("vcpu.rip: {:x}", regs.rip);

    assert!(run.exit_reason == Exit::Io);
    let io = unsafe { *run.io() };
    assert!(io.direction == IoDirection::In);
    assert!(io.size == 1);
    assert!(io.port == 0x1);
    unsafe {
        println!("{:#?}", *run.io());
    }*/
}
<|endoftext|>


<commit_before>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;
// ratsio_stream is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::manager::{service::{HealthCheckResult,
                               Service},
                     sys::Sys};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info = EventConnectionInfo::new(config.token, config.url,
                                                     event_core.supervisor_id.clone());
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    site:        Option<String>,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl<'a> From<&'a ArgMatches<'a>> for EventStreamConfig {
    fn from(m: &ArgMatches) -> Self {
        EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            application: m.value_of("EVENT_STREAM_APPLICATION")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            site:        m.value_of("EVENT_STREAM_SITE").map(str::to_string),
                            meta:        EventStreamMetadata::from(m),
                            token:       AutomateAuthToken::from(m),
                            url:         m.value_of("EVENT_STREAM_URL")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"), }
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String, supervisor_id: String) -> Self {
        EventConnectionInfo { name: format!("hab_client_{}", supervisor_id),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    fqdn: String,
    application: String,
    environment: String,
    site: Option<String>,
    meta: EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys, fqdn: String) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address: sys.gossip_listen(),
                    fqdn,
                    environment: config.environment.clone(),
                    application: config.application.clone(),
                    site: config.site.clone(),
                    meta: config.meta.clone() }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
/// `execution` will be `Some` if the service had a hook to run, and
/// records how long it took that hook to execute completely.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    execution: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheckResult = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   execution:        execution.map(Duration::into), });
    }
}

////////////////////////////////////////////////////////////////////////

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { occurred_at:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait {
    secs: u64,
}

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self { secs: 5 } }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self { secs: s.parse()? }) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.secs) }
}
<commit_msg>Fix lint:format errors<commit_after>//! Main interface for a stream of events the Supervisor can send out
//! in the course of its operations.
//!
//! Currently, the Supervisor is able to send events to a [NATS
//! Streaming][1] server. The `init_stream` function must be called
//! before sending events to initialize the publishing thread in the
//! background. Thereafter, you can pass "event" structs to the
//! `event` function, which will publish the event to the stream.
//!
//! All events are published under the "habitat" subject.
//!
//! [1]:https://github.com/nats-io/nats-streaming-server

mod error;
// ratsio_stream is the default, but setting it as a default in Cargo.toml
// makes it trickier to use nitox instead.
#[cfg(feature = "nitox_stream")]
#[path = "event/nitox.rs"]
mod stream_impl;
#[cfg(any(feature = "ratsio_stream", not(feature = "nitox_stream")))]
#[path = "event/ratsio.rs"]
mod stream_impl;
mod types;

pub(crate) use self::types::ServiceMetadata;
use self::types::{EventMessage,
                  EventMetadata,
                  HealthCheckEvent,
                  ServiceStartedEvent,
                  ServiceStoppedEvent};
use crate::manager::{service::{HealthCheckResult,
                               Service},
                     sys::Sys};
use clap::ArgMatches;
pub use error::{Error,
                Result};
use futures::sync::mpsc::UnboundedSender;
use habitat_common::types::{AutomateAuthToken,
                            EventStreamMetadata};
use habitat_core::env::Config as EnvConfig;
use state::Container;
use std::{net::SocketAddr,
          num::ParseIntError,
          str::FromStr,
          sync::Once,
          time::Duration};

static INIT: Once = Once::new();
lazy_static! {
    // TODO (CM): When const fn support lands in stable, we can ditch
    // this lazy_static call.

    /// Reference to the event stream.
    static ref EVENT_STREAM: Container = Container::new();
    /// Core information that is shared between all events.
    static ref EVENT_CORE: Container = Container::new();
}

/// Starts a new thread for sending events to a NATS Streaming
/// server. Stashes the handle to the stream, as well as the core
/// event information that will be a part of all events, in a global
/// static reference for access later.
pub fn init_stream(config: EventStreamConfig, event_core: EventCore) -> Result<()> {
    // call_once can't return a Result (or anything), so we'll fake it
    // by hanging onto any error we might receive.
    let mut return_value: Result<()> = Ok(());

    INIT.call_once(|| {
            let conn_info =
                EventConnectionInfo::new(config.token, config.url, &event_core.supervisor_id);
            match stream_impl::init_stream(conn_info) {
                Ok(event_stream) => {
                    EVENT_STREAM.set(event_stream);
                    EVENT_CORE.set(event_core);
                }
                Err(e) => return_value = Err(e),
            }
        });

    return_value
}

/// Captures all event stream-related configuration options that would
/// be passed in by a user
#[derive(Clone, Debug)]
pub struct EventStreamConfig {
    environment: String,
    application: String,
    site:        Option<String>,
    meta:        EventStreamMetadata,
    token:       AutomateAuthToken,
    url:         String,
}

impl<'a> From<&'a ArgMatches<'a>> for EventStreamConfig {
    fn from(m: &ArgMatches) -> Self {
        EventStreamConfig { environment: m.value_of("EVENT_STREAM_ENVIRONMENT")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            application: m.value_of("EVENT_STREAM_APPLICATION")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"),
                            site:        m.value_of("EVENT_STREAM_SITE").map(str::to_string),
                            meta:        EventStreamMetadata::from(m),
                            token:       AutomateAuthToken::from(m),
                            url:         m.value_of("EVENT_STREAM_URL")
                                          .map(str::to_string)
                                          .expect("Required option for EventStream feature"), }
    }
}

/// All the information needed to establish a connection to a NATS
/// Streaming server.
pub struct EventConnectionInfo {
    pub name:        String,
    pub verbose:     bool,
    pub cluster_uri: String,
    pub cluster_id:  String,
    pub auth_token:  AutomateAuthToken,
}

impl EventConnectionInfo {
    pub fn new(auth_token: AutomateAuthToken, cluster_uri: String, supervisor_id: &str) -> Self {
        EventConnectionInfo { name: format!("hab_client_{}", supervisor_id),
                              verbose: true,
                              cluster_uri,
                              cluster_id: "event-service".to_string(),
                              auth_token }
    }
}

/// A collection of data that will be present in all events. Rather
/// than baking this into the structure of each event, we represent it
/// once and merge the information into the final rendered form of the
/// event.
///
/// This prevents us from having to thread information throughout the
/// system, just to get it to the places where the events are
/// generated (e.g., not all code has direct access to the
/// Supervisor's ID).
#[derive(Clone, Debug)]
pub struct EventCore {
    /// The unique identifier of the Supervisor sending the event.
    supervisor_id: String,
    ip_address: SocketAddr,
    fqdn: String,
    application: String,
    environment: String,
    site: Option<String>,
    meta: EventStreamMetadata,
}

impl EventCore {
    pub fn new(config: &EventStreamConfig, sys: &Sys, fqdn: String) -> Self {
        EventCore { supervisor_id: sys.member_id.clone(),
                    ip_address: sys.gossip_listen(),
                    fqdn,
                    environment: config.environment.clone(),
                    application: config.application.clone(),
                    site: config.site.clone(),
                    meta: config.meta.clone() }
    }
}

/// Send an event for the start of a Service.
pub fn service_started(service: &Service) {
    if stream_initialized() {
        publish(ServiceStartedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

/// Send an event for the stop of a Service.
pub fn service_stopped(service: &Service) {
    if stream_initialized() {
        publish(ServiceStoppedEvent { service_metadata: Some(service.to_service_metadata()),
                                      event_metadata:   None, });
    }
}

// Takes metadata directly, rather than a `&Service` like other event
// functions, because of how the asynchronous health checking
// currently works. Revisit when async/await + Pin is all stabilized.
/// `execution` will be `Some` if the service had a hook to run, and
/// records how long it took that hook to execute completely.
pub fn health_check(metadata: ServiceMetadata,
                    check_result: HealthCheckResult,
                    execution: Option<Duration>) {
    if stream_initialized() {
        let check_result: types::HealthCheckResult = check_result.into();
        publish(HealthCheckEvent { service_metadata: Some(metadata),
                                   event_metadata:   None,
                                   result:           i32::from(check_result),
                                   execution:        execution.map(Duration::into), });
    }
}

////////////////////////////////////////////////////////////////////////

/// Internal helper function to know whether or not to go to the trouble of
/// creating event structures. If the event stream hasn't been
/// initialized, then we shouldn't need to do anything.
fn stream_initialized() -> bool { EVENT_STREAM.try_get::<EventStream>().is_some() }

/// Publish an event. This is the main interface that client code will
/// use.
///
/// If `init_stream` has not been called already, this function will
/// be a no-op.
fn publish(mut event: impl EventMessage) {
    if let Some(e) = EVENT_STREAM.try_get::<EventStream>() {
        // TODO (CM): Yeah... this is looking pretty gross. The
        // intention is to be able to timestamp the events right as
        // they go out.
        //
        // We *could* set the time when we convert the EventCore to a
        // EventMetadata struct, but that seems odd.
        //
        // It probably means that this structure just isn't the right
        // one.
        //
        // The ugliness is at least contained, though.
        event.event_metadata(EventMetadata { occurred_at:
                                                 Some(std::time::SystemTime::now().into()),
                                             ..EVENT_CORE.get::<EventCore>().to_event_metadata() });

        e.send(event.to_bytes());
    }
}

/// A lightweight handle for the event stream. All events get to the
/// event stream through this.
struct EventStream(UnboundedSender<Vec<u8>>);

impl EventStream {
    /// Queues an event to be sent out.
    fn send(&self, event: Vec<u8>) {
        trace!("About to queue an event: {:?}", event);
        if let Err(e) = self.0.unbounded_send(event) {
            error!("Failed to queue event: {:?}", e);
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// How long should we for the event thread to start up before
/// abandoning it and shutting down?
#[derive(Clone, Debug)]
struct EventThreadStartupWait {
    secs: u64,
}

impl Default for EventThreadStartupWait {
    fn default() -> Self { Self { secs: 5 } }
}

impl FromStr for EventThreadStartupWait {
    type Err = ParseIntError;

    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> { Ok(Self { secs: s.parse()? }) }
}

impl EnvConfig for EventThreadStartupWait {
    const ENVVAR: &'static str = "HAB_EVENT_THREAD_STARTUP_WAIT_SEC";
}

impl Into<Duration> for EventThreadStartupWait {
    fn into(self) -> Duration { Duration::from_secs(self.secs) }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "unstable", feature(core))]

extern crate itertools;

use std::fmt::Debug;
#[cfg(feature = "unstable")]
use std::iter::RandomAccessIterator;
use itertools::Itertools;
use itertools::EitherOrBoth::{Both, Left};
use itertools::{
    Zip,
};
#[cfg(feature = "unstable")]
use itertools::ZipTrusted;

#[cfg(feature = "unstable")]
#[test]
fn test_zip_longest_size_hint() {
    let c = 0..;
    let v: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let v2 = &[10, 11, 12];

    assert_eq!(c.zip_longest(v.iter()).size_hint(), (std::usize::MAX, None));

    assert_eq!(v.iter().zip_longest(v2.iter()).size_hint(), (10, Some(10)));
}

#[test]
fn test_double_ended_zip_longest() {
    let xs = [1, 2, 3, 4, 5, 6];
    let ys = [1, 2, 3, 7];
    let a = xs.iter().map(|&x| x);
    let b = ys.iter().map(|&x| x);
    let mut it = a.zip_longest(b);
    assert_eq!(it.next(), Some(Both(1, 1)));
    assert_eq!(it.next(), Some(Both(2, 2)));
    assert_eq!(it.next_back(), Some(Left(6)));
    assert_eq!(it.next_back(), Some(Left(5)));
    assert_eq!(it.next_back(), Some(Both(4, 7)));
    assert_eq!(it.next(), Some(Both(3, 3)));
    assert_eq!(it.next(), None);
}

#[cfg(feature = "unstable")]
// This function copied from std::iter in rust-lang/rust
#[cfg(test)]
fn check_randacc_iter<A, T>(a: T, len: usize) where
    A: PartialEq,
    T: Clone + Iterator<Item=A> + RandomAccessIterator
{
    let mut b = a.clone();
    assert_eq!(len, b.indexable());
    let mut n = 0;
    for (i, elt) in a.enumerate() {
        assert!(Some(elt) == b.idx(i));
        n += 1;
    }
    assert_eq!(n, len);
    assert!(None == b.idx(n));
    // call recursively to check after picking off an element
    if len > 0 {
        b.next();
        check_randacc_iter(b, len-1);
    }
}

#[cfg(feature = "unstable")]
#[test]
fn test_random_access_zip_longest() {
    let xs = [1, 2, 3, 4, 5];
    let ys = [7, 9, 11];
    check_randacc_iter(xs.iter().zip_longest(ys.iter()), std::cmp::max(xs.len(), ys.len()));
}

#[test]
fn zip_tuple() {
    let xs = [1, 2, 3];
    let ys = b"ab";
    let mut it = Zip::new((xs.iter().cloned(), ));
    assert!(it.next() != None);
    let mut jt = Zip::new((xs.iter().cloned(), ys.iter().cloned()));
    assert_eq!(jt.next(), Some((1, b'a')));
    assert_eq!(jt.next(), Some((2, b'b')));
    assert_eq!(jt.next(), None);

    let mut jt = Zip::new((xs.iter().cloned(), xs.iter().cloned(), xs.iter().cloned()));
    assert_eq!(jt.next(), Some((1, 1, 1)));
}

fn assert_iters_equal<
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>>(mut it: I, mut jt: J)
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[cfg(feature = "unstable")]
#[test]
fn ziptrusted_1() {
    let mut xs = [0; 6];
    let mut ys = [0; 8];
    let mut zs = [0; 7];

    xs.iter_mut().set_from(0..);
    ys.iter_mut().set_from(0..);
    zs.iter_mut().set_from(0..);

    let it = ZipTrusted::new((xs.iter(), ys.iter()));
    assert_eq!(it.size_hint(), (6, Some(6)));
    assert_iters_equal(it, xs.iter().zip(ys.iter()));

    let it = ZipTrusted::new((xs.iter(), ys.iter(), zs.iter()));
    assert_eq!(it.size_hint(), (6, Some(6)));
    assert_iters_equal(it, xs.iter()
                             .zip(ys.iter())
                             .zip(zs.iter())
                             .map(|((a, b), c)| (a, b, c)));
}
<commit_msg>Fix zip length test<commit_after>#![cfg_attr(feature = "unstable", feature(core))]

extern crate itertools;

use std::fmt::Debug;
#[cfg(feature = "unstable")]
use std::iter::RandomAccessIterator;
use itertools::Itertools;
use itertools::EitherOrBoth::{Both, Left};
use itertools::{
    Zip,
};
#[cfg(feature = "unstable")]
use itertools::ZipTrusted;

#[cfg(feature = "unstable")]
#[test]
fn test_zip_longest_size_hint() {
    let c = (1..10).cycle();
    let v: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let v2 = &[10, 11, 12];

    assert_eq!(c.zip_longest(v.iter()).size_hint(), (std::usize::MAX, None));

    assert_eq!(v.iter().zip_longest(v2.iter()).size_hint(), (10, Some(10)));
}

#[test]
fn test_double_ended_zip_longest() {
    let xs = [1, 2, 3, 4, 5, 6];
    let ys = [1, 2, 3, 7];
    let a = xs.iter().map(|&x| x);
    let b = ys.iter().map(|&x| x);
    let mut it = a.zip_longest(b);
    assert_eq!(it.next(), Some(Both(1, 1)));
    assert_eq!(it.next(), Some(Both(2, 2)));
    assert_eq!(it.next_back(), Some(Left(6)));
    assert_eq!(it.next_back(), Some(Left(5)));
    assert_eq!(it.next_back(), Some(Both(4, 7)));
    assert_eq!(it.next(), Some(Both(3, 3)));
    assert_eq!(it.next(), None);
}

#[cfg(feature = "unstable")]
// This function copied from std::iter in rust-lang/rust
#[cfg(test)]
fn check_randacc_iter<A, T>(a: T, len: usize) where
    A: PartialEq,
    T: Clone + Iterator<Item=A> + RandomAccessIterator
{
    let mut b = a.clone();
    assert_eq!(len, b.indexable());
    let mut n = 0;
    for (i, elt) in a.enumerate() {
        assert!(Some(elt) == b.idx(i));
        n += 1;
    }
    assert_eq!(n, len);
    assert!(None == b.idx(n));
    // call recursively to check after picking off an element
    if len > 0 {
        b.next();
        check_randacc_iter(b, len-1);
    }
}

#[cfg(feature = "unstable")]
#[test]
fn test_random_access_zip_longest() {
    let xs = [1, 2, 3, 4, 5];
    let ys = [7, 9, 11];
    check_randacc_iter(xs.iter().zip_longest(ys.iter()), std::cmp::max(xs.len(), ys.len()));
}

#[test]
fn zip_tuple() {
    let xs = [1, 2, 3];
    let ys = b"ab";
    let mut it = Zip::new((xs.iter().cloned(), ));
    assert!(it.next() != None);
    let mut jt = Zip::new((xs.iter().cloned(), ys.iter().cloned()));
    assert_eq!(jt.next(), Some((1, b'a')));
    assert_eq!(jt.next(), Some((2, b'b')));
    assert_eq!(jt.next(), None);

    let mut jt = Zip::new((xs.iter().cloned(), xs.iter().cloned(), xs.iter().cloned()));
    assert_eq!(jt.next(), Some((1, 1, 1)));
}

fn assert_iters_equal<
    A: PartialEq + Debug,
    I: Iterator<Item=A>,
    J: Iterator<Item=A>>(mut it: I, mut jt: J)
{
    loop {
        let elti = it.next();
        let eltj = jt.next();
        assert_eq!(elti, eltj);
        if elti.is_none() { break; }
    }
}

#[cfg(feature = "unstable")]
#[test]
fn ziptrusted_1() {
    let mut xs = [0; 6];
    let mut ys = [0; 8];
    let mut zs = [0; 7];

    xs.iter_mut().set_from(0..);
    ys.iter_mut().set_from(0..);
    zs.iter_mut().set_from(0..);

    let it = ZipTrusted::new((xs.iter(), ys.iter()));
    assert_eq!(it.size_hint(), (6, Some(6)));
    assert_iters_equal(it, xs.iter().zip(ys.iter()));

    let it = ZipTrusted::new((xs.iter(), ys.iter(), zs.iter()));
    assert_eq!(it.size_hint(), (6, Some(6)));
    assert_iters_equal(it, xs.iter()
                             .zip(ys.iter())
                             .zip(zs.iter())
                             .map(|((a, b), c)| (a, b, c)));
}
<|endoftext|>


<commit_before>// Copyright 2019 Tyler Neely
//
// 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.

extern crate rocksdb;
extern crate libc;

mod util;

use libc::{size_t};

use rocksdb::{DB, DBVector, Error, IteratorMode, Options, WriteBatch};
use std::str;
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { mem::transmute(libc::calloc(len, mem::size_of::<u8>())) };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
  let path = DBPath::new("_rust_rocksdb_externaltest");
  let db = DB::open_default(&path).unwrap();
  let p = db.put(b"k1", b"v1111");
  assert!(p.is_ok());
  let r: Result<Option<DBVector>, Error> = db.get(b"k1");
  assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
  assert!(db.delete(b"k1").is_ok());
  assert!(db.get(b"k1").unwrap().is_none());
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            let p = db.write(batch);
            assert!(p.is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            let p = db.write(batch);
            assert!(p.is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let db = DB::open_default(&path).unwrap();
        let p = db.put(b"k1", b"v1111");
        assert!(p.is_ok());
        let p = db.put(b"k2", b"v2222");
        assert!(p.is_ok());
        let p = db.put(b"k3", b"v3333");
        assert!(p.is_ok());
        let iter = db.iterator(IteratorMode::Start);
        for (k, v) in iter {
            println!(
                "Hello {}: {}",
                str::from_utf8(&*k).unwrap(),
                str::from_utf8(&*v).unwrap()
            );
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();
        let p = db.put(b"k1", b"v1111");
        assert!(p.is_ok());

        let snap = db.snapshot();
        let r: Result<Option<DBVector>, Error> = snap.get(b"k1");
        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");

        let p = db.put(b"k2", b"v2222");
        assert!(p.is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<commit_msg>Some test cleanup...<commit_after>// Copyright 2019 Tyler Neely
//
// 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.

extern crate rocksdb;
extern crate libc;

mod util;

use libc::{size_t};

use rocksdb::{DB, DBVector, Error, IteratorMode, Options, WriteBatch};
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { mem::transmute(libc::calloc(len, mem::size_of::<u8>())) };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
    let path = DBPath::new("_rust_rocksdb_externaltest");

    {
        let db = DB::open_default(&path).unwrap();

        let p = db.put(b"k1", b"v1111");
        assert!(p.is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");

        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            let p = db.write(batch);
            assert!(p.is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            let p = db.write(batch);
            assert!(p.is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let data = [(b"k1", b"v1111"), (b"k2", b"v2222"), (b"k3", b"v3333")];
        let db = DB::open_default(&path).unwrap();

        for (key, value) in &data {
            assert!(db.put(key, value).is_ok());
        }

        let iter = db.iterator(IteratorMode::Start);

        for (idx, (db_key, db_value)) in iter.enumerate() {
            let (key, value) = data[idx];
            assert_eq!((&key[..], &value[..]), (db_key.as_ref(), db_value.as_ref()));
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let snap = db.snapshot();
        let r: Result<Option<DBVector>, Error> = snap.get(b"k1");
        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");

        assert!(db.put(b"k2", b"v2222").is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fmt;
use std::fs;
use std::io::{self, Read};
use std::path::{Path, PathBuf};
use std::process;
use std::str::FromStr;
use std::sync::atomic;
use std::time::Duration;

use csv;

use Csv;

static XSV_INTEGRATION_TEST_DIR: &'static str = "xit";

static NEXT_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;

pub struct Workdir {
    root: PathBuf,
    dir: PathBuf,
    flexible: bool,
}

impl Workdir {
    pub fn new(name: &str) -> Workdir {
        let id = NEXT_ID.fetch_add(1, atomic::Ordering::SeqCst);
        let root = env::current_exe().unwrap()
                       .parent()
                       .expect("executable's directory")
                       .to_path_buf();
        let dir = root.join(XSV_INTEGRATION_TEST_DIR)
                      .join(name)
                      .join(&format!("test-{}", id));
        if let Err(err) = create_dir_all(&dir) {
            panic!("Could not create '{:?}': {}", dir, err);
        }
        Workdir { root: root, dir: dir, flexible: false }
    }

    pub fn flexible(mut self, yes: bool) -> Workdir {
        self.flexible = yes;
        self
    }

    pub fn create<T: Csv>(&self, name: &str, rows: T) {
        let mut wtr = csv::Writer::from_file(&self.path(name))
                                  .unwrap().flexible(self.flexible);
        for row in rows.to_vecs().into_iter() {
            wtr.write(row.iter()).unwrap();
        }
        wtr.flush().unwrap();
    }

    pub fn create_indexed<T: Csv>(&self, name: &str, rows: T) {
        self.create(name, rows);

        let mut cmd = self.command("index");
        cmd.arg(name);
        self.run(&mut cmd);
    }

    pub fn read_stdout<T: Csv>(&self, cmd: &mut process::Command) -> T {
        let stdout: String = self.stdout(cmd);
        let mut rdr = csv::Reader::from_string(stdout).has_headers(false);
        Csv::from_vecs(rdr.records().collect::<Result<_, _>>().unwrap())
    }

    pub fn command(&self, sub_command: &str) -> process::Command {
        let mut cmd = process::Command::new(&self.xsv_bin());
        cmd.current_dir(&self.dir).arg(sub_command);
        cmd
    }

    pub fn output(&self, cmd: &mut process::Command) -> process::Output {
        debug!("[{}]: {:?}", self.dir.display(), cmd);
        let o = cmd.output().unwrap();
        if !o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command failed but expected success!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr))
        }
        o
    }

    pub fn run(&self, cmd: &mut process::Command) {
        self.output(cmd);
    }

    pub fn stdout<T: FromStr>(&self, cmd: &mut process::Command) -> T {
        let o = self.output(cmd);
        let stdout = String::from_utf8_lossy(&o.stdout);
        stdout.trim_matches(&['\r', '\n'][..]).parse().ok().expect(
            &format!("Could not convert from string: '{}'", stdout))
    }

    pub fn assert_err(&self, cmd: &mut process::Command) {
        let o = cmd.output().unwrap();
        if o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command succeeded but expected failure!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr));
        }
    }

    pub fn from_str<T: FromStr>(&self, name: &Path) -> T {
        let mut o = String::new();
        fs::File::open(name).unwrap().read_to_string(&mut o).unwrap();
        o.parse().ok().expect("fromstr")
    }

    pub fn path(&self, name: &str) -> PathBuf {
        self.dir.join(name)
    }

    pub fn xsv_bin(&self) -> PathBuf {
        self.root.join("xsv")
    }
}

impl fmt::Debug for Workdir {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "path={}", self.dir.display())
    }
}

// For whatever reason, `fs::create_dir_all` fails intermittently on Travis
// with a weird "file exists" error. Despite my best efforts to get to the
// bottom of it, I've decided a try-wait-and-retry hack is good enough.
fn create_dir_all<P: AsRef<Path>>(p: P) -> io::Result<()> {
    let mut last_err = None;
    for _ in 0..10 {
        if let Err(err) = fs::create_dir_all(&p) {
            last_err = Some(err);
            ::std::thread::sleep(Duration::from_millis(500));
        } else {
            return Ok(())
        }
    }
    Err(last_err.unwrap())
}
<commit_msg>Fix tests on beta<commit_after>use std::env;
use std::fmt;
use std::fs;
use std::io::{self, Read};
use std::path::{Path, PathBuf};
use std::process;
use std::str::FromStr;
use std::sync::atomic;
use std::time::Duration;

use csv;

use Csv;

static XSV_INTEGRATION_TEST_DIR: &'static str = "xit";

static NEXT_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;

pub struct Workdir {
    root: PathBuf,
    dir: PathBuf,
    flexible: bool,
}

impl Workdir {
    pub fn new(name: &str) -> Workdir {
        let id = NEXT_ID.fetch_add(1, atomic::Ordering::SeqCst);
        let mut root = env::current_exe().unwrap()
                           .parent()
                           .expect("executable's directory")
                           .to_path_buf();
        if root.ends_with("deps") {
            root.pop();
        }
        let dir = root.join(XSV_INTEGRATION_TEST_DIR)
                      .join(name)
                      .join(&format!("test-{}", id));
        if let Err(err) = create_dir_all(&dir) {
            panic!("Could not create '{:?}': {}", dir, err);
        }
        Workdir { root: root, dir: dir, flexible: false }
    }

    pub fn flexible(mut self, yes: bool) -> Workdir {
        self.flexible = yes;
        self
    }

    pub fn create<T: Csv>(&self, name: &str, rows: T) {
        let mut wtr = csv::Writer::from_file(&self.path(name))
                                  .unwrap().flexible(self.flexible);
        for row in rows.to_vecs().into_iter() {
            wtr.write(row.iter()).unwrap();
        }
        wtr.flush().unwrap();
    }

    pub fn create_indexed<T: Csv>(&self, name: &str, rows: T) {
        self.create(name, rows);

        let mut cmd = self.command("index");
        cmd.arg(name);
        self.run(&mut cmd);
    }

    pub fn read_stdout<T: Csv>(&self, cmd: &mut process::Command) -> T {
        let stdout: String = self.stdout(cmd);
        let mut rdr = csv::Reader::from_string(stdout).has_headers(false);
        Csv::from_vecs(rdr.records().collect::<Result<_, _>>().unwrap())
    }

    pub fn command(&self, sub_command: &str) -> process::Command {
        let mut cmd = process::Command::new(&self.xsv_bin());
        cmd.current_dir(&self.dir).arg(sub_command);
        cmd
    }

    pub fn output(&self, cmd: &mut process::Command) -> process::Output {
        debug!("[{}]: {:?}", self.dir.display(), cmd);
        let o = cmd.output().unwrap();
        if !o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command failed but expected success!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr))
        }
        o
    }

    pub fn run(&self, cmd: &mut process::Command) {
        self.output(cmd);
    }

    pub fn stdout<T: FromStr>(&self, cmd: &mut process::Command) -> T {
        let o = self.output(cmd);
        let stdout = String::from_utf8_lossy(&o.stdout);
        stdout.trim_matches(&['\r', '\n'][..]).parse().ok().expect(
            &format!("Could not convert from string: '{}'", stdout))
    }

    pub fn assert_err(&self, cmd: &mut process::Command) {
        let o = cmd.output().unwrap();
        if o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command succeeded but expected failure!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr));
        }
    }

    pub fn from_str<T: FromStr>(&self, name: &Path) -> T {
        let mut o = String::new();
        fs::File::open(name).unwrap().read_to_string(&mut o).unwrap();
        o.parse().ok().expect("fromstr")
    }

    pub fn path(&self, name: &str) -> PathBuf {
        self.dir.join(name)
    }

    pub fn xsv_bin(&self) -> PathBuf {
        self.root.join("xsv")
    }
}

impl fmt::Debug for Workdir {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "path={}", self.dir.display())
    }
}

// For whatever reason, `fs::create_dir_all` fails intermittently on Travis
// with a weird "file exists" error. Despite my best efforts to get to the
// bottom of it, I've decided a try-wait-and-retry hack is good enough.
fn create_dir_all<P: AsRef<Path>>(p: P) -> io::Result<()> {
    let mut last_err = None;
    for _ in 0..10 {
        if let Err(err) = fs::create_dir_all(&p) {
            last_err = Some(err);
            ::std::thread::sleep(Duration::from_millis(500));
        } else {
            return Ok(())
        }
    }
    Err(last_err.unwrap())
}
<|endoftext|>


<commit_before>//! This module provides the Torrance Sparrow microfacet BRDF, see
//! [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
//! for details.

use std::f32;
use enum_set::EnumSet;

use linalg::{self, Vector};
use film::Colorf;
use bxdf::{self, BxDF, BxDFType};
use bxdf::fresnel::Fresnel;
use bxdf::microfacet::{MicrofacetDistribution};

/// Struct providing the Torrance Sparrow BRDF, implemented as described in
/// [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
#[derive(Copy, Clone)]
pub struct TorranceSparrow<'a> {
    reflectance: Colorf,
    fresnel: &'a Fresnel,
    /// Microfacet distribution describing the structure of the microfacets of
    /// the material
    microfacet: &'a MicrofacetDistribution,
}

impl<'a> TorranceSparrow<'a> {
    /// Create a new Torrance Sparrow microfacet BRDF
    pub fn new(c: &Colorf, fresnel: &'a Fresnel, microfacet: &'a MicrofacetDistribution)
        -> TorranceSparrow<'a> {
        TorranceSparrow { reflectance: *c, fresnel: fresnel, microfacet: microfacet }
    }
}

impl<'a> BxDF for TorranceSparrow<'a> {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Glossy);
        e.insert(BxDFType::Reflection);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        let cos_to = f32::abs(bxdf::cos_theta(w_o));
        let cos_ti = f32::abs(bxdf::cos_theta(w_i));
        if cos_to == 0.0 || cos_ti == 0.0 {
            return Colorf::new(0.0, 0.0, 0.0)
        }
        let mut w_h = *w_i + *w_o;
        if w_h.x == 0.0 && w_h.y == 0.0 && w_h.z == 0.0 {
            return Colorf::new(0.0, 0.0, 0.0)
        }
        w_h = w_h.normalized();
        let d = self.microfacet.normal_distribution(&w_h);
        let f = self.fresnel.fresnel(linalg::dot(w_i, &w_h));
        let g = self.microfacet.shadowing_masking(w_i, w_o, &w_h);
        (self.reflectance * f * d * g / (4.0 * cos_ti * cos_to))
    }
    fn sample(&self, w_o: &Vector, samples: &(f32, f32)) -> (Colorf, Vector, f32) {
        let mut w_h = self.microfacet.sample(w_o, samples);
        if !bxdf::same_hemisphere(w_o, &w_h) {
            w_h = -w_h;
        }
        let w_i = linalg::reflect(w_o, &w_h);
        if !bxdf::same_hemisphere(w_o, &w_i) {
            (Colorf::black(), Vector::broadcast(0.0), 0.0)
        } else {
            // This term is p_o(o) in eq. 38 of Walter et al's 07 paper and is for reflection so
            // we use the Jacobian for reflection, eq. 14
            let jacobian = 1.0 / (4.0 * f32::abs(linalg::dot(w_o, &w_h)));
            let pdf = self.microfacet.pdf(&w_h) * jacobian;
            (self.eval(w_o, &w_i), w_i, pdf)
        }
    }
    fn pdf(&self, w_o: &Vector, w_i: &Vector) -> f32 {
        if !bxdf::same_hemisphere(w_o, w_i) {
            0.0
        } else {
            let mut w_h = *w_o + *w_i;
            if w_h.x == 0.0 && w_h.y == 0.0 && w_h.z == 0.0 {
                0.0
            } else {
                w_h = w_h.normalized();
                // This term is p_o(o) in eq. 38 of Walter et al's 07 paper and is for reflection so
                // we use the Jacobian for reflection, eq. 14
                let jacobian = 1.0 / (4.0 * f32::abs(linalg::dot(w_o, &w_h)));
                self.microfacet.pdf(&w_h.normalized()) * jacobian
            }
        }
    }
}

<commit_msg>some minor clean-ups<commit_after>//! This module provides the Torrance Sparrow microfacet BRDF, see
//! [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
//! for details.

use std::f32;
use enum_set::EnumSet;

use linalg::{self, Vector};
use film::Colorf;
use bxdf::{self, BxDF, BxDFType};
use bxdf::fresnel::Fresnel;
use bxdf::microfacet::{MicrofacetDistribution};

/// Struct providing the Torrance Sparrow BRDF, implemented as described in
/// [Walter et al. 07](https://www.cs.cornell.edu/~srm/publications/EGSR07-btdf.pdf)
#[derive(Copy, Clone)]
pub struct TorranceSparrow<'a> {
    reflectance: Colorf,
    fresnel: &'a Fresnel,
    /// Microfacet distribution describing the structure of the microfacets of
    /// the material
    microfacet: &'a MicrofacetDistribution,
}

impl<'a> TorranceSparrow<'a> {
    /// Create a new Torrance Sparrow microfacet BRDF
    pub fn new(c: &Colorf, fresnel: &'a Fresnel, microfacet: &'a MicrofacetDistribution)
        -> TorranceSparrow<'a> {
        TorranceSparrow { reflectance: *c, fresnel: fresnel, microfacet: microfacet }
    }
}

impl<'a> BxDF for TorranceSparrow<'a> {
    fn bxdf_type(&self) -> EnumSet<BxDFType> {
        let mut e = EnumSet::new();
        e.insert(BxDFType::Glossy);
        e.insert(BxDFType::Reflection);
        e
    }
    fn eval(&self, w_o: &Vector, w_i: &Vector) -> Colorf {
        let cos_to = f32::abs(bxdf::cos_theta(w_o));
        let cos_ti = f32::abs(bxdf::cos_theta(w_i));
        if cos_to == 0.0 || cos_ti == 0.0 {
            return Colorf::new(0.0, 0.0, 0.0)
        }
        let mut w_h = *w_i + *w_o;
        if w_h == Vector::broadcast(0.0) {
            return Colorf::new(0.0, 0.0, 0.0)
        }
        w_h = w_h.normalized();
        let d = self.microfacet.normal_distribution(&w_h);
        let f = self.fresnel.fresnel(linalg::dot(w_i, &w_h));
        let g = self.microfacet.shadowing_masking(w_i, w_o, &w_h);
        (self.reflectance * f * d * g / (4.0 * cos_ti * cos_to))
    }
    fn sample(&self, w_o: &Vector, samples: &(f32, f32)) -> (Colorf, Vector, f32) {
        if w_o.z == 0.0 {
            return (Colorf::black(), Vector::broadcast(0.0), 0.0)
        }
        let mut w_h = self.microfacet.sample(w_o, samples);
        if !bxdf::same_hemisphere(w_o, &w_h) {
            w_h = -w_h;
        }
        let w_i = linalg::reflect(w_o, &w_h);
        if !bxdf::same_hemisphere(w_o, &w_i) {
            (Colorf::black(), Vector::broadcast(0.0), 0.0)
        } else {
            (self.eval(w_o, &w_i), w_i, self.pdf(w_o, &w_i))
        }
    }
    fn pdf(&self, w_o: &Vector, w_i: &Vector) -> f32 {
        if !bxdf::same_hemisphere(w_o, w_i) {
            0.0
        } else {
            let w_h = (*w_o + *w_i).normalized();
            // This term is p_o(o) in eq. 38 of Walter et al's 07 paper and is for reflection so
            // we use the Jacobian for reflection, eq. 14
            let jacobian = 1.0 / (4.0 * f32::abs(linalg::dot(w_o, &w_h)));
            self.microfacet.pdf(&w_h) * jacobian
        }
    }
}

<|endoftext|>