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

use crate::{
    graph::DotIdentifier,
    message::Message,
    proto::oak::introspection_events::{event::EventDetails, ChannelDestroyed},
    Runtime,
};
use log::{debug, error};
use oak_abi::OakStatus;
use std::{
    collections::{HashMap, VecDeque},
    sync::{
        atomic::{AtomicU64, Ordering::SeqCst},
        Arc, Mutex, RwLock, RwLockReadGuard, Weak,
    },
    thread::{Thread, ThreadId},
};

type Messages = VecDeque<Message>;

/// A `HashMap` of waiting threads, keyed by `ThreadId`.
///
/// Implemented as a `HashMap` to prevent build up of stale `Weak<Thread>`s.
///
/// That is: If a thread waiting/blocked on a channel is woken by a different channel, its
/// `Weak<Thread>` will remain in the first channel's waiting_thread member. If a thread keeps
/// waiting on this first channel, and keeps being woken by other channels, it will keep re-adding
/// itself. We use a `HashMap` and insert at the current `ThreadId` so that we replace any stale
/// `Weak<Thread>`s which will have gone out of scope. (`wait_on_channels` drops the underlying arc
/// as soon as it is resumed.)
type WaitingThreads = Mutex<HashMap<ThreadId, Weak<Thread>>>;

/// The internal implementation of a channel representation backed by a `VecDeque<Message>`.
///
/// Channels are reference counted using `Arc<Channel>`, which are always in the form of a
/// [`ChannelHalf`] object that references one end of the [`Channel`] (read or write) and which
/// is included in the `reader_count` or `writer_count`.
pub struct Channel {
    // An internal identifier for this channel. Purely for disambiguation in debugging output.
    id: ChannelId,

    pub messages: RwLock<Messages>,

    // Counts of the numbers of `ChannelHalf` objects that refer to this channel.
    writer_count: AtomicU64,
    reader_count: AtomicU64,

    /// A HashMap of `ThreadId`s to `Weak<Thread>`s. This allows threads to insert a weak reference
    /// to themselves to be woken when a new message is available. Weak references are used so that
    /// if the thread is woken by a different channel, it can deallocate the underlying `Arc`
    /// instead of removing itself from all the `Channel`s it subscribed to.
    /// Threads can be woken up spuriously without issue.
    waiting_threads: WaitingThreads,

    /// The `Label` associated with this channel.
    ///
    /// This is set at channel creation time and does not change after that.
    ///
    /// See https://github.com/project-oak/oak/blob/main/docs/concepts.md#labels
    pub label: oak_abi::label::Label,

    /// The name for the channel.
    ///
    /// The name does not have to be unique and can be empty. It is used in logs to help with
    /// identifying channels during debugging.
    pub name: String,

    /// Weak reference to the Runtime used for sending introspection events.
    runtime_weak: Weak<Runtime>,
}

impl std::fmt::Debug for Channel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Channel {{ id={}, #readers={}, #writers={}, name={:?}, label={:?} }}",
            self.id,
            self.reader_count.load(SeqCst),
            self.writer_count.load(SeqCst),
            self.name,
            self.label,
        )
    }
}

/// A reference to one half of a `Channel`
pub struct ChannelHalf {
    channel: Arc<Channel>,
    pub direction: ChannelHalfDirection,
}

impl ChannelHalf {
    /// Constructor for [`ChannelHalf`] keeps the underlying `Channel`'s reader/writer count
    /// up-to-date.
    pub fn new(channel: Arc<Channel>, direction: ChannelHalfDirection) -> Self {
        match direction {
            ChannelHalfDirection::Write => channel.inc_writer_count(),
            ChannelHalfDirection::Read => channel.inc_reader_count(),
        };
        ChannelHalf { channel, direction }
    }

    /// Get the ID of the underlying channel.  For debugging/introspection
    /// purposes.
    pub fn get_channel_id(&self) -> u64 {
        self.channel.id
    }

    /// Get the ID of the underlying channel.  For debugging/introspection
    /// purposes.
    pub fn get_channel_name(&self) -> &str {
        &self.channel.name
    }

    /// Get read-only access to the channel's messages.  For debugging/introspection
    /// purposes.
    pub fn get_messages(&self) -> RwLockReadGuard<Messages> {
        self.channel.messages.read().unwrap()
    }

    /// Visit all channel halves that are reachable via this `ChannelHalf`, starting with
    /// `self`. The `visitor` function should return a boolean indicating whether the provided half
    /// needs to be further explored.  For debugging/introspection purposes.
    pub fn visit_halves<F>(&self, visitor: &mut F)
    where
        F: FnMut(&ChannelHalf) -> bool,
    {
        if visitor(self) {
            let mut to_visit = Vec::new();
            {
                let messages = self.get_messages();
                for msg in messages.iter() {
                    for half in &msg.channels {
                        to_visit.push(half.clone())
                    }
                }
            }
            // Now the lock is not held, let us go and make our visit.
            for half in to_visit {
                half.visit_halves(visitor);
            }
        }
    }

    /// Wake any threads waiting on the underlying channel.
    pub fn wake_waiters(&self) {
        self.channel.wake_waiters();
    }
}

/// Manual implementation of the [`Clone`] trait to keep the counts for the underlying `Channel` in
/// sync.
impl Clone for ChannelHalf {
    fn clone(&self) -> Self {
        ChannelHalf::new(self.channel.clone(), self.direction)
    }
}

/// Manual implementation of the [`Drop`] trait to keep the counts for the underlying `Channel` in
/// sync.
impl Drop for ChannelHalf {
    fn drop(&mut self) {
        match self.direction {
            ChannelHalfDirection::Write => self.channel.dec_writer_count(),
            ChannelHalfDirection::Read => self.channel.dec_reader_count(),
        };
        if self.direction == ChannelHalfDirection::Write && !self.channel.has_writers() {
            // This was the last writer to the channel: wake any waiters so they
            // can be aware that the channel is orphaned.
            debug!(
                "last writer for channel {} gone, wake waiters",
                self.channel.id
            );
            self.wake_waiters();
        }
    }
}

impl std::fmt::Debug for ChannelHalf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Channel {} {}",
            self.channel.id,
            match self.direction {
                ChannelHalfDirection::Write => "WRITE",
                ChannelHalfDirection::Read => "READ",
            },
        )
    }
}

impl DotIdentifier for ChannelHalf {
    fn dot_id(&self) -> String {
        self.channel.id.dot_id()
    }
}

/// Perform an operation on a [`Channel`] associated with a [`ChannelId`].
fn with_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    f(half.channel.clone())
}

/// Perform an operation on a [`Channel`] associated with a reader [`ChannelHalf`].
pub fn with_reader_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    if half.direction != ChannelHalfDirection::Read {
        return Err(OakStatus::ErrBadHandle);
    }
    with_channel(half, f)
}

/// Perform an operation on a [`Channel`] associated with a writer [`ChannelHalf`].
pub fn with_writer_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    if half.direction != ChannelHalfDirection::Write {
        return Err(OakStatus::ErrBadHandle);
    }
    with_channel(half, f)
}

/// The direction of a [`ChannelHalf`].
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum ChannelHalfDirection {
    Read,
    Write,
}

/// An internal identifier to track a [`Channel`].
type ChannelId = u64;

impl DotIdentifier for ChannelId {
    fn dot_id(&self) -> String {
        format!("channel{}", self)
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        debug!("dropping Channel object {:?}", self);
        match self.runtime_weak.upgrade() {
            Some(runtime) => {
                runtime.introspection_event(EventDetails::ChannelDestroyed(ChannelDestroyed {
                    channel_id: self.id,
                }));
            }
            None => {
                error!("Couldn't send ChannelDestroyed event, since the runtime has been dropped. {:?}", self);
            }
        }

        // There should be no waiters for this channel (a waiting Node would have
        // to have a `Handle` to wait on, which would be a reference that pins this
        // channel to existence) so no need to `wake_waiters()`.
        // Deliberately clear the HashMap under the lock to avoid a TSAN report.
        self.waiting_threads.lock().unwrap().clear();
    }
}

impl Channel {
    pub fn new(
        id: ChannelId,
        name: &str,
        label: &oak_abi::label::Label,
        runtime_weak: Weak<Runtime>,
    ) -> Arc<Channel> {
        debug!("create new Channel object with ID {}", id);
        Arc::new(Channel {
            id,
            messages: RwLock::new(Messages::new()),
            writer_count: AtomicU64::new(0),
            reader_count: AtomicU64::new(0),
            waiting_threads: Mutex::new(HashMap::new()),
            label: label.clone(),
            name: format!("{}({})", name, id),
            runtime_weak,
        })
    }

    /// Determine whether there are any readers of the channel.
    pub fn has_readers(&self) -> bool {
        self.reader_count.load(SeqCst) > 0
    }

    /// Determine whether there are any writers of the channel.
    pub fn has_writers(&self) -> bool {
        self.writer_count.load(SeqCst) > 0
    }

    /// Decrement the [`Channel`] writer counter.
    fn dec_writer_count(&self) {
        if self.writer_count.fetch_sub(1, SeqCst) == 0 {
            panic!("remove_reader: Writer count was already 0, something is very wrong!")
        }
    }

    /// Decrement the [`Channel`] reader counter.
    fn dec_reader_count(&self) {
        if self.reader_count.fetch_sub(1, SeqCst) == 0 {
            panic!("remove_reader: Reader count was already 0, something is very wrong!")
        }
    }

    /// Increment the [`Channel`] writer counter.
    fn inc_writer_count(&self) -> u64 {
        self.writer_count.fetch_add(1, SeqCst)
    }

    /// Increment the [`Channel`] reader counter.
    fn inc_reader_count(&self) -> u64 {
        self.reader_count.fetch_add(1, SeqCst)
    }

    /// Add the given [`Thread`] reference into the collection of [`Thread`]s waiting on this
    /// [`Channel`]'s readability.  Threads waiting on the [`Channel`] will be woken when
    /// data is available, or if the [`Channel`] becomes orphaned (no writers left).
    pub fn add_waiter(&self, thread: &Arc<Thread>) {
        self.waiting_threads
            .lock()
            .unwrap()
            .insert(thread.id(), Arc::downgrade(thread));
    }

    /// Wake any [`Thread`]s that are waiting on the [`Channel`].
    pub fn wake_waiters(&self) {
        // Unpark (wake up) all waiting threads that still have live references. The first thread
        // woken can immediately read the message, and others might find `messages` is empty before
        // they are even woken. This should not be an issue (being woken does not guarantee a
        // message is available), but it could potentially result in some particular thread always
        // getting first chance to read the message.
        //
        // If a thread is woken and finds no message it will take the `waiting_threads` lock and
        // add itself again. Note that since that lock is currently held, the woken thread will add
        // itself to waiting_threads *after* we call clear below as we release the lock implicilty
        // on leaving this function.
        let mut waiting_threads = self.waiting_threads.lock().unwrap();
        for thread in waiting_threads.values() {
            if let Some(thread) = thread.upgrade() {
                thread.unpark();
            }
        }
        waiting_threads.clear();
    }
}
<commit_msg>oak_runtime: show names in introspection pages<commit_after>//
// Copyright 2020 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use crate::{
    graph::DotIdentifier,
    message::Message,
    proto::oak::introspection_events::{event::EventDetails, ChannelDestroyed},
    Runtime,
};
use log::{debug, error};
use oak_abi::OakStatus;
use std::{
    collections::{HashMap, VecDeque},
    sync::{
        atomic::{AtomicU64, Ordering::SeqCst},
        Arc, Mutex, RwLock, RwLockReadGuard, Weak,
    },
    thread::{Thread, ThreadId},
};

type Messages = VecDeque<Message>;

/// A `HashMap` of waiting threads, keyed by `ThreadId`.
///
/// Implemented as a `HashMap` to prevent build up of stale `Weak<Thread>`s.
///
/// That is: If a thread waiting/blocked on a channel is woken by a different channel, its
/// `Weak<Thread>` will remain in the first channel's waiting_thread member. If a thread keeps
/// waiting on this first channel, and keeps being woken by other channels, it will keep re-adding
/// itself. We use a `HashMap` and insert at the current `ThreadId` so that we replace any stale
/// `Weak<Thread>`s which will have gone out of scope. (`wait_on_channels` drops the underlying arc
/// as soon as it is resumed.)
type WaitingThreads = Mutex<HashMap<ThreadId, Weak<Thread>>>;

/// The internal implementation of a channel representation backed by a `VecDeque<Message>`.
///
/// Channels are reference counted using `Arc<Channel>`, which are always in the form of a
/// [`ChannelHalf`] object that references one end of the [`Channel`] (read or write) and which
/// is included in the `reader_count` or `writer_count`.
pub struct Channel {
    // An internal identifier for this channel. Purely for disambiguation in debugging output.
    id: ChannelId,

    pub messages: RwLock<Messages>,

    // Counts of the numbers of `ChannelHalf` objects that refer to this channel.
    writer_count: AtomicU64,
    reader_count: AtomicU64,

    /// A HashMap of `ThreadId`s to `Weak<Thread>`s. This allows threads to insert a weak reference
    /// to themselves to be woken when a new message is available. Weak references are used so that
    /// if the thread is woken by a different channel, it can deallocate the underlying `Arc`
    /// instead of removing itself from all the `Channel`s it subscribed to.
    /// Threads can be woken up spuriously without issue.
    waiting_threads: WaitingThreads,

    /// The `Label` associated with this channel.
    ///
    /// This is set at channel creation time and does not change after that.
    ///
    /// See https://github.com/project-oak/oak/blob/main/docs/concepts.md#labels
    pub label: oak_abi::label::Label,

    /// The name for the channel.
    ///
    /// The name does not have to be unique and can be empty. It is used in logs to help with
    /// identifying channels during debugging.
    pub name: String,

    /// Weak reference to the Runtime used for sending introspection events.
    runtime_weak: Weak<Runtime>,
}

impl std::fmt::Debug for Channel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Channel {{ id={}, #readers={}, #writers={}, name={:?}, label={:?} }}",
            self.id,
            self.reader_count.load(SeqCst),
            self.writer_count.load(SeqCst),
            self.name,
            self.label,
        )
    }
}

/// A reference to one half of a `Channel`
pub struct ChannelHalf {
    channel: Arc<Channel>,
    pub direction: ChannelHalfDirection,
}

impl ChannelHalf {
    /// Constructor for [`ChannelHalf`] keeps the underlying `Channel`'s reader/writer count
    /// up-to-date.
    pub fn new(channel: Arc<Channel>, direction: ChannelHalfDirection) -> Self {
        match direction {
            ChannelHalfDirection::Write => channel.inc_writer_count(),
            ChannelHalfDirection::Read => channel.inc_reader_count(),
        };
        ChannelHalf { channel, direction }
    }

    /// Get the ID of the underlying channel.  For debugging/introspection
    /// purposes.
    pub fn get_channel_id(&self) -> u64 {
        self.channel.id
    }

    /// Get the ID of the underlying channel.  For debugging/introspection
    /// purposes.
    pub fn get_channel_name(&self) -> &str {
        &self.channel.name
    }

    /// Get read-only access to the channel's messages.  For debugging/introspection
    /// purposes.
    pub fn get_messages(&self) -> RwLockReadGuard<Messages> {
        self.channel.messages.read().unwrap()
    }

    /// Visit all channel halves that are reachable via this `ChannelHalf`, starting with
    /// `self`. The `visitor` function should return a boolean indicating whether the provided half
    /// needs to be further explored.  For debugging/introspection purposes.
    pub fn visit_halves<F>(&self, visitor: &mut F)
    where
        F: FnMut(&ChannelHalf) -> bool,
    {
        if visitor(self) {
            let mut to_visit = Vec::new();
            {
                let messages = self.get_messages();
                for msg in messages.iter() {
                    for half in &msg.channels {
                        to_visit.push(half.clone())
                    }
                }
            }
            // Now the lock is not held, let us go and make our visit.
            for half in to_visit {
                half.visit_halves(visitor);
            }
        }
    }

    /// Wake any threads waiting on the underlying channel.
    pub fn wake_waiters(&self) {
        self.channel.wake_waiters();
    }
}

/// Manual implementation of the [`Clone`] trait to keep the counts for the underlying `Channel` in
/// sync.
impl Clone for ChannelHalf {
    fn clone(&self) -> Self {
        ChannelHalf::new(self.channel.clone(), self.direction)
    }
}

/// Manual implementation of the [`Drop`] trait to keep the counts for the underlying `Channel` in
/// sync.
impl Drop for ChannelHalf {
    fn drop(&mut self) {
        match self.direction {
            ChannelHalfDirection::Write => self.channel.dec_writer_count(),
            ChannelHalfDirection::Read => self.channel.dec_reader_count(),
        };
        if self.direction == ChannelHalfDirection::Write && !self.channel.has_writers() {
            // This was the last writer to the channel: wake any waiters so they
            // can be aware that the channel is orphaned.
            debug!(
                "last writer for channel {} gone, wake waiters",
                self.channel.id
            );
            self.wake_waiters();
        }
    }
}

impl std::fmt::Debug for ChannelHalf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Channel {} ('{}') {}",
            self.channel.id,
            self.channel.name,
            match self.direction {
                ChannelHalfDirection::Write => "WRITE",
                ChannelHalfDirection::Read => "READ",
            },
        )
    }
}

impl DotIdentifier for ChannelHalf {
    fn dot_id(&self) -> String {
        self.channel.id.dot_id()
    }
}

/// Perform an operation on a [`Channel`] associated with a [`ChannelId`].
fn with_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    f(half.channel.clone())
}

/// Perform an operation on a [`Channel`] associated with a reader [`ChannelHalf`].
pub fn with_reader_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    if half.direction != ChannelHalfDirection::Read {
        return Err(OakStatus::ErrBadHandle);
    }
    with_channel(half, f)
}

/// Perform an operation on a [`Channel`] associated with a writer [`ChannelHalf`].
pub fn with_writer_channel<U, F: FnOnce(Arc<Channel>) -> Result<U, OakStatus>>(
    half: &ChannelHalf,
    f: F,
) -> Result<U, OakStatus> {
    if half.direction != ChannelHalfDirection::Write {
        return Err(OakStatus::ErrBadHandle);
    }
    with_channel(half, f)
}

/// The direction of a [`ChannelHalf`].
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum ChannelHalfDirection {
    Read,
    Write,
}

/// An internal identifier to track a [`Channel`].
type ChannelId = u64;

impl DotIdentifier for ChannelId {
    fn dot_id(&self) -> String {
        format!("channel{}", self)
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        debug!("dropping Channel object {:?}", self);
        match self.runtime_weak.upgrade() {
            Some(runtime) => {
                runtime.introspection_event(EventDetails::ChannelDestroyed(ChannelDestroyed {
                    channel_id: self.id,
                }));
            }
            None => {
                error!("Couldn't send ChannelDestroyed event, since the runtime has been dropped. {:?}", self);
            }
        }

        // There should be no waiters for this channel (a waiting Node would have
        // to have a `Handle` to wait on, which would be a reference that pins this
        // channel to existence) so no need to `wake_waiters()`.
        // Deliberately clear the HashMap under the lock to avoid a TSAN report.
        self.waiting_threads.lock().unwrap().clear();
    }
}

impl Channel {
    pub fn new(
        id: ChannelId,
        name: &str,
        label: &oak_abi::label::Label,
        runtime_weak: Weak<Runtime>,
    ) -> Arc<Channel> {
        debug!("create new Channel object with ID {}", id);
        Arc::new(Channel {
            id,
            messages: RwLock::new(Messages::new()),
            writer_count: AtomicU64::new(0),
            reader_count: AtomicU64::new(0),
            waiting_threads: Mutex::new(HashMap::new()),
            label: label.clone(),
            name: format!("{}({})", name, id),
            runtime_weak,
        })
    }

    /// Determine whether there are any readers of the channel.
    pub fn has_readers(&self) -> bool {
        self.reader_count.load(SeqCst) > 0
    }

    /// Determine whether there are any writers of the channel.
    pub fn has_writers(&self) -> bool {
        self.writer_count.load(SeqCst) > 0
    }

    /// Decrement the [`Channel`] writer counter.
    fn dec_writer_count(&self) {
        if self.writer_count.fetch_sub(1, SeqCst) == 0 {
            panic!("remove_reader: Writer count was already 0, something is very wrong!")
        }
    }

    /// Decrement the [`Channel`] reader counter.
    fn dec_reader_count(&self) {
        if self.reader_count.fetch_sub(1, SeqCst) == 0 {
            panic!("remove_reader: Reader count was already 0, something is very wrong!")
        }
    }

    /// Increment the [`Channel`] writer counter.
    fn inc_writer_count(&self) -> u64 {
        self.writer_count.fetch_add(1, SeqCst)
    }

    /// Increment the [`Channel`] reader counter.
    fn inc_reader_count(&self) -> u64 {
        self.reader_count.fetch_add(1, SeqCst)
    }

    /// Add the given [`Thread`] reference into the collection of [`Thread`]s waiting on this
    /// [`Channel`]'s readability.  Threads waiting on the [`Channel`] will be woken when
    /// data is available, or if the [`Channel`] becomes orphaned (no writers left).
    pub fn add_waiter(&self, thread: &Arc<Thread>) {
        self.waiting_threads
            .lock()
            .unwrap()
            .insert(thread.id(), Arc::downgrade(thread));
    }

    /// Wake any [`Thread`]s that are waiting on the [`Channel`].
    pub fn wake_waiters(&self) {
        // Unpark (wake up) all waiting threads that still have live references. The first thread
        // woken can immediately read the message, and others might find `messages` is empty before
        // they are even woken. This should not be an issue (being woken does not guarantee a
        // message is available), but it could potentially result in some particular thread always
        // getting first chance to read the message.
        //
        // If a thread is woken and finds no message it will take the `waiting_threads` lock and
        // add itself again. Note that since that lock is currently held, the woken thread will add
        // itself to waiting_threads *after* we call clear below as we release the lock implicilty
        // on leaving this function.
        let mut waiting_threads = self.waiting_threads.lock().unwrap();
        for thread in waiting_threads.values() {
            if let Some(thread) = thread.upgrade() {
                thread.unpark();
            }
        }
        waiting_threads.clear();
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    collections::VecDeque,
    mem,
    sync::Arc,
    thread::{self, JoinHandle},
    time::Duration,
};

use async_task::{Runnable, Task};
use slab::Slab;
use sync::{Condvar, Mutex};

#[derive(Default)]
struct State {
    tasks: VecDeque<Runnable>,
    num_threads: usize,
    num_idle: usize,
    worker_threads: Slab<JoinHandle<()>>,
    last_exited_thread: Option<JoinHandle<()>>,
    shutting_down: bool,
}

fn run_blocking_thread(idx: usize, inner: Arc<Inner>) {
    let mut state = inner.state.lock();
    while !state.shutting_down {
        if let Some(runnable) = state.tasks.pop_front() {
            drop(state);
            runnable.run();
            state = inner.state.lock();
            continue;
        }

        // No more tasks so wait for more work.
        state.num_idle += 1;

        let (guard, result) = inner
            .condvar
            .wait_timeout_while(state, inner.keepalive, |s| {
                !s.shutting_down && s.tasks.is_empty()
            });
        state = guard;

        // Only decrement the idle count if we timed out. Otherwise, it was decremented when new
        // work was added to `state.tasks`.
        if result.timed_out() {
            state.num_idle -= 1;
            break;
        }
    }

    state.num_threads -= 1;

    // If we're shutting down then the BlockingPool will take care of joining all the threads.
    // Otherwise, we need to join the last worker thread that exited here.
    let last_exited_thread = if !state.shutting_down {
        let this_thread = state.worker_threads.remove(idx);
        mem::replace(&mut state.last_exited_thread, Some(this_thread))
    } else {
        None
    };

    // Drop the lock before trying to join the last exited thread.
    drop(state);

    if let Some(handle) = last_exited_thread {
        let _ = handle.join();
    }
}

struct Inner {
    state: Mutex<State>,
    condvar: Condvar,
    max_threads: usize,
    keepalive: Duration,
}

impl Inner {
    fn schedule(self: &Arc<Inner>, runnable: Runnable) {
        let mut state = self.state.lock();

        // If we're shutting down then nothing is going to run this task.
        if state.shutting_down {
            return;
        }

        state.tasks.push_back(runnable);

        if state.num_idle == 0 {
            // There are no idle threads.  Spawn a new one if possible.
            if state.num_threads < self.max_threads {
                state.num_threads += 1;
                let entry = state.worker_threads.vacant_entry();
                let idx = entry.key();
                let inner = self.clone();
                entry.insert(
                    thread::Builder::new()
                        .name(format!("blockingPool{}", idx))
                        .spawn(move || run_blocking_thread(idx, inner))
                        .unwrap(),
                );
            }
        } else {
            // We have idle threads, wake one up.
            state.num_idle -= 1;
            self.condvar.notify_one();
        }
    }
}

/// A thread pool for running work that may block.
///
/// It is generally discouraged to do any blocking work inside an async function. However, this is
/// sometimes unavoidable when dealing with interfaces that don't provide async variants. In this
/// case callers may use the `BlockingPool` to run the blocking work on a different thread and
/// `await` for its result to finish, which will prevent blocking the main thread of the
/// application.
///
/// Since the blocking work is sent to another thread, users should be careful when using the
/// `BlockingPool` for latency-sensitive operations. Additionally, the `BlockingPool` is intended to
/// be used for work that will eventually complete on its own. Users who want to spawn a thread
/// should just use `thread::spawn` directly.
///
/// There is no way to cancel work once it has been picked up by one of the worker threads in the
/// `BlockingPool` and dropping or shutting down the pool will block until all worker threads finish
/// their current task.
///
/// # Examples
///
/// Spawn a task to run in the `BlockingPool` and await on its result.
///
/// ```edition2018
/// use cros_async::BlockingPool;
///
/// # async fn do_it() {
///     let pool = BlockingPool::default();
///
///     let res = pool.spawn(move || {
///         // Do some CPU-intensive or blocking work here.
///
///         42
///     }).await;
///
///     assert_eq!(res, 42);
/// # }
/// # cros_async::block_on(do_it());
/// ```
pub struct BlockingPool {
    inner: Arc<Inner>,
}

impl BlockingPool {
    /// Create a new `BlockingPool`.
    ///
    /// The `BlockingPool` will never spawn more than `max_threads` threads to do work, regardless
    /// of the number of tasks that are added to it. This value should be set relatively low (for
    /// example, the number of CPUs on the machine) if the pool is intended to run CPU intensive
    /// work or it should be set relatively high (128 or more) if the pool is intended to be used
    /// for various IO operations that cannot be completed asynchronously. The default value is 256.
    ///
    /// Worker threads are spawned on demand when new work is added to the pool and will
    /// automatically exit after being idle for some time so there is no overhead for setting
    /// `max_threads` to a large value when there is little to no work assigned to the
    /// `BlockingPool`. `keepalive` determines the idle duration after which the worker thread will
    /// exit. The default value is 10 seconds.
    pub fn new(max_threads: usize, keepalive: Duration) -> BlockingPool {
        BlockingPool {
            inner: Arc::new(Inner {
                state: Default::default(),
                condvar: Condvar::new(),
                max_threads,
                keepalive,
            }),
        }
    }

    /// Like new but with pre-allocating capacity for up to `max_threads`.
    pub fn with_capacity(max_threads: usize, keepalive: Duration) -> BlockingPool {
        BlockingPool {
            inner: Arc::new(Inner {
                state: Mutex::new(State {
                    tasks: VecDeque::new(),
                    num_threads: 0,
                    num_idle: 0,
                    worker_threads: Slab::with_capacity(max_threads),
                    last_exited_thread: None,
                    shutting_down: false,
                }),
                condvar: Condvar::new(),
                max_threads,
                keepalive,
            }),
        }
    }

    /// Spawn a task to run in the `BlockingPool`.
    ///
    /// Callers may `await` the returned `Task` to be notified when the work is completed.
    ///
    /// # Panics
    ///
    /// `await`ing a `Task` after dropping the `BlockingPool` or calling `BlockingPool::shutdown`
    /// will panic if the work was not completed before the pool was shut down.
    pub fn spawn<F, R>(&self, f: F) -> Task<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let raw = Arc::downgrade(&self.inner);
        let schedule = move |runnable| {
            if let Some(i) = raw.upgrade() {
                i.schedule(runnable);
            }
        };

        let (runnable, task) = async_task::spawn(async move { f() }, schedule);
        runnable.schedule();

        task
    }

    /// Shut down the `BlockingPool`.
    ///
    /// This will block until all work that has been started by the worker threads is finished. Any
    /// work that was added to the `BlockingPool` but not yet picked up by a worker thread will not
    /// complete and `await`ing on the `Task` for that work will panic.
    pub fn shutdown(&self) {
        let mut state = self.inner.state.lock();

        if state.shutting_down {
            // We've already shut down this BlockingPool.
            return;
        }

        state.shutting_down = true;
        let last_exited_thread = state.last_exited_thread.take();
        let unfinished_tasks = mem::replace(&mut state.tasks, VecDeque::new());
        let mut worker_threads = mem::replace(&mut state.worker_threads, Slab::new());
        drop(state);

        self.inner.condvar.notify_all();

        // Cancel any unfinished work after releasing the lock.
        drop(unfinished_tasks);

        // Now wait for all worker threads to exit.
        if let Some(handle) = last_exited_thread {
            let _ = handle.join();
        }

        for handle in worker_threads.drain() {
            let _ = handle.join();
        }
    }
}

impl Default for BlockingPool {
    fn default() -> BlockingPool {
        BlockingPool::new(256, Duration::from_secs(10))
    }
}

impl Drop for BlockingPool {
    fn drop(&mut self) {
        self.shutdown()
    }
}

#[cfg(test)]
mod test {
    use std::{
        sync::Arc,
        thread,
        time::{Duration, Instant},
    };

    use futures::{stream::FuturesUnordered, StreamExt};
    use sync::{Condvar, Mutex};

    use crate::{block_on, BlockingPool};

    #[test]
    fn blocking_sleep() {
        let pool = BlockingPool::default();

        let res = block_on(pool.spawn(|| 42));
        assert_eq!(res, 42);
    }

    #[test]
    fn more_tasks_than_threads() {
        let pool = BlockingPool::new(4, Duration::from_secs(10));

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);
    }

    #[test]
    fn shutdown() {
        let pool = BlockingPool::default();

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);

        pool.shutdown();
        let state = pool.inner.state.lock();
        assert_eq!(state.num_threads, 0);
    }

    #[test]
    fn keepalive_timeout() {
        // Set the keepalive to a very low value so that threads will exit soon after they run out
        // of work.
        let pool = BlockingPool::new(7, Duration::from_millis(1));

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);

        // Wait for all threads to exit.
        let deadline = Instant::now() + Duration::from_secs(10);
        while Instant::now() < deadline {
            thread::sleep(Duration::from_millis(100));
            let state = pool.inner.state.lock();
            if state.num_threads == 0 {
                break;
            }
        }

        {
            let state = pool.inner.state.lock();
            assert_eq!(state.num_threads, 0);
            assert_eq!(state.num_idle, 0);
        }
    }

    #[test]
    #[should_panic]
    fn shutdown_with_pending_work() {
        let pool = BlockingPool::new(1, Duration::from_secs(10));

        let mu = Arc::new(Mutex::new(false));
        let cv = Arc::new(Condvar::new());

        // First spawn a thread that blocks the pool.
        let task_mu = mu.clone();
        let task_cv = cv.clone();
        pool.spawn(move || {
            let mut ready = task_mu.lock();
            while !*ready {
                ready = task_cv.wait(ready);
            }
        })
        .detach();

        // This task will never finish because we will shut down the pool first.
        let unfinished = pool.spawn(|| 5);

        // Spawn a thread to unblock the work we started earlier once it sees that the pool is
        // shutting down.
        let inner = pool.inner.clone();
        thread::spawn(move || {
            let mut state = inner.state.lock();
            while !state.shutting_down {
                state = inner.condvar.wait(state);
            }

            *mu.lock() = true;
            cv.notify_all();
        });
        pool.shutdown();

        // This should panic.
        assert_eq!(block_on(unfinished), 5);
    }
}
<commit_msg>cros_async: Add a deadline to BlockingPool::shutdown<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    collections::VecDeque,
    mem,
    sync::{
        mpsc::{channel, Receiver, Sender},
        Arc,
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

use async_task::{Runnable, Task};
use slab::Slab;
use sync::{Condvar, Mutex};
use sys_util::{error, warn};

const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);

struct State {
    tasks: VecDeque<Runnable>,
    num_threads: usize,
    num_idle: usize,
    worker_threads: Slab<JoinHandle<()>>,
    exited_threads: Option<Receiver<usize>>,
    exit: Sender<usize>,
    shutting_down: bool,
}

fn run_blocking_thread(idx: usize, inner: Arc<Inner>, exit: Sender<usize>) {
    let mut state = inner.state.lock();
    while !state.shutting_down {
        if let Some(runnable) = state.tasks.pop_front() {
            drop(state);
            runnable.run();
            state = inner.state.lock();
            continue;
        }

        // No more tasks so wait for more work.
        state.num_idle += 1;

        let (guard, result) = inner
            .condvar
            .wait_timeout_while(state, inner.keepalive, |s| {
                !s.shutting_down && s.tasks.is_empty()
            });
        state = guard;

        // Only decrement the idle count if we timed out. Otherwise, it was decremented when new
        // work was added to `state.tasks`.
        if result.timed_out() {
            state.num_idle -= 1;
            break;
        }
    }

    state.num_threads -= 1;

    // If we're shutting down then the BlockingPool will take care of joining all the threads.
    // Otherwise, we need to join the last worker thread that exited here.
    let last_exited_thread = if let Some(exited_threads) = state.exited_threads.as_mut() {
        exited_threads
            .try_recv()
            .map(|idx| state.worker_threads.remove(idx))
            .ok()
    } else {
        None
    };

    // Drop the lock before trying to join the last exited thread.
    drop(state);

    if let Some(handle) = last_exited_thread {
        let _ = handle.join();
    }

    if let Err(e) = exit.send(idx) {
        error!("Failed to send thread exit event on channel: {}", e);
    }
}

struct Inner {
    state: Mutex<State>,
    condvar: Condvar,
    max_threads: usize,
    keepalive: Duration,
}

impl Inner {
    fn schedule(self: &Arc<Inner>, runnable: Runnable) {
        let mut state = self.state.lock();

        // If we're shutting down then nothing is going to run this task.
        if state.shutting_down {
            return;
        }

        state.tasks.push_back(runnable);

        if state.num_idle == 0 {
            // There are no idle threads.  Spawn a new one if possible.
            if state.num_threads < self.max_threads {
                state.num_threads += 1;
                let exit = state.exit.clone();
                let entry = state.worker_threads.vacant_entry();
                let idx = entry.key();
                let inner = self.clone();
                entry.insert(
                    thread::Builder::new()
                        .name(format!("blockingPool{}", idx))
                        .spawn(move || run_blocking_thread(idx, inner, exit))
                        .unwrap(),
                );
            }
        } else {
            // We have idle threads, wake one up.
            state.num_idle -= 1;
            self.condvar.notify_one();
        }
    }
}

#[derive(Debug, thiserror::Error)]
#[error("{0} BlockingPool threads did not exit in time and will be detached")]
pub struct ShutdownTimedOut(usize);

/// A thread pool for running work that may block.
///
/// It is generally discouraged to do any blocking work inside an async function. However, this is
/// sometimes unavoidable when dealing with interfaces that don't provide async variants. In this
/// case callers may use the `BlockingPool` to run the blocking work on a different thread and
/// `await` for its result to finish, which will prevent blocking the main thread of the
/// application.
///
/// Since the blocking work is sent to another thread, users should be careful when using the
/// `BlockingPool` for latency-sensitive operations. Additionally, the `BlockingPool` is intended to
/// be used for work that will eventually complete on its own. Users who want to spawn a thread
/// should just use `thread::spawn` directly.
///
/// There is no way to cancel work once it has been picked up by one of the worker threads in the
/// `BlockingPool`. Dropping or shutting down the pool will block up to a timeout (default 10
/// seconds) to wait for any active blocking work to finish. Any threads running tasks that have not
/// completed by that time will be detached.
///
/// # Examples
///
/// Spawn a task to run in the `BlockingPool` and await on its result.
///
/// ```edition2018
/// use cros_async::BlockingPool;
///
/// # async fn do_it() {
///     let pool = BlockingPool::default();
///
///     let res = pool.spawn(move || {
///         // Do some CPU-intensive or blocking work here.
///
///         42
///     }).await;
///
///     assert_eq!(res, 42);
/// # }
/// # cros_async::block_on(do_it());
/// ```
pub struct BlockingPool {
    inner: Arc<Inner>,
}

impl BlockingPool {
    /// Create a new `BlockingPool`.
    ///
    /// The `BlockingPool` will never spawn more than `max_threads` threads to do work, regardless
    /// of the number of tasks that are added to it. This value should be set relatively low (for
    /// example, the number of CPUs on the machine) if the pool is intended to run CPU intensive
    /// work or it should be set relatively high (128 or more) if the pool is intended to be used
    /// for various IO operations that cannot be completed asynchronously. The default value is 256.
    ///
    /// Worker threads are spawned on demand when new work is added to the pool and will
    /// automatically exit after being idle for some time so there is no overhead for setting
    /// `max_threads` to a large value when there is little to no work assigned to the
    /// `BlockingPool`. `keepalive` determines the idle duration after which the worker thread will
    /// exit. The default value is 10 seconds.
    pub fn new(max_threads: usize, keepalive: Duration) -> BlockingPool {
        let (exit, exited_threads) = channel();
        BlockingPool {
            inner: Arc::new(Inner {
                state: Mutex::new(State {
                    tasks: VecDeque::new(),
                    num_threads: 0,
                    num_idle: 0,
                    worker_threads: Slab::new(),
                    exited_threads: Some(exited_threads),
                    exit,
                    shutting_down: false,
                }),
                condvar: Condvar::new(),
                max_threads,
                keepalive,
            }),
        }
    }

    /// Like new but with pre-allocating capacity for up to `max_threads`.
    pub fn with_capacity(max_threads: usize, keepalive: Duration) -> BlockingPool {
        let (exit, exited_threads) = channel();
        BlockingPool {
            inner: Arc::new(Inner {
                state: Mutex::new(State {
                    tasks: VecDeque::new(),
                    num_threads: 0,
                    num_idle: 0,
                    worker_threads: Slab::with_capacity(max_threads),
                    exited_threads: Some(exited_threads),
                    exit,
                    shutting_down: false,
                }),
                condvar: Condvar::new(),
                max_threads,
                keepalive,
            }),
        }
    }

    /// Spawn a task to run in the `BlockingPool`.
    ///
    /// Callers may `await` the returned `Task` to be notified when the work is completed.
    ///
    /// # Panics
    ///
    /// `await`ing a `Task` after dropping the `BlockingPool` or calling `BlockingPool::shutdown`
    /// will panic if the work was not completed before the pool was shut down.
    pub fn spawn<F, R>(&self, f: F) -> Task<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let raw = Arc::downgrade(&self.inner);
        let schedule = move |runnable| {
            if let Some(i) = raw.upgrade() {
                i.schedule(runnable);
            }
        };

        let (runnable, task) = async_task::spawn(async move { f() }, schedule);
        runnable.schedule();

        task
    }

    /// Shut down the `BlockingPool`.
    ///
    /// If `deadline` is provided then this will block until either all worker threads exit or the
    /// deadline is exceeded. If `deadline` is not given then this will block indefinitely until all
    /// worker threads exit. Any work that was added to the `BlockingPool` but not yet picked up by
    /// a worker thread will not complete and `await`ing on the `Task` for that work will panic.
    pub fn shutdown(&self, deadline: Option<Instant>) -> Result<(), ShutdownTimedOut> {
        let mut state = self.inner.state.lock();

        if state.shutting_down {
            // We've already shut down this BlockingPool.
            return Ok(());
        }

        state.shutting_down = true;
        let exited_threads = state.exited_threads.take().expect("exited_threads missing");
        let unfinished_tasks = mem::replace(&mut state.tasks, VecDeque::new());
        let mut worker_threads = mem::replace(&mut state.worker_threads, Slab::new());
        drop(state);

        self.inner.condvar.notify_all();

        // Cancel any unfinished work after releasing the lock.
        drop(unfinished_tasks);

        // Now wait for all worker threads to exit.
        if let Some(deadline) = deadline {
            let mut now = Instant::now();
            while now < deadline && !worker_threads.is_empty() {
                if let Ok(idx) = exited_threads.recv_timeout(deadline - now) {
                    let _ = worker_threads.remove(idx).join();
                }
                now = Instant::now();
            }

            // Any threads that have not yet joined will just be detached.
            if !worker_threads.is_empty() {
                return Err(ShutdownTimedOut(worker_threads.len()));
            }

            Ok(())
        } else {
            // Block indefinitely until all worker threads exit.
            for handle in worker_threads.drain() {
                let _ = handle.join();
            }

            Ok(())
        }
    }
}

impl Default for BlockingPool {
    fn default() -> BlockingPool {
        BlockingPool::new(256, Duration::from_secs(10))
    }
}

impl Drop for BlockingPool {
    fn drop(&mut self) {
        if let Err(e) = self.shutdown(Some(Instant::now() + DEFAULT_SHUTDOWN_TIMEOUT)) {
            warn!("{}", e);
        }
    }
}

#[cfg(test)]
mod test {
    use std::{
        sync::{Arc, Barrier},
        thread,
        time::{Duration, Instant},
    };

    use futures::{stream::FuturesUnordered, StreamExt};
    use sync::{Condvar, Mutex};

    use crate::{block_on, BlockingPool};

    #[test]
    fn blocking_sleep() {
        let pool = BlockingPool::default();

        let res = block_on(pool.spawn(|| 42));
        assert_eq!(res, 42);
    }

    #[test]
    fn more_tasks_than_threads() {
        let pool = BlockingPool::new(4, Duration::from_secs(10));

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);
    }

    #[test]
    fn shutdown() {
        let pool = BlockingPool::default();

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);

        pool.shutdown(Some(Instant::now() + Duration::from_secs(10)))
            .unwrap();
        let state = pool.inner.state.lock();
        assert_eq!(state.num_threads, 0);
    }

    #[test]
    fn keepalive_timeout() {
        // Set the keepalive to a very low value so that threads will exit soon after they run out
        // of work.
        let pool = BlockingPool::new(7, Duration::from_millis(1));

        let stream = (0..19)
            .map(|_| pool.spawn(|| thread::sleep(Duration::from_millis(5))))
            .collect::<FuturesUnordered<_>>();

        let results = block_on(stream.collect::<Vec<_>>());
        assert_eq!(results.len(), 19);

        // Wait for all threads to exit.
        let deadline = Instant::now() + Duration::from_secs(10);
        while Instant::now() < deadline {
            thread::sleep(Duration::from_millis(100));
            let state = pool.inner.state.lock();
            if state.num_threads == 0 {
                break;
            }
        }

        {
            let state = pool.inner.state.lock();
            assert_eq!(state.num_threads, 0);
            assert_eq!(state.num_idle, 0);
        }
    }

    #[test]
    #[should_panic]
    fn shutdown_with_pending_work() {
        let pool = BlockingPool::new(1, Duration::from_secs(10));

        let mu = Arc::new(Mutex::new(false));
        let cv = Arc::new(Condvar::new());

        // First spawn a thread that blocks the pool.
        let task_mu = mu.clone();
        let task_cv = cv.clone();
        pool.spawn(move || {
            let mut ready = task_mu.lock();
            while !*ready {
                ready = task_cv.wait(ready);
            }
        })
        .detach();

        // This task will never finish because we will shut down the pool first.
        let unfinished = pool.spawn(|| 5);

        // Spawn a thread to unblock the work we started earlier once it sees that the pool is
        // shutting down.
        let inner = pool.inner.clone();
        thread::spawn(move || {
            let mut state = inner.state.lock();
            while !state.shutting_down {
                state = inner.condvar.wait(state);
            }

            *mu.lock() = true;
            cv.notify_all();
        });
        pool.shutdown(None).unwrap();

        // This should panic.
        assert_eq!(block_on(unfinished), 5);
    }

    #[test]
    fn unfinished_worker_thread() {
        let pool = BlockingPool::default();

        let ready = Arc::new(Mutex::new(false));
        let cv = Arc::new(Condvar::new());
        let barrier = Arc::new(Barrier::new(2));

        let thread_ready = ready.clone();
        let thread_barrier = barrier.clone();
        let thread_cv = cv.clone();

        let task = pool.spawn(move || {
            thread_barrier.wait();
            let mut ready = thread_ready.lock();
            while !*ready {
                ready = thread_cv.wait(ready);
            }
        });

        // Wait to shut down the pool until after the worker thread has started.
        barrier.wait();
        pool.shutdown(Some(Instant::now() + Duration::from_millis(5)))
            .unwrap_err();

        let num_threads = pool.inner.state.lock().num_threads;
        assert_eq!(num_threads, 1);

        // Now wake up the blocked task so we don't leak the thread.
        *ready.lock() = true;
        cv.notify_all();

        block_on(task);

        let deadline = Instant::now() + Duration::from_secs(10);
        while Instant::now() < deadline {
            thread::sleep(Duration::from_millis(100));
            let state = pool.inner.state.lock();
            if state.num_threads == 0 {
                break;
            }
        }

        {
            let state = pool.inner.state.lock();
            assert_eq!(state.num_threads, 0);
            assert_eq!(state.num_idle, 0);
        }
    }
}
<|endoftext|>


<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

use rand::distributions::{Sample, Range};
use rand::Rng;

use std::usize;

use super::MCMove;
use super::select_molecule;

use types::Vector3D;
use sys::{System, EnergyCache};

/// Monte-Carlo move for translating a molecule
pub struct Translate {
    /// Type of molecule to translate. `None` means all molecules.
    moltype: Option<u64>,
    /// Index of the molecule to translate
    molid: usize,
    /// New positions of the atom in the translated molecule
    newpos: Vec<Vector3D>,
    /// Maximum displacement value
    dr: f64,
    /// Translation range for random number generation
    range: Range<f64>,
}

impl Translate {
    /// Create a new `Translate` move, with maximum displacement of `dr`.
    /// Translating all the molecules in the system.
    pub fn new(dr: f64) -> Translate {
        Translate::create(dr, None)
    }

    /// Create a new `Translate` move, with maximum displacement of `dr`.
    /// Translating only molecules with `moltype` type.
    pub fn with_moltype(dr: f64, moltype: u64) -> Translate {
        Translate::create(dr, Some(moltype))
    }

    /// Factorizing the constructors
    fn create(dr: f64, moltype: Option<u64>) -> Translate {
        assert!(dr > 0.0, "dr must be positive in Translate move");
        let dr = dr / f64::sqrt(3.0);
        Translate {
            moltype: moltype,
            molid: usize::MAX,
            newpos: Vec::new(),
            dr: dr,
            range: Range::new(-dr, dr),
        }
    }
}

impl Default for Translate {
    fn default() -> Translate {
        Translate::new(1.0)
    }
}

impl MCMove for Translate {
    fn describe(&self) -> &str {
        "molecular translation"
    }

    fn setup(&mut self, _: &System) {
        // TODO: introduce a limitation for dr, dr_max, such that it can not
        // be larger than the largest cutoff in the system? We would initialize
        // the value here.
    }

    fn prepare(&mut self, system: &mut System, rng: &mut Box<Rng>) -> bool {
        if let Some(id) = select_molecule(system, self.moltype, rng) {
            self.molid = id;
        } else {
            warn!("Can not translate molecule: no molecule of this type in the system.");
            return false;
        }

        // Create random displacement vector.
        let delta = Vector3D::new(
            self.range.sample(rng),
            self.range.sample(rng),
            self.range.sample(rng)
        );

        self.newpos.clear();
        // Create new position.
        // Note that this may move a particles' center-
        // of-mass (com) out of the cell.
        // If the move is accepted, we have to wrap the
        // com such that it lies inside the cell.
        for i in system.molecule(self.molid) {
            self.newpos.push(system[i].position + delta);
        }
        return true;
    }

    fn cost(&self, system: &System, beta: f64, cache: &mut EnergyCache) -> f64 {
        let idxes = system.molecule(self.molid).iter().collect::<Vec<_>>();
        let cost = cache.move_particles_cost(system, idxes, &self.newpos);
        return cost*beta;
    }

    fn apply(&mut self, system: &mut System) {
        // Update positions.
        for (i, pi) in system.molecule(self.molid).iter().enumerate() {
            system[pi].position = self.newpos[i];
        }
        // Move molecule such that its center-of-mass
        // is inside the simulation cell.
        // Note that particles of the molecule may still be
        // outside the cell, but that is not important.
        system.wrap_molecule(self.molid)
    }

    fn restore(&mut self, _: &mut System) {
        // Nothing to do.
    }

    fn update_amplitude(&mut self, scaling_factor: Option<f64>) {
        if let Some(s) = scaling_factor {
            self.dr *= s;
            self.range = Range::new(-self.dr, self.dr);
        }
    }
}
<commit_msg>Add restriction to the amplitude for translations in setup and update_amplitude.<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

use rand::distributions::{Sample, Range};
use rand::Rng;

use std::usize;
use std::f64;

use super::MCMove;
use super::select_molecule;

use types::Vector3D;
use sys::{System, EnergyCache};

/// Monte-Carlo move for translating a molecule
pub struct Translate {
    /// Type of molecule to translate. `None` means all molecules.
    moltype: Option<u64>,
    /// Index of the molecule to translate
    molid: usize,
    /// New positions of the atom in the translated molecule
    newpos: Vec<Vector3D>,
    /// Maximum displacement value
    dr: f64,
    /// The maximum value must not exceed this value
    dr_max: f64,
    /// Translation range for random number generation
    range: Range<f64>,
}

impl Translate {
    /// Create a new `Translate` move, with maximum displacement of `dr`.
    /// Translating all the molecules in the system.
    pub fn new(dr: f64) -> Translate {
        Translate::create(dr, None)
    }

    /// Create a new `Translate` move, with maximum displacement of `dr`.
    /// Translating only molecules with `moltype` type.
    pub fn with_moltype(dr: f64, moltype: u64) -> Translate {
        Translate::create(dr, Some(moltype))
    }

    /// Factorizing the constructors
    fn create(dr: f64, moltype: Option<u64>) -> Translate {
        assert!(dr > 0.0, "dr must be positive in Translate move");
        let dr = dr / f64::sqrt(3.0);
        Translate {
            moltype: moltype,
            molid: usize::MAX,
            newpos: Vec::new(),
            dr: dr,
            dr_max: 0.0,
            range: Range::new(-dr, dr),
        }
    }
}

impl Default for Translate {
    fn default() -> Translate {
        Translate::new(1.0)
    }
}

impl MCMove for Translate {
    fn describe(&self) -> &str {
        "molecular translation"
    }

    fn setup(&mut self, system: &System) {
        // Limit the displacement range to the maximum cutoff
        self.dr_max = system.interactions()
                            .all_pairs()
                            .iter()
                            .map(|i| i.get_cutoff())
                            .fold(f64::NAN, f64::max);
        if self.dr > self.dr_max {
            self.dr = self.dr_max
        }
    }

    fn prepare(&mut self, system: &mut System, rng: &mut Box<Rng>) -> bool {
        if let Some(id) = select_molecule(system, self.moltype, rng) {
            self.molid = id;
        } else {
            warn!("Can not translate molecule: no molecule of this type in the system.");
            return false;
        }

        // Create random displacement vector.
        let delta = Vector3D::new(
            self.range.sample(rng),
            self.range.sample(rng),
            self.range.sample(rng)
        );

        self.newpos.clear();
        // Create new position.
        // Note that this may move a particles' center-
        // of-mass (com) out of the cell.
        // If the move is accepted, we have to wrap the
        // com such that it lies inside the cell.
        for i in system.molecule(self.molid) {
            self.newpos.push(system[i].position + delta);
        }
        return true;
    }

    fn cost(&self, system: &System, beta: f64, cache: &mut EnergyCache) -> f64 {
        let idxes = system.molecule(self.molid).iter().collect::<Vec<_>>();
        let cost = cache.move_particles_cost(system, idxes, &self.newpos);
        return cost*beta;
    }

    fn apply(&mut self, system: &mut System) {
        // Update positions.
        for (i, pi) in system.molecule(self.molid).iter().enumerate() {
            system[pi].position = self.newpos[i];
        }
        // Move molecule such that its center-of-mass
        // is inside the simulation cell.
        // Note that particles of the molecule may still be
        // outside the cell, but that is not important.
        system.wrap_molecule(self.molid)
    }

    fn restore(&mut self, _: &mut System) {
        // Nothing to do.
    }

    fn update_amplitude(&mut self, scaling_factor: Option<f64>) {
        if let Some(s) = scaling_factor {
            if (self.dr * s) < self.dr_max {
                self.dr *= s;
                self.range = Range::new(-self.dr, self.dr);
            } else {
                // FIXME: fix the warn_once macro formatting usage
                #[allow(useless_format)] {
                    warn_once!(
                        "Tried to increase the maximum amplitude for translations to more than the set maximum."
                    );
                }
            }

        };
    }
}
<|endoftext|>


<commit_before>#![no_std]

use core::fmt;
use core::marker::PhantomData;
use libtock_platform as platform;
use libtock_platform::allow_ro::AllowRo;
use libtock_platform::allow_rw::AllowRw;
use libtock_platform::share;
use libtock_platform::subscribe::Subscribe;
use libtock_platform::{DefaultConfig, ErrorCode, Syscalls};

/// The console driver.
///
/// It allows libraries to pass strings to the kernel's console driver.
///
/// # Example
/// ```ignore
/// use libtock2::Console;
///
/// // Writes "foo", followed by a newline, to the console
/// let mut writer = Console::writer();
/// writeln!(writer, foo).unwrap();
/// ```
pub struct Console<
    S: Syscalls,
    C: platform::allow_ro::Config
        + platform::allow_rw::Config
        + platform::subscribe::Config
        = DefaultConfig,
>(S, C);

impl<
        S: Syscalls,
        C: platform::allow_ro::Config + platform::allow_rw::Config + platform::subscribe::Config,
    > Console<S, C>
{
    /// Run a check against the console capsule to ensure it is present.
    ///
    /// Returns `true` if the driver was present. This does not necessarily mean
    /// that the driver is working, as it may still fail to allocate grant
    /// memory.
    #[inline(always)]
    pub fn driver_check() -> bool {
        S::command(DRIVER_NUM, command::DRIVER_CHECK, 0, 0).is_success()
    }

    /// Writes bytes.
    /// This is an alternative to `fmt::Write::write`
    /// because this can actually return an error code.
    pub fn write(s: &[u8]) -> Result<(), ErrorCode> {
        let called = core::cell::Cell::new(Option::<(u32,)>::None);
        share::scope::<
            (
                AllowRo<_, DRIVER_NUM, { allow_ro::WRITE }>,
                Subscribe<_, DRIVER_NUM, { subscribe::WRITE }>,
            ),
            _,
            _,
        >(|handle| {
            let (allow_ro, subscribe) = handle.split();

            S::allow_ro::<C, DRIVER_NUM, { allow_ro::WRITE }>(allow_ro, s)?;

            S::subscribe::<_, _, C, DRIVER_NUM, { subscribe::WRITE }>(subscribe, &called)?;

            S::command(DRIVER_NUM, command::WRITE, s.len() as u32, 0).to_result()?;

            loop {
                S::yield_wait();
                if let Some((_,)) = called.get() {
                    return Ok(());
                }
            }
        })
    }

    /// Reads bytes
    /// Reads from the device and writes to `buf`, starting from index 0.
    /// No special guarantees about when the read stops.
    /// Returns count of bytes written to `buf`.
    pub fn read(buf: &mut [u8]) -> (usize, Result<(), ErrorCode>) {
        let called = core::cell::Cell::new(Option::<(u32, u32)>::None);
        let mut bytes_received = 0;
        let r = share::scope::<
            (
                AllowRw<_, DRIVER_NUM, { allow_rw::READ }>,
                Subscribe<_, DRIVER_NUM, { subscribe::READ }>,
            ),
            _,
            _,
        >(|handle| {
            let (allow_rw, subscribe) = handle.split();
            let len = buf.len();
            S::allow_rw::<C, DRIVER_NUM, { allow_rw::READ }>(allow_rw, buf)?;
            S::subscribe::<_, _, C, DRIVER_NUM, { subscribe::READ }>(subscribe, &called)?;

            // When this fails, `called` is guaranteed unmodified,
            // because upcalls are never processed within a `command` call.
            S::command(DRIVER_NUM, command::READ, len as u32, 0).to_result()?;

            loop {
                S::yield_wait();
                if let Some((status, bytes_pushed_count)) = called.get() {
                    bytes_received = bytes_pushed_count as usize;
                    return match status {
                        0 => Ok(()),
                        e_status => Err(e_status.try_into().unwrap_or(ErrorCode::Fail)),
                    };
                }
            }
        });
        (bytes_received, r)
    }

    pub fn writer() -> ConsoleWriter<S> {
        ConsoleWriter {
            syscalls: Default::default(),
        }
    }
}

pub struct ConsoleWriter<S: Syscalls> {
    syscalls: PhantomData<S>,
}

impl<S: Syscalls> fmt::Write for ConsoleWriter<S> {
    fn write_str(&mut self, s: &str) -> Result<(), fmt::Error> {
        Console::<S>::write(s.as_bytes()).map_err(|_e| fmt::Error)
    }
}

#[cfg(test)]
mod tests;

// -----------------------------------------------------------------------------
// Driver number and command IDs
// -----------------------------------------------------------------------------

const DRIVER_NUM: u32 = 1;

// Command IDs
#[allow(unused)]
mod command {
    pub const DRIVER_CHECK: u32 = 0;
    pub const WRITE: u32 = 1;
    pub const READ: u32 = 2;
    pub const ABORT: u32 = 3;
}

#[allow(unused)]
mod subscribe {
    pub const WRITE: u32 = 1;
    pub const READ: u32 = 2;
}

mod allow_ro {
    pub const WRITE: u32 = 1;
}

mod allow_rw {
    pub const READ: u32 = 1;
}
<commit_msg>Update apis/console/src/lib.rs<commit_after>#![no_std]

use core::fmt;
use core::marker::PhantomData;
use libtock_platform as platform;
use libtock_platform::allow_ro::AllowRo;
use libtock_platform::allow_rw::AllowRw;
use libtock_platform::share;
use libtock_platform::subscribe::Subscribe;
use libtock_platform::{DefaultConfig, ErrorCode, Syscalls};

/// The console driver.
///
/// It allows libraries to pass strings to the kernel's console driver.
///
/// # Example
/// ```ignore
/// use libtock2::Console;
///
/// // Writes "foo", followed by a newline, to the console
/// let mut writer = Console::writer();
/// writeln!(writer, foo).unwrap();
/// ```
pub struct Console<
    S: Syscalls,
    C: platform::allow_ro::Config
        + platform::allow_rw::Config
        + platform::subscribe::Config
        = DefaultConfig,
>(S, C);

impl<
        S: Syscalls,
        C: platform::allow_ro::Config + platform::allow_rw::Config + platform::subscribe::Config,
    > Console<S, C>
{
    /// Run a check against the console capsule to ensure it is present.
    ///
    /// Returns `true` if the driver was present. This does not necessarily mean
    /// that the driver is working, as it may still fail to allocate grant
    /// memory.
    #[inline(always)]
    pub fn driver_check() -> bool {
        S::command(DRIVER_NUM, command::DRIVER_CHECK, 0, 0).is_success()
    }

    /// Writes bytes.
    /// This is an alternative to `fmt::Write::write`
    /// because this can actually return an error code.
    pub fn write(s: &[u8]) -> Result<(), ErrorCode> {
        let called = core::cell::Cell::new(Option::<(u32,)>::None);
        share::scope::<
            (
                AllowRo<_, DRIVER_NUM, { allow_ro::WRITE }>,
                Subscribe<_, DRIVER_NUM, { subscribe::WRITE }>,
            ),
            _,
            _,
        >(|handle| {
            let (allow_ro, subscribe) = handle.split();

            S::allow_ro::<C, DRIVER_NUM, { allow_ro::WRITE }>(allow_ro, s)?;

            S::subscribe::<_, _, C, DRIVER_NUM, { subscribe::WRITE }>(subscribe, &called)?;

            S::command(DRIVER_NUM, command::WRITE, s.len() as u32, 0).to_result()?;

            loop {
                S::yield_wait();
                if let Some((_,)) = called.get() {
                    return Ok(());
                }
            }
        })
    }

    /// Reads bytes
    /// Reads from the device and writes to `buf`, starting from index 0.
    /// No special guarantees about when the read stops.
    /// Returns count of bytes written to `buf`.
    pub fn read(buf: &mut [u8]) -> (usize, Result<(), ErrorCode>) {
        let called = core::cell::Cell::new(Option::<(u32, u32)>::None);
        let mut bytes_received = 0;
        let r = share::scope::<
            (
                AllowRw<_, DRIVER_NUM, { allow_rw::READ }>,
                Subscribe<_, DRIVER_NUM, { subscribe::READ }>,
            ),
            _,
            _,
        >(|handle| {
            let (allow_rw, subscribe) = handle.split();
            let len = buf.len();
            S::allow_rw::<C, DRIVER_NUM, { allow_rw::READ }>(allow_rw, buf)?;
            S::subscribe::<_, _, C, DRIVER_NUM, { subscribe::READ }>(subscribe, &called)?;

            // When this fails, `called` is guaranteed unmodified,
            // because upcalls are never processed until we call `yield`.
            S::command(DRIVER_NUM, command::READ, len as u32, 0).to_result()?;

            loop {
                S::yield_wait();
                if let Some((status, bytes_pushed_count)) = called.get() {
                    bytes_received = bytes_pushed_count as usize;
                    return match status {
                        0 => Ok(()),
                        e_status => Err(e_status.try_into().unwrap_or(ErrorCode::Fail)),
                    };
                }
            }
        });
        (bytes_received, r)
    }

    pub fn writer() -> ConsoleWriter<S> {
        ConsoleWriter {
            syscalls: Default::default(),
        }
    }
}

pub struct ConsoleWriter<S: Syscalls> {
    syscalls: PhantomData<S>,
}

impl<S: Syscalls> fmt::Write for ConsoleWriter<S> {
    fn write_str(&mut self, s: &str) -> Result<(), fmt::Error> {
        Console::<S>::write(s.as_bytes()).map_err(|_e| fmt::Error)
    }
}

#[cfg(test)]
mod tests;

// -----------------------------------------------------------------------------
// Driver number and command IDs
// -----------------------------------------------------------------------------

const DRIVER_NUM: u32 = 1;

// Command IDs
#[allow(unused)]
mod command {
    pub const DRIVER_CHECK: u32 = 0;
    pub const WRITE: u32 = 1;
    pub const READ: u32 = 2;
    pub const ABORT: u32 = 3;
}

#[allow(unused)]
mod subscribe {
    pub const WRITE: u32 = 1;
    pub const READ: u32 = 2;
}

mod allow_ro {
    pub const WRITE: u32 = 1;
}

mod allow_rw {
    pub const READ: u32 = 1;
}
<|endoftext|>


<commit_before>/// Return the Hamming distance between the strings,
/// or an error if the lengths are mismatched.
pub fn hamming_distance(s1: &str, s2: &str) -> Result<u64, ()> {
    unimplemented!("What is the Hamming Distance between {:?} and {:?}",
                   s1, s2);
}<commit_msg>hamming: `Result` should be `usize`<commit_after>/// Return the Hamming distance between the strings,
/// or an error if the lengths are mismatched.
pub fn hamming_distance(s1: &str, s2: &str) -> Result<usize, ()> {
    unimplemented!("What is the Hamming Distance between {:?} and {:?}",
                   s1, s2);
}
<|endoftext|>


<commit_before>/*!
 * Helpers for headers - Tags
 */

/**
 * Spec helpers for header-tags
 */
pub mod spec {
    use storage::file::header::spec::FileHeaderSpec as FHS;
    use module::helpers::spec::{named_text, named_text_array};

    /**
     * helper for a Header spec for
     *
     *  { "URL": "<Text>" }
     */
    pub fn url_key() -> FHS {
        named_text("URL")
    }

    /**
     * helper for a Header spec for
     *
     *  { "TAGS": [ "<Text>", ... ] }
     */
    pub fn tags_key() -> FHS {
        named_text_array("TAGS")
    }

}

/**
 * Data helpers for header-tags
 */
pub mod data {
    use std::ops::Deref;
    use storage::file::header::data::FileHeaderData as FHD;
    use module::Module;
    use clap::ArgMatches;
    use storage::parser::Parser;
    use storage::parser::FileHeaderParser;

    /**
     * Use a Vec<String> to build a Tag-Array:
     *
     *  [ "<Text>", ... ]
     */
    pub fn build_tag_array(tags: Vec<String>) -> FHD {
        let texttags = tags.into_iter().map(|t| FHD::Text(t.clone())).collect();
        FHD::Array { values: Box::new(texttags) }
    }

    /**
     * Fetch tags from a header, whereas the header looks like this:
     *
     *   { ...,
     *     "TAGS": [ "<Text>", ... ],
     *     ...
     *   }
     *
     * Does no spec verification.
     */
    pub fn get_tags_from_header(header: &FHD) -> Vec<String> {
        let mut tags : Vec<String> = vec![];

        fn match_array(a: &Box<FHD>) -> Vec<String> {
            let mut tags : Vec<String> = vec![];

            match a.deref() {
                &FHD::Array{values: ref vs} => {
                    let values : Vec<FHD> = vs.deref().clone();
                    for value in values {
                        match value {
                            FHD::Text(t) => tags.push(t),
                            _ => warn!("Malformed Header Data: Expected Text, found non-Text"),
                        }
                    }
                }
                _ => warn!("Malformed Header Data: Expected Array, found non-Array"),
            }

            tags
        }

        match header {
            &FHD::Map{keys: ref ks} => {
                let keys : Vec<FHD> = ks.clone();
                for key in keys {
                    match key {
                        FHD::Key{name: _, value: ref v} => return match_array(v),
                        _ => warn!("Malformed Header Data: Expected Key, found non-Key"),
                    }
                }
            },
            _ => warn!("Malformed Header Data: Expected Map, found non-Map"),
        }

        tags
    }

    /**
     * Helper function to alter the tags in a file
     */
    pub fn alter_tags_in_files<HP, F, R>(m: &Module,
                                     matches: &ArgMatches,
                                     parser: &Parser<HP>,
                                     generate_new_tags: F,
                                     rebuild_header: R) -> bool
        where HP: FileHeaderParser,
              F:  Fn(Vec<String>, &Vec<String>) -> Vec<String>,
              R:  Fn(&FHD, Vec<String>) -> Option<FHD>
    {
        use std::process::exit;
        use module::helpers::cli::create_tag_filter;
        use module::helpers::cli::create_hash_filter;
        use module::helpers::cli::create_text_header_field_grep_filter;
        use module::helpers::cli::create_content_grep_filter;
        use module::helpers::cli::CliFileFilter;

        let cli_tags = matches.value_of("tags")
                          .map(|ts| {
                            ts.split(",")
                              .map(String::from)
                              .collect::<Vec<String>>()
                          })
                          .unwrap_or(vec![]);

        let filter = {
            let hash_filter = create_hash_filter(matches, "with:id", false);
            let text_filter = create_text_header_field_grep_filter(matches, "with_match", "URL", false);
            let tags_filter = create_tag_filter(matches, "with_tags", false);
            hash_filter.or(Box::new(text_filter)).or(Box::new(tags_filter))
        };

        m.runtime()
            .store()
            .load_for_module(m, &parser)
            .into_iter()
            .filter(|file| filter.filter_file(file))
            .map(|file| {
                debug!("Alter tags in file: {:?}", file);

                let hdr = {
                    let f = file.deref().borrow();
                    f.header().clone()
                };

                debug!("Tags:...");
                let old_tags = get_tags_from_header(&hdr);
                debug!("    old_tags = {:?}", &old_tags);
                debug!("    cli_tags = {:?}", &cli_tags);

                let new_tags = generate_new_tags(old_tags, &cli_tags);
                debug!("    new_tags = {:?}", &new_tags);

                let new_header = rebuild_header(&hdr, new_tags)
                    .unwrap_or_else(|| {
                        error!("Could not rebuild header for file");
                        exit(1);
                    });
                {
                    let mut f_mut = file.deref().borrow_mut();
                    f_mut.set_header(new_header);
                }

                m.runtime().store().persist(&parser, file);
                true
            })
            .all(|x| x)
    }

}

<commit_msg>Auto merge of #55 - matthiasbeyer:fix-warn_header_data_malformed, r=matthiasbeyer<commit_after>/*!
 * Helpers for headers - Tags
 */

/**
 * Spec helpers for header-tags
 */
pub mod spec {
    use storage::file::header::spec::FileHeaderSpec as FHS;
    use module::helpers::spec::{named_text, named_text_array};

    /**
     * helper for a Header spec for
     *
     *  { "URL": "<Text>" }
     */
    pub fn url_key() -> FHS {
        named_text("URL")
    }

    /**
     * helper for a Header spec for
     *
     *  { "TAGS": [ "<Text>", ... ] }
     */
    pub fn tags_key() -> FHS {
        named_text_array("TAGS")
    }

}

/**
 * Data helpers for header-tags
 */
pub mod data {
    use std::ops::Deref;
    use storage::file::header::data::FileHeaderData as FHD;
    use module::Module;
    use clap::ArgMatches;
    use storage::parser::Parser;
    use storage::parser::FileHeaderParser;

    /**
     * Use a Vec<String> to build a Tag-Array:
     *
     *  [ "<Text>", ... ]
     */
    pub fn build_tag_array(tags: Vec<String>) -> FHD {
        let texttags = tags.into_iter().map(|t| FHD::Text(t.clone())).collect();
        FHD::Array { values: Box::new(texttags) }
    }

    /**
     * Fetch tags from a header, whereas the header looks like this:
     *
     *   { ...,
     *     "TAGS": [ "<Text>", ... ],
     *     ...
     *   }
     *
     * Does no spec verification.
     */
    pub fn get_tags_from_header(header: &FHD) -> Vec<String> {
        let mut tags : Vec<String> = vec![];

        fn match_array(a: &Box<FHD>) -> Vec<String> {
            let mut tags : Vec<String> = vec![];

            match a.deref() {
                &FHD::Array{values: ref vs} => {
                    let values : Vec<FHD> = vs.deref().clone();
                    for value in values {
                        match value {
                            FHD::Text(t) => tags.push(t),
                            _ => warn!("Malformed Header Data: Expected Text, found non-Text"),
                        }
                    }
                }
                _ => warn!("Malformed Header Data: Expected Array, found non-Array"),
            }

            tags
        }

        match header {
            &FHD::Map{keys: ref ks} => {
                let keys : Vec<FHD> = ks.clone();
                for key in keys {
                    match key {
                        FHD::Key{name: ref name, value: ref v} => {
                            if name == "TAGS" {
                                return match_array(v)
                            }
                        },
                        _ => warn!("Malformed Header Data: Expected Key, found non-Key"),
                    }
                }
            },
            _ => warn!("Malformed Header Data: Expected Map, found non-Map"),
        }

        tags
    }

    /**
     * Helper function to alter the tags in a file
     */
    pub fn alter_tags_in_files<HP, F, R>(m: &Module,
                                     matches: &ArgMatches,
                                     parser: &Parser<HP>,
                                     generate_new_tags: F,
                                     rebuild_header: R) -> bool
        where HP: FileHeaderParser,
              F:  Fn(Vec<String>, &Vec<String>) -> Vec<String>,
              R:  Fn(&FHD, Vec<String>) -> Option<FHD>
    {
        use std::process::exit;
        use module::helpers::cli::create_tag_filter;
        use module::helpers::cli::create_hash_filter;
        use module::helpers::cli::create_text_header_field_grep_filter;
        use module::helpers::cli::create_content_grep_filter;
        use module::helpers::cli::CliFileFilter;

        let cli_tags = matches.value_of("tags")
                          .map(|ts| {
                            ts.split(",")
                              .map(String::from)
                              .collect::<Vec<String>>()
                          })
                          .unwrap_or(vec![]);

        let filter = {
            let hash_filter = create_hash_filter(matches, "with:id", false);
            let text_filter = create_text_header_field_grep_filter(matches, "with_match", "URL", false);
            let tags_filter = create_tag_filter(matches, "with_tags", false);
            hash_filter.or(Box::new(text_filter)).or(Box::new(tags_filter))
        };

        m.runtime()
            .store()
            .load_for_module(m, &parser)
            .into_iter()
            .filter(|file| filter.filter_file(file))
            .map(|file| {
                debug!("Alter tags in file: {:?}", file);

                let hdr = {
                    let f = file.deref().borrow();
                    f.header().clone()
                };

                debug!("Tags:...");
                let old_tags = get_tags_from_header(&hdr);
                debug!("    old_tags = {:?}", &old_tags);
                debug!("    cli_tags = {:?}", &cli_tags);

                let new_tags = generate_new_tags(old_tags, &cli_tags);
                debug!("    new_tags = {:?}", &new_tags);

                let new_header = rebuild_header(&hdr, new_tags)
                    .unwrap_or_else(|| {
                        error!("Could not rebuild header for file");
                        exit(1);
                    });
                {
                    let mut f_mut = file.deref().borrow_mut();
                    f_mut.set_header(new_header);
                }

                m.runtime().store().persist(&parser, file);
                true
            })
            .all(|x| x)
    }

}

<|endoftext|>


<commit_before>use lsp_types::{Location, Position, Range, TextDocumentIdentifier, Url};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use url_serde;

pub use lsp_types::{
    notification::*, request::*, ApplyWorkspaceEditParams, CodeActionParams, CodeLens, CodeLensParams,
    CompletionParams, CompletionResponse, DocumentOnTypeFormattingParams, DocumentSymbolParams,
    DocumentSymbolResponse, ExecuteCommandParams, Hover, InitializeResult,
    PublishDiagnosticsParams, ReferenceParams, SignatureHelp, TextDocumentEdit,
    TextDocumentPositionParams, TextEdit, WorkspaceEdit, WorkspaceSymbolParams,
    MessageType, ShowMessageParams,
};

pub enum AnalyzerStatus {}

impl Request for AnalyzerStatus {
    type Params = ();
    type Result = String;
    const METHOD: &'static str = "rust-analyzer/analyzerStatus";
}

pub enum CollectGarbage {}

impl Request for CollectGarbage {
    type Params = ();
    type Result = ();
    const METHOD: &'static str = "rust-analyzer/collectGarbage";
}

pub enum SyntaxTree {}

impl Request for SyntaxTree {
    type Params = SyntaxTreeParams;
    type Result = String;
    const METHOD: &'static str = "rust-analyzer/syntaxTree";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SyntaxTreeParams {
    pub text_document: TextDocumentIdentifier,
    pub range: Option<Range>,
}

pub enum ExtendSelection {}

impl Request for ExtendSelection {
    type Params = ExtendSelectionParams;
    type Result = ExtendSelectionResult;
    const METHOD: &'static str = "rust-analyzer/extendSelection";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExtendSelectionParams {
    pub text_document: TextDocumentIdentifier,
    pub selections: Vec<Range>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExtendSelectionResult {
    pub selections: Vec<Range>,
}

pub enum FindMatchingBrace {}

impl Request for FindMatchingBrace {
    type Params = FindMatchingBraceParams;
    type Result = Vec<Position>;
    const METHOD: &'static str = "rust-analyzer/findMatchingBrace";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct FindMatchingBraceParams {
    pub text_document: TextDocumentIdentifier,
    pub offsets: Vec<Position>,
}

pub enum DecorationsRequest {}

impl Request for DecorationsRequest {
    type Params = TextDocumentIdentifier;
    type Result = Vec<Decoration>;
    const METHOD: &'static str = "rust-analyzer/decorationsRequest";
}

pub enum PublishDecorations {}

impl Notification for PublishDecorations {
    type Params = PublishDecorationsParams;
    const METHOD: &'static str = "rust-analyzer/publishDecorations";
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct PublishDecorationsParams {
    #[serde(with = "url_serde")]
    pub uri: Url,
    pub decorations: Vec<Decoration>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Decoration {
    pub range: Range,
    pub tag: &'static str,
}

pub enum ParentModule {}

impl Request for ParentModule {
    type Params = TextDocumentPositionParams;
    type Result = Vec<Location>;
    const METHOD: &'static str = "rust-analyzer/parentModule";
}

pub enum JoinLines {}

impl Request for JoinLines {
    type Params = JoinLinesParams;
    type Result = SourceChange;
    const METHOD: &'static str = "rust-analyzer/joinLines";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct JoinLinesParams {
    pub text_document: TextDocumentIdentifier,
    pub range: Range,
}

pub enum OnEnter {}

impl Request for OnEnter {
    type Params = TextDocumentPositionParams;
    type Result = Option<SourceChange>;
    const METHOD: &'static str = "rust-analyzer/onEnter";
}

pub enum Runnables {}

impl Request for Runnables {
    type Params = RunnablesParams;
    type Result = Vec<Runnable>;
    const METHOD: &'static str = "rust-analyzer/runnables";
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct RunnablesParams {
    pub text_document: TextDocumentIdentifier,
    pub position: Option<Position>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Runnable {
    pub range: Range,
    pub label: String,
    pub bin: String,
    pub args: Vec<String>,
    pub env: FxHashMap<String, String>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SourceChange {
    pub label: String,
    pub workspace_edit: WorkspaceEdit,
    pub cursor_position: Option<TextDocumentPositionParams>,
}

pub enum InternalFeedback {}

impl Notification for InternalFeedback {
    const METHOD: &'static str = "internalFeedback";
    type Params = String;
}
<commit_msg>Remove InternalFeedback<commit_after>use lsp_types::{Location, Position, Range, TextDocumentIdentifier, Url};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use url_serde;

pub use lsp_types::{
    notification::*, request::*, ApplyWorkspaceEditParams, CodeActionParams, CodeLens, CodeLensParams,
    CompletionParams, CompletionResponse, DocumentOnTypeFormattingParams, DocumentSymbolParams,
    DocumentSymbolResponse, ExecuteCommandParams, Hover, InitializeResult,
    PublishDiagnosticsParams, ReferenceParams, SignatureHelp, TextDocumentEdit,
    TextDocumentPositionParams, TextEdit, WorkspaceEdit, WorkspaceSymbolParams,
    MessageType, ShowMessageParams,
};

pub enum AnalyzerStatus {}

impl Request for AnalyzerStatus {
    type Params = ();
    type Result = String;
    const METHOD: &'static str = "rust-analyzer/analyzerStatus";
}

pub enum CollectGarbage {}

impl Request for CollectGarbage {
    type Params = ();
    type Result = ();
    const METHOD: &'static str = "rust-analyzer/collectGarbage";
}

pub enum SyntaxTree {}

impl Request for SyntaxTree {
    type Params = SyntaxTreeParams;
    type Result = String;
    const METHOD: &'static str = "rust-analyzer/syntaxTree";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SyntaxTreeParams {
    pub text_document: TextDocumentIdentifier,
    pub range: Option<Range>,
}

pub enum ExtendSelection {}

impl Request for ExtendSelection {
    type Params = ExtendSelectionParams;
    type Result = ExtendSelectionResult;
    const METHOD: &'static str = "rust-analyzer/extendSelection";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExtendSelectionParams {
    pub text_document: TextDocumentIdentifier,
    pub selections: Vec<Range>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ExtendSelectionResult {
    pub selections: Vec<Range>,
}

pub enum FindMatchingBrace {}

impl Request for FindMatchingBrace {
    type Params = FindMatchingBraceParams;
    type Result = Vec<Position>;
    const METHOD: &'static str = "rust-analyzer/findMatchingBrace";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct FindMatchingBraceParams {
    pub text_document: TextDocumentIdentifier,
    pub offsets: Vec<Position>,
}

pub enum DecorationsRequest {}

impl Request for DecorationsRequest {
    type Params = TextDocumentIdentifier;
    type Result = Vec<Decoration>;
    const METHOD: &'static str = "rust-analyzer/decorationsRequest";
}

pub enum PublishDecorations {}

impl Notification for PublishDecorations {
    type Params = PublishDecorationsParams;
    const METHOD: &'static str = "rust-analyzer/publishDecorations";
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct PublishDecorationsParams {
    #[serde(with = "url_serde")]
    pub uri: Url,
    pub decorations: Vec<Decoration>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Decoration {
    pub range: Range,
    pub tag: &'static str,
}

pub enum ParentModule {}

impl Request for ParentModule {
    type Params = TextDocumentPositionParams;
    type Result = Vec<Location>;
    const METHOD: &'static str = "rust-analyzer/parentModule";
}

pub enum JoinLines {}

impl Request for JoinLines {
    type Params = JoinLinesParams;
    type Result = SourceChange;
    const METHOD: &'static str = "rust-analyzer/joinLines";
}

#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct JoinLinesParams {
    pub text_document: TextDocumentIdentifier,
    pub range: Range,
}

pub enum OnEnter {}

impl Request for OnEnter {
    type Params = TextDocumentPositionParams;
    type Result = Option<SourceChange>;
    const METHOD: &'static str = "rust-analyzer/onEnter";
}

pub enum Runnables {}

impl Request for Runnables {
    type Params = RunnablesParams;
    type Result = Vec<Runnable>;
    const METHOD: &'static str = "rust-analyzer/runnables";
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct RunnablesParams {
    pub text_document: TextDocumentIdentifier,
    pub position: Option<Position>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Runnable {
    pub range: Range,
    pub label: String,
    pub bin: String,
    pub args: Vec<String>,
    pub env: FxHashMap<String, String>,
}

#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct SourceChange {
    pub label: String,
    pub workspace_edit: WorkspaceEdit,
    pub cursor_position: Option<TextDocumentPositionParams>,
}
<|endoftext|>


<commit_before>//! Types for the *m.space* events.
//!
//! See [MSC2758] and [MSC1772].
//! [MSC2758]: https://github.com/matrix-org/matrix-doc/blob/master/proposals/2758-textual-id-grammar.md
//! [MSC1772]: https://github.com/matrix-org/matrix-doc/blob/master/proposals/1772-groups-as-rooms.md

pub mod child;
pub mod parent;
<commit_msg>events: Add missing newline<commit_after>//! Types for the *m.space* events.
//!
//! See [MSC2758] and [MSC1772].
//!
//! [MSC2758]: https://github.com/matrix-org/matrix-doc/blob/master/proposals/2758-textual-id-grammar.md
//! [MSC1772]: https://github.com/matrix-org/matrix-doc/blob/master/proposals/1772-groups-as-rooms.md

pub mod child;
pub mod parent;
<|endoftext|>