<commit_before>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! ARM PL031 Real Time Clock
//!
//! This module implements a PL031 Real Time Clock (RTC) that provides to provides long time base counter.
//! This is achieved by generating an interrupt signal after counting for a programmed number of cycles of
//! a real-time clock input.
//!
use std::fmt;
use std::time::Instant;
use std::{io, result};

use crate::BusDevice;
use logger::{warn, IncMetric, METRICS};
use utils::byte_order;
use utils::eventfd::EventFd;
//use bus::Error;

// As you can see in https://static.docs.arm.com/ddi0224/c/real_time_clock_pl031_r1p3_technical_reference_manual_DDI0224C.pdf
// at section 3.2 Summary of RTC registers, the total size occupied by this device is 0x000 -> 0xFFC + 4 = 0x1000.
// From 0x0 to 0x1C we have following registers:
const RTCDR: u64 = 0x0; // Data Register.
const RTCMR: u64 = 0x4; // Match Register.
const RTCLR: u64 = 0x8; // Load Regiser.
const RTCCR: u64 = 0xc; // Control Register.
const RTCIMSC: u64 = 0x10; // Interrupt Mask Set or Clear Register.
const RTCRIS: u64 = 0x14; // Raw Interrupt Status.
const RTCMIS: u64 = 0x18; // Masked Interrupt Status.
const RTCICR: u64 = 0x1c; // Interrupt Clear Register.
                          // From 0x020 to 0xFDC => reserved space.
                          // From 0xFE0 to 0x1000 => Peripheral and PrimeCell Identification Registers which are Read Only registers.
                          // AMBA standard devices have CIDs (Cell IDs) and PIDs (Peripheral IDs). The linux kernel will look for these in order to assert the identity
                          // of these devices (i.e look at the `amba_device_try_add` function).
                          // We are putting the expected values (look at 'Reset value' column from above mentioned document) in an array.
const PL031_ID: [u8; 8] = [0x31, 0x10, 0x14, 0x00, 0x0d, 0xf0, 0x05, 0xb1];
// We are only interested in the margins.
const AMBA_ID_LOW: u64 = 0xFE0;
const AMBA_ID_HIGH: u64 = 0x1000;

#[derive(Debug)]
pub enum Error {
    BadWriteOffset(u64),
    InterruptFailure(io::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::BadWriteOffset(offset) => write!(f, "Bad Write Offset: {}", offset),
            Error::InterruptFailure(e) => write!(f, "Failed to trigger interrupt: {}", e),
        }
    }
}
type Result<T> = result::Result<T, Error>;

/// A RTC device following the PL031 specification..
pub struct RTC {
    previous_now: Instant,
    tick_offset: i64,
    // This is used for implementing the RTC alarm. However, in Firecracker we do not need it.
    match_value: u32,
    // Writes to this register load an update value into the RTC.
    load: u32,
    imsc: u32,
    ris: u32,
    interrupt_evt: EventFd,
}

impl RTC {
    /// Constructs an AMBA PL031 RTC device.
    pub fn new(interrupt_evt: EventFd) -> RTC {
        RTC {
            // This is used only for duration measuring purposes.
            previous_now: Instant::now(),
            tick_offset: utils::time::get_time_ns(utils::time::ClockType::Real) as i64,
            match_value: 0,
            load: 0,
            imsc: 0,
            ris: 0,
            interrupt_evt,
        }
    }

    fn trigger_interrupt(&mut self) -> Result<()> {
        self.interrupt_evt.write(1).map_err(Error::InterruptFailure)
    }

    fn get_time(&self) -> u32 {
        let ts = (self.tick_offset as i128)
            + (Instant::now().duration_since(self.previous_now).as_nanos() as i128);
        (ts / utils::time::NANOS_PER_SECOND as i128) as u32
    }

    fn handle_write(&mut self, offset: u64, val: u32) -> Result<()> {
        match offset {
            RTCMR => {
                // The MR register is used for implementing the RTC alarm. A real time clock alarm is
                // a feature that can be used to allow a computer to 'wake up' after shut down to execute
                // tasks every day or on a certain day. It can sometimes be found in the 'Power Management'
                // section of a motherboard's BIOS setup. This is functionality that extends beyond
                // Firecracker intended use. However, we increment a metric just in case.
                self.match_value = val;
                METRICS.rtc.missed_write_count.inc();
            }
            RTCLR => {
                self.load = val;
                self.previous_now = Instant::now();
                // If the unwrap fails, then the internal value of the clock has been corrupted and
                // we want to terminate the execution of the process.
                self.tick_offset = utils::time::seconds_to_nanoseconds(i64::from(val))
                    .expect("Time conversion overflow");
            }
            RTCIMSC => {
                self.imsc = val & 1;
                self.trigger_interrupt()?;
            }
            RTCICR => {
                // As per above mentioned doc, the interrupt is cleared by writing any data value to
                // the Interrupt Clear Register.
                self.ris = 0;
                self.trigger_interrupt()?;
            }
            RTCCR => (), // ignore attempts to turn off the timer.
            o => {
                return Err(Error::BadWriteOffset(o));
            }
        }
        Ok(())
    }
}

impl BusDevice for RTC {
    fn read(&mut self, offset: u64, data: &mut [u8]) {
        let v;
        let mut read_ok = true;

        if offset < AMBA_ID_HIGH && offset >= AMBA_ID_LOW {
            let index = ((offset - AMBA_ID_LOW) >> 2) as usize;
            v = u32::from(PL031_ID[index]);
        } else {
            v = match offset {
                RTCDR => self.get_time(),
                RTCMR => {
                    METRICS.rtc.missed_read_count.inc();
                    // Even though we are not implementing RTC alarm we return the last value
                    self.match_value
                }
                RTCLR => self.load,
                RTCCR => 1, // RTC is always enabled.
                RTCIMSC => self.imsc,
                RTCRIS => self.ris,
                RTCMIS => self.ris & self.imsc,
                _ => {
                    read_ok = false;
                    0
                }
            };
        }
        if read_ok && data.len() <= 4 {
            byte_order::write_le_u32(data, v);
        } else {
            warn!(
                "Invalid RTC PL031 read: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }

    fn write(&mut self, offset: u64, data: &[u8]) {
        if data.len() <= 4 {
            let v = byte_order::read_le_u32(&data[..]);
            if let Err(e) = self.handle_write(offset, v) {
                warn!("Failed to write to RTC PL031 device: {}", e);
                METRICS.rtc.error_count.inc();
            }
        } else {
            warn!(
                "Invalid RTC PL031 write: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }
}

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

    #[test]
    fn test_rtc_read_write_and_event() {
        let mut rtc = RTC::new(EventFd::new(libc::EFD_NONBLOCK).unwrap());
        let mut data = [0; 4];

        // Read and write to the MR register.
        byte_order::write_le_u32(&mut data, 123);
        rtc.write(RTCMR, &data);
        rtc.read(RTCMR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 123);

        // Read and write to the LR register.
        let v = utils::time::get_time_ns(utils::time::ClockType::Real);
        byte_order::write_le_u32(&mut data, (v / utils::time::NANOS_PER_SECOND) as u32);
        let previous_now_before = rtc.previous_now;
        rtc.write(RTCLR, &data);

        assert!(rtc.previous_now > previous_now_before);

        rtc.read(RTCLR, &mut data);
        let v_read = byte_order::read_le_u32(&data[..]);
        assert_eq!((v / utils::time::NANOS_PER_SECOND) as u32, v_read);

        // Read and write to IMSC register.
        // Test with non zero value.
        let non_zero = 1;
        byte_order::write_le_u32(&mut data, non_zero);
        rtc.write(RTCIMSC, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() == 1);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(non_zero & 1, v);

        // Now test with 0.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCIMSC, &data);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(0, v);

        // Read and write to the ICR register.
        byte_order::write_le_u32(&mut data, 1);
        rtc.write(RTCICR, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() > 1);
        let v_before = byte_order::read_le_u32(&data[..]);
        let no_errors_before = METRICS.rtc.error_count.count();

        rtc.read(RTCICR, &mut data);
        let no_errors_after = METRICS.rtc.error_count.count();
        let v = byte_order::read_le_u32(&data[..]);
        // ICR is a  write only register. Data received should stay equal to data sent.
        assert_eq!(v, v_before);
        assert_eq!(no_errors_after - no_errors_before, 1);

        // Attempts to turn off the RTC should not go through.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCCR, &data);
        rtc.read(RTCCR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 1);

        // Attempts to write beyond the writable space. Using here the space used to read
        // the CID and PID from.
        byte_order::write_le_u32(&mut data, 0);
        let no_errors_before = METRICS.rtc.error_count.count();
        rtc.write(AMBA_ID_LOW, &data);
        let no_errors_after = METRICS.rtc.error_count.count();
        assert_eq!(no_errors_after - no_errors_before, 1);
        // However, reading from the AMBA_ID_LOW should succeed upon read.

        let mut data = [0; 4];
        rtc.read(AMBA_ID_LOW, &mut data);
        let index = AMBA_ID_LOW + 3;
        assert_eq!(data[0], PL031_ID[((index - AMBA_ID_LOW) >> 2) as usize]);
    }
}
<commit_msg>rtc: API for obtaining interrupt event<commit_after>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! ARM PL031 Real Time Clock
//!
//! This module implements a PL031 Real Time Clock (RTC) that provides to provides long time base counter.
//! This is achieved by generating an interrupt signal after counting for a programmed number of cycles of
//! a real-time clock input.
//!
use std::fmt;
use std::time::Instant;
use std::{io, result};

use crate::BusDevice;
use logger::{warn, IncMetric, METRICS};
use utils::byte_order;
use utils::eventfd::EventFd;
//use bus::Error;

// As you can see in https://static.docs.arm.com/ddi0224/c/real_time_clock_pl031_r1p3_technical_reference_manual_DDI0224C.pdf
// at section 3.2 Summary of RTC registers, the total size occupied by this device is 0x000 -> 0xFFC + 4 = 0x1000.
// From 0x0 to 0x1C we have following registers:
const RTCDR: u64 = 0x0; // Data Register.
const RTCMR: u64 = 0x4; // Match Register.
const RTCLR: u64 = 0x8; // Load Regiser.
const RTCCR: u64 = 0xc; // Control Register.
const RTCIMSC: u64 = 0x10; // Interrupt Mask Set or Clear Register.
const RTCRIS: u64 = 0x14; // Raw Interrupt Status.
const RTCMIS: u64 = 0x18; // Masked Interrupt Status.
const RTCICR: u64 = 0x1c; // Interrupt Clear Register.
                          // From 0x020 to 0xFDC => reserved space.
                          // From 0xFE0 to 0x1000 => Peripheral and PrimeCell Identification Registers which are Read Only registers.
                          // AMBA standard devices have CIDs (Cell IDs) and PIDs (Peripheral IDs). The linux kernel will look for these in order to assert the identity
                          // of these devices (i.e look at the `amba_device_try_add` function).
                          // We are putting the expected values (look at 'Reset value' column from above mentioned document) in an array.
const PL031_ID: [u8; 8] = [0x31, 0x10, 0x14, 0x00, 0x0d, 0xf0, 0x05, 0xb1];
// We are only interested in the margins.
const AMBA_ID_LOW: u64 = 0xFE0;
const AMBA_ID_HIGH: u64 = 0x1000;

#[derive(Debug)]
pub enum Error {
    BadWriteOffset(u64),
    InterruptFailure(io::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::BadWriteOffset(offset) => write!(f, "Bad Write Offset: {}", offset),
            Error::InterruptFailure(e) => write!(f, "Failed to trigger interrupt: {}", e),
        }
    }
}
type Result<T> = result::Result<T, Error>;

/// A RTC device following the PL031 specification..
pub struct RTC {
    previous_now: Instant,
    tick_offset: i64,
    // This is used for implementing the RTC alarm. However, in Firecracker we do not need it.
    match_value: u32,
    // Writes to this register load an update value into the RTC.
    load: u32,
    imsc: u32,
    ris: u32,
    interrupt_evt: EventFd,
}

impl RTC {
    /// Constructs an AMBA PL031 RTC device.
    pub fn new(interrupt_evt: EventFd) -> RTC {
        RTC {
            // This is used only for duration measuring purposes.
            previous_now: Instant::now(),
            tick_offset: utils::time::get_time_ns(utils::time::ClockType::Real) as i64,
            match_value: 0,
            load: 0,
            imsc: 0,
            ris: 0,
            interrupt_evt,
        }
    }

    fn trigger_interrupt(&mut self) -> Result<()> {
        self.interrupt_evt.write(1).map_err(Error::InterruptFailure)
    }

    fn get_time(&self) -> u32 {
        let ts = (self.tick_offset as i128)
            + (Instant::now().duration_since(self.previous_now).as_nanos() as i128);
        (ts / utils::time::NANOS_PER_SECOND as i128) as u32
    }

    fn handle_write(&mut self, offset: u64, val: u32) -> Result<()> {
        match offset {
            RTCMR => {
                // The MR register is used for implementing the RTC alarm. A real time clock alarm is
                // a feature that can be used to allow a computer to 'wake up' after shut down to execute
                // tasks every day or on a certain day. It can sometimes be found in the 'Power Management'
                // section of a motherboard's BIOS setup. This is functionality that extends beyond
                // Firecracker intended use. However, we increment a metric just in case.
                self.match_value = val;
                METRICS.rtc.missed_write_count.inc();
            }
            RTCLR => {
                self.load = val;
                self.previous_now = Instant::now();
                // If the unwrap fails, then the internal value of the clock has been corrupted and
                // we want to terminate the execution of the process.
                self.tick_offset = utils::time::seconds_to_nanoseconds(i64::from(val))
                    .expect("Time conversion overflow");
            }
            RTCIMSC => {
                self.imsc = val & 1;
                self.trigger_interrupt()?;
            }
            RTCICR => {
                // As per above mentioned doc, the interrupt is cleared by writing any data value to
                // the Interrupt Clear Register.
                self.ris = 0;
                self.trigger_interrupt()?;
            }
            RTCCR => (), // ignore attempts to turn off the timer.
            o => {
                return Err(Error::BadWriteOffset(o));
            }
        }
        Ok(())
    }
    /// Provides a reference to the interrupt event fd.
    pub fn interrupt_evt(&self) -> &EventFd {
        &self.interrupt_evt
    }
}

impl BusDevice for RTC {
    fn read(&mut self, offset: u64, data: &mut [u8]) {
        let v;
        let mut read_ok = true;

        if offset < AMBA_ID_HIGH && offset >= AMBA_ID_LOW {
            let index = ((offset - AMBA_ID_LOW) >> 2) as usize;
            v = u32::from(PL031_ID[index]);
        } else {
            v = match offset {
                RTCDR => self.get_time(),
                RTCMR => {
                    METRICS.rtc.missed_read_count.inc();
                    // Even though we are not implementing RTC alarm we return the last value
                    self.match_value
                }
                RTCLR => self.load,
                RTCCR => 1, // RTC is always enabled.
                RTCIMSC => self.imsc,
                RTCRIS => self.ris,
                RTCMIS => self.ris & self.imsc,
                _ => {
                    read_ok = false;
                    0
                }
            };
        }
        if read_ok && data.len() <= 4 {
            byte_order::write_le_u32(data, v);
        } else {
            warn!(
                "Invalid RTC PL031 read: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }

    fn write(&mut self, offset: u64, data: &[u8]) {
        if data.len() <= 4 {
            let v = byte_order::read_le_u32(&data[..]);
            if let Err(e) = self.handle_write(offset, v) {
                warn!("Failed to write to RTC PL031 device: {}", e);
                METRICS.rtc.error_count.inc();
            }
        } else {
            warn!(
                "Invalid RTC PL031 write: offset {}, data length {}",
                offset,
                data.len()
            );
            METRICS.rtc.error_count.inc();
        }
    }
}

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

    #[test]
    fn test_rtc_read_write_and_event() {
        let mut rtc = RTC::new(EventFd::new(libc::EFD_NONBLOCK).unwrap());
        let mut data = [0; 4];

        // Read and write to the MR register.
        byte_order::write_le_u32(&mut data, 123);
        rtc.write(RTCMR, &data);
        rtc.read(RTCMR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 123);

        // Read and write to the LR register.
        let v = utils::time::get_time_ns(utils::time::ClockType::Real);
        byte_order::write_le_u32(&mut data, (v / utils::time::NANOS_PER_SECOND) as u32);
        let previous_now_before = rtc.previous_now;
        rtc.write(RTCLR, &data);

        assert!(rtc.previous_now > previous_now_before);

        rtc.read(RTCLR, &mut data);
        let v_read = byte_order::read_le_u32(&data[..]);
        assert_eq!((v / utils::time::NANOS_PER_SECOND) as u32, v_read);

        // Read and write to IMSC register.
        // Test with non zero value.
        let non_zero = 1;
        byte_order::write_le_u32(&mut data, non_zero);
        rtc.write(RTCIMSC, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() == 1);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(non_zero & 1, v);

        // Now test with 0.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCIMSC, &data);
        rtc.read(RTCIMSC, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(0, v);

        // Read and write to the ICR register.
        byte_order::write_le_u32(&mut data, 1);
        rtc.write(RTCICR, &data);
        // The interrupt line should be on.
        assert!(rtc.interrupt_evt.read().unwrap() > 1);
        let v_before = byte_order::read_le_u32(&data[..]);
        let no_errors_before = METRICS.rtc.error_count.count();

        rtc.read(RTCICR, &mut data);
        let no_errors_after = METRICS.rtc.error_count.count();
        let v = byte_order::read_le_u32(&data[..]);
        // ICR is a  write only register. Data received should stay equal to data sent.
        assert_eq!(v, v_before);
        assert_eq!(no_errors_after - no_errors_before, 1);

        // Attempts to turn off the RTC should not go through.
        byte_order::write_le_u32(&mut data, 0);
        rtc.write(RTCCR, &data);
        rtc.read(RTCCR, &mut data);
        let v = byte_order::read_le_u32(&data[..]);
        assert_eq!(v, 1);

        // Attempts to write beyond the writable space. Using here the space used to read
        // the CID and PID from.
        byte_order::write_le_u32(&mut data, 0);
        let no_errors_before = METRICS.rtc.error_count.count();
        rtc.write(AMBA_ID_LOW, &data);
        let no_errors_after = METRICS.rtc.error_count.count();
        assert_eq!(no_errors_after - no_errors_before, 1);
        // However, reading from the AMBA_ID_LOW should succeed upon read.

        let mut data = [0; 4];
        rtc.read(AMBA_ID_LOW, &mut data);
        let index = AMBA_ID_LOW + 3;
        assert_eq!(data[0], PL031_ID[((index - AMBA_ID_LOW) >> 2) as usize]);
    }
}
<|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--async]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --async            Use a threadpool to perform disk IO in the background.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
    flag_async: bool,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
    flag_async: false,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                if let Err(e) = tree.cas(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= scan_max => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if v % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| d.argv(std::env::args().into_iter()).deserialize())
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(16)
        .io_buf_size(8_000_000)
        .async_io(args.flag_async)
        .async_io_threads(3)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(8)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(200))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::Builder::new()
                .name(format!("t({})", i))
                .spawn(move || run(tree, shutdown))
                .unwrap()
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(
            !args.flag_burn_in,
            "don't set both --burn-in and --total-ops"
        );
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<commit_msg>get rid of removed io threads configurable<commit_after>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--async]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --async            Use a threadpool to perform disk IO in the background.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
    flag_async: bool,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
    flag_async: false,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                if let Err(e) = tree.cas(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= scan_max => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if v % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| d.argv(std::env::args().into_iter()).deserialize())
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(16)
        .io_buf_size(8_000_000)
        .async_io(args.flag_async)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(8)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(200))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::Builder::new()
                .name(format!("t({})", i))
                .spawn(move || run(tree, shutdown))
                .unwrap()
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(
            !args.flag_burn_in,
            "don't set both --burn-in and --total-ops"
        );
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use chan_signal::Signal;
use docopt::Docopt;
use rand::{thread_rng, Rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --kv-len=<l>       The length of both keys and values [default: 2].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_kv_len: usize,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
    kv_len: usize,
) {
    let byte = || {
        thread_rng()
            .gen_iter::<u8>()
            .take(kv_len)
            .collect::<Vec<_>>()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(byte(), byte()).unwrap();
            }
            1 => {
                tree.get(&*byte()).unwrap();
            }
            2 => {
                tree.del(&*byte()).unwrap();
            }
            3 => match tree.cas(byte(), Some(byte()), Some(byte())) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            4 => {
                let _ = tree
                    .scan(&*byte())
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| {
            d.argv(std::env::args().into_iter()).deserialize()
        }).unwrap_or_else(|e| e.exit());

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .print_profile_on_drop(true)
        .build();

    let tree = Arc::new(sled::Tree::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;
    let kv_len = args.flag_kv_len;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total, kv_len))
        };

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(
            args.flag_duration,
        ));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<commit_msg>Fix stress2 to deal with fewer copies<commit_after>#[macro_use]
extern crate serde_derive;
extern crate chan_signal;
extern crate docopt;
extern crate rand;
extern crate sled;

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use chan_signal::Signal;
use docopt::Docopt;
use rand::{thread_rng, Rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --kv-len=<l>       The length of both keys and values [default: 2].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_kv_len: usize,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
    kv_len: usize,
) {
    let bytes = || {
        thread_rng()
            .gen_iter::<u8>()
            .take(kv_len)
            .collect::<Vec<_>>()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(bytes(), bytes()).unwrap();
            }
            1 => {
                tree.get(&*bytes()).unwrap();
            }
            2 => {
                tree.del(&*bytes()).unwrap();
            }
            3 => match tree.cas(
                bytes(),
                Some(&*bytes()),
                Some(bytes()),
            ) {
                Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                other => panic!("operational error: {:?}", other),
            },
            4 => {
                let _ = tree
                    .scan(&*bytes())
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| {
            d.argv(std::env::args().into_iter()).deserialize()
        }).unwrap_or_else(|e| e.exit());

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .print_profile_on_drop(true)
        .build();

    let tree = Arc::new(sled::Tree::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;
    let kv_len = args.flag_kv_len;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total, kv_len))
        };

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(
            args.flag_duration,
        ));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                match tree.cas(&key, old, new) {
                    Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                    other => panic!("operational error: {:?}", other),
                }
            }
            v if v > cas_max && v <= scan_max => {
                let _ = tree
                    .scan(&key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            })
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(4)
        .io_buf_size(8_000_000)
        .async_io(false)
        .async_io_threads(8)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(200))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown))
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(
            !args.flag_burn_in,
            "don't set both --burn-in and --total-ops"
        );
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record
                    .module_path()
                    .unwrap()
                    .split("::")
                    .last()
                    .unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<commit_msg>Make async configurable on stress2<commit_after>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use docopt::Docopt;
use rand::{thread_rng, Rng};
use serde::Deserialize;

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no_jemalloc"),
    ),
    global_allocator
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);
static SEQ: AtomicUsize = AtomicUsize::new(0);

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--async]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --async            Use a threadpool to perform disk IO in the background.
";

#[derive(Deserialize, Clone)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
    flag_key_len: usize,
    flag_val_len: usize,
    flag_get_prop: usize,
    flag_set_prop: usize,
    flag_del_prop: usize,
    flag_cas_prop: usize,
    flag_scan_prop: usize,
    flag_merge_prop: usize,
    flag_entries: usize,
    flag_sequential: bool,
    flag_total_ops: Option<usize>,
    flag_async: bool,
}

// defaults will be applied later based on USAGE above
static mut ARGS: Args = Args {
    flag_threads: 0,
    flag_burn_in: false,
    flag_duration: 0,
    flag_key_len: 0,
    flag_val_len: 0,
    flag_get_prop: 0,
    flag_set_prop: 0,
    flag_del_prop: 0,
    flag_cas_prop: 0,
    flag_scan_prop: 0,
    flag_merge_prop: 0,
    flag_entries: 0,
    flag_sequential: false,
    flag_total_ops: None,
    flag_async: false,
};

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret =
        old_value.map(|ov| ov.to_vec()).unwrap_or_else(|| vec![]);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let args = unsafe { ARGS.clone() };

    let get_max = args.flag_get_prop;
    let set_max = get_max + args.flag_set_prop;
    let del_max = set_max + args.flag_del_prop;
    let cas_max = del_max + args.flag_cas_prop;
    let scan_max = cas_max + args.flag_scan_prop;
    let merge_max = scan_max + args.flag_merge_prop;

    let bytes = |len| -> Vec<u8> {
        let i = if args.flag_sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.flag_entries;

        let i_bytes: [u8; std::mem::size_of::<usize>()] =
            unsafe { std::mem::transmute(i) };

        i_bytes.into_iter().cycle().take(len).cloned().collect()
    };
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        TOTAL.fetch_add(1, Ordering::Release);
        let key = bytes(args.flag_key_len);
        let choice = rng.gen_range(0, merge_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                tree.set(&key, bytes(args.flag_val_len)).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.del(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old_k = bytes(args.flag_val_len);

                let old = if rng.gen::<bool>() {
                    Some(old_k.as_slice())
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    Some(bytes(args.flag_val_len))
                } else {
                    None
                };

                match tree.cas(&key, old, new) {
                    Ok(_) | Err(sled::Error::CasFailed(_)) => {}
                    other => panic!("operational error: {:?}", other),
                }
            }
            v if v > cas_max && v <= scan_max => {
                let _ = tree
                    .scan(&key)
                    .take(rng.gen_range(0, 15))
                    .map(|res| res.unwrap())
                    .collect::<Vec<_>>();
            }
            _ => {
                tree.merge(&key, bytes(args.flag_val_len)).unwrap();
            }
        }
    }
}

fn main() {
    setup_logger();

    let args = unsafe {
        ARGS = Docopt::new(USAGE)
            .and_then(|d| {
                d.argv(std::env::args().into_iter()).deserialize()
            })
            .unwrap_or_else(|e| e.exit());
        ARGS.clone()
    };

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(16)
        .io_buf_size(8_000_000)
        .async_io(args.flag_async)
        .async_io_threads(3)
        .blink_node_split_size(4096)
        .page_consolidation_threshold(10)
        .cache_bits(0)
        .cache_capacity(1_000_000_000)
        .flush_every_ms(Some(200))
        .snapshot_after_ops(100_000_000_000)
        .print_profile_on_drop(true)
        .merge_operator(concatenate_merge)
        .build();

    let tree = Arc::new(sled::Db::start(config).unwrap());

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown))
        } else {
            thread::spawn(move || run(tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.flag_total_ops {
        assert!(
            !args.flag_burn_in,
            "don't set both --burn-in and --total-ops"
        );
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.flag_burn_in {
        thread::sleep(std::time::Duration::from_secs(unsafe {
            ARGS.flag_duration
        }));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}

pub fn setup_logger() {
    use std::io::Write;

    fn tn() -> String {
        std::thread::current()
            .name()
            .unwrap_or("unknown")
            .to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record
                    .module_path()
                    .unwrap()
                    .split("::")
                    .last()
                    .unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if std::env::var("RUST_LOG").is_ok() {
        builder.parse(&std::env::var("RUST_LOG").unwrap());
    }

    let _r = builder.try_init();
}
<|endoftext|>


<commit_before>//! Render to a window created by Glutin, using Glium's OpenGL functions

// FIXME: Support resizing the window, like the SDL renderer does

use frontend_api::{FrontendAction, Renderer};
use frontend_api::ppu::{SCREEN_WIDTH, SCREEN_HEIGHT};

use glium::{DisplayBuild, Surface, Rect};
use glium::backend::glutin_backend::GlutinFacade;
use glium::index::{NoIndices, PrimitiveType};
use glium::glutin::WindowBuilder;
use glium::program::Program;
use glium::texture::{ClientFormat, RawImage2d, Texture2d};
use glium::uniforms::MagnifySamplerFilter;
use glium::vertex::VertexBuffer;

use std::borrow::Cow;

/// Our vertices are extremely simple: There's no need for more than 2 dimensions and the texture
/// coordinates can be calculated from the position, so we don't need to store them.
#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 2],
}

impl Vertex {
    fn new(x: f32, y: f32) -> Self { Vertex { position: [x, y] } }
}

implement_vertex!(Vertex, position);

const VERTEX_SHADER_SRC: &'static str = r#"
    #version 140

    in vec2 position;
    out vec2 tex_coords;

    void main() {
        tex_coords = vec2((position.x + 1) * 0.5, 1 - (position.y + 1) * 0.5);
        gl_Position = vec4(position, 0.0, 1.0);
    }
"#;

const FRAGMENT_SHADER_SRC: &'static str = r#"
    #version 140

    in vec2 tex_coords;
    out vec4 color;

    uniform sampler2D tex;

    void main() {
        color = texture(tex, tex_coords);
    }
"#;

pub struct GliumRenderer {
    display: GlutinFacade,
    /// This vertex buffer will only ever store 4 vertices spanning the whole window
    vbuf: VertexBuffer<Vertex>,
    /// A simple shader that maps our texture onto the window
    program: Program,
    /// This texture is updated with the PPU's data every frame
    texture: Texture2d,
}

impl Default for GliumRenderer {
    fn default() -> Self {
        let display = WindowBuilder::new()
            .with_dimensions(SCREEN_WIDTH * 3, SCREEN_HEIGHT * 3)
            .with_title("breeze".to_owned())
            .build_glium().unwrap();
        // Create a rectangle spanning the whole viewport/window. This is a triangle strip.
        let shape = [
            Vertex::new(-1.0, 1.0),
            Vertex::new(1.0, 1.0),
            Vertex::new(-1.0, -1.0),
            Vertex::new(1.0, -1.0),
        ];

        GliumRenderer {
            vbuf: VertexBuffer::new(&display, &shape).unwrap(),
            program: Program::from_source(&display, VERTEX_SHADER_SRC, FRAGMENT_SHADER_SRC, None)
                .unwrap(),
            texture: Texture2d::empty(&display, SCREEN_WIDTH, SCREEN_HEIGHT)
                .unwrap(),
            display: display,
        }
    }
}

impl GliumRenderer {
    fn handle_events(&mut self) -> Option<FrontendAction> {
        use glium::glutin::Event::*;

        for ev in self.display.poll_events() {
            match ev {
                Closed => {
                    info!("quit event -> exiting");
                    return Some(FrontendAction::Exit)
                }
                _ => {}
            }
        }
        None
    }
}

impl Renderer for GliumRenderer {
    fn render(&mut self, frame_data: &[u8]) -> Option<FrontendAction> {
        // upload new texture data
        self.texture.write(Rect {
            left: 0,
            bottom: 0,
            width: SCREEN_WIDTH,
            height: SCREEN_HEIGHT,
        }, RawImage2d {
            data: Cow::Borrowed(frame_data),
            width: SCREEN_WIDTH,
            height: SCREEN_HEIGHT,
            format: ClientFormat::U8U8U8,
        });

        let mut target = self.display.draw();
        target.draw(
            &self.vbuf,
            &NoIndices(PrimitiveType::TriangleStrip),
            &self.program,
            &uniform! {
                tex: self.texture.sampled()
                    .magnify_filter(MagnifySamplerFilter::Nearest),
            },
            &Default::default()).unwrap();
        target.finish().unwrap();

        self.handle_events()
    }

    fn set_rom_title(&mut self, title: &str) {
        match self.display.get_window() {
            Some(win_ref) => win_ref.set_title(title),
            None => {},     // Headless? Ignore.
        }
    }
}
<commit_msg>Use SrgbTexture2d (Fixes #10)<commit_after>//! Render to a window created by Glutin, using Glium's OpenGL functions

// FIXME: Support resizing the window, like the SDL renderer does

use frontend_api::{FrontendAction, Renderer};
use frontend_api::ppu::{SCREEN_WIDTH, SCREEN_HEIGHT};

use glium::{DisplayBuild, Surface, Rect};
use glium::backend::glutin_backend::GlutinFacade;
use glium::index::{NoIndices, PrimitiveType};
use glium::glutin::WindowBuilder;
use glium::program::Program;
use glium::texture::{ClientFormat, RawImage2d, SrgbTexture2d};
use glium::uniforms::MagnifySamplerFilter;
use glium::vertex::VertexBuffer;

use std::borrow::Cow;

/// Our vertices are extremely simple: There's no need for more than 2 dimensions and the texture
/// coordinates can be calculated from the position, so we don't need to store them.
#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 2],
}

impl Vertex {
    fn new(x: f32, y: f32) -> Self { Vertex { position: [x, y] } }
}

implement_vertex!(Vertex, position);

const VERTEX_SHADER_SRC: &'static str = r#"
    #version 140

    in vec2 position;
    out vec2 tex_coords;

    void main() {
        tex_coords = vec2((position.x + 1) * 0.5, 1 - (position.y + 1) * 0.5);
        gl_Position = vec4(position, 0.0, 1.0);
    }
"#;

const FRAGMENT_SHADER_SRC: &'static str = r#"
    #version 140

    in vec2 tex_coords;
    out vec4 color;

    uniform sampler2D tex;

    void main() {
        color = texture(tex, tex_coords);
    }
"#;

pub struct GliumRenderer {
    display: GlutinFacade,
    /// This vertex buffer will only ever store 4 vertices spanning the whole window
    vbuf: VertexBuffer<Vertex>,
    /// A simple shader that maps our texture onto the window
    program: Program,
    /// This texture is updated with the PPU's data every frame
    texture: SrgbTexture2d,
}

impl Default for GliumRenderer {
    fn default() -> Self {
        let display = WindowBuilder::new()
            .with_dimensions(SCREEN_WIDTH * 3, SCREEN_HEIGHT * 3)
            .with_title("breeze".to_owned())
            .build_glium().unwrap();
        // Create a rectangle spanning the whole viewport/window. This is a triangle strip.
        let shape = [
            Vertex::new(-1.0, 1.0),
            Vertex::new(1.0, 1.0),
            Vertex::new(-1.0, -1.0),
            Vertex::new(1.0, -1.0),
        ];

        GliumRenderer {
            vbuf: VertexBuffer::new(&display, &shape).unwrap(),
            program: Program::from_source(&display, VERTEX_SHADER_SRC, FRAGMENT_SHADER_SRC, None)
                .unwrap(),
            texture: SrgbTexture2d::empty(&display, SCREEN_WIDTH, SCREEN_HEIGHT)
                .unwrap(),
            display: display,
        }
    }
}

impl GliumRenderer {
    fn handle_events(&mut self) -> Option<FrontendAction> {
        use glium::glutin::Event::*;

        for ev in self.display.poll_events() {
            match ev {
                Closed => {
                    info!("quit event -> exiting");
                    return Some(FrontendAction::Exit)
                }
                _ => {}
            }
        }
        None
    }
}

impl Renderer for GliumRenderer {
    fn render(&mut self, frame_data: &[u8]) -> Option<FrontendAction> {
        // upload new texture data
        self.texture.write(Rect {
            left: 0,
            bottom: 0,
            width: SCREEN_WIDTH,
            height: SCREEN_HEIGHT,
        }, RawImage2d {
            data: Cow::Borrowed(frame_data),
            width: SCREEN_WIDTH,
            height: SCREEN_HEIGHT,
            format: ClientFormat::U8U8U8,
        });

        let mut target = self.display.draw();
        target.draw(
            &self.vbuf,
            &NoIndices(PrimitiveType::TriangleStrip),
            &self.program,
            &uniform! {
                tex: self.texture.sampled()
                    .magnify_filter(MagnifySamplerFilter::Nearest),
            },
            &Default::default()).unwrap();
        target.finish().unwrap();

        self.handle_events()
    }

    fn set_rom_title(&mut self, title: &str) {
        match self.display.get_window() {
            Some(win_ref) => win_ref.set_title(title),
            None => {},     // Headless? Ignore.
        }
    }
}
<|endoftext|>


<commit_before>use super::{StringReader, UnmatchedBrace};

use rustc_ast::token::{self, Token};
use rustc_ast::tokenstream::{
    DelimSpan,
    IsJoint::{self, *},
    TokenStream, TokenTree, TreeAndJoint,
};
use rustc_ast_pretty::pprust::token_to_string;
use rustc_data_structures::fx::FxHashMap;
use rustc_errors::PResult;
use rustc_span::Span;

impl<'a> StringReader<'a> {
    crate fn into_token_trees(self) -> (PResult<'a, TokenStream>, Vec<UnmatchedBrace>) {
        let mut tt_reader = TokenTreesReader {
            string_reader: self,
            token: Token::dummy(),
            joint_to_prev: Joint,
            open_braces: Vec::new(),
            unmatched_braces: Vec::new(),
            matching_delim_spans: Vec::new(),
            last_unclosed_found_span: None,
            last_delim_empty_block_spans: FxHashMap::default(),
        };
        let res = tt_reader.parse_all_token_trees();
        (res, tt_reader.unmatched_braces)
    }
}

struct TokenTreesReader<'a> {
    string_reader: StringReader<'a>,
    token: Token,
    joint_to_prev: IsJoint,
    /// Stack of open delimiters and their spans. Used for error message.
    open_braces: Vec<(token::DelimToken, Span)>,
    unmatched_braces: Vec<UnmatchedBrace>,
    /// The type and spans for all braces
    ///
    /// Used only for error recovery when arriving to EOF with mismatched braces.
    matching_delim_spans: Vec<(token::DelimToken, Span, Span)>,
    last_unclosed_found_span: Option<Span>,
    /// Collect empty block spans that might have been auto-inserted by editors.
    last_delim_empty_block_spans: FxHashMap<token::DelimToken, Span>,
}

impl<'a> TokenTreesReader<'a> {
    // Parse a stream of tokens into a list of `TokenTree`s, up to an `Eof`.
    fn parse_all_token_trees(&mut self) -> PResult<'a, TokenStream> {
        let mut buf = TokenStreamBuilder::default();

        self.real_token();
        while self.token != token::Eof {
            buf.push(self.parse_token_tree()?);
        }

        Ok(buf.into_token_stream())
    }

    // Parse a stream of tokens into a list of `TokenTree`s, up to a `CloseDelim`.
    fn parse_token_trees_until_close_delim(&mut self) -> TokenStream {
        let mut buf = TokenStreamBuilder::default();
        loop {
            if let token::CloseDelim(..) = self.token.kind {
                return buf.into_token_stream();
            }

            match self.parse_token_tree() {
                Ok(tree) => buf.push(tree),
                Err(mut e) => {
                    e.emit();
                    return buf.into_token_stream();
                }
            }
        }
    }

    fn parse_token_tree(&mut self) -> PResult<'a, TreeAndJoint> {
        let sm = self.string_reader.sess.source_map();

        match self.token.kind {
            token::Eof => {
                let msg = "this file contains an unclosed delimiter";
                let mut err =
                    self.string_reader.sess.span_diagnostic.struct_span_err(self.token.span, msg);
                for &(_, sp) in &self.open_braces {
                    err.span_label(sp, "unclosed delimiter");
                    self.unmatched_braces.push(UnmatchedBrace {
                        expected_delim: token::DelimToken::Brace,
                        found_delim: None,
                        found_span: self.token.span,
                        unclosed_span: Some(sp),
                        candidate_span: None,
                    });
                }

                if let Some((delim, _)) = self.open_braces.last() {
                    if let Some((_, open_sp, close_sp)) =
                        self.matching_delim_spans.iter().find(|(d, open_sp, close_sp)| {
                            if let Some(close_padding) = sm.span_to_margin(*close_sp) {
                                if let Some(open_padding) = sm.span_to_margin(*open_sp) {
                                    return delim == d && close_padding != open_padding;
                                }
                            }
                            false
                        })
                    // these are in reverse order as they get inserted on close, but
                    {
                        // we want the last open/first close
                        err.span_label(*open_sp, "this delimiter might not be properly closed...");
                        err.span_label(
                            *close_sp,
                            "...as it matches this but it has different indentation",
                        );
                    }
                }
                Err(err)
            }
            token::OpenDelim(delim) => {
                // The span for beginning of the delimited section
                let pre_span = self.token.span;

                // Parse the open delimiter.
                self.open_braces.push((delim, self.token.span));
                self.real_token();

                // Parse the token trees within the delimiters.
                // We stop at any delimiter so we can try to recover if the user
                // uses an incorrect delimiter.
                let tts = self.parse_token_trees_until_close_delim();

                // Expand to cover the entire delimited token tree
                let delim_span = DelimSpan::from_pair(pre_span, self.token.span);

                match self.token.kind {
                    // Correct delimiter.
                    token::CloseDelim(d) if d == delim => {
                        let (open_brace, open_brace_span) = self.open_braces.pop().unwrap();
                        let close_brace_span = self.token.span;

                        if tts.is_empty() {
                            let empty_block_span = open_brace_span.to(close_brace_span);
                            if !sm.is_multiline(empty_block_span) {
                                // Only track if the block is in the form of `{}`, otherwise it is
                                // likely that it was written on purpose.
                                self.last_delim_empty_block_spans.insert(delim, empty_block_span);
                            }
                        }

                        if self.open_braces.is_empty() {
                            // Clear up these spans to avoid suggesting them as we've found
                            // properly matched delimiters so far for an entire block.
                            self.matching_delim_spans.clear();
                        } else {
                            self.matching_delim_spans.push((
                                open_brace,
                                open_brace_span,
                                close_brace_span,
                            ));
                        }
                        // Parse the closing delimiter.
                        self.real_token();
                    }
                    // Incorrect delimiter.
                    token::CloseDelim(other) => {
                        let mut unclosed_delimiter = None;
                        let mut candidate = None;
                        if self.last_unclosed_found_span != Some(self.token.span) {
                            // do not complain about the same unclosed delimiter multiple times
                            self.last_unclosed_found_span = Some(self.token.span);
                            // This is a conservative error: only report the last unclosed
                            // delimiter. The previous unclosed delimiters could actually be
                            // closed! The parser just hasn't gotten to them yet.
                            if let Some(&(_, sp)) = self.open_braces.last() {
                                unclosed_delimiter = Some(sp);
                            };
                            if let Some(current_padding) = sm.span_to_margin(self.token.span) {
                                for (brace, brace_span) in &self.open_braces {
                                    if let Some(padding) = sm.span_to_margin(*brace_span) {
                                        // high likelihood of these two corresponding
                                        if current_padding == padding && brace == &other {
                                            candidate = Some(*brace_span);
                                        }
                                    }
                                }
                            }
                            let (tok, _) = self.open_braces.pop().unwrap();
                            self.unmatched_braces.push(UnmatchedBrace {
                                expected_delim: tok,
                                found_delim: Some(other),
                                found_span: self.token.span,
                                unclosed_span: unclosed_delimiter,
                                candidate_span: candidate,
                            });
                        } else {
                            self.open_braces.pop();
                        }

                        // If the incorrect delimiter matches an earlier opening
                        // delimiter, then don't consume it (it can be used to
                        // close the earlier one). Otherwise, consume it.
                        // E.g., we try to recover from:
                        // fn foo() {
                        //     bar(baz(
                        // }  // Incorrect delimiter but matches the earlier `{`
                        if !self.open_braces.iter().any(|&(b, _)| b == other) {
                            self.real_token();
                        }
                    }
                    token::Eof => {
                        // Silently recover, the EOF token will be seen again
                        // and an error emitted then. Thus we don't pop from
                        // self.open_braces here.
                    }
                    _ => {}
                }

                Ok(TokenTree::Delimited(delim_span, delim, tts).into())
            }
            token::CloseDelim(_delim) => {
                // An unexpected closing delimiter (i.e., there is no
                // matching opening delimiter).
                let token_str = token_to_string(&self.token);
                let msg = format!("unexpected closing delimiter: `{}`", token_str);
                let mut err =
                    self.string_reader.sess.span_diagnostic.struct_span_err(self.token.span, &msg);

                err.span_label(self.token.span, "unexpected closing delimiter");

                Err(err)
            }
            _ => {
                let tt = TokenTree::Token(self.token.take());
                self.real_token();
                let is_joint = self.joint_to_prev == Joint && self.token.is_op();
                Ok((tt, if is_joint { Joint } else { NonJoint }))
            }
        }
    }

    fn real_token(&mut self) {
        self.joint_to_prev = Joint;
        loop {
            let token = self.string_reader.next_token();
            match token.kind {
                token::Whitespace | token::Comment | token::Shebang(_) | token::Unknown(_) => {
                    self.joint_to_prev = NonJoint;
                }
                _ => {
                    self.token = token;
                    return;
                }
            }
        }
    }
}

#[derive(Default)]
struct TokenStreamBuilder {
    buf: Vec<TreeAndJoint>,
}

impl TokenStreamBuilder {
    fn push(&mut self, (tree, joint): TreeAndJoint) {
        if let Some((TokenTree::Token(prev_token), Joint)) = self.buf.last() {
            if let TokenTree::Token(token) = &tree {
                if let Some(glued) = prev_token.glue(token) {
                    self.buf.pop();
                    self.buf.push((TokenTree::Token(glued), joint));
                    return;
                }
            }
        }
        self.buf.push((tree, joint))
    }

    fn into_token_stream(self) -> TokenStream {
        TokenStream::new(self.buf)
    }
}
<commit_msg>Revert "Dropping "block is empty..." error message on incorrect closing delimiter #70583"<commit_after>use super::{StringReader, UnmatchedBrace};

use rustc_ast::token::{self, Token};
use rustc_ast::tokenstream::{
    DelimSpan,
    IsJoint::{self, *},
    TokenStream, TokenTree, TreeAndJoint,
};
use rustc_ast_pretty::pprust::token_to_string;
use rustc_data_structures::fx::FxHashMap;
use rustc_errors::PResult;
use rustc_span::Span;

impl<'a> StringReader<'a> {
    crate fn into_token_trees(self) -> (PResult<'a, TokenStream>, Vec<UnmatchedBrace>) {
        let mut tt_reader = TokenTreesReader {
            string_reader: self,
            token: Token::dummy(),
            joint_to_prev: Joint,
            open_braces: Vec::new(),
            unmatched_braces: Vec::new(),
            matching_delim_spans: Vec::new(),
            last_unclosed_found_span: None,
            last_delim_empty_block_spans: FxHashMap::default(),
        };
        let res = tt_reader.parse_all_token_trees();
        (res, tt_reader.unmatched_braces)
    }
}

struct TokenTreesReader<'a> {
    string_reader: StringReader<'a>,
    token: Token,
    joint_to_prev: IsJoint,
    /// Stack of open delimiters and their spans. Used for error message.
    open_braces: Vec<(token::DelimToken, Span)>,
    unmatched_braces: Vec<UnmatchedBrace>,
    /// The type and spans for all braces
    ///
    /// Used only for error recovery when arriving to EOF with mismatched braces.
    matching_delim_spans: Vec<(token::DelimToken, Span, Span)>,
    last_unclosed_found_span: Option<Span>,
    /// Collect empty block spans that might have been auto-inserted by editors.
    last_delim_empty_block_spans: FxHashMap<token::DelimToken, Span>,
}

impl<'a> TokenTreesReader<'a> {
    // Parse a stream of tokens into a list of `TokenTree`s, up to an `Eof`.
    fn parse_all_token_trees(&mut self) -> PResult<'a, TokenStream> {
        let mut buf = TokenStreamBuilder::default();

        self.real_token();
        while self.token != token::Eof {
            buf.push(self.parse_token_tree()?);
        }

        Ok(buf.into_token_stream())
    }

    // Parse a stream of tokens into a list of `TokenTree`s, up to a `CloseDelim`.
    fn parse_token_trees_until_close_delim(&mut self) -> TokenStream {
        let mut buf = TokenStreamBuilder::default();
        loop {
            if let token::CloseDelim(..) = self.token.kind {
                return buf.into_token_stream();
            }

            match self.parse_token_tree() {
                Ok(tree) => buf.push(tree),
                Err(mut e) => {
                    e.emit();
                    return buf.into_token_stream();
                }
            }
        }
    }

    fn parse_token_tree(&mut self) -> PResult<'a, TreeAndJoint> {
        let sm = self.string_reader.sess.source_map();
        match self.token.kind {
            token::Eof => {
                let msg = "this file contains an unclosed delimiter";
                let mut err =
                    self.string_reader.sess.span_diagnostic.struct_span_err(self.token.span, msg);
                for &(_, sp) in &self.open_braces {
                    err.span_label(sp, "unclosed delimiter");
                    self.unmatched_braces.push(UnmatchedBrace {
                        expected_delim: token::DelimToken::Brace,
                        found_delim: None,
                        found_span: self.token.span,
                        unclosed_span: Some(sp),
                        candidate_span: None,
                    });
                }

                if let Some((delim, _)) = self.open_braces.last() {
                    if let Some((_, open_sp, close_sp)) =
                        self.matching_delim_spans.iter().find(|(d, open_sp, close_sp)| {
                            if let Some(close_padding) = sm.span_to_margin(*close_sp) {
                                if let Some(open_padding) = sm.span_to_margin(*open_sp) {
                                    return delim == d && close_padding != open_padding;
                                }
                            }
                            false
                        })
                    // these are in reverse order as they get inserted on close, but
                    {
                        // we want the last open/first close
                        err.span_label(*open_sp, "this delimiter might not be properly closed...");
                        err.span_label(
                            *close_sp,
                            "...as it matches this but it has different indentation",
                        );
                    }
                }
                Err(err)
            }
            token::OpenDelim(delim) => {
                // The span for beginning of the delimited section
                let pre_span = self.token.span;

                // Parse the open delimiter.
                self.open_braces.push((delim, self.token.span));
                self.real_token();

                // Parse the token trees within the delimiters.
                // We stop at any delimiter so we can try to recover if the user
                // uses an incorrect delimiter.
                let tts = self.parse_token_trees_until_close_delim();

                // Expand to cover the entire delimited token tree
                let delim_span = DelimSpan::from_pair(pre_span, self.token.span);

                match self.token.kind {
                    // Correct delimiter.
                    token::CloseDelim(d) if d == delim => {
                        let (open_brace, open_brace_span) = self.open_braces.pop().unwrap();
                        let close_brace_span = self.token.span;

                        if tts.is_empty() {
                            let empty_block_span = open_brace_span.to(close_brace_span);
                            if !sm.is_multiline(empty_block_span) {
                                // Only track if the block is in the form of `{}`, otherwise it is
                                // likely that it was written on purpose.
                                self.last_delim_empty_block_spans.insert(delim, empty_block_span);
                            }
                        }

                        if self.open_braces.is_empty() {
                            // Clear up these spans to avoid suggesting them as we've found
                            // properly matched delimiters so far for an entire block.
                            self.matching_delim_spans.clear();
                        } else {
                            self.matching_delim_spans.push((
                                open_brace,
                                open_brace_span,
                                close_brace_span,
                            ));
                        }
                        // Parse the closing delimiter.
                        self.real_token();
                    }
                    // Incorrect delimiter.
                    token::CloseDelim(other) => {
                        let mut unclosed_delimiter = None;
                        let mut candidate = None;
                        if self.last_unclosed_found_span != Some(self.token.span) {
                            // do not complain about the same unclosed delimiter multiple times
                            self.last_unclosed_found_span = Some(self.token.span);
                            // This is a conservative error: only report the last unclosed
                            // delimiter. The previous unclosed delimiters could actually be
                            // closed! The parser just hasn't gotten to them yet.
                            if let Some(&(_, sp)) = self.open_braces.last() {
                                unclosed_delimiter = Some(sp);
                            };
                            if let Some(current_padding) = sm.span_to_margin(self.token.span) {
                                for (brace, brace_span) in &self.open_braces {
                                    if let Some(padding) = sm.span_to_margin(*brace_span) {
                                        // high likelihood of these two corresponding
                                        if current_padding == padding && brace == &other {
                                            candidate = Some(*brace_span);
                                        }
                                    }
                                }
                            }
                            let (tok, _) = self.open_braces.pop().unwrap();
                            self.unmatched_braces.push(UnmatchedBrace {
                                expected_delim: tok,
                                found_delim: Some(other),
                                found_span: self.token.span,
                                unclosed_span: unclosed_delimiter,
                                candidate_span: candidate,
                            });
                        } else {
                            self.open_braces.pop();
                        }

                        // If the incorrect delimiter matches an earlier opening
                        // delimiter, then don't consume it (it can be used to
                        // close the earlier one). Otherwise, consume it.
                        // E.g., we try to recover from:
                        // fn foo() {
                        //     bar(baz(
                        // }  // Incorrect delimiter but matches the earlier `{`
                        if !self.open_braces.iter().any(|&(b, _)| b == other) {
                            self.real_token();
                        }
                    }
                    token::Eof => {
                        // Silently recover, the EOF token will be seen again
                        // and an error emitted then. Thus we don't pop from
                        // self.open_braces here.
                    }
                    _ => {}
                }

                Ok(TokenTree::Delimited(delim_span, delim, tts).into())
            }
            token::CloseDelim(delim) => {
                // An unexpected closing delimiter (i.e., there is no
                // matching opening delimiter).
                let token_str = token_to_string(&self.token);
                let msg = format!("unexpected closing delimiter: `{}`", token_str);
                let mut err =
                    self.string_reader.sess.span_diagnostic.struct_span_err(self.token.span, &msg);

                if let Some(span) = self.last_delim_empty_block_spans.remove(&delim) {
                    err.span_label(
                        span,
                        "this block is empty, you might have not meant to close it",
                    );
                }
                err.span_label(self.token.span, "unexpected closing delimiter");
                Err(err)
            }
            _ => {
                let tt = TokenTree::Token(self.token.take());
                self.real_token();
                let is_joint = self.joint_to_prev == Joint && self.token.is_op();
                Ok((tt, if is_joint { Joint } else { NonJoint }))
            }
        }
    }

    fn real_token(&mut self) {
        self.joint_to_prev = Joint;
        loop {
            let token = self.string_reader.next_token();
            match token.kind {
                token::Whitespace | token::Comment | token::Shebang(_) | token::Unknown(_) => {
                    self.joint_to_prev = NonJoint;
                }
                _ => {
                    self.token = token;
                    return;
                }
            }
        }
    }
}

#[derive(Default)]
struct TokenStreamBuilder {
    buf: Vec<TreeAndJoint>,
}

impl TokenStreamBuilder {
    fn push(&mut self, (tree, joint): TreeAndJoint) {
        if let Some((TokenTree::Token(prev_token), Joint)) = self.buf.last() {
            if let TokenTree::Token(token) = &tree {
                if let Some(glued) = prev_token.glue(token) {
                    self.buf.pop();
                    self.buf.push((TokenTree::Token(glued), joint));
                    return;
                }
            }
        }
        self.buf.push((tree, joint))
    }

    fn into_token_stream(self) -> TokenStream {
        TokenStream::new(self.buf)
    }
}
<|endoftext|>