<commit_before>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::io::{Reader, Writer};

pub fn read_at_least<R : Reader>(reader : &mut R,
                                 buf: &mut [u8],
                                 min_bytes : uint) -> uint {
    let mut pos = 0;
    let bufLen = buf.len();
    while pos < min_bytes {
        let buf1 = buf.mut_slice(pos, bufLen);
        match reader.read(buf1) {
            Err(e) => fail!("premature EOF? {}", e),
            Ok(n) => pos += n
        }
    }
    return pos;
}

pub trait BufferedInputStream : Reader {
    fn skip(&mut self, bytes : uint);
    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8);
}

pub struct BufferedInputStreamWrapper<'a, R> {
    priv inner : &'a mut R,
    priv buf : ~[u8],
    priv pos : uint,
    priv cap : uint
}

impl <'a, R> BufferedInputStreamWrapper<'a, R> {
    pub fn new<'a> (r : &'a mut R) -> BufferedInputStreamWrapper<'a, R> {
        let mut result = BufferedInputStreamWrapper {
            inner : r,
            buf : std::vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            result.buf.set_len(8192)
        }
        return result;
    }
}

impl<'a, R: Reader> BufferedInputStream for BufferedInputStreamWrapper<'a, R> {

   fn skip(&mut self, mut bytes : uint) {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, bytes);
                self.pos = bytes;
                self.cap = n;

            } else {
                //# Forward large skip to the underlying stream.
                fail!("TODO")
            }
        }
    }

    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8) {
        if self.cap - self.pos == 0 {
            let n = read_at_least(self.inner, self.buf, 1);
            self.cap = n;
            self.pos = 0;
        }
        (std::ptr::to_unsafe_ptr(self.buf.unsafe_ref(self.pos)),
         std::ptr::to_unsafe_ptr(self.buf.unsafe_ref(self.cap)))
    }
}

impl<'a, R: Reader> Reader for BufferedInputStreamWrapper<'a, R> {
    fn read(&mut self, dst: &mut [u8]) -> Result<uint, std::io::IoError> {
        let mut num_bytes = dst.len();
        if num_bytes <= self.cap - self.pos {
            //# Serve from the current buffer.
            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.pos + num_bytes));
            self.pos += num_bytes;
            return Ok(num_bytes);
        } else {
            //# Copy current available into destination.

            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap));
            let fromFirstBuffer = self.cap - self.pos;

            let dst1 = dst.mut_slice(fromFirstBuffer, num_bytes);
            num_bytes -= fromFirstBuffer;
            if num_bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, num_bytes);
                std::vec::bytes::copy_memory(dst1,
                                             self.buf.slice(0, num_bytes));
                self.cap = n;
                self.pos = num_bytes;
                return Ok(fromFirstBuffer + num_bytes);
            } else {
                //# Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Ok(fromFirstBuffer + read_at_least(self.inner, dst1, num_bytes));
            }
        }
    }
}

pub struct ArrayInputStream<'a> {
    priv array : &'a [u8]
}

impl <'a> ArrayInputStream<'a> {
    pub fn new<'b>(array : &'b [u8]) -> ArrayInputStream<'b> {
        ArrayInputStream { array : array }
    }
}

impl <'a> Reader for ArrayInputStream<'a> {
    fn read(&mut self, dst: &mut [u8]) -> Result<uint, std::io::IoError> {
        let n = std::cmp::min(dst.len(), self.array.len());
        unsafe { dst.copy_memory(self.array.slice_to(n)); }
        self.array = self.array.slice_from(n);
        Ok(n)
    }
}

impl <'a> BufferedInputStream for ArrayInputStream<'a> {
    fn skip(&mut self, bytes : uint) {
        assert!(self.array.len() >= bytes,
                "ArrayInputStream ended prematurely.");
        self.array = self.array.slice_from(bytes);
    }
    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8){
        let len = self.array.len();
        (self.array.as_ptr(),
         std::ptr::to_unsafe_ptr(self.array.unsafe_ref(len)))
    }
}

pub trait BufferedOutputStream : Writer {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8);
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> std::io::IoResult<()>;
}

pub struct BufferedOutputStreamWrapper<'a, W> {
    priv inner: &'a mut W,
    priv buf: ~[u8],
    priv pos: uint
}

impl <'a, W> BufferedOutputStreamWrapper<'a, W> {
    pub fn new<'b> (w : &'b mut W) -> BufferedOutputStreamWrapper<'b, W> {
        let mut result = BufferedOutputStreamWrapper {
            inner: w,
            buf : std::vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            result.buf.set_len(8192);
        }
        return result;
    }
}

impl<'a, W: Writer> BufferedOutputStream for BufferedOutputStreamWrapper<'a, W> {
    #[inline]
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        (std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(self.pos)),
         std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(len)))
    }

    #[inline]
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> std::io::IoResult<()> {
        let easyCase = ptr == std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(self.pos));
        if easyCase {
            self.pos += size;
            Ok(())
        } else {
            std::vec::raw::mut_buf_as_slice::<u8,std::io::IoResult<()>>(ptr, size, |buf| {
                self.write(buf)
            })
        }
    }

}

impl<'a, W: Writer> Writer for BufferedOutputStreamWrapper<'a, W> {
    fn write(&mut self, buf: &[u8]) -> std::io::IoResult<()> {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = self.buf.mut_slice_from(self.pos);
            std::vec::bytes::copy_memory(dst, buf);
            self.pos += size;
        } else if size <= self.buf.len() {
            //# Too much for this buffer, but not a full buffer's
            //# worth, so we'll go ahead and copy.
            {
                let dst = self.buf.mut_slice_from(self.pos);
                std::vec::bytes::copy_memory(dst, buf.slice(0, available));
            }
            if_ok!(self.inner.write(self.buf));

            size -= available;
            let src = buf.slice_from(available);
            let dst = self.buf.mut_slice_from(0);
            std::vec::bytes::copy_memory(dst, src);
            self.pos = size;
        } else {
            //# Writing so much data that we might as well write
            //# directly to avoid a copy.
            self.inner.write(self.buf.slice(0, self.pos)).unwrap();
            self.pos = 0;
            if_ok!(self.inner.write(buf));
        }
        return Ok(());
    }

    fn flush(&mut self) -> std::io::IoResult<()> {
        if self.pos > 0 {
            if_ok!(self.inner.write(self.buf.slice(0, self.pos)));
            self.pos = 0;
        }
        return Ok(())
    }
}

pub struct ArrayOutputStream<'a> {
    priv array : &'a mut [u8],
    priv fill_pos : uint,
}

impl <'a> ArrayOutputStream<'a> {
    pub fn new<'b>(array : &'b mut [u8]) -> ArrayOutputStream<'b> {
        ArrayOutputStream {
            array : array,
            fill_pos : 0
        }
    }
}

impl <'a> Writer for ArrayOutputStream<'a> {
    fn write(&mut self, buf: &[u8]) -> std::io::IoResult<()> {
        assert!(buf.len() <= self.array.len() - self.fill_pos,
                "ArrayOutputStream's backing array was not large enough for the data written.");
        unsafe { self.array.mut_slice_from(self.fill_pos).copy_memory(buf); }
        self.fill_pos += buf.len();
        Ok(())
    }
}

impl <'a> BufferedOutputStream for ArrayOutputStream<'a> {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.array.len();
        (std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(self.fill_pos)),
         std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(len)))
    }
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> std::io::IoResult<()> {
        let easyCase = ptr == std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(self.fill_pos));
        if easyCase {
            self.fill_pos += size;
            Ok(())
        } else {
            std::vec::raw::mut_buf_as_slice::<u8,std::io::IoResult<()>>(ptr, size, |buf| {
                self.write(buf)
            })
        }
    }
}
<commit_msg>new IO API cleanup<commit_after>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::io::{Reader, Writer, IoResult};

pub fn read_at_least<R : Reader>(reader : &mut R,
                                 buf: &mut [u8],
                                 min_bytes : uint) -> uint {
    let mut pos = 0;
    let bufLen = buf.len();
    while pos < min_bytes {
        let buf1 = buf.mut_slice(pos, bufLen);
        match reader.read(buf1) {
            Err(e) => fail!("premature EOF? {}", e),
            Ok(n) => pos += n
        }
    }
    return pos;
}

pub trait BufferedInputStream : Reader {
    fn skip(&mut self, bytes : uint);
    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8);
}

pub struct BufferedInputStreamWrapper<'a, R> {
    priv inner : &'a mut R,
    priv buf : ~[u8],
    priv pos : uint,
    priv cap : uint
}

impl <'a, R> BufferedInputStreamWrapper<'a, R> {
    pub fn new<'a> (r : &'a mut R) -> BufferedInputStreamWrapper<'a, R> {
        let mut result = BufferedInputStreamWrapper {
            inner : r,
            buf : std::vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            result.buf.set_len(8192)
        }
        return result;
    }
}

impl<'a, R: Reader> BufferedInputStream for BufferedInputStreamWrapper<'a, R> {

   fn skip(&mut self, mut bytes : uint) {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, bytes);
                self.pos = bytes;
                self.cap = n;

            } else {
                //# Forward large skip to the underlying stream.
                fail!("TODO")
            }
        }
    }

    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8) {
        if self.cap - self.pos == 0 {
            let n = read_at_least(self.inner, self.buf, 1);
            self.cap = n;
            self.pos = 0;
        }
        (std::ptr::to_unsafe_ptr(self.buf.unsafe_ref(self.pos)),
         std::ptr::to_unsafe_ptr(self.buf.unsafe_ref(self.cap)))
    }
}

impl<'a, R: Reader> Reader for BufferedInputStreamWrapper<'a, R> {
    fn read(&mut self, dst: &mut [u8]) -> IoResult<uint> {
        let mut num_bytes = dst.len();
        if num_bytes <= self.cap - self.pos {
            //# Serve from the current buffer.
            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.pos + num_bytes));
            self.pos += num_bytes;
            return Ok(num_bytes);
        } else {
            //# Copy current available into destination.

            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap));
            let fromFirstBuffer = self.cap - self.pos;

            let dst1 = dst.mut_slice(fromFirstBuffer, num_bytes);
            num_bytes -= fromFirstBuffer;
            if num_bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, num_bytes);
                std::vec::bytes::copy_memory(dst1,
                                             self.buf.slice(0, num_bytes));
                self.cap = n;
                self.pos = num_bytes;
                return Ok(fromFirstBuffer + num_bytes);
            } else {
                //# Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Ok(fromFirstBuffer + read_at_least(self.inner, dst1, num_bytes));
            }
        }
    }
}

pub struct ArrayInputStream<'a> {
    priv array : &'a [u8]
}

impl <'a> ArrayInputStream<'a> {
    pub fn new<'b>(array : &'b [u8]) -> ArrayInputStream<'b> {
        ArrayInputStream { array : array }
    }
}

impl <'a> Reader for ArrayInputStream<'a> {
    fn read(&mut self, dst: &mut [u8]) -> Result<uint, std::io::IoError> {
        let n = std::cmp::min(dst.len(), self.array.len());
        unsafe { dst.copy_memory(self.array.slice_to(n)); }
        self.array = self.array.slice_from(n);
        Ok(n)
    }
}

impl <'a> BufferedInputStream for ArrayInputStream<'a> {
    fn skip(&mut self, bytes : uint) {
        assert!(self.array.len() >= bytes,
                "ArrayInputStream ended prematurely.");
        self.array = self.array.slice_from(bytes);
    }
    unsafe fn get_read_buffer(&mut self) -> (*u8, *u8){
        let len = self.array.len();
        (self.array.as_ptr(),
         std::ptr::to_unsafe_ptr(self.array.unsafe_ref(len)))
    }
}

pub trait BufferedOutputStream : Writer {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8);
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> IoResult<()>;
}

pub struct BufferedOutputStreamWrapper<'a, W> {
    priv inner: &'a mut W,
    priv buf: ~[u8],
    priv pos: uint
}

impl <'a, W> BufferedOutputStreamWrapper<'a, W> {
    pub fn new<'b> (w : &'b mut W) -> BufferedOutputStreamWrapper<'b, W> {
        let mut result = BufferedOutputStreamWrapper {
            inner: w,
            buf : std::vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            result.buf.set_len(8192);
        }
        return result;
    }
}

impl<'a, W: Writer> BufferedOutputStream for BufferedOutputStreamWrapper<'a, W> {
    #[inline]
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        (std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(self.pos)),
         std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(len)))
    }

    #[inline]
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> IoResult<()> {
        let easyCase = ptr == std::ptr::to_mut_unsafe_ptr(self.buf.unsafe_mut_ref(self.pos));
        if easyCase {
            self.pos += size;
            Ok(())
        } else {
            std::vec::raw::mut_buf_as_slice::<u8,IoResult<()>>(ptr, size, |buf| {
                self.write(buf)
            })
        }
    }

}

impl<'a, W: Writer> Writer for BufferedOutputStreamWrapper<'a, W> {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = self.buf.mut_slice_from(self.pos);
            std::vec::bytes::copy_memory(dst, buf);
            self.pos += size;
        } else if size <= self.buf.len() {
            //# Too much for this buffer, but not a full buffer's
            //# worth, so we'll go ahead and copy.
            {
                let dst = self.buf.mut_slice_from(self.pos);
                std::vec::bytes::copy_memory(dst, buf.slice(0, available));
            }
            if_ok!(self.inner.write(self.buf));

            size -= available;
            let src = buf.slice_from(available);
            let dst = self.buf.mut_slice_from(0);
            std::vec::bytes::copy_memory(dst, src);
            self.pos = size;
        } else {
            //# Writing so much data that we might as well write
            //# directly to avoid a copy.
            if_ok!(self.inner.write(self.buf.slice(0, self.pos)));
            self.pos = 0;
            if_ok!(self.inner.write(buf));
        }
        return Ok(());
    }

    fn flush(&mut self) -> IoResult<()> {
        if self.pos > 0 {
            if_ok!(self.inner.write(self.buf.slice(0, self.pos)));
            self.pos = 0;
        }
        return Ok(())
    }
}

pub struct ArrayOutputStream<'a> {
    priv array : &'a mut [u8],
    priv fill_pos : uint,
}

impl <'a> ArrayOutputStream<'a> {
    pub fn new<'b>(array : &'b mut [u8]) -> ArrayOutputStream<'b> {
        ArrayOutputStream {
            array : array,
            fill_pos : 0
        }
    }
}

impl <'a> Writer for ArrayOutputStream<'a> {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        assert!(buf.len() <= self.array.len() - self.fill_pos,
                "ArrayOutputStream's backing array was not large enough for the data written.");
        unsafe { self.array.mut_slice_from(self.fill_pos).copy_memory(buf); }
        self.fill_pos += buf.len();
        Ok(())
    }
}

impl <'a> BufferedOutputStream for ArrayOutputStream<'a> {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.array.len();
        (std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(self.fill_pos)),
         std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(len)))
    }
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) -> IoResult<()> {
        let easyCase = ptr == std::ptr::to_mut_unsafe_ptr(self.array.unsafe_mut_ref(self.fill_pos));
        if easyCase {
            self.fill_pos += size;
            Ok(())
        } else {
            std::vec::raw::mut_buf_as_slice::<u8,IoResult<()>>(ptr, size, |buf| {
                self.write(buf)
            })
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

use crate::{
    metrics::{
        DIEM_STORAGE_OTHER_TIMERS_SECONDS, DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION,
        DIEM_STORAGE_PRUNE_WINDOW,
    },
    schema::{
        jellyfish_merkle_node::JellyfishMerkleNodeSchema, stale_node_index::StaleNodeIndexSchema,
    },
};
use anyhow::Result;
use diem_infallible::Mutex;
use diem_jellyfish_merkle::StaleNodeIndex;
use diem_logger::prelude::*;
use diem_types::transaction::Version;
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::{
    iter::Peekable,
    sync::{
        atomic::{AtomicU64, Ordering},
        mpsc::{channel, Receiver, Sender},
        Arc,
    },
    thread::{sleep, JoinHandle},
    time::{Duration, Instant},
};

/// The `Pruner` is meant to be part of a `DiemDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// Other than the latest version, how many historical versions to keep being readable. For
    /// example, this being 0 means keep only the latest version.
    historical_versions_to_keep: u64,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// (For tests) A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets this atomic value to `V`, all versions before `V` can no longer be accessed.
    #[allow(dead_code)]
    worker_progress: Arc<AtomicU64>,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(db: Arc<DB>, historical_versions_to_keep: u64) -> Self {
        let (command_sender, command_receiver) = channel();

        let worker_progress = Arc::new(AtomicU64::new(0));
        let worker_progress_clone = Arc::clone(&worker_progress);

        DIEM_STORAGE_PRUNE_WINDOW.set(historical_versions_to_keep as i64);
        let worker_thread = std::thread::Builder::new()
            .name("diemdb_pruner".into())
            .spawn(move || Worker::new(db, command_receiver, worker_progress_clone).work())
            .expect("Creating pruner thread should succeed.");

        Self {
            historical_versions_to_keep,
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            worker_progress,
        }
    }

    /// Sends pruning command to the worker thread when necessary.
    pub fn wake(&self, latest_version: Version) {
        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            self.command_sender
                .lock()
                .send(Command::Prune {
                    least_readable_version,
                })
                .expect("Receiver should not destruct prematurely.");
        }
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(&self, latest_version: Version) -> Result<()> {
        self.wake(latest_version);

        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if self.worker_progress.load(Ordering::Relaxed) >= least_readable_version {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}

enum Command {
    Quit,
    Prune { least_readable_version: Version },
}

struct Worker {
    db: Arc<DB>,
    command_receiver: Receiver<Command>,
    target_least_readable_version: Version,
    /// Keeps a record of the pruning progress. If this equals to version `V`, we know versions
    /// smaller than `V` are no longer readable.
    /// This being an atomic value is to communicate the info with the Pruner thread (for tests).
    least_readable_version: Arc<AtomicU64>,
    /// Indicates if there's NOT any pending work to do currently, to hint
    /// `Self::receive_commands()` to `recv()` blocking-ly.
    blocking_recv: bool,
    index_min_nonpurged_version: Version,
    index_purged_at: Instant,
}

impl Worker {
    const MAX_VERSIONS_TO_PRUNE_PER_BATCH: usize = 100;

    fn new(
        db: Arc<DB>,
        command_receiver: Receiver<Command>,
        least_readable_version: Arc<AtomicU64>,
    ) -> Self {
        Self {
            db,
            command_receiver,
            least_readable_version,
            target_least_readable_version: 0,
            blocking_recv: true,
            index_min_nonpurged_version: 0,
            index_purged_at: Instant::now(),
        }
    }

    fn work(mut self) {
        self.initialize();

        while self.receive_commands() {
            // Process a reasonably small batch of work before trying to receive commands again,
            // in case `Command::Quit` is received (that's when we should quit.)
            match prune_state(
                Arc::clone(&self.db),
                self.least_readable_version.load(Ordering::Relaxed),
                self.target_least_readable_version,
                Self::MAX_VERSIONS_TO_PRUNE_PER_BATCH,
            ) {
                Ok(least_readable_version) => {
                    self.record_progress(least_readable_version);

                    // Make next recv() blocking if all done.
                    self.blocking_recv =
                        least_readable_version == self.target_least_readable_version;

                    // Try to purge the log.
                    if let Err(e) = self.maybe_purge_index() {
                        warn!(
                            error = ?e,
                            "Failed purging state node index, ignored.",
                        );
                    }
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "Error pruning stale state nodes.",
                    );
                    // On error, stop retrying vigorously by making next recv() blocking.
                    self.blocking_recv = true;
                }
            }
        }
    }

    /// Find out the first undeleted item in the stale node index.
    ///
    /// Seeking from the beginning (version 0) is potentially costly, we do it once upon worker
    /// thread start, record the progress and seek from that position afterwards.
    fn initialize(&mut self) {
        loop {
            match self.get_least_readable_version() {
                Ok(least_readable_version) => {
                    info!(
                        least_readable_version = least_readable_version,
                        "[state pruner worker] initialized."
                    );
                    self.target_least_readable_version = least_readable_version;
                    self.record_progress(least_readable_version);
                    return;
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "[state pruner worker] Error on first seek. Retrying in 1 second.",
                    );
                    sleep(Duration::from_secs(1));
                }
            }
        }
    }

    fn get_least_readable_version(&self) -> Result<Version> {
        let mut iter = self
            .db
            .iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek_to_first();
        Ok(iter.next().transpose()?.map_or(0, |(index, _)| {
            index
                .stale_since_version
                .checked_sub(1)
                .expect("Nothing is stale since version 0.")
        }))
    }

    /// Log the progress.
    fn record_progress(&mut self, least_readable_version: Version) {
        self.least_readable_version
            .store(least_readable_version, Ordering::Relaxed);
        DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION.set(least_readable_version as i64);
    }

    /// Tries to receive all pending commands, blocking waits for the next command if no work needs
    /// to be done, otherwise quits with `true` to allow the outer loop to do some work before
    /// getting back here.
    ///
    /// Returns `false` if `Command::Quit` is received, to break the outer loop and let
    /// `work_loop()` return.
    fn receive_commands(&mut self) -> bool {
        loop {
            let command = if self.blocking_recv {
                // Worker has nothing to do, blocking wait for the next command.
                self.command_receiver
                    .recv()
                    .expect("Sender should not destruct prematurely.")
            } else {
                // Worker has pending work to do, non-blocking recv.
                match self.command_receiver.try_recv() {
                    Ok(command) => command,
                    // Channel has drained, yield control to the outer loop.
                    Err(_) => return true,
                }
            };

            match command {
                // On `Command::Quit` inform the outer loop to quit by returning `false`.
                Command::Quit => return false,
                Command::Prune {
                    least_readable_version,
                } => {
                    if least_readable_version > self.target_least_readable_version {
                        self.target_least_readable_version = least_readable_version;
                        // Switch to non-blocking to allow some work to be done after the
                        // channel has drained.
                        self.blocking_recv = false;
                    }
                }
            }
        }
    }

    /// Purge the stale node index so that after restart not too much already pruned stuff is dealt
    /// with again (although no harm is done deleting those then non-existent things.)
    ///
    /// We issue (range) deletes on the index only periodically instead of after every pruning batch
    /// to avoid sending too many deletions to the DB, which takes disk space and slows it down.
    fn maybe_purge_index(&mut self) -> Result<()> {
        const MIN_INTERVAL: Duration = Duration::from_secs(60);
        const MIN_VERSIONS: u64 = 60000;

        // A deletion is issued at most once in one minute and when the pruner has progressed by at
        // least 60000 versions (assuming the pruner deletes as slow as 1000 versions per second,
        // this imposes at most one minute of work in vain after restarting.)
        let now = Instant::now();
        if now - self.index_purged_at > MIN_INTERVAL {
            let least_readable_version = self.least_readable_version.load(Ordering::Relaxed);

            if least_readable_version - self.index_min_nonpurged_version + 1 > MIN_VERSIONS {
                let new_min_non_purged_version = least_readable_version + 1;
                self.db.range_delete::<StaleNodeIndexSchema, Version>(
                    &self.index_min_nonpurged_version,
                    &new_min_non_purged_version, // end is exclusive
                )?;
                self.index_min_nonpurged_version = new_min_non_purged_version;
                self.index_purged_at = now;
            }
        }

        Ok(())
    }
}

struct StaleNodeIndicesByVersionIterator<'a> {
    inner: Peekable<SchemaIterator<'a, StaleNodeIndexSchema>>,
    target_least_readable_version: Version,
}

impl<'a> StaleNodeIndicesByVersionIterator<'a> {
    fn new(
        db: &'a DB,
        least_readable_version: Version,
        target_least_readable_version: Version,
    ) -> Result<Self> {
        let mut iter = db.iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek(&least_readable_version)?;

        Ok(Self {
            inner: iter.peekable(),
            target_least_readable_version,
        })
    }

    fn next_result(&mut self) -> Result<Option<Vec<StaleNodeIndex>>> {
        match self.inner.next().transpose()? {
            None => Ok(None),
            Some((index, _)) => {
                let version = index.stale_since_version;
                if version > self.target_least_readable_version {
                    return Ok(None);
                }

                let mut indices = vec![index];
                while let Some(res) = self.inner.peek() {
                    if let Ok((index_ref, _)) = res {
                        if index_ref.stale_since_version != version {
                            break;
                        }
                    }

                    let (index, _) = self.inner.next().transpose()?.expect("Should be Some.");
                    indices.push(index);
                }

                Ok(Some(indices))
            }
        }
    }
}

impl<'a> Iterator for StaleNodeIndicesByVersionIterator<'a> {
    type Item = Result<Vec<StaleNodeIndex>>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_result().transpose()
    }
}

pub fn prune_state(
    db: Arc<DB>,
    least_readable_version: Version,
    target_least_readable_version: Version,
    max_versions: usize,
) -> Result<Version> {
    let indices = StaleNodeIndicesByVersionIterator::new(
        &db,
        least_readable_version,
        target_least_readable_version,
    )?
    .take(max_versions) // Iterator<Item = Result<Vec<StaleNodeIndex>>>
    .collect::<Result<Vec<_>>>()? // now Vec<Vec<StaleNodeIndex>>
    .into_iter()
    .flatten()
    .collect::<Vec<_>>();

    if indices.is_empty() {
        Ok(least_readable_version)
    } else {
        let _timer = DIEM_STORAGE_OTHER_TIMERS_SECONDS
            .with_label_values(&["pruner_commit"])
            .start_timer();
        let new_least_readable_version = indices.last().expect("Should exist.").stale_since_version;
        let mut batch = SchemaBatch::new();
        indices
            .into_iter()
            .try_for_each(|index| batch.delete::<JellyfishMerkleNodeSchema>(&index.node_key))?;
        db.write_schemas(batch)?;
        Ok(new_least_readable_version)
    }
}

#[cfg(test)]
mod test;
<commit_msg>[pruner] fix busy looping after db-restore<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

use crate::{
    metrics::{
        DIEM_STORAGE_OTHER_TIMERS_SECONDS, DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION,
        DIEM_STORAGE_PRUNE_WINDOW,
    },
    schema::{
        jellyfish_merkle_node::JellyfishMerkleNodeSchema, stale_node_index::StaleNodeIndexSchema,
    },
};
use anyhow::Result;
use diem_infallible::Mutex;
use diem_jellyfish_merkle::StaleNodeIndex;
use diem_logger::prelude::*;
use diem_types::transaction::Version;
use schemadb::{ReadOptions, SchemaBatch, SchemaIterator, DB};
use std::{
    iter::Peekable,
    sync::{
        atomic::{AtomicU64, Ordering},
        mpsc::{channel, Receiver, Sender},
        Arc,
    },
    thread::{sleep, JoinHandle},
    time::{Duration, Instant},
};

/// The `Pruner` is meant to be part of a `DiemDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// Other than the latest version, how many historical versions to keep being readable. For
    /// example, this being 0 means keep only the latest version.
    historical_versions_to_keep: u64,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// (For tests) A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets this atomic value to `V`, all versions before `V` can no longer be accessed.
    #[allow(dead_code)]
    worker_progress: Arc<AtomicU64>,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(db: Arc<DB>, historical_versions_to_keep: u64) -> Self {
        let (command_sender, command_receiver) = channel();

        let worker_progress = Arc::new(AtomicU64::new(0));
        let worker_progress_clone = Arc::clone(&worker_progress);

        DIEM_STORAGE_PRUNE_WINDOW.set(historical_versions_to_keep as i64);
        let worker_thread = std::thread::Builder::new()
            .name("diemdb_pruner".into())
            .spawn(move || Worker::new(db, command_receiver, worker_progress_clone).work())
            .expect("Creating pruner thread should succeed.");

        Self {
            historical_versions_to_keep,
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            worker_progress,
        }
    }

    /// Sends pruning command to the worker thread when necessary.
    pub fn wake(&self, latest_version: Version) {
        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            self.command_sender
                .lock()
                .send(Command::Prune {
                    least_readable_version,
                })
                .expect("Receiver should not destruct prematurely.");
        }
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(&self, latest_version: Version) -> Result<()> {
        self.wake(latest_version);

        if latest_version > self.historical_versions_to_keep {
            let least_readable_version = latest_version - self.historical_versions_to_keep;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if self.worker_progress.load(Ordering::Relaxed) >= least_readable_version {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}

enum Command {
    Quit,
    Prune { least_readable_version: Version },
}

struct Worker {
    db: Arc<DB>,
    command_receiver: Receiver<Command>,
    target_least_readable_version: Version,
    /// Keeps a record of the pruning progress. If this equals to version `V`, we know versions
    /// smaller than `V` are no longer readable.
    /// This being an atomic value is to communicate the info with the Pruner thread (for tests).
    least_readable_version: Arc<AtomicU64>,
    /// Indicates if there's NOT any pending work to do currently, to hint
    /// `Self::receive_commands()` to `recv()` blocking-ly.
    blocking_recv: bool,
    index_min_nonpurged_version: Version,
    index_purged_at: Instant,
}

impl Worker {
    const MAX_VERSIONS_TO_PRUNE_PER_BATCH: usize = 100;

    fn new(
        db: Arc<DB>,
        command_receiver: Receiver<Command>,
        least_readable_version: Arc<AtomicU64>,
    ) -> Self {
        Self {
            db,
            command_receiver,
            least_readable_version,
            target_least_readable_version: 0,
            blocking_recv: true,
            index_min_nonpurged_version: 0,
            index_purged_at: Instant::now(),
        }
    }

    fn work(mut self) {
        self.initialize();

        while self.receive_commands() {
            // Process a reasonably small batch of work before trying to receive commands again,
            // in case `Command::Quit` is received (that's when we should quit.)
            let least_readable_version = self.least_readable_version.load(Ordering::Relaxed);
            match prune_state(
                Arc::clone(&self.db),
                least_readable_version,
                self.target_least_readable_version,
                Self::MAX_VERSIONS_TO_PRUNE_PER_BATCH,
            ) {
                Ok(new_least_readable_version) => {
                    self.record_progress(new_least_readable_version);

                    // Make next recv() blocking if nothing left to do.
                    self.blocking_recv = new_least_readable_version == least_readable_version // did nothing
                        || new_least_readable_version == self.target_least_readable_version; // did all

                    // Try to purge the log.
                    if let Err(e) = self.maybe_purge_index() {
                        warn!(
                            error = ?e,
                            "Failed purging state node index, ignored.",
                        );
                    }
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "Error pruning stale state nodes.",
                    );
                    // On error, stop retrying vigorously by making next recv() blocking.
                    self.blocking_recv = true;
                }
            }
        }
    }

    /// Find out the first undeleted item in the stale node index.
    ///
    /// Seeking from the beginning (version 0) is potentially costly, we do it once upon worker
    /// thread start, record the progress and seek from that position afterwards.
    fn initialize(&mut self) {
        loop {
            match self.get_least_readable_version() {
                Ok(least_readable_version) => {
                    info!(
                        least_readable_version = least_readable_version,
                        "[state pruner worker] initialized."
                    );
                    self.target_least_readable_version = least_readable_version;
                    self.record_progress(least_readable_version);
                    return;
                }
                Err(e) => {
                    error!(
                        error = ?e,
                        "[state pruner worker] Error on first seek. Retrying in 1 second.",
                    );
                    sleep(Duration::from_secs(1));
                }
            }
        }
    }

    fn get_least_readable_version(&self) -> Result<Version> {
        let mut iter = self
            .db
            .iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek_to_first();
        Ok(iter.next().transpose()?.map_or(0, |(index, _)| {
            index
                .stale_since_version
                .checked_sub(1)
                .expect("Nothing is stale since version 0.")
        }))
    }

    /// Log the progress.
    fn record_progress(&mut self, least_readable_version: Version) {
        self.least_readable_version
            .store(least_readable_version, Ordering::Relaxed);
        DIEM_STORAGE_PRUNER_LEAST_READABLE_STATE_VERSION.set(least_readable_version as i64);
    }

    /// Tries to receive all pending commands, blocking waits for the next command if no work needs
    /// to be done, otherwise quits with `true` to allow the outer loop to do some work before
    /// getting back here.
    ///
    /// Returns `false` if `Command::Quit` is received, to break the outer loop and let
    /// `work_loop()` return.
    fn receive_commands(&mut self) -> bool {
        loop {
            let command = if self.blocking_recv {
                // Worker has nothing to do, blocking wait for the next command.
                self.command_receiver
                    .recv()
                    .expect("Sender should not destruct prematurely.")
            } else {
                // Worker has pending work to do, non-blocking recv.
                match self.command_receiver.try_recv() {
                    Ok(command) => command,
                    // Channel has drained, yield control to the outer loop.
                    Err(_) => return true,
                }
            };

            match command {
                // On `Command::Quit` inform the outer loop to quit by returning `false`.
                Command::Quit => return false,
                Command::Prune {
                    least_readable_version,
                } => {
                    if least_readable_version > self.target_least_readable_version {
                        self.target_least_readable_version = least_readable_version;
                        // Switch to non-blocking to allow some work to be done after the
                        // channel has drained.
                        self.blocking_recv = false;
                    }
                }
            }
        }
    }

    /// Purge the stale node index so that after restart not too much already pruned stuff is dealt
    /// with again (although no harm is done deleting those then non-existent things.)
    ///
    /// We issue (range) deletes on the index only periodically instead of after every pruning batch
    /// to avoid sending too many deletions to the DB, which takes disk space and slows it down.
    fn maybe_purge_index(&mut self) -> Result<()> {
        const MIN_INTERVAL: Duration = Duration::from_secs(60);
        const MIN_VERSIONS: u64 = 60000;

        // A deletion is issued at most once in one minute and when the pruner has progressed by at
        // least 60000 versions (assuming the pruner deletes as slow as 1000 versions per second,
        // this imposes at most one minute of work in vain after restarting.)
        let now = Instant::now();
        if now - self.index_purged_at > MIN_INTERVAL {
            let least_readable_version = self.least_readable_version.load(Ordering::Relaxed);

            if least_readable_version - self.index_min_nonpurged_version + 1 > MIN_VERSIONS {
                let new_min_non_purged_version = least_readable_version + 1;
                self.db.range_delete::<StaleNodeIndexSchema, Version>(
                    &self.index_min_nonpurged_version,
                    &new_min_non_purged_version, // end is exclusive
                )?;
                self.index_min_nonpurged_version = new_min_non_purged_version;
                self.index_purged_at = now;
            }
        }

        Ok(())
    }
}

struct StaleNodeIndicesByVersionIterator<'a> {
    inner: Peekable<SchemaIterator<'a, StaleNodeIndexSchema>>,
    target_least_readable_version: Version,
}

impl<'a> StaleNodeIndicesByVersionIterator<'a> {
    fn new(
        db: &'a DB,
        least_readable_version: Version,
        target_least_readable_version: Version,
    ) -> Result<Self> {
        let mut iter = db.iter::<StaleNodeIndexSchema>(ReadOptions::default())?;
        iter.seek(&least_readable_version)?;

        Ok(Self {
            inner: iter.peekable(),
            target_least_readable_version,
        })
    }

    fn next_result(&mut self) -> Result<Option<Vec<StaleNodeIndex>>> {
        match self.inner.next().transpose()? {
            None => Ok(None),
            Some((index, _)) => {
                let version = index.stale_since_version;
                if version > self.target_least_readable_version {
                    return Ok(None);
                }

                let mut indices = vec![index];
                while let Some(res) = self.inner.peek() {
                    if let Ok((index_ref, _)) = res {
                        if index_ref.stale_since_version != version {
                            break;
                        }
                    }

                    let (index, _) = self.inner.next().transpose()?.expect("Should be Some.");
                    indices.push(index);
                }

                Ok(Some(indices))
            }
        }
    }
}

impl<'a> Iterator for StaleNodeIndicesByVersionIterator<'a> {
    type Item = Result<Vec<StaleNodeIndex>>;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_result().transpose()
    }
}

pub fn prune_state(
    db: Arc<DB>,
    least_readable_version: Version,
    target_least_readable_version: Version,
    max_versions: usize,
) -> Result<Version> {
    let indices = StaleNodeIndicesByVersionIterator::new(
        &db,
        least_readable_version,
        target_least_readable_version,
    )?
    .take(max_versions) // Iterator<Item = Result<Vec<StaleNodeIndex>>>
    .collect::<Result<Vec<_>>>()? // now Vec<Vec<StaleNodeIndex>>
    .into_iter()
    .flatten()
    .collect::<Vec<_>>();

    if indices.is_empty() {
        Ok(least_readable_version)
    } else {
        let _timer = DIEM_STORAGE_OTHER_TIMERS_SECONDS
            .with_label_values(&["pruner_commit"])
            .start_timer();
        let new_least_readable_version = indices.last().expect("Should exist.").stale_since_version;
        let mut batch = SchemaBatch::new();
        indices
            .into_iter()
            .try_for_each(|index| batch.delete::<JellyfishMerkleNodeSchema>(&index.node_key))?;
        db.write_schemas(batch)?;
        Ok(new_least_readable_version)
    }
}

#[cfg(test)]
mod test;
<|endoftext|>


<commit_before>/*
 * libgit2 "init" example - shows how to initialize a new repo
 *
 * Written by the libgit2 contributors
 *
 * To the extent possible under law, the author(s) have dedicated all copyright
 * and related and neighboring rights to this software to the public domain
 * worldwide. This software is distributed without any warranty.
 *
 * You should have received a copy of the CC0 Public Domain Dedication along
 * with this software. If not, see
 * <http://creativecommons.org/publicdomain/zero/1.0/>.
 */

#![deny(warnings)]
#![feature(core, convert)]

extern crate git2;
extern crate docopt;
extern crate rustc_serialize;

use docopt::Docopt;
use git2::{Repository, RepositoryInitOptions, RepositoryInitMode, Error};
use std::num::FromStrRadix;
use std::path::{PathBuf, Path};

#[derive(RustcDecodable)]
struct Args {
    arg_directory: String,
    flag_quiet: bool,
    flag_bare: bool,
    flag_template: Option<String>,
    flag_separate_git_dir: Option<String>,
    flag_initial_commit: bool,
    flag_shared: Option<String>,
}

fn run(args: &Args) -> Result<(), Error> {
    let mut path = PathBuf::from(&args.arg_directory);
    let repo = if !args.flag_bare && args.flag_template.is_none() &&
                  args.flag_shared.is_none() &&
                  args.flag_separate_git_dir.is_none() {
        try!(Repository::init(&path))
    } else {
        let mut opts = RepositoryInitOptions::new();
        opts.bare(args.flag_bare);
        if let Some(ref s) = args.flag_template {
            opts.template_path(Path::new(s));
        }

        // If you specified a separate git directory, then initialize
        // the repository at that path and use the second path as the
        // working directory of the repository (with a git-link file)
        if let Some(ref s) = args.flag_separate_git_dir {
            opts.workdir_path(&path);
            path = PathBuf::from(s);
        }

        if let Some(ref s) = args.flag_shared {
            opts.mode(try!(parse_shared(&s)));
        }
        try!(Repository::init_opts(&path, &opts))
    };

    // Print a message to stdout like "git init" does
    if !args.flag_quiet {
        if args.flag_bare || args.flag_separate_git_dir.is_some() {
            path = repo.path().to_path_buf();
        } else {
            path = repo.workdir().unwrap().to_path_buf();
        }
        println!("Initialized empty Git repository in {}", path.display());
    }

    if args.flag_initial_commit {
        try!(create_initial_commit(&repo));
        println!("Created empty initial commit");
    }

    Ok(())
}

/// Unlike regular "git init", this example shows how to create an initial empty
/// commit in the repository. This is the helper function that does that.
fn create_initial_commit(repo: &Repository) -> Result<(), Error> {
    // First use the config to initialize a commit signature for the user.
    let sig = try!(repo.signature());

    // Now let's create an empty tree for this commit
    let tree_id = {
        let mut index = try!(repo.index());

        // Outside of this example, you could call index.add_path()
        // here to put actual files into the index. For our purposes, we'll
        // leave it empty for now.

        try!(index.write_tree())
    };

    let tree = try!(repo.find_tree(tree_id));

    // Ready to create the initial commit.
    //
    // Normally creating a commit would involve looking up the current HEAD
    // commit and making that be the parent of the initial commit, but here this
    // is the first commit so there will be no parent.
    try!(repo.commit(Some("HEAD"), &sig, &sig, "Initial commit", &tree, &[]));

    Ok(())
}

fn parse_shared(shared: &str) -> Result<RepositoryInitMode, Error> {
    match shared {
        "false" | "umask" => Ok(git2::REPOSITORY_INIT_SHARED_UMASK),
        "true" | "group" => Ok(git2::REPOSITORY_INIT_SHARED_GROUP),
        "all" | "world" => Ok(git2::REPOSITORY_INIT_SHARED_ALL),
        _ => {
            if shared.starts_with("0") {
                match FromStrRadix::from_str_radix(&shared[1..], 8).ok() {
                    Some(n) => {
                        return Ok(RepositoryInitMode::from_bits_truncate(n))
                    }
                    None => {
                        Err(Error::from_str("invalid octal value for --shared"))
                    }
                }
            } else {
                Err(Error::from_str("unknown value for --shared"))
            }
        }
    }
}

fn main() {
    const USAGE: &'static str = "
usage: init [options] <directory>

Options:
    -q, --quiet                 don't print information to stdout
    --bare                      initialize a new bare repository
    --template <dir>            use <dir> as an initialization template
    --separate-git-dir <dir>    use <dir> as the .git directory
    --initial-commit            create an initial empty commit
    --shared <perms>            permissions to create the repository with
";

    let args = Docopt::new(USAGE).and_then(|d| d.decode())
                                 .unwrap_or_else(|e| e.exit());
    match run(&args) {
        Ok(()) => {}
        Err(e) => println!("error: {}", e),
    }
}
<commit_msg>Update to rust master<commit_after>/*
 * libgit2 "init" example - shows how to initialize a new repo
 *
 * Written by the libgit2 contributors
 *
 * To the extent possible under law, the author(s) have dedicated all copyright
 * and related and neighboring rights to this software to the public domain
 * worldwide. This software is distributed without any warranty.
 *
 * You should have received a copy of the CC0 Public Domain Dedication along
 * with this software. If not, see
 * <http://creativecommons.org/publicdomain/zero/1.0/>.
 */

#![deny(warnings)]
#![feature(convert)]

extern crate git2;
extern crate docopt;
extern crate rustc_serialize;

use docopt::Docopt;
use git2::{Repository, RepositoryInitOptions, RepositoryInitMode, Error};
use std::path::{PathBuf, Path};

#[derive(RustcDecodable)]
struct Args {
    arg_directory: String,
    flag_quiet: bool,
    flag_bare: bool,
    flag_template: Option<String>,
    flag_separate_git_dir: Option<String>,
    flag_initial_commit: bool,
    flag_shared: Option<String>,
}

fn run(args: &Args) -> Result<(), Error> {
    let mut path = PathBuf::from(&args.arg_directory);
    let repo = if !args.flag_bare && args.flag_template.is_none() &&
                  args.flag_shared.is_none() &&
                  args.flag_separate_git_dir.is_none() {
        try!(Repository::init(&path))
    } else {
        let mut opts = RepositoryInitOptions::new();
        opts.bare(args.flag_bare);
        if let Some(ref s) = args.flag_template {
            opts.template_path(Path::new(s));
        }

        // If you specified a separate git directory, then initialize
        // the repository at that path and use the second path as the
        // working directory of the repository (with a git-link file)
        if let Some(ref s) = args.flag_separate_git_dir {
            opts.workdir_path(&path);
            path = PathBuf::from(s);
        }

        if let Some(ref s) = args.flag_shared {
            opts.mode(try!(parse_shared(&s)));
        }
        try!(Repository::init_opts(&path, &opts))
    };

    // Print a message to stdout like "git init" does
    if !args.flag_quiet {
        if args.flag_bare || args.flag_separate_git_dir.is_some() {
            path = repo.path().to_path_buf();
        } else {
            path = repo.workdir().unwrap().to_path_buf();
        }
        println!("Initialized empty Git repository in {}", path.display());
    }

    if args.flag_initial_commit {
        try!(create_initial_commit(&repo));
        println!("Created empty initial commit");
    }

    Ok(())
}

/// Unlike regular "git init", this example shows how to create an initial empty
/// commit in the repository. This is the helper function that does that.
fn create_initial_commit(repo: &Repository) -> Result<(), Error> {
    // First use the config to initialize a commit signature for the user.
    let sig = try!(repo.signature());

    // Now let's create an empty tree for this commit
    let tree_id = {
        let mut index = try!(repo.index());

        // Outside of this example, you could call index.add_path()
        // here to put actual files into the index. For our purposes, we'll
        // leave it empty for now.

        try!(index.write_tree())
    };

    let tree = try!(repo.find_tree(tree_id));

    // Ready to create the initial commit.
    //
    // Normally creating a commit would involve looking up the current HEAD
    // commit and making that be the parent of the initial commit, but here this
    // is the first commit so there will be no parent.
    try!(repo.commit(Some("HEAD"), &sig, &sig, "Initial commit", &tree, &[]));

    Ok(())
}

fn parse_shared(shared: &str) -> Result<RepositoryInitMode, Error> {
    match shared {
        "false" | "umask" => Ok(git2::REPOSITORY_INIT_SHARED_UMASK),
        "true" | "group" => Ok(git2::REPOSITORY_INIT_SHARED_GROUP),
        "all" | "world" => Ok(git2::REPOSITORY_INIT_SHARED_ALL),
        _ => {
            if shared.starts_with("0") {
                match u32::from_str_radix(&shared[1..], 8).ok() {
                    Some(n) => {
                        return Ok(RepositoryInitMode::from_bits_truncate(n))
                    }
                    None => {
                        Err(Error::from_str("invalid octal value for --shared"))
                    }
                }
            } else {
                Err(Error::from_str("unknown value for --shared"))
            }
        }
    }
}

fn main() {
    const USAGE: &'static str = "
usage: init [options] <directory>

Options:
    -q, --quiet                 don't print information to stdout
    --bare                      initialize a new bare repository
    --template <dir>            use <dir> as an initialization template
    --separate-git-dir <dir>    use <dir> as the .git directory
    --initial-commit            create an initial empty commit
    --shared <perms>            permissions to create the repository with
";

    let args = Docopt::new(USAGE).and_then(|d| d.decode())
                                 .unwrap_or_else(|e| e.exit());
    match run(&args) {
        Ok(()) => {}
        Err(e) => println!("error: {}", e),
    }
}
<|endoftext|>


<commit_before>#![feature(io)]

extern crate tcl;
use std::io;
use std::io::{BufRead, Write};

/**
 * A simple program that passes commands directly to the interpreter
 */
pub fn main() {

    let input = io::stdin();
    let output = io::stdout();
    let mut buf = String::new();

    // Initialize Tcl
    let env = tcl::init();
    // Get an interpreter
    let mut interp = env.interpreter();

    loop {
        {
            let mut in_lock = input.lock();
            in_lock.read_line(&mut buf).ok().expect("Couldn't read stdin!");
        }

        let result = match interp.eval(buf.as_slice(), tcl::EvalScope::Local) {
            tcl::TclResult::Error(string) => {
                string
            },
            _ => interp.string_result()
        };

        {
            let mut out_lock = output.lock();
            out_lock.write_all(result.as_bytes()).ok().expect("Couldn't write stdout!");
        }

    }
}
<commit_msg>Repl now actually works<commit_after>#![feature(io)]

extern crate tcl;
use std::io;
use std::io::{BufRead, Write};

/**
 * A simple program that passes commands directly to the interpreter
 */
pub fn main() {

    let input = io::stdin();

    // Initialize Tcl
    let env = tcl::init();
    // Get an interpreter
    let mut interp = env.interpreter();

    loop {
        let mut buf = String::new();
        {
            let mut in_lock = input.lock();
            in_lock.read_line(&mut buf).ok().expect("Couldn't read stdin!");
        }

        let result = interp.eval(buf.as_slice(), tcl::EvalScope::Local);

        match result {
            tcl::TclResult::Error(string) => {
                println!("Error: {}", string);
            },
            _ => println!("{}", interp.string_result())
        };

    }
}
<|endoftext|>


<commit_before>//! See https://github.com/matklad/cargo-xtask/.
//!
//! This binary defines various auxiliary build commands, which are not
//! expressible with just `cargo`. Notably, it provides `cargo xtask codegen`
//! for code generation and `cargo xtask install` for installation of
//! rust-analyzer server and client.
//!
//! This binary is integrated into the `cargo` command line by using an alias in
//! `.cargo/config`.
mod help;

use std::{env, fmt::Write, path::PathBuf, str};

use anyhow::Context;
use pico_args::Arguments;
use xtask::{
    codegen::{self, Mode},
    install_pre_commit_hook, reformat_staged_files, run, run_clippy, run_fuzzer, run_rustfmt,
    run_with_output, Cmd, Result,
};

// Latest stable, feel free to send a PR if this lags behind.
const REQUIRED_RUST_VERSION: u32 = 39;

struct InstallOpt {
    client: Option<ClientOpt>,
    server: Option<ServerOpt>,
}

enum ClientOpt {
    VsCode,
}

struct ServerOpt {
    jemalloc: bool,
}

fn main() -> Result<()> {
    if env::args().next().map(|it| it.contains("pre-commit")) == Some(true) {
        return reformat_staged_files();
    }

    let subcommand = match std::env::args_os().nth(1) {
        None => {
            eprintln!("{}", help::GLOBAL_HELP);
            return Ok(());
        }
        Some(s) => s,
    };
    let mut matches = Arguments::from_vec(std::env::args_os().skip(2).collect());
    let subcommand = &*subcommand.to_string_lossy();
    match subcommand {
        "install" => {
            if matches.contains(["-h", "--help"]) {
                eprintln!("{}", help::INSTALL_HELP);
                return Ok(());
            }
            let server = matches.contains("--server");
            let client_code = matches.contains("--client-code");
            if server && client_code {
                eprintln!("{}", help::INSTALL_RA_CONFLICT);
                return Ok(());
            }
            let jemalloc = matches.contains("--jemalloc");
            matches.finish().or_else(handle_extra_flags)?;
            let opts = InstallOpt {
                client: if server { None } else { Some(ClientOpt::VsCode) },
                server: if client_code { None } else { Some(ServerOpt { jemalloc }) },
            };
            install(opts)?
        }
        "codegen" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            codegen::generate_syntax(Mode::Overwrite)?;
            codegen::generate_parser_tests(Mode::Overwrite)?;
            codegen::generate_assists_docs(Mode::Overwrite)?;
        }
        "format" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_rustfmt(Mode::Overwrite)?
        }
        "install-pre-commit-hook" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            install_pre_commit_hook()?
        }
        "lint" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_clippy()?
        }
        "fuzz-tests" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_fuzzer()?
        }
        _ => eprintln!("{}", help::GLOBAL_HELP),
    }
    Ok(())
}

fn handle_extra_flags(e: pico_args::Error) -> Result<()> {
    if let pico_args::Error::UnusedArgsLeft(flags) = e {
        let mut invalid_flags = String::new();
        for flag in flags {
            write!(&mut invalid_flags, "{}, ", flag)?;
        }
        let (invalid_flags, _) = invalid_flags.split_at(invalid_flags.len() - 2);
        anyhow::bail!("Invalid flags: {}", invalid_flags)
    } else {
        anyhow::bail!(e.to_string())
    }
}

fn install(opts: InstallOpt) -> Result<()> {
    if cfg!(target_os = "macos") {
        fix_path_for_mac().context("Fix path for mac")?
    }
    if let Some(server) = opts.server {
        install_server(server).context("install server")?;
    }
    if let Some(client) = opts.client {
        install_client(client).context("install client")?;
    }
    Ok(())
}

fn fix_path_for_mac() -> Result<()> {
    let mut vscode_path: Vec<PathBuf> = {
        const COMMON_APP_PATH: &str =
            r"/Applications/Visual Studio Code.app/Contents/Resources/app/bin";
        const ROOT_DIR: &str = "";
        let home_dir = match env::var("HOME") {
            Ok(home) => home,
            Err(e) => anyhow::bail!("Failed getting HOME from environment with error: {}.", e),
        };

        [ROOT_DIR, &home_dir]
            .iter()
            .map(|dir| String::from(*dir) + COMMON_APP_PATH)
            .map(PathBuf::from)
            .filter(|path| path.exists())
            .collect()
    };

    if !vscode_path.is_empty() {
        let vars = match env::var_os("PATH") {
            Some(path) => path,
            None => anyhow::bail!("Could not get PATH variable from env."),
        };

        let mut paths = env::split_paths(&vars).collect::<Vec<_>>();
        paths.append(&mut vscode_path);
        let new_paths = env::join_paths(paths).context("build env PATH")?;
        env::set_var("PATH", &new_paths);
    }

    Ok(())
}

fn install_client(ClientOpt::VsCode: ClientOpt) -> Result<()> {
    let npm_version = Cmd {
        unix: r"npm --version",
        windows: r"cmd.exe /c npm.cmd --version",
        work_dir: "./editors/code",
    }
    .run();

    if npm_version.is_err() {
        eprintln!("\nERROR: `npm --version` failed, `npm` is required to build the VS Code plugin")
    }

    Cmd { unix: r"npm ci", windows: r"cmd.exe /c npm.cmd ci", work_dir: "./editors/code" }.run()?;
    Cmd {
        unix: r"npm run package --scripts-prepend-node-path",
        windows: r"cmd.exe /c npm.cmd run package",
        work_dir: "./editors/code",
    }
    .run()?;

    let code_binary = ["code", "code-insiders", "codium"].iter().find(|bin| {
        Cmd {
            unix: &format!("{} --version", bin),
            windows: &format!("cmd.exe /c {}.cmd --version", bin),
            work_dir: "./editors/code",
        }
        .run()
        .is_ok()
    });

    let code_binary = match code_binary {
        Some(it) => it,
        None => anyhow::bail!("Can't execute `code --version`. Perhaps it is not in $PATH?"),
    };

    Cmd {
        unix: &format!(r"{} --install-extension ./ra-lsp-0.0.1.vsix --force", code_binary),
        windows: &format!(
            r"cmd.exe /c {}.cmd --install-extension ./ra-lsp-0.0.1.vsix --force",
            code_binary
        ),
        work_dir: "./editors/code",
    }
    .run()?;

    let output = Cmd {
        unix: &format!(r"{} --list-extensions", code_binary),
        windows: &format!(r"cmd.exe /c {}.cmd --list-extensions", code_binary),
        work_dir: ".",
    }
    .run_with_output()?;

    if !str::from_utf8(&output.stdout)?.contains("ra-lsp") {
        anyhow::bail!(
            "Could not install the Visual Studio Code extension. \
             Please make sure you have at least NodeJS 10.x installed and try again."
        );
    }

    Ok(())
}

fn install_server(opts: ServerOpt) -> Result<()> {
    let mut old_rust = false;
    if let Ok(output) = run_with_output("cargo --version", ".") {
        if let Ok(stdout) = String::from_utf8(output.stdout) {
            println!("{}", stdout);
            if !check_version(&stdout, REQUIRED_RUST_VERSION) {
                old_rust = true;
            }
        }
    }

    if old_rust {
        eprintln!(
            "\nWARNING: at least rust 1.{}.0 is required to compile rust-analyzer\n",
            REQUIRED_RUST_VERSION,
        )
    }

    let res = if opts.jemalloc {
        run("cargo install --path crates/ra_lsp_server --locked --force --features jemalloc", ".")
    } else {
        run("cargo install --path crates/ra_lsp_server --locked --force", ".")
    };

    if res.is_err() && old_rust {
        eprintln!(
            "\nWARNING: at least rust 1.{}.0 is required to compile rust-analyzer\n",
            REQUIRED_RUST_VERSION,
        )
    }

    res
}

fn check_version(version_output: &str, min_minor_version: u32) -> bool {
    // Parse second the number out of
    //      cargo 1.39.0-beta (1c6ec66d5 2019-09-30)
    let minor: Option<u32> = version_output.split('.').nth(1).and_then(|it| it.parse().ok());
    match minor {
        None => true,
        Some(minor) => minor >= min_minor_version,
    }
}
<commit_msg>Remove npm.cmd<commit_after>//! See https://github.com/matklad/cargo-xtask/.
//!
//! This binary defines various auxiliary build commands, which are not
//! expressible with just `cargo`. Notably, it provides `cargo xtask codegen`
//! for code generation and `cargo xtask install` for installation of
//! rust-analyzer server and client.
//!
//! This binary is integrated into the `cargo` command line by using an alias in
//! `.cargo/config`.
mod help;

use std::{env, fmt::Write, path::PathBuf, str};

use anyhow::Context;
use pico_args::Arguments;
use xtask::{
    codegen::{self, Mode},
    install_pre_commit_hook, reformat_staged_files, run, run_clippy, run_fuzzer, run_rustfmt,
    run_with_output, Cmd, Result,
};

// Latest stable, feel free to send a PR if this lags behind.
const REQUIRED_RUST_VERSION: u32 = 39;

struct InstallOpt {
    client: Option<ClientOpt>,
    server: Option<ServerOpt>,
}

enum ClientOpt {
    VsCode,
}

struct ServerOpt {
    jemalloc: bool,
}

fn main() -> Result<()> {
    if env::args().next().map(|it| it.contains("pre-commit")) == Some(true) {
        return reformat_staged_files();
    }

    let subcommand = match std::env::args_os().nth(1) {
        None => {
            eprintln!("{}", help::GLOBAL_HELP);
            return Ok(());
        }
        Some(s) => s,
    };
    let mut matches = Arguments::from_vec(std::env::args_os().skip(2).collect());
    let subcommand = &*subcommand.to_string_lossy();
    match subcommand {
        "install" => {
            if matches.contains(["-h", "--help"]) {
                eprintln!("{}", help::INSTALL_HELP);
                return Ok(());
            }
            let server = matches.contains("--server");
            let client_code = matches.contains("--client-code");
            if server && client_code {
                eprintln!("{}", help::INSTALL_RA_CONFLICT);
                return Ok(());
            }
            let jemalloc = matches.contains("--jemalloc");
            matches.finish().or_else(handle_extra_flags)?;
            let opts = InstallOpt {
                client: if server { None } else { Some(ClientOpt::VsCode) },
                server: if client_code { None } else { Some(ServerOpt { jemalloc }) },
            };
            install(opts)?
        }
        "codegen" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            codegen::generate_syntax(Mode::Overwrite)?;
            codegen::generate_parser_tests(Mode::Overwrite)?;
            codegen::generate_assists_docs(Mode::Overwrite)?;
        }
        "format" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_rustfmt(Mode::Overwrite)?
        }
        "install-pre-commit-hook" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            install_pre_commit_hook()?
        }
        "lint" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_clippy()?
        }
        "fuzz-tests" => {
            if matches.contains(["-h", "--help"]) {
                help::print_no_param_subcommand_help(&subcommand);
                return Ok(());
            }
            run_fuzzer()?
        }
        _ => eprintln!("{}", help::GLOBAL_HELP),
    }
    Ok(())
}

fn handle_extra_flags(e: pico_args::Error) -> Result<()> {
    if let pico_args::Error::UnusedArgsLeft(flags) = e {
        let mut invalid_flags = String::new();
        for flag in flags {
            write!(&mut invalid_flags, "{}, ", flag)?;
        }
        let (invalid_flags, _) = invalid_flags.split_at(invalid_flags.len() - 2);
        anyhow::bail!("Invalid flags: {}", invalid_flags)
    } else {
        anyhow::bail!(e.to_string())
    }
}

fn install(opts: InstallOpt) -> Result<()> {
    if cfg!(target_os = "macos") {
        fix_path_for_mac().context("Fix path for mac")?
    }
    if let Some(server) = opts.server {
        install_server(server).context("install server")?;
    }
    if let Some(client) = opts.client {
        install_client(client).context("install client")?;
    }
    Ok(())
}

fn fix_path_for_mac() -> Result<()> {
    let mut vscode_path: Vec<PathBuf> = {
        const COMMON_APP_PATH: &str =
            r"/Applications/Visual Studio Code.app/Contents/Resources/app/bin";
        const ROOT_DIR: &str = "";
        let home_dir = match env::var("HOME") {
            Ok(home) => home,
            Err(e) => anyhow::bail!("Failed getting HOME from environment with error: {}.", e),
        };

        [ROOT_DIR, &home_dir]
            .iter()
            .map(|dir| String::from(*dir) + COMMON_APP_PATH)
            .map(PathBuf::from)
            .filter(|path| path.exists())
            .collect()
    };

    if !vscode_path.is_empty() {
        let vars = match env::var_os("PATH") {
            Some(path) => path,
            None => anyhow::bail!("Could not get PATH variable from env."),
        };

        let mut paths = env::split_paths(&vars).collect::<Vec<_>>();
        paths.append(&mut vscode_path);
        let new_paths = env::join_paths(paths).context("build env PATH")?;
        env::set_var("PATH", &new_paths);
    }

    Ok(())
}

fn install_client(ClientOpt::VsCode: ClientOpt) -> Result<()> {
    let npm_version = Cmd {
        unix: r"npm --version",
        windows: r"cmd.exe /c npm --version",
        work_dir: "./editors/code",
    }
    .run();

    if npm_version.is_err() {
        eprintln!("\nERROR: `npm --version` failed, `npm` is required to build the VS Code plugin")
    }

    Cmd { unix: r"npm ci", windows: r"cmd.exe /c npm ci", work_dir: "./editors/code" }.run()?;
    Cmd {
        unix: r"npm run package --scripts-prepend-node-path",
        windows: r"cmd.exe /c npm run package",
        work_dir: "./editors/code",
    }
    .run()?;

    let code_binary = ["code", "code-insiders", "codium"].iter().find(|bin| {
        Cmd {
            unix: &format!("{} --version", bin),
            windows: &format!("cmd.exe /c {}.cmd --version", bin),
            work_dir: "./editors/code",
        }
        .run()
        .is_ok()
    });

    let code_binary = match code_binary {
        Some(it) => it,
        None => anyhow::bail!("Can't execute `code --version`. Perhaps it is not in $PATH?"),
    };

    Cmd {
        unix: &format!(r"{} --install-extension ./ra-lsp-0.0.1.vsix --force", code_binary),
        windows: &format!(
            r"cmd.exe /c {}.cmd --install-extension ./ra-lsp-0.0.1.vsix --force",
            code_binary
        ),
        work_dir: "./editors/code",
    }
    .run()?;

    let output = Cmd {
        unix: &format!(r"{} --list-extensions", code_binary),
        windows: &format!(r"cmd.exe /c {}.cmd --list-extensions", code_binary),
        work_dir: ".",
    }
    .run_with_output()?;

    if !str::from_utf8(&output.stdout)?.contains("ra-lsp") {
        anyhow::bail!(
            "Could not install the Visual Studio Code extension. \
             Please make sure you have at least NodeJS 10.x installed and try again."
        );
    }

    Ok(())
}

fn install_server(opts: ServerOpt) -> Result<()> {
    let mut old_rust = false;
    if let Ok(output) = run_with_output("cargo --version", ".") {
        if let Ok(stdout) = String::from_utf8(output.stdout) {
            println!("{}", stdout);
            if !check_version(&stdout, REQUIRED_RUST_VERSION) {
                old_rust = true;
            }
        }
    }

    if old_rust {
        eprintln!(
            "\nWARNING: at least rust 1.{}.0 is required to compile rust-analyzer\n",
            REQUIRED_RUST_VERSION,
        )
    }

    let res = if opts.jemalloc {
        run("cargo install --path crates/ra_lsp_server --locked --force --features jemalloc", ".")
    } else {
        run("cargo install --path crates/ra_lsp_server --locked --force", ".")
    };

    if res.is_err() && old_rust {
        eprintln!(
            "\nWARNING: at least rust 1.{}.0 is required to compile rust-analyzer\n",
            REQUIRED_RUST_VERSION,
        )
    }

    res
}

fn check_version(version_output: &str, min_minor_version: u32) -> bool {
    // Parse second the number out of
    //      cargo 1.39.0-beta (1c6ec66d5 2019-09-30)
    let minor: Option<u32> = version_output.split('.').nth(1).and_then(|it| it.parse().ok());
    match minor {
        None => true,
        Some(minor) => minor >= min_minor_version,
    }
}
<|endoftext|>


<commit_before>use roaring::{RoaringBitmap, Iter};

use std::collections::{BTreeMap, HashMap};
use std::collections::btree_map::Keys;

use document::Document;
use schema::{Schema, SchemaRead, SchemaWrite, FieldType, FieldRef, AddFieldError};
use store::{IndexStore, IndexReader, DocRefIterator};

#[derive(Debug)]
pub struct MemoryIndexStoreFieldTerm {
    pub docs: RoaringBitmap<u64>,
}

impl MemoryIndexStoreFieldTerm {
    pub fn new() -> MemoryIndexStoreFieldTerm {
        MemoryIndexStoreFieldTerm {
            docs: RoaringBitmap::new(),
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStoreField {
    pub docs: RoaringBitmap<u64>,
    pub terms: BTreeMap<Vec<u8>, MemoryIndexStoreFieldTerm>,
    pub num_tokens: u64,
}

impl MemoryIndexStoreField {
    pub fn new() -> MemoryIndexStoreField {
        MemoryIndexStoreField {
            docs: RoaringBitmap::new(),
            terms: BTreeMap::new(),
            num_tokens: 0,
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStore {
    schema: Schema,
    docs: BTreeMap<u64, Document>,
    fields: HashMap<FieldRef, MemoryIndexStoreField>,
    next_doc_id: u64,
    doc_key2id_map: HashMap<String, u64>,
}

impl MemoryIndexStore {
    pub fn new() -> MemoryIndexStore {
        MemoryIndexStore {
            schema: Schema::new(),
            docs: BTreeMap::new(),
            fields: HashMap::new(),
            next_doc_id: 1,
            doc_key2id_map: HashMap::new(),
        }
    }
}

impl<'a> IndexStore<'a> for MemoryIndexStore {
    type Reader = MemoryIndexStoreReader<'a>;

    fn reader(&'a self) -> MemoryIndexStoreReader<'a> {
        MemoryIndexStoreReader {
            store: self,
        }
    }

    fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let field_ref = try!(self.schema.add_field(name, field_type));
        self.fields.insert(field_ref, MemoryIndexStoreField::new());

        Ok(field_ref)
    }

    fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        self.schema.remove_field(field_ref);

        match self.fields.remove(field_ref) {
            Some(field_ref) => true,
            None => false,
        }
    }

    fn insert_or_update_document(&mut self, doc: Document) {
        let doc_id = self.next_doc_id;
        self.next_doc_id += 1;

        // Put field contents in inverted index
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                // Silently ignore unrecognised fields
                // TODO: Review this
                if let Some(field) = self.fields.get_mut(field_ref) {
                    field.docs.insert(doc_id);
                    field.num_tokens += 1;

                    let term_bytes = token.term.to_bytes();
                    if !field.terms.contains_key(&term_bytes) {
                        // TODO: We shouldn't need to clone here
                        field.terms.insert(term_bytes.clone(), MemoryIndexStoreFieldTerm::new());
                    }

                    let mut term = field.terms.get_mut(&term_bytes).unwrap();
                    term.docs.insert(doc_id);
                }
            }
        }

        self.doc_key2id_map.insert(doc.key.clone(), doc_id);
        self.docs.insert(doc_id, doc);
    }

    fn remove_document_by_key(&mut self, doc_key: &str) -> bool {
        match self.doc_key2id_map.remove(doc_key) {
            Some(doc_id) => {
                self.docs.remove(&doc_id);

                true
            }
            None => false
        }
    }
}

pub struct MemoryIndexStoreReader<'a> {
    store: &'a MemoryIndexStore,
}

impl<'a> IndexReader<'a> for MemoryIndexStoreReader<'a> {
    type AllDocRefIterator = MemoryIndexStoreAllDocRefIterator<'a>;
    type TermDocRefIterator = MemoryIndexStoreTermDocRefIterator<'a>;

    fn schema(&self) -> &Schema {
        &self.store.schema
    }

    fn get_document_by_key(&self, doc_key: &str) -> Option<&Document> {
        match self.store.doc_key2id_map.get(doc_key) {
            Some(doc_id) => self.store.docs.get(doc_id),
            None => None,
        }
    }

    fn get_document_by_id(&self, doc_id: &u64) -> Option<&Document> {
        self.store.docs.get(doc_id)
    }

    fn contains_document_key(&self, doc_key: &str) -> bool {
        self.store.doc_key2id_map.contains_key(doc_key)
    }

    fn num_docs(&self) -> usize {
        self.store.docs.len()
    }

    fn iter_all_docs(&'a self) -> MemoryIndexStoreAllDocRefIterator<'a> {
        MemoryIndexStoreAllDocRefIterator {
            keys: self.store.docs.keys(),
        }
    }

    fn iter_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> Option<MemoryIndexStoreTermDocRefIterator<'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return None,
        };

        Some(MemoryIndexStoreTermDocRefIterator {
            iterator: term.docs.iter(),
        })
    }

    fn iter_all_terms(&'a self, field_ref: &FieldRef) -> Option<Box<Iterator<Item=&'a [u8]> + 'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        Some(Box::new(field.terms.keys().map(|t| &t[..])))
    }

    fn num_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return 0,
        };

        term.docs.len()
    }

    fn total_tokens(&'a self, field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        field.num_tokens
    }
}

pub struct MemoryIndexStoreAllDocRefIterator<'a> {
    keys: Keys<'a, u64, Document>,
}

impl<'a> Iterator for MemoryIndexStoreAllDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.keys.next().cloned()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreAllDocRefIterator<'a> {

}

pub struct MemoryIndexStoreTermDocRefIterator<'a> {
    iterator: Iter<'a, u64>,
}

impl<'a> Iterator for MemoryIndexStoreTermDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.iterator.next()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreTermDocRefIterator<'a> {

}

#[cfg(test)]
mod tests {
    use super::{MemoryIndexStore, MemoryIndexStoreReader};

    use term::Term;
    use token::Token;
    use document::Document;
    use schema::{Schema, SchemaRead, FieldType, FieldRef};
    use store::{IndexStore, IndexReader};

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    #[test]
    fn test_num_docs() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.num_docs(), 2);
    }

    #[test]
    fn test_all_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.iter_all_docs().count(), 2);
    }

    #[test]
    fn test_term_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();
        let title_field = reader.schema().get_field_by_name("title").unwrap();

        assert_eq!(reader.iter_docs_with_term(&Term::String("hello".to_string()).to_bytes(), &title_field).unwrap().count(), 1);
    }
}
<commit_msg>Improve performance of insert_or_update_document<commit_after>use roaring::{RoaringBitmap, Iter};

use std::collections::{BTreeMap, HashMap};
use std::collections::btree_map::Keys;

use document::Document;
use schema::{Schema, SchemaRead, SchemaWrite, FieldType, FieldRef, AddFieldError};
use store::{IndexStore, IndexReader, DocRefIterator};

#[derive(Debug)]
pub struct MemoryIndexStoreFieldTerm {
    pub docs: RoaringBitmap<u64>,
}

impl MemoryIndexStoreFieldTerm {
    pub fn new() -> MemoryIndexStoreFieldTerm {
        MemoryIndexStoreFieldTerm {
            docs: RoaringBitmap::new(),
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStoreField {
    pub docs: RoaringBitmap<u64>,
    pub terms: BTreeMap<Vec<u8>, MemoryIndexStoreFieldTerm>,
    pub num_tokens: u64,
}

impl MemoryIndexStoreField {
    pub fn new() -> MemoryIndexStoreField {
        MemoryIndexStoreField {
            docs: RoaringBitmap::new(),
            terms: BTreeMap::new(),
            num_tokens: 0,
        }
    }
}

#[derive(Debug)]
pub struct MemoryIndexStore {
    schema: Schema,
    docs: BTreeMap<u64, Document>,
    fields: HashMap<FieldRef, MemoryIndexStoreField>,
    next_doc_id: u64,
    doc_key2id_map: HashMap<String, u64>,
}

impl MemoryIndexStore {
    pub fn new() -> MemoryIndexStore {
        MemoryIndexStore {
            schema: Schema::new(),
            docs: BTreeMap::new(),
            fields: HashMap::new(),
            next_doc_id: 1,
            doc_key2id_map: HashMap::new(),
        }
    }
}

impl<'a> IndexStore<'a> for MemoryIndexStore {
    type Reader = MemoryIndexStoreReader<'a>;

    fn reader(&'a self) -> MemoryIndexStoreReader<'a> {
        MemoryIndexStoreReader {
            store: self,
        }
    }

    fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let field_ref = try!(self.schema.add_field(name, field_type));
        self.fields.insert(field_ref, MemoryIndexStoreField::new());

        Ok(field_ref)
    }

    fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        self.schema.remove_field(field_ref);

        match self.fields.remove(field_ref) {
            Some(field_ref) => true,
            None => false,
        }
    }

    fn insert_or_update_document(&mut self, doc: Document) {
        let doc_id = self.next_doc_id;
        self.next_doc_id += 1;

        // Put field contents in inverted index
        for (field_ref, tokens) in doc.fields.iter() {
            // Silently ignore unrecognised fields
            // TODO: Review this
            if let Some(field) = self.fields.get_mut(field_ref) {
                field.docs.insert(doc_id);

                for token in tokens.iter() {
                    field.num_tokens += 1;

                    let term_bytes = token.term.to_bytes();
                    let mut term = field.terms.entry(term_bytes).or_insert_with(|| MemoryIndexStoreFieldTerm::new());
                    term.docs.insert(doc_id);
                }
            }
        }

        self.doc_key2id_map.insert(doc.key.clone(), doc_id);
        self.docs.insert(doc_id, doc);
    }

    fn remove_document_by_key(&mut self, doc_key: &str) -> bool {
        match self.doc_key2id_map.remove(doc_key) {
            Some(doc_id) => {
                self.docs.remove(&doc_id);

                true
            }
            None => false
        }
    }
}

pub struct MemoryIndexStoreReader<'a> {
    store: &'a MemoryIndexStore,
}

impl<'a> IndexReader<'a> for MemoryIndexStoreReader<'a> {
    type AllDocRefIterator = MemoryIndexStoreAllDocRefIterator<'a>;
    type TermDocRefIterator = MemoryIndexStoreTermDocRefIterator<'a>;

    fn schema(&self) -> &Schema {
        &self.store.schema
    }

    fn get_document_by_key(&self, doc_key: &str) -> Option<&Document> {
        match self.store.doc_key2id_map.get(doc_key) {
            Some(doc_id) => self.store.docs.get(doc_id),
            None => None,
        }
    }

    fn get_document_by_id(&self, doc_id: &u64) -> Option<&Document> {
        self.store.docs.get(doc_id)
    }

    fn contains_document_key(&self, doc_key: &str) -> bool {
        self.store.doc_key2id_map.contains_key(doc_key)
    }

    fn num_docs(&self) -> usize {
        self.store.docs.len()
    }

    fn iter_all_docs(&'a self) -> MemoryIndexStoreAllDocRefIterator<'a> {
        MemoryIndexStoreAllDocRefIterator {
            keys: self.store.docs.keys(),
        }
    }

    fn iter_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> Option<MemoryIndexStoreTermDocRefIterator<'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return None,
        };

        Some(MemoryIndexStoreTermDocRefIterator {
            iterator: term.docs.iter(),
        })
    }

    fn iter_all_terms(&'a self, field_ref: &FieldRef) -> Option<Box<Iterator<Item=&'a [u8]> + 'a>> {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return None,
        };

        Some(Box::new(field.terms.keys().map(|t| &t[..])))
    }

    fn num_docs_with_term(&'a self, term: &[u8], field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        let term = match field.terms.get(term) {
            Some(term) => term,
            None => return 0,
        };

        term.docs.len()
    }

    fn total_tokens(&'a self, field_ref: &FieldRef) -> u64 {
        let field = match self.store.fields.get(field_ref) {
            Some(field) => field,
            None => return 0,
        };

        field.num_tokens
    }
}

pub struct MemoryIndexStoreAllDocRefIterator<'a> {
    keys: Keys<'a, u64, Document>,
}

impl<'a> Iterator for MemoryIndexStoreAllDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.keys.next().cloned()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreAllDocRefIterator<'a> {

}

pub struct MemoryIndexStoreTermDocRefIterator<'a> {
    iterator: Iter<'a, u64>,
}

impl<'a> Iterator for MemoryIndexStoreTermDocRefIterator<'a> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        self.iterator.next()
    }
}

impl<'a> DocRefIterator<'a> for MemoryIndexStoreTermDocRefIterator<'a> {

}

#[cfg(test)]
mod tests {
    use super::{MemoryIndexStore, MemoryIndexStoreReader};

    use term::Term;
    use token::Token;
    use document::Document;
    use schema::{Schema, SchemaRead, FieldType, FieldRef};
    use store::{IndexStore, IndexReader};

    fn make_test_store() -> MemoryIndexStore {
        let mut store = MemoryIndexStore::new();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    #[test]
    fn test_num_docs() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.num_docs(), 2);
    }

    #[test]
    fn test_all_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();

        assert_eq!(reader.iter_all_docs().count(), 2);
    }

    #[test]
    fn test_term_docs_iterator() {
        let store = make_test_store();
        let reader = store.reader();
        let title_field = reader.schema().get_field_by_name("title").unwrap();

        assert_eq!(reader.iter_docs_with_term(&Term::String("hello".to_string()).to_bytes(), &title_field).unwrap().count(), 1);
    }
}
<|endoftext|>