<commit_before>use core::fmt;
use core::cell::UnsafeCell;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering::{Acquire, Release, AcqRel};
use crate::task::Waker;

/// A synchronization primitive for task wakeup.
///
/// Sometimes the task interested in a given event will change over time.
/// An `AtomicWaker` can coordinate concurrent notifications with the consumer
/// potentially "updating" the underlying task to wake up. This is useful in
/// scenarios where a computation completes in another thread and wants to
/// notify the consumer, but the consumer is in the process of being migrated to
/// a new logical task.
///
/// Consumers should call `register` before checking the result of a computation
/// and producers should call `wake` after producing the computation (this
/// differs from the usual `thread::park` pattern). It is also permitted for
/// `wake` to be called **before** `register`. This results in a no-op.
///
/// A single `AtomicWaker` may be reused for any number of calls to `register` or
/// `wake`.
///
/// # Memory ordering
///
/// Calling `register` "acquires" all memory "released" by calls to `wake`
/// before the call to `register`.  Later calls to `wake` will wake the
/// registered waker (on contention this wake might be triggered in `register`).
///
/// For concurrent calls to `register` (should be avoided) the ordering is only
/// guaranteed for the winning call.
///
/// # Examples
///
/// Here is a simple example providing a `Flag` that can be signalled manually
/// when it is ready.
///
/// ```
/// use futures::future::Future;
/// use futures::task::{Context, Poll, AtomicWaker};
/// use std::sync::Arc;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::Relaxed;
/// use std::pin::Pin;
///
/// struct Inner {
///     waker: AtomicWaker,
///     set: AtomicBool,
/// }
///
/// #[derive(Clone)]
/// struct Flag(Arc<Inner>);
///
/// impl Flag {
///     pub fn new() -> Self {
///         Flag(Arc::new(Inner {
///             waker: AtomicWaker::new(),
///             set: AtomicBool::new(false),
///         }))
///     }
///
///     pub fn signal(&self) {
///         self.0.set.store(true, Relaxed);
///         self.0.waker.wake();
///     }
/// }
///
/// impl Future for Flag {
///     type Output = ();
///
///     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
///         // quick check to avoid registration if already done.
///         if self.0.set.load(Relaxed) {
///             return Poll::Ready(());
///         }
///
///         self.0.waker.register(cx.waker());
///
///         // Need to check condition **after** `register` to avoid a race
///         // condition that would result in lost notifications.
///         if self.0.set.load(Relaxed) {
///             Poll::Ready(())
///         } else {
///             Poll::Pending
///         }
///     }
/// }
/// ```
pub struct AtomicWaker {
    state: AtomicUsize,
    waker: UnsafeCell<Option<Waker>>,
}

// `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell
// stores a `Waker` value produced by calls to `register` and many threads can
// race to take the waker (to wake it) by calling `wake`.
//
// If a new `Waker` instance is produced by calling `register` before an
// existing one is consumed, then the existing one is overwritten.
//
// While `AtomicWaker` is single-producer, the implementation ensures memory
// safety. In the event of concurrent calls to `register`, there will be a
// single winner whose waker will get stored in the cell. The losers will not
// have their tasks woken. As such, callers should ensure to add synchronization
// to calls to `register`.
//
// The implementation uses a single `AtomicUsize` value to coordinate access to
// the `Waker` cell. There are two bits that are operated on independently.
// These are represented by `REGISTERING` and `WAKING`.
//
// The `REGISTERING` bit is set when a producer enters the critical section. The
// `WAKING` bit is set when a consumer enters the critical section. Neither bit
// being set is represented by `WAITING`.
//
// A thread obtains an exclusive lock on the waker cell by transitioning the
// state from `WAITING` to `REGISTERING` or `WAKING`, depending on the operation
// the thread wishes to perform. When this transition is made, it is guaranteed
// that no other thread will access the waker cell.
//
// # Registering
//
// On a call to `register`, an attempt to transition the state from WAITING to
// REGISTERING is made. On success, the caller obtains a lock on the waker cell.
//
// If the lock is obtained, then the thread sets the waker cell to the waker
// provided as an argument. Then it attempts to transition the state back from
// `REGISTERING` -> `WAITING`.
//
// If this transition is successful, then the registering process is complete
// and the next call to `wake` will observe the waker.
//
// If the transition fails, then there was a concurrent call to `wake` that was
// unable to access the waker cell (due to the registering thread holding the
// lock). To handle this, the registering thread removes the waker it just set
// from the cell and calls `wake` on it. This call to wake represents the
// attempt to wake by the other thread (that set the `WAKING` bit). The state is
// then transitioned from `REGISTERING | WAKING` back to `WAITING`.  This
// transition must succeed because, at this point, the state cannot be
// transitioned by another thread.
//
// # Waking
//
// On a call to `wake`, an attempt to transition the state from `WAITING` to
// `WAKING` is made. On success, the caller obtains a lock on the waker cell.
//
// If the lock is obtained, then the thread takes ownership of the current value
// in the waker cell, and calls `wake` on it. The state is then transitioned
// back to `WAITING`. This transition must succeed as, at this point, the state
// cannot be transitioned by another thread.
//
// If the thread is unable to obtain the lock, the `WAKING` bit is still.  This
// is because it has either been set by the current thread but the previous
// value included the `REGISTERING` bit **or** a concurrent thread is in the
// `WAKING` critical section. Either way, no action must be taken.
//
// If the current thread is the only concurrent call to `wake` and another
// thread is in the `register` critical section, when the other thread **exits**
// the `register` critical section, it will observe the `WAKING` bit and handle
// the wake itself.
//
// If another thread is in the `wake` critical section, then it will handle
// waking the task.
//
// # A potential race (is safely handled).
//
// Imagine the following situation:
//
// * Thread A obtains the `wake` lock and wakes a task.
//
// * Before thread A releases the `wake` lock, the woken task is scheduled.
//
// * Thread B attempts to wake the task. In theory this should result in the
//   task being woken, but it cannot because thread A still holds the wake lock.
//
// This case is handled by requiring users of `AtomicWaker` to call `register`
// **before** attempting to observe the application state change that resulted
// in the task being awoken. The wakers also change the application state before
// calling wake.
//
// Because of this, the waker will do one of two things.
//
// 1) Observe the application state change that Thread B is woken for. In this
//    case, it is OK for Thread B's wake to be lost.
//
// 2) Call register before attempting to observe the application state. Since
//    Thread A still holds the `wake` lock, the call to `register` will result
//    in the task waking itself and get scheduled again.

/// Idle state
const WAITING: usize = 0;

/// A new waker value is being registered with the `AtomicWaker` cell.
const REGISTERING: usize = 0b01;

/// The waker currently registered with the `AtomicWaker` cell is being woken.
const WAKING: usize = 0b10;

impl AtomicWaker {
    /// Create an `AtomicWaker`.
    pub fn new() -> AtomicWaker {
        // Make sure that task is Sync
        trait AssertSync: Sync {}
        impl AssertSync for Waker {}

        AtomicWaker {
            state: AtomicUsize::new(WAITING),
            waker: UnsafeCell::new(None),
        }
    }

    /// Registers the waker to be notified on calls to `wake`.
    ///
    /// The new task will take place of any previous tasks that were registered
    /// by previous calls to `register`. Any calls to `wake` that happen after
    /// a call to `register` (as defined by the memory ordering rules), will
    /// notify the `register` caller's task and deregister the waker from future
    /// notifications. Because of this, callers should ensure `register` gets
    /// invoked with a new `Waker` **each** time they require a wakeup.
    ///
    /// It is safe to call `register` with multiple other threads concurrently
    /// calling `wake`. This will result in the `register` caller's current
    /// task being notified once.
    ///
    /// This function is safe to call concurrently, but this is generally a bad
    /// idea. Concurrent calls to `register` will attempt to register different
    /// tasks to be notified. One of the callers will win and have its task set,
    /// but there is no guarantee as to which caller will succeed.
    ///
    /// # Examples
    ///
    /// Here is how `register` is used when implementing a flag.
    ///
    /// ```
    /// use futures::future::Future;
    /// use futures::task::{Context, Poll, AtomicWaker};
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::Relaxed;
    /// use std::pin::Pin;
    ///
    /// struct Flag {
    ///     waker: AtomicWaker,
    ///     set: AtomicBool,
    /// }
    ///
    /// impl Future for Flag {
    ///     type Output = ();
    ///
    ///     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
    ///         // Register **before** checking `set` to avoid a race condition
    ///         // that would result in lost notifications.
    ///         self.waker.register(cx.waker());
    ///
    ///         if self.set.load(Relaxed) {
    ///             Poll::Ready(())
    ///         } else {
    ///             Poll::Pending
    ///         }
    ///     }
    /// }
    /// ```
    pub fn register(&self, waker: &Waker) {
        match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) {
            WAITING => {
                unsafe {
                    // Locked acquired, update the waker cell
                    *self.waker.get() = Some(waker.clone());

                    // Release the lock. If the state transitioned to include
                    // the `WAKING` bit, this means that at least one wake has
                    // been called concurrently.
                    //
                    // Start by assuming that the state is `REGISTERING` as this
                    // is what we jut set it to.
                    let res = self.state.compare_exchange(
                        REGISTERING, WAITING, AcqRel, Acquire);

                    match res {
                        Ok(_) => {
                            // memory ordering: acquired self.state
                            // - if previous wakes went through it syncs with
                            //   their final release (`fetch_and`)
                            // - if there was no previous wake the next wake
                            //   will wake us, no sync needed.
                        }
                        Err(actual) => {
                            // This branch can only be reached if at least one
                            // concurrent thread called `wake`. In this
                            // case, `actual` **must** be `REGISTERING |
                            // `WAKING`.
                            debug_assert_eq!(actual, REGISTERING | WAKING);

                            // Take the waker to wake once the atomic operation has
                            // completed.
                            let waker = (*self.waker.get()).take().unwrap();

                            // Just swap, because no one could change state while state == `REGISTERING` | `WAKING`.
                            self.state.swap(WAITING, AcqRel);

                            // memory ordering: we acquired the state for all
                            // concurrent wakes, but future wakes might still
                            // need to wake us in case we can't make progress
                            // from the pending wakes.
                            //
                            // So we simply schedule to come back later (we could
                            // also simply leave the registration in place above).
                            waker.wake();
                        }
                    }
                }
            }
            WAKING => {
                // Currently in the process of waking the task, i.e.,
                // `wake` is currently being called on the old task handle.
                //
                // memory ordering: we acquired the state for all
                // concurrent wakes, but future wakes might still
                // need to wake us in case we can't make progress
                // from the pending wakes.
                //
                // So we simply schedule to come back later (we
                // could also spin here trying to acquire the lock
                // to register).
                waker.wake_by_ref();
            }
            state => {
                // In this case, a concurrent thread is holding the
                // "registering" lock. This probably indicates a bug in the
                // caller's code as racing to call `register` doesn't make much
                // sense.
                //
                // memory ordering: don't care. a concurrent register() is going
                // to succeed and provide proper memory ordering.
                //
                // We just want to maintain memory safety. It is ok to drop the
                // call to `register`.
                debug_assert!(
                    state == REGISTERING ||
                    state == REGISTERING | WAKING);
            }
        }
    }

    /// Calls `wake` on the last `Waker` passed to `register`.
    ///
    /// If `register` has not been called yet, then this does nothing.
    pub fn wake(&self) {
        if let Some(waker) = self.take() {
            waker.wake();
        }
    }

    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
    ///
    ///
    /// Sometimes, just waking the AtomicWaker is not fine grained enough. This allows the user
    /// to take the waker and then wake it separately, rather than performing both steps in one
    /// atomic action.
    ///
    /// If a waker has not been registered, this returns `None`.
    pub fn take(&self) -> Option<Waker> {
        // AcqRel ordering is used in order to acquire the value of the `task`
        // cell as well as to establish a `release` ordering with whatever
        // memory the `AtomicWaker` is associated with.
        match self.state.fetch_or(WAKING, AcqRel) {
            WAITING => {
                // The waking lock has been acquired.
                let waker = unsafe { (*self.waker.get()).take() };

                // Release the lock
                self.state.fetch_and(!WAKING, Release);

                waker
            }
            state => {
                // There is a concurrent thread currently updating the
                // associated task.
                //
                // Nothing more to do as the `WAKING` bit has been set. It
                // doesn't matter if there are concurrent registering threads or
                // not.
                //
                debug_assert!(
                    state == REGISTERING ||
                    state == REGISTERING | WAKING ||
                    state == WAKING);
                None
            }
        }
    }
}

impl Default for AtomicWaker {
    fn default() -> Self {
        AtomicWaker::new()
    }
}

impl fmt::Debug for AtomicWaker {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "AtomicWaker")
    }
}

unsafe impl Send for AtomicWaker {}
unsafe impl Sync for AtomicWaker {}
<commit_msg>improve comments and relax one cmpxch in AtomicWaker<commit_after>use core::fmt;
use core::cell::UnsafeCell;
use core::sync::atomic::AtomicUsize;
use core::sync::atomic::Ordering::{Acquire, Release, AcqRel};
use crate::task::Waker;

/// A synchronization primitive for task wakeup.
///
/// Sometimes the task interested in a given event will change over time.
/// An `AtomicWaker` can coordinate concurrent notifications with the consumer
/// potentially "updating" the underlying task to wake up. This is useful in
/// scenarios where a computation completes in another thread and wants to
/// notify the consumer, but the consumer is in the process of being migrated to
/// a new logical task.
///
/// Consumers should call `register` before checking the result of a computation
/// and producers should call `wake` after producing the computation (this
/// differs from the usual `thread::park` pattern). It is also permitted for
/// `wake` to be called **before** `register`. This results in a no-op.
///
/// A single `AtomicWaker` may be reused for any number of calls to `register` or
/// `wake`.
///
/// # Memory ordering
///
/// Calling `register` "acquires" all memory "released" by calls to `wake`
/// before the call to `register`.  Later calls to `wake` will wake the
/// registered waker (on contention this wake might be triggered in `register`).
///
/// For concurrent calls to `register` (should be avoided) the ordering is only
/// guaranteed for the winning call.
///
/// # Examples
///
/// Here is a simple example providing a `Flag` that can be signalled manually
/// when it is ready.
///
/// ```
/// use futures::future::Future;
/// use futures::task::{Context, Poll, AtomicWaker};
/// use std::sync::Arc;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::Relaxed;
/// use std::pin::Pin;
///
/// struct Inner {
///     waker: AtomicWaker,
///     set: AtomicBool,
/// }
///
/// #[derive(Clone)]
/// struct Flag(Arc<Inner>);
///
/// impl Flag {
///     pub fn new() -> Self {
///         Flag(Arc::new(Inner {
///             waker: AtomicWaker::new(),
///             set: AtomicBool::new(false),
///         }))
///     }
///
///     pub fn signal(&self) {
///         self.0.set.store(true, Relaxed);
///         self.0.waker.wake();
///     }
/// }
///
/// impl Future for Flag {
///     type Output = ();
///
///     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
///         // quick check to avoid registration if already done.
///         if self.0.set.load(Relaxed) {
///             return Poll::Ready(());
///         }
///
///         self.0.waker.register(cx.waker());
///
///         // Need to check condition **after** `register` to avoid a race
///         // condition that would result in lost notifications.
///         if self.0.set.load(Relaxed) {
///             Poll::Ready(())
///         } else {
///             Poll::Pending
///         }
///     }
/// }
/// ```
pub struct AtomicWaker {
    state: AtomicUsize,
    waker: UnsafeCell<Option<Waker>>,
}

// `AtomicWaker` is a multi-consumer, single-producer transfer cell. The cell
// stores a `Waker` value produced by calls to `register` and many threads can
// race to take the waker (to wake it) by calling `wake`.
//
// If a new `Waker` instance is produced by calling `register` before an
// existing one is consumed, then the existing one is overwritten.
//
// While `AtomicWaker` is single-producer, the implementation ensures memory
// safety. In the event of concurrent calls to `register`, there will be a
// single winner whose waker will get stored in the cell. The losers will not
// have their tasks woken. As such, callers should ensure to add synchronization
// to calls to `register`.
//
// The implementation uses a single `AtomicUsize` value to coordinate access to
// the `Waker` cell. There are two bits that are operated on independently.
// These are represented by `REGISTERING` and `WAKING`.
//
// The `REGISTERING` bit is set when a producer enters the critical section. The
// `WAKING` bit is set when a consumer enters the critical section. Neither bit
// being set is represented by `WAITING`.
//
// A thread obtains an exclusive lock on the waker cell by transitioning the
// state from `WAITING` to `REGISTERING` or `WAKING`, depending on the operation
// the thread wishes to perform. When this transition is made, it is guaranteed
// that no other thread will access the waker cell.
//
// # Registering
//
// On a call to `register`, an attempt to transition the state from WAITING to
// REGISTERING is made. On success, the caller obtains a lock on the waker cell.
//
// If the lock is obtained, then the thread sets the waker cell to the waker
// provided as an argument. Then it attempts to transition the state back from
// `REGISTERING` -> `WAITING`.
//
// If this transition is successful, then the registering process is complete
// and the next call to `wake` will observe the waker.
//
// If the transition fails, then there was a concurrent call to `wake` that was
// unable to access the waker cell (due to the registering thread holding the
// lock). To handle this, the registering thread removes the waker it just set
// from the cell and calls `wake` on it. This call to wake represents the
// attempt to wake by the other thread (that set the `WAKING` bit). The state is
// then transitioned from `REGISTERING | WAKING` back to `WAITING`.  This
// transition must succeed because, at this point, the state cannot be
// transitioned by another thread.
//
// # Waking
//
// On a call to `wake`, an attempt to transition the state from `WAITING` to
// `WAKING` is made. On success, the caller obtains a lock on the waker cell.
//
// If the lock is obtained, then the thread takes ownership of the current value
// in the waker cell, and calls `wake` on it. The state is then transitioned
// back to `WAITING`. This transition must succeed as, at this point, the state
// cannot be transitioned by another thread.
//
// If the thread is unable to obtain the lock, the `WAKING` bit is still.  This
// is because it has either been set by the current thread but the previous
// value included the `REGISTERING` bit **or** a concurrent thread is in the
// `WAKING` critical section. Either way, no action must be taken.
//
// If the current thread is the only concurrent call to `wake` and another
// thread is in the `register` critical section, when the other thread **exits**
// the `register` critical section, it will observe the `WAKING` bit and handle
// the wake itself.
//
// If another thread is in the `wake` critical section, then it will handle
// waking the task.
//
// # A potential race (is safely handled).
//
// Imagine the following situation:
//
// * Thread A obtains the `wake` lock and wakes a task.
//
// * Before thread A releases the `wake` lock, the woken task is scheduled.
//
// * Thread B attempts to wake the task. In theory this should result in the
//   task being woken, but it cannot because thread A still holds the wake lock.
//
// This case is handled by requiring users of `AtomicWaker` to call `register`
// **before** attempting to observe the application state change that resulted
// in the task being awoken. The wakers also change the application state before
// calling wake.
//
// Because of this, the waker will do one of two things.
//
// 1) Observe the application state change that Thread B is woken for. In this
//    case, it is OK for Thread B's wake to be lost.
//
// 2) Call register before attempting to observe the application state. Since
//    Thread A still holds the `wake` lock, the call to `register` will result
//    in the task waking itself and get scheduled again.

/// Idle state
const WAITING: usize = 0;

/// A new waker value is being registered with the `AtomicWaker` cell.
const REGISTERING: usize = 0b01;

/// The waker currently registered with the `AtomicWaker` cell is being woken.
const WAKING: usize = 0b10;

impl AtomicWaker {
    /// Create an `AtomicWaker`.
    pub fn new() -> AtomicWaker {
        // Make sure that task is Sync
        trait AssertSync: Sync {}
        impl AssertSync for Waker {}

        AtomicWaker {
            state: AtomicUsize::new(WAITING),
            waker: UnsafeCell::new(None),
        }
    }

    /// Registers the waker to be notified on calls to `wake`.
    ///
    /// The new task will take place of any previous tasks that were registered
    /// by previous calls to `register`. Any calls to `wake` that happen after
    /// a call to `register` (as defined by the memory ordering rules), will
    /// notify the `register` caller's task and deregister the waker from future
    /// notifications. Because of this, callers should ensure `register` gets
    /// invoked with a new `Waker` **each** time they require a wakeup.
    ///
    /// It is safe to call `register` with multiple other threads concurrently
    /// calling `wake`. This will result in the `register` caller's current
    /// task being notified once.
    ///
    /// This function is safe to call concurrently, but this is generally a bad
    /// idea. Concurrent calls to `register` will attempt to register different
    /// tasks to be notified. One of the callers will win and have its task set,
    /// but there is no guarantee as to which caller will succeed.
    ///
    /// # Examples
    ///
    /// Here is how `register` is used when implementing a flag.
    ///
    /// ```
    /// use futures::future::Future;
    /// use futures::task::{Context, Poll, AtomicWaker};
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::Relaxed;
    /// use std::pin::Pin;
    ///
    /// struct Flag {
    ///     waker: AtomicWaker,
    ///     set: AtomicBool,
    /// }
    ///
    /// impl Future for Flag {
    ///     type Output = ();
    ///
    ///     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
    ///         // Register **before** checking `set` to avoid a race condition
    ///         // that would result in lost notifications.
    ///         self.waker.register(cx.waker());
    ///
    ///         if self.set.load(Relaxed) {
    ///             Poll::Ready(())
    ///         } else {
    ///             Poll::Pending
    ///         }
    ///     }
    /// }
    /// ```
    pub fn register(&self, waker: &Waker) {
        match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) {
            WAITING => {
                unsafe {
                    // Locked acquired, update the waker cell
                    *self.waker.get() = Some(waker.clone());

                    // Release the lock. If the state transitioned to include
                    // the `WAKING` bit, this means that at least one wake has
                    // been called concurrently.
                    //
                    // Start by assuming that the state is `REGISTERING` as this
                    // is what we just set it to. If this holds, we know that no
                    // other writes were performed in the meantime, so there is
                    // nothing to acquire, but release the update. In case of
                    // concurrent wakers, we need to acquire their releases.
                    let res = self.state.compare_exchange(
                        REGISTERING, WAITING, Release, Acquire);

                    match res {
                        Ok(_) => {
                            // memory ordering: acquired self.state during CAS
                            // - if previous wakes went through it syncs with
                            //   their final release (`fetch_and`)
                            // - if there was no previous wake the next wake
                            //   will wake us, no sync needed.
                        }
                        Err(actual) => {
                            // This branch can only be reached if at least one
                            // concurrent thread called `wake`. In this
                            // case, `actual` **must** be `REGISTERING |
                            // `WAKING`.
                            debug_assert_eq!(actual, REGISTERING | WAKING);

                            // Take the waker to wake once the atomic operation has
                            // completed.
                            let waker = (*self.waker.get()).take().unwrap();

                            // We need to return to WAITING state (clear our lock and
                            // concurrent WAKING flag). This needs to acquire all
                            // WAKING fetch_or releases and it needs to release our
                            // update to self.waker, so we need a `swap` operation.
                            self.state.swap(WAITING, AcqRel);

                            // memory ordering: we acquired the state for all
                            // concurrent wakes, but future wakes might still
                            // need to wake us in case we can't make progress
                            // from the pending wakes.
                            //
                            // So we simply schedule to come back later (we could
                            // also simply leave the registration in place above).
                            waker.wake();
                        }
                    }
                }
            }
            WAKING => {
                // Currently in the process of waking the task, i.e.,
                // `wake` is currently being called on the old task handle.
                //
                // memory ordering: we acquired the state for all
                // concurrent wakes, but future wakes might still
                // need to wake us in case we can't make progress
                // from the pending wakes.
                //
                // So we simply schedule to come back later (we
                // could also spin here trying to acquire the lock
                // to register).
                waker.wake_by_ref();
            }
            state => {
                // In this case, a concurrent thread is holding the
                // "registering" lock. This probably indicates a bug in the
                // caller's code as racing to call `register` doesn't make much
                // sense.
                //
                // memory ordering: don't care. a concurrent register() is going
                // to succeed and provide proper memory ordering.
                //
                // We just want to maintain memory safety. It is ok to drop the
                // call to `register`.
                debug_assert!(
                    state == REGISTERING ||
                    state == REGISTERING | WAKING);
            }
        }
    }

    /// Calls `wake` on the last `Waker` passed to `register`.
    ///
    /// If `register` has not been called yet, then this does nothing.
    pub fn wake(&self) {
        if let Some(waker) = self.take() {
            waker.wake();
        }
    }

    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
    ///
    ///
    /// Sometimes, just waking the AtomicWaker is not fine grained enough. This allows the user
    /// to take the waker and then wake it separately, rather than performing both steps in one
    /// atomic action.
    ///
    /// If a waker has not been registered, this returns `None`.
    pub fn take(&self) -> Option<Waker> {
        // AcqRel ordering is used in order to acquire the value of the `task`
        // cell as well as to establish a `release` ordering with whatever
        // memory the `AtomicWaker` is associated with.
        match self.state.fetch_or(WAKING, AcqRel) {
            WAITING => {
                // The waking lock has been acquired.
                let waker = unsafe { (*self.waker.get()).take() };

                // Release the lock
                self.state.fetch_and(!WAKING, Release);

                waker
            }
            state => {
                // There is a concurrent thread currently updating the
                // associated task.
                //
                // Nothing more to do as the `WAKING` bit has been set. It
                // doesn't matter if there are concurrent registering threads or
                // not.
                //
                debug_assert!(
                    state == REGISTERING ||
                    state == REGISTERING | WAKING ||
                    state == WAKING);
                None
            }
        }
    }
}

impl Default for AtomicWaker {
    fn default() -> Self {
        AtomicWaker::new()
    }
}

impl fmt::Debug for AtomicWaker {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "AtomicWaker")
    }
}

unsafe impl Send for AtomicWaker {}
unsafe impl Sync for AtomicWaker {}
<|endoftext|>


<commit_before>#![feature(test)]
#![deny(warnings)]

extern crate test;

use std::net::SocketAddr;

use futures_util::future::join_all;
use tokio::runtime::current_thread::Runtime;

use hyper::{Body, Method, Request, Response, Server};
use hyper::client::HttpConnector;

// HTTP1

#[bench]
fn http1_get(b: &mut test::Bencher) {
    opts()
        .bench(b)
}

#[bench]
fn http1_post(b: &mut test::Bencher) {
    opts()
        .method(Method::POST)
        .request_body(b"foo bar baz quux")
        .bench(b)
}

#[bench]
fn http1_body_both_100kb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 100];
    opts()
        .method(Method::POST)
        .request_body(body)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_body_both_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .method(Method::POST)
        .request_body(body)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_empty(b: &mut test::Bencher) {
    opts()
        .parallel(10)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_req_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .parallel(10)
        .method(Method::POST)
        .request_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_res_1mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 1];
    opts()
        .parallel(10)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_res_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .parallel(10)
        .response_body(body)
        .bench(b)
}

// HTTP2

const HTTP2_MAX_WINDOW: u32 = std::u32::MAX >> 1;

#[bench]
fn http2_get(b: &mut test::Bencher) {
    opts()
        .http2()
        .bench(b)
}

#[bench]
fn http2_post(b: &mut test::Bencher) {
    opts()
        .http2()
        .method(Method::POST)
        .request_body(b"foo bar baz quux")
        .bench(b)
}

#[bench]
fn http2_req_100kb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 100];
    opts()
        .http2()
        .method(Method::POST)
        .request_body(body)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_empty(b: &mut test::Bencher) {
    opts()
        .http2()
        .parallel(10)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_req_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .http2()
        .parallel(10)
        .method(Method::POST)
        .request_body(body)
        //.http2_stream_window(HTTP2_MAX_WINDOW)
        //.http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_res_1mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 1];
    opts()
        .http2()
        .parallel(10)
        .response_body(body)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_res_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .http2()
        .parallel(10)
        .response_body(body)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

// ==== Benchmark Options =====

struct Opts {
    http2: bool,
    http2_stream_window: Option<u32>,
    http2_conn_window: Option<u32>,
    parallel_cnt: u32,
    request_method: Method,
    request_body: Option<&'static [u8]>,
    response_body: &'static [u8],
}

fn opts() -> Opts {
    Opts {
        http2: false,
        http2_stream_window: None,
        http2_conn_window: None,
        parallel_cnt: 1,
        request_method: Method::GET,
        request_body: None,
        response_body: b"Hello",
    }
}

impl Opts {
    fn http2(mut self) -> Self {
        self.http2 = true;
        self
    }

    fn http2_stream_window(mut self, sz: impl Into<Option<u32>>) -> Self {
        self.http2_stream_window = sz.into();
        self
    }

    fn http2_conn_window(mut self, sz: impl Into<Option<u32>>) -> Self {
        self.http2_conn_window = sz.into();
        self
    }

    fn method(mut self, m: Method) -> Self {
        self.request_method = m;
        self
    }

    fn request_body(mut self, body: &'static [u8]) -> Self {
        self.request_body = Some(body);
        self
    }

    fn response_body(mut self, body: &'static [u8]) -> Self {
        self.response_body = body;
        self
    }

    fn parallel(mut self, cnt: u32) -> Self {
        assert!(cnt > 0, "parallel count must be larger than 0");
        self.parallel_cnt = cnt;
        self
    }

    fn bench(self, b: &mut test::Bencher) {
        let _ = pretty_env_logger::try_init();
        // Create a runtime of current thread.
        let mut rt = Runtime::new().unwrap();

        b.bytes = self.response_body.len() as u64 + self.request_body.map(|b| b.len()).unwrap_or(0) as u64;

        let addr = spawn_server(&mut rt, &self);

        let connector = HttpConnector::new();
        let client = hyper::Client::builder()
            .http2_only(self.http2)
            .http2_initial_stream_window_size(self.http2_stream_window)
            .http2_initial_connection_window_size(self.http2_conn_window)
            .build::<_, Body>(connector);

        let url: hyper::Uri = format!("http://{}/hello", addr).parse().unwrap();

        let make_request = || {
            let body = self
                .request_body
                .map(Body::from)
                .unwrap_or_else(|| Body::empty());
            let mut req = Request::new(body);
            *req.method_mut() = self.request_method.clone();
            *req.uri_mut() = url.clone();
            req
        };

        let send_request = |req: Request<Body>| {
            let fut = client.request(req);
            async {
                let res = fut.await.expect("client wait");
                let mut body = res.into_body();
                while let Some(_chunk) = body.next().await {}
            }
        };

        if self.parallel_cnt == 1 {
            b.iter(|| {
                let req = make_request();
                rt.block_on(send_request(req));
            });

        } else {
            b.iter(|| {
                let futs = (0..self.parallel_cnt)
                    .map(|_| {
                        let req = make_request();
                        send_request(req)
                    });
                // Await all spawned futures becoming completed.
                rt.block_on(join_all(futs));
            });
        }
    }
}

fn spawn_server(rt: &mut Runtime, opts: &Opts) -> SocketAddr {
    use hyper::service::{make_service_fn, service_fn};
    let addr = "127.0.0.1:0".parse().unwrap();

    let body = opts.response_body;
    let srv = Server::bind(&addr)
        .http2_only(opts.http2)
        .http2_initial_stream_window_size(opts.http2_stream_window)
        .http2_initial_connection_window_size(opts.http2_conn_window)
        .serve(make_service_fn( move |_| async move {
            Ok::<_, hyper::Error>(service_fn(move |req: Request<Body>| async move {
                let mut req_body = req.into_body();
                while let Some(_chunk) = req_body.next().await {}
                Ok::<_, hyper::Error>(Response::new(Body::from(body)))
            }))
        }));
    let addr = srv.local_addr();
    rt.spawn(async {
        if let Err(err) = srv.await {
            panic!("server error: {}", err);
        }
    });
    return addr
}
<commit_msg>test(benches): add end-to-end chunked benchmarks<commit_after>#![feature(test)]
#![deny(warnings)]

extern crate test;

use std::net::SocketAddr;

use futures_util::future::join_all;
use tokio::runtime::current_thread::Runtime;

use hyper::{Body, Method, Request, Response, Server};
use hyper::client::HttpConnector;

// HTTP1

#[bench]
fn http1_get(b: &mut test::Bencher) {
    opts()
        .bench(b)
}

#[bench]
fn http1_post(b: &mut test::Bencher) {
    opts()
        .method(Method::POST)
        .request_body(b"foo bar baz quux")
        .bench(b)
}

#[bench]
fn http1_body_both_100kb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 100];
    opts()
        .method(Method::POST)
        .request_body(body)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_body_both_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .method(Method::POST)
        .request_body(body)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_empty(b: &mut test::Bencher) {
    opts()
        .parallel(10)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_req_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .parallel(10)
        .method(Method::POST)
        .request_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_req_10kb_100_chunks(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 10];
    opts()
        .parallel(10)
        .method(Method::POST)
        .request_chunks(body, 100)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_res_1mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 1];
    opts()
        .parallel(10)
        .response_body(body)
        .bench(b)
}

#[bench]
fn http1_parallel_x10_res_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .parallel(10)
        .response_body(body)
        .bench(b)
}

// HTTP2

const HTTP2_MAX_WINDOW: u32 = std::u32::MAX >> 1;

#[bench]
fn http2_get(b: &mut test::Bencher) {
    opts()
        .http2()
        .bench(b)
}

#[bench]
fn http2_post(b: &mut test::Bencher) {
    opts()
        .http2()
        .method(Method::POST)
        .request_body(b"foo bar baz quux")
        .bench(b)
}

#[bench]
fn http2_req_100kb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 100];
    opts()
        .http2()
        .method(Method::POST)
        .request_body(body)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_empty(b: &mut test::Bencher) {
    opts()
        .http2()
        .parallel(10)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_req_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .http2()
        .parallel(10)
        .method(Method::POST)
        .request_body(body)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_req_10kb_100_chunks(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 10];
    opts()
        .http2()
        .parallel(10)
        .method(Method::POST)
        .request_chunks(body, 100)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_res_1mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 1];
    opts()
        .http2()
        .parallel(10)
        .response_body(body)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

#[bench]
fn http2_parallel_x10_res_10mb(b: &mut test::Bencher) {
    let body = &[b'x'; 1024 * 1024 * 10];
    opts()
        .http2()
        .parallel(10)
        .response_body(body)
        .http2_stream_window(HTTP2_MAX_WINDOW)
        .http2_conn_window(HTTP2_MAX_WINDOW)
        .bench(b)
}

// ==== Benchmark Options =====

struct Opts {
    http2: bool,
    http2_stream_window: Option<u32>,
    http2_conn_window: Option<u32>,
    parallel_cnt: u32,
    request_method: Method,
    request_body: Option<&'static [u8]>,
    request_chunks: usize,
    response_body: &'static [u8],
}

fn opts() -> Opts {
    Opts {
        http2: false,
        http2_stream_window: None,
        http2_conn_window: None,
        parallel_cnt: 1,
        request_method: Method::GET,
        request_body: None,
        request_chunks: 0,
        response_body: b"Hello",
    }
}

impl Opts {
    fn http2(mut self) -> Self {
        self.http2 = true;
        self
    }

    fn http2_stream_window(mut self, sz: impl Into<Option<u32>>) -> Self {
        self.http2_stream_window = sz.into();
        self
    }

    fn http2_conn_window(mut self, sz: impl Into<Option<u32>>) -> Self {
        self.http2_conn_window = sz.into();
        self
    }

    fn method(mut self, m: Method) -> Self {
        self.request_method = m;
        self
    }

    fn request_body(mut self, body: &'static [u8]) -> Self {
        self.request_body = Some(body);
        self
    }

    fn request_chunks(mut self, chunk: &'static [u8], cnt: usize) -> Self {
        assert!(cnt > 0);
        self.request_body = Some(chunk);
        self.request_chunks = cnt;
        self
    }

    fn response_body(mut self, body: &'static [u8]) -> Self {
        self.response_body = body;
        self
    }

    fn parallel(mut self, cnt: u32) -> Self {
        assert!(cnt > 0, "parallel count must be larger than 0");
        self.parallel_cnt = cnt;
        self
    }

    fn bench(self, b: &mut test::Bencher) {
        let _ = pretty_env_logger::try_init();
        // Create a runtime of current thread.
        let mut rt = Runtime::new().unwrap();
        let exec = rt.handle();

        let req_len = self.request_body.map(|b| b.len()).unwrap_or(0) as u64;
        let req_len = if self.request_chunks > 0 {
            req_len * self.request_chunks as u64
        } else {
            req_len
        };
        let bytes_per_iter = (req_len + self.response_body.len() as u64) * self.parallel_cnt as u64;
        b.bytes = bytes_per_iter;

        let addr = spawn_server(&mut rt, &self);

        let connector = HttpConnector::new();
        let client = hyper::Client::builder()
            .http2_only(self.http2)
            .http2_initial_stream_window_size(self.http2_stream_window)
            .http2_initial_connection_window_size(self.http2_conn_window)
            .build::<_, Body>(connector);

        let url: hyper::Uri = format!("http://{}/hello", addr).parse().unwrap();

        let make_request = || {
            let chunk_cnt = self.request_chunks;
            let body = if chunk_cnt > 0 {

                let (mut tx, body) = Body::channel();
                let chunk = self.request_body.expect("request_chunks means request_body");
                exec.spawn(async move {
                    for _ in 0..chunk_cnt {
                        tx.send_data(chunk.into()).await.expect("send_data");
                    }
                }).expect("body tx spawn");
                body
            } else {
                self
                    .request_body
                    .map(Body::from)
                    .unwrap_or_else(|| Body::empty())
            };
            let mut req = Request::new(body);
            *req.method_mut() = self.request_method.clone();
            *req.uri_mut() = url.clone();
            req
        };

        let send_request = |req: Request<Body>| {
            let fut = client.request(req);
            async {
                let res = fut.await.expect("client wait");
                let mut body = res.into_body();
                while let Some(_chunk) = body.next().await {}
            }
        };

        if self.parallel_cnt == 1 {
            b.iter(|| {
                let req = make_request();
                rt.block_on(send_request(req));
            });

        } else {
            b.iter(|| {
                let futs = (0..self.parallel_cnt)
                    .map(|_| {
                        let req = make_request();
                        send_request(req)
                    });
                // Await all spawned futures becoming completed.
                rt.block_on(join_all(futs));
            });
        }
    }
}

fn spawn_server(rt: &mut Runtime, opts: &Opts) -> SocketAddr {
    use hyper::service::{make_service_fn, service_fn};
    let addr = "127.0.0.1:0".parse().unwrap();

    let body = opts.response_body;
    let srv = Server::bind(&addr)
        .http2_only(opts.http2)
        .http2_initial_stream_window_size(opts.http2_stream_window)
        .http2_initial_connection_window_size(opts.http2_conn_window)
        .serve(make_service_fn( move |_| async move {
            Ok::<_, hyper::Error>(service_fn(move |req: Request<Body>| async move {
                let mut req_body = req.into_body();
                while let Some(_chunk) = req_body.next().await {}
                Ok::<_, hyper::Error>(Response::new(Body::from(body)))
            }))
        }));
    let addr = srv.local_addr();
    rt.spawn(async {
        if let Err(err) = srv.await {
            panic!("server error: {}", err);
        }
    });
    return addr
}
<|endoftext|>


<commit_before>#[derive(Clone, Copy)]
pub struct True;
#[derive(Clone, Copy)]
pub struct False;

const TRUE: True = True;
const FALSE: False = False;

impl True {
    pub const fn not<'a>(&'a self) -> &'a False { &FALSE }
    pub const fn and<'a, T>(&'a self, other: &'a T) -> &'a T { other }
    pub const fn or<'a, T>(&'a self, _: &'a T) -> &'a True { &TRUE }
}

impl False {
    pub const fn not<'a>(&'a self) -> &'a True { &TRUE }
    pub const fn and<'a, T>(&'a self, _: &'a T) -> &'a False { &FALSE }
    pub const fn or<'a, T>(&'a self, other: &'a T) -> &'a T { other }
}

pub trait ToBool: Sized {
    type Bool: Sized;
    const TO_BOOL: Self::Bool;
}

impl ToBool for [(); 0] {
    type Bool = False;
    const TO_BOOL: Self::Bool = False;
}

impl ToBool for [(); 1] {
    type Bool = True;
    const TO_BOOL: Self::Bool = True;
}

/// Converts a `const bool` to a type-level boolean.
#[doc(hidden)]
#[macro_export]
macro_rules! _to_bool {
    ($x:expr) => {
        {
            const X: bool = $x;
            <[(); X as usize] as $crate::_bool::ToBool>::TO_BOOL
        }
    };
}
<commit_msg>Remove `TRUE` and `FALSE` constants<commit_after>#[derive(Clone, Copy)]
pub struct True;
#[derive(Clone, Copy)]
pub struct False;

impl True {
    pub const fn not<'a>(&'a self) -> &'a False { &False }
    pub const fn and<'a, T>(&'a self, other: &'a T) -> &'a T { other }
    pub const fn or<'a, T>(&'a self, _: &'a T) -> &'a True { &True }
}

impl False {
    pub const fn not<'a>(&'a self) -> &'a True { &True }
    pub const fn and<'a, T>(&'a self, _: &'a T) -> &'a False { &False }
    pub const fn or<'a, T>(&'a self, other: &'a T) -> &'a T { other }
}

pub trait ToBool: Sized {
    type Bool: Sized;
    const TO_BOOL: Self::Bool;
}

impl ToBool for [(); 0] {
    type Bool = False;
    const TO_BOOL: Self::Bool = False;
}

impl ToBool for [(); 1] {
    type Bool = True;
    const TO_BOOL: Self::Bool = True;
}

/// Converts a `const bool` to a type-level boolean.
#[doc(hidden)]
#[macro_export]
macro_rules! _to_bool {
    ($x:expr) => {
        {
            const X: bool = $x;
            <[(); X as usize] as $crate::_bool::ToBool>::TO_BOOL
        }
    };
}
<|endoftext|>


<commit_before>///! Busy handler (when the database is locked)
use std::mem;
use std::os::raw::{c_int, c_void};
use std::panic::catch_unwind;
use std::ptr;
use std::time::Duration;

use crate::ffi;
use crate::{Connection, InnerConnection, Result};

impl Connection {
    /// Set a busy handler that sleeps for a specified amount of time when a
    /// table is locked. The handler will sleep multiple times until at
    /// least "ms" milliseconds of sleeping have accumulated.
    ///
    /// Calling this routine with an argument equal to zero turns off all busy
    /// handlers.
    ///
    /// There can only be a single busy handler for a particular database
    /// connection at any given moment. If another busy handler was defined
    /// (using `busy_handler`) prior to calling this routine, that other
    /// busy handler is cleared.
    pub fn busy_timeout(&self, timeout: Duration) -> Result<()> {
        let ms = timeout
            .as_secs()
            .checked_mul(1000)
            .and_then(|t| t.checked_add(timeout.subsec_millis().into()))
            .expect("too big");
        self.db.borrow_mut().busy_timeout(ms as i32)
    }

    /// Register a callback to handle `SQLITE_BUSY` errors.
    ///
    /// If the busy callback is `None`, then `SQLITE_BUSY is returned
    /// immediately upon encountering the lock.` The argument to the busy
    /// handler callback is the number of times that the
    /// busy handler has been invoked previously for the
    /// same locking event. If the busy callback returns `false`, then no
    /// additional attempts are made to access the
    /// database and `SQLITE_BUSY` is returned to the
    /// application. If the callback returns `true`, then another attempt
    /// is made to access the database and the cycle repeats.
    ///
    /// There can only be a single busy handler defined for each database
    /// connection. Setting a new busy handler clears any previously set
    /// handler. Note that calling `busy_timeout()` or evaluating `PRAGMA
    /// busy_timeout=N` will change the busy handler and thus
    /// clear any previously set busy handler.
    pub fn busy_handler(&self, callback: Option<fn(i32) -> bool>) -> Result<()> {
        unsafe extern "C" fn busy_handler_callback(p_arg: *mut c_void, count: c_int) -> c_int {
            let handler_fn: fn(i32) -> bool = mem::transmute(p_arg);
            if let Ok(true) = catch_unwind(|| handler_fn(count)) {
                1
            } else {
                0
            }
        }
        let mut c = self.db.borrow_mut();
        let r = match callback {
            Some(f) => unsafe {
                ffi::sqlite3_busy_handler(c.db(), Some(busy_handler_callback), mem::transmute(f))
            },
            None => unsafe { ffi::sqlite3_busy_handler(c.db(), None, ptr::null_mut()) },
        };
        c.decode_result(r)
    }
}

impl InnerConnection {
    fn busy_timeout(&mut self, timeout: c_int) -> Result<()> {
        let r = unsafe { ffi::sqlite3_busy_timeout(self.db, timeout) };
        self.decode_result(r)
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::mpsc::sync_channel;
    use std::thread;
    use std::time::Duration;
    use tempfile;

    use crate::{Connection, Error, ErrorCode, Result, TransactionBehavior, NO_PARAMS};

    #[test]
    fn test_default_busy() {
        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let mut db1 = Connection::open(&path).unwrap();
        let tx1 = db1
            .transaction_with_behavior(TransactionBehavior::Exclusive)
            .unwrap();
        let db2 = Connection::open(&path).unwrap();
        let r: Result<()> = db2.query_row("PRAGMA schema_version", NO_PARAMS, |_| unreachable!());
        match r.unwrap_err() {
            Error::SqliteFailure(err, _) => {
                assert_eq!(err.code, ErrorCode::DatabaseBusy);
            }
            err => panic!("Unexpected error {}", err),
        }
        tx1.rollback().unwrap();
    }

    #[test]
    #[ignore] // FIXME: unstable
    fn test_busy_timeout() {
        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let db2 = Connection::open(&path).unwrap();
        db2.busy_timeout(Duration::from_secs(1)).unwrap();

        let (rx, tx) = sync_channel(0);
        let child = thread::spawn(move || {
            let mut db1 = Connection::open(&path).unwrap();
            let tx1 = db1
                .transaction_with_behavior(TransactionBehavior::Exclusive)
                .unwrap();
            rx.send(1).unwrap();
            thread::sleep(Duration::from_millis(100));
            tx1.rollback().unwrap();
        });

        assert_eq!(tx.recv().unwrap(), 1);
        let _ = db2
            .query_row("PRAGMA schema_version", NO_PARAMS, |row| {
                row.get::<_, i32>(0)
            })
            .expect("unexpected error");

        child.join().unwrap();
    }

    #[test]
    #[ignore] // FIXME: unstable
    fn test_busy_handler() {
        lazy_static::lazy_static! {
            static ref CALLED: AtomicBool = AtomicBool::new(false);
        }
        fn busy_handler(_: i32) -> bool {
            CALLED.store(true, Ordering::Relaxed);
            thread::sleep(Duration::from_millis(100));
            true
        }

        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let db2 = Connection::open(&path).unwrap();
        db2.busy_handler(Some(busy_handler)).unwrap();

        let (rx, tx) = sync_channel(0);
        let child = thread::spawn(move || {
            let mut db1 = Connection::open(&path).unwrap();
            let tx1 = db1
                .transaction_with_behavior(TransactionBehavior::Exclusive)
                .unwrap();
            rx.send(1).unwrap();
            thread::sleep(Duration::from_millis(100));
            tx1.rollback().unwrap();
        });

        assert_eq!(tx.recv().unwrap(), 1);
        let _ = db2
            .query_row("PRAGMA schema_version", NO_PARAMS, |row| {
                row.get::<_, i32>(0)
            })
            .expect("unexpected error");
        assert_eq!(CALLED.load(Ordering::Relaxed), true);

        child.join().unwrap();
    }
}
<commit_msg>Fix i32 overflow in Connection::busy_timeout<commit_after>///! Busy handler (when the database is locked)
use std::convert::TryInto;
use std::mem;
use std::os::raw::{c_int, c_void};
use std::panic::catch_unwind;
use std::ptr;
use std::time::Duration;

use crate::ffi;
use crate::{Connection, InnerConnection, Result};

impl Connection {
    /// Set a busy handler that sleeps for a specified amount of time when a
    /// table is locked. The handler will sleep multiple times until at
    /// least "ms" milliseconds of sleeping have accumulated.
    ///
    /// Calling this routine with an argument equal to zero turns off all busy
    /// handlers.
    ///
    /// There can only be a single busy handler for a particular database
    /// connection at any given moment. If another busy handler was defined
    /// (using `busy_handler`) prior to calling this routine, that other
    /// busy handler is cleared.
    pub fn busy_timeout(&self, timeout: Duration) -> Result<()> {
        let ms: i32 = timeout
            .as_secs()
            .checked_mul(1000)
            .and_then(|t| t.checked_add(timeout.subsec_millis().into()))
            .and_then(|t| t.try_into().ok())
            .expect("too big");
        self.db.borrow_mut().busy_timeout(ms)
    }

    /// Register a callback to handle `SQLITE_BUSY` errors.
    ///
    /// If the busy callback is `None`, then `SQLITE_BUSY is returned
    /// immediately upon encountering the lock.` The argument to the busy
    /// handler callback is the number of times that the
    /// busy handler has been invoked previously for the
    /// same locking event. If the busy callback returns `false`, then no
    /// additional attempts are made to access the
    /// database and `SQLITE_BUSY` is returned to the
    /// application. If the callback returns `true`, then another attempt
    /// is made to access the database and the cycle repeats.
    ///
    /// There can only be a single busy handler defined for each database
    /// connection. Setting a new busy handler clears any previously set
    /// handler. Note that calling `busy_timeout()` or evaluating `PRAGMA
    /// busy_timeout=N` will change the busy handler and thus
    /// clear any previously set busy handler.
    pub fn busy_handler(&self, callback: Option<fn(i32) -> bool>) -> Result<()> {
        unsafe extern "C" fn busy_handler_callback(p_arg: *mut c_void, count: c_int) -> c_int {
            let handler_fn: fn(i32) -> bool = mem::transmute(p_arg);
            if let Ok(true) = catch_unwind(|| handler_fn(count)) {
                1
            } else {
                0
            }
        }
        let mut c = self.db.borrow_mut();
        let r = match callback {
            Some(f) => unsafe {
                ffi::sqlite3_busy_handler(c.db(), Some(busy_handler_callback), mem::transmute(f))
            },
            None => unsafe { ffi::sqlite3_busy_handler(c.db(), None, ptr::null_mut()) },
        };
        c.decode_result(r)
    }
}

impl InnerConnection {
    fn busy_timeout(&mut self, timeout: c_int) -> Result<()> {
        let r = unsafe { ffi::sqlite3_busy_timeout(self.db, timeout) };
        self.decode_result(r)
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::mpsc::sync_channel;
    use std::thread;
    use std::time::Duration;
    use tempfile;

    use crate::{Connection, Error, ErrorCode, Result, TransactionBehavior, NO_PARAMS};

    #[test]
    fn test_default_busy() {
        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let mut db1 = Connection::open(&path).unwrap();
        let tx1 = db1
            .transaction_with_behavior(TransactionBehavior::Exclusive)
            .unwrap();
        let db2 = Connection::open(&path).unwrap();
        let r: Result<()> = db2.query_row("PRAGMA schema_version", NO_PARAMS, |_| unreachable!());
        match r.unwrap_err() {
            Error::SqliteFailure(err, _) => {
                assert_eq!(err.code, ErrorCode::DatabaseBusy);
            }
            err => panic!("Unexpected error {}", err),
        }
        tx1.rollback().unwrap();
    }

    #[test]
    #[ignore] // FIXME: unstable
    fn test_busy_timeout() {
        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let db2 = Connection::open(&path).unwrap();
        db2.busy_timeout(Duration::from_secs(1)).unwrap();

        let (rx, tx) = sync_channel(0);
        let child = thread::spawn(move || {
            let mut db1 = Connection::open(&path).unwrap();
            let tx1 = db1
                .transaction_with_behavior(TransactionBehavior::Exclusive)
                .unwrap();
            rx.send(1).unwrap();
            thread::sleep(Duration::from_millis(100));
            tx1.rollback().unwrap();
        });

        assert_eq!(tx.recv().unwrap(), 1);
        let _ = db2
            .query_row("PRAGMA schema_version", NO_PARAMS, |row| {
                row.get::<_, i32>(0)
            })
            .expect("unexpected error");

        child.join().unwrap();
    }

    #[test]
    #[ignore] // FIXME: unstable
    fn test_busy_handler() {
        lazy_static::lazy_static! {
            static ref CALLED: AtomicBool = AtomicBool::new(false);
        }
        fn busy_handler(_: i32) -> bool {
            CALLED.store(true, Ordering::Relaxed);
            thread::sleep(Duration::from_millis(100));
            true
        }

        let temp_dir = tempfile::tempdir().unwrap();
        let path = temp_dir.path().join("test.db3");

        let db2 = Connection::open(&path).unwrap();
        db2.busy_handler(Some(busy_handler)).unwrap();

        let (rx, tx) = sync_channel(0);
        let child = thread::spawn(move || {
            let mut db1 = Connection::open(&path).unwrap();
            let tx1 = db1
                .transaction_with_behavior(TransactionBehavior::Exclusive)
                .unwrap();
            rx.send(1).unwrap();
            thread::sleep(Duration::from_millis(100));
            tx1.rollback().unwrap();
        });

        assert_eq!(tx.recv().unwrap(), 1);
        let _ = db2
            .query_row("PRAGMA schema_version", NO_PARAMS, |row| {
                row.get::<_, i32>(0)
            })
            .expect("unexpected error");
        assert_eq!(CALLED.load(Ordering::Relaxed), true);

        child.join().unwrap();
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::time::{Duration, SystemTime};

use serde::{Deserialize, Serialize};
use serde_json;
use thiserror::Error;

use crate::feed;

#[derive(Error, Debug)]
pub enum DataError {
    #[error("io error")]
    Io(#[from] std::io::Error),
    #[error("json error")]
    Json(#[from] serde_json::Error),
}

fn gen_hash<T: Hash>(t: &T) -> u64 {
    let mut hasher = std::collections::hash_map::DefaultHasher::default();
    t.hash(&mut hasher);
    hasher.finish()
}

type FeedID = u64;
type SubscriberID = i64;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Feed {
    pub link: String,
    pub title: String,
    pub down_time: Option<SystemTime>,
    pub subscribers: HashSet<SubscriberID>,
    pub ttl: Option<u32>,
    hash_list: Vec<u64>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Hub {
    pub callback: String,
    pub secret: String,
}

#[derive(Debug)]
pub struct Database {
    path: PathBuf,
    feeds: HashMap<FeedID, Feed>,
    subscribers: HashMap<SubscriberID, HashSet<FeedID>>,
}

impl Database {
    pub fn create(path: PathBuf) -> Result<Database, DataError> {
        let feeds: HashMap<FeedID, Feed> = HashMap::new();
        let subscribers: HashMap<SubscriberID, HashSet<FeedID>> = HashMap::new();
        let result = Database {
            path: path,
            feeds: feeds,
            subscribers: subscribers,
        };

        result.save()?;

        Ok(result)
    }

    pub fn open(path: PathBuf) -> Result<Database, DataError> {
        if path.exists() {
            let f = File::open(&path)?;
            let feeds_list: Vec<Feed> = serde_json::from_reader(&f)?;

            let mut feeds: HashMap<FeedID, Feed> = HashMap::with_capacity(feeds_list.len());
            let mut subscribers: HashMap<SubscriberID, HashSet<FeedID>> = HashMap::new();

            for feed in feeds_list {
                let feed_id = gen_hash(&feed.link);
                for subscriber in &feed.subscribers {
                    let subscribed_feeds = subscribers
                        .entry(subscriber.to_owned())
                        .or_insert_with(HashSet::new);
                    subscribed_feeds.insert(feed_id);
                }
                feeds.insert(feed_id, feed);
            }

            Ok(Database {
                path: path,
                feeds: feeds,
                subscribers: subscribers,
            })
        } else {
            Database::create(path)
        }
    }

    pub fn all_feeds(&self) -> Vec<Feed> {
        self.feeds.iter().map(|(_, v)| v.clone()).collect()
    }

    pub fn all_subscribers(&self) -> Vec<SubscriberID> {
        self.subscribers.iter().map(|(k, _)| *k).collect()
    }

    pub fn subscribed_feeds(&self, subscriber: SubscriberID) -> Option<Vec<Feed>> {
        self.subscribers.get(&subscriber).map(|feeds| {
            feeds
                .iter()
                .map(|feed_id| &self.feeds[feed_id])
                .cloned()
                .collect()
        })
    }

    pub fn get_or_update_down_time(&mut self, rss_link: &str) -> Duration {
        let feed_id = gen_hash(&rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();
        let now = SystemTime::now();
        if let Some(t) = feed.down_time {
            now.duration_since(t).unwrap_or_default()
        } else {
            feed.down_time = Some(now);
            Duration::default()
        }
    }

    pub fn reset_down_time(&mut self, rss_link: &str) {
        let feed_id = gen_hash(&rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();
        feed.down_time = None;
    }

    pub fn is_subscribed(&self, subscriber: SubscriberID, rss_link: &str) -> bool {
        self.subscribers
            .get(&subscriber)
            .map(|feeds| feeds.contains(&gen_hash(&rss_link)))
            .unwrap_or(false)
    }

    pub fn subscribe(&mut self, subscriber: SubscriberID, rss_link: &str, rss: &feed::Rss) -> bool {
        let feed_id = gen_hash(&rss_link);
        {
            let subscribed_feeds = self
                .subscribers
                .entry(subscriber)
                .or_insert_with(HashSet::new);
            if !subscribed_feeds.insert(feed_id) {
                return false;
            }
        }
        {
            let feed = self.feeds.entry(feed_id).or_insert_with(|| Feed {
                link: rss_link.to_owned(),
                title: rss.title.to_owned(),
                down_time: None,
                ttl: rss.ttl,
                hash_list: rss.items.iter().map(gen_item_hash).collect(),
                subscribers: HashSet::new(),
            });
            feed.subscribers.insert(subscriber);
        }
        self.save().unwrap_or_default();
        true
    }

    pub fn unsubscribe(&mut self, subscriber: SubscriberID, rss_link: &str) -> Option<Feed> {
        let feed_id = gen_hash(&rss_link);

        let clear_subscriber;
        if let Some(subscribed_feeds) = self.subscribers.get_mut(&subscriber) {
            if subscribed_feeds.remove(&feed_id) {
                clear_subscriber = subscribed_feeds.is_empty();
            } else {
                return None;
            }
        } else {
            return None;
        }
        if clear_subscriber {
            self.subscribers.remove(&subscriber);
        }

        let result;
        let clear_feed;
        if let Some(feed) = self.feeds.get_mut(&feed_id) {
            if feed.subscribers.remove(&subscriber) {
                clear_feed = feed.subscribers.is_empty();
                result = feed.clone();
            } else {
                return None;
            }
        } else {
            return None;
        };
        if clear_feed {
            self.feeds.remove(&feed_id);
        }
        self.save().unwrap_or_default();
        Some(result)
    }

    pub fn delete_subscriber(&mut self, subscriber: SubscriberID) {
        self.subscribed_feeds(subscriber)
            .map(|feeds| {
                for feed in feeds {
                    let _ = self.unsubscribe(subscriber, &feed.link);
                }
            })
            .unwrap_or_default();
    }

    pub fn update_subscriber(&mut self, from: SubscriberID, to: SubscriberID) {
        let feeds = self.subscribers.remove(&from).unwrap();
        for feed_id in &feeds {
            let feed = self.feeds.get_mut(&feed_id).unwrap();
            feed.subscribers.remove(&from);
            feed.subscribers.insert(to);
        }
        self.subscribers.insert(to, feeds);
    }

    /// Update the feed in database, return updates
    pub fn update(&mut self, rss_link: &str, new_feed: feed::Rss) -> Vec<FeedUpdate> {
        let feed_id = gen_hash(&rss_link);
        if self.feeds.get(&feed_id).is_none() {
            return Vec::new();
        }

        self.reset_down_time(rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();

        let mut updates = Vec::new();
        let mut new_items = Vec::new();
        let mut new_hash_list = Vec::new();
        let items_len = new_feed.items.len();
        for item in new_feed.items {
            let hash = gen_item_hash(&item);
            if !feed.hash_list.contains(&hash) {
                new_hash_list.push(hash);
                new_items.push(item);
            }
        }
        if !new_items.is_empty() {
            updates.push(FeedUpdate::Items(new_items));

            let max_size = items_len * 2;
            let mut append: Vec<u64> = feed
                .hash_list
                .iter()
                .take(max_size - new_hash_list.len())
                .cloned()
                .collect();
            new_hash_list.append(&mut append);
            feed.hash_list = new_hash_list;
        }
        if new_feed.title != feed.title {
            updates.push(FeedUpdate::Title(new_feed.title.clone()));
            feed.title = new_feed.title;
        }
        feed.ttl = new_feed.ttl;
        if !updates.is_empty() {
            self.save().unwrap_or_default();
        }
        updates
    }

    pub fn save(&self) -> Result<(), DataError> {
        let feeds_list: Vec<&Feed> = self.feeds.iter().map(|(_id, feed)| feed).collect();
        let mut file = File::create(&self.path)?;
        if let Err(e) = serde_json::to_writer(&mut file, &feeds_list) {
            if e.is_io() {
                return Err(DataError::Io(e.into()));
            } else {
                unreachable!(e);
            };
        }
        Ok(())
    }
}

pub enum FeedUpdate {
    Items(Vec<feed::Item>),
    Title(String),
}

fn gen_item_hash(item: &feed::Item) -> u64 {
    item.id.as_ref().map(|id| gen_hash(&id)).unwrap_or_else(|| {
        let title = item.title.as_ref().map(|s| s.as_str()).unwrap_or_default();
        let link = item.link.as_ref().map(|s| s.as_str()).unwrap_or_default();
        gen_hash(&format!("{}{}", title, link))
    })
}
<commit_msg>Add a specialized hasher for database<commit_after>use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::hash::{BuildHasherDefault, Hash, Hasher};
use std::path::PathBuf;
use std::time::{Duration, SystemTime};

use serde::{Deserialize, Serialize};
use serde_json;
use thiserror::Error;

use crate::feed;

#[derive(Error, Debug)]
pub enum DataError {
    #[error("io error")]
    Io(#[from] std::io::Error),
    #[error("json error")]
    Json(#[from] serde_json::Error),
}

fn gen_hash<T: Hash>(t: &T) -> u64 {
    let mut hasher = std::collections::hash_map::DefaultHasher::default();
    t.hash(&mut hasher);
    hasher.finish()
}

type FeedId = u64;
type SubscriberId = i64;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Feed {
    pub link: String,
    pub title: String,
    pub down_time: Option<SystemTime>,
    pub subscribers: HashSet<SubscriberId, Size64>,
    pub ttl: Option<u32>,
    hash_list: Vec<u64>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Hub {
    pub callback: String,
    pub secret: String,
}

#[derive(Debug)]
pub struct Database {
    path: PathBuf,
    feeds: HashMap<FeedId, Feed, Size64>,
    subscribers: HashMap<SubscriberId, HashSet<FeedId, Size64>, Size64>,
}

impl Database {
    pub fn create(path: PathBuf) -> Result<Database, DataError> {
        let result = Database {
            path,
            feeds: HashMap::with_hasher(Size64::default()),
            subscribers: HashMap::with_hasher(Size64::default()),
        };

        result.save()?;

        Ok(result)
    }

    pub fn open(path: PathBuf) -> Result<Database, DataError> {
        if path.exists() {
            let f = File::open(&path)?;
            let feeds_list: Vec<Feed> = serde_json::from_reader(&f)?;

            let mut feeds = HashMap::with_capacity_and_hasher(feeds_list.len(), Size64::default());
            let mut subscribers = HashMap::with_hasher(Size64::default());

            for feed in feeds_list {
                let feed_id = gen_hash(&feed.link);
                for subscriber in &feed.subscribers {
                    let subscribed_feeds = subscribers
                        .entry(subscriber.to_owned())
                        .or_insert_with(HashSet::default);
                    subscribed_feeds.insert(feed_id);
                }
                feeds.insert(feed_id, feed);
            }

            Ok(Database {
                path,
                feeds,
                subscribers,
            })
        } else {
            Database::create(path)
        }
    }

    pub fn all_feeds(&self) -> Vec<Feed> {
        self.feeds.iter().map(|(_, v)| v.clone()).collect()
    }

    pub fn all_subscribers(&self) -> Vec<SubscriberId> {
        self.subscribers.iter().map(|(k, _)| *k).collect()
    }

    pub fn subscribed_feeds(&self, subscriber: SubscriberId) -> Option<Vec<Feed>> {
        self.subscribers.get(&subscriber).map(|feeds| {
            feeds
                .iter()
                .map(|feed_id| &self.feeds[feed_id])
                .cloned()
                .collect()
        })
    }

    pub fn get_or_update_down_time(&mut self, rss_link: &str) -> Duration {
        let feed_id = gen_hash(&rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();
        let now = SystemTime::now();
        if let Some(t) = feed.down_time {
            now.duration_since(t).unwrap_or_default()
        } else {
            feed.down_time = Some(now);
            Duration::default()
        }
    }

    pub fn reset_down_time(&mut self, rss_link: &str) {
        let feed_id = gen_hash(&rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();
        feed.down_time = None;
    }

    pub fn is_subscribed(&self, subscriber: SubscriberId, rss_link: &str) -> bool {
        self.subscribers
            .get(&subscriber)
            .map(|feeds| feeds.contains(&gen_hash(&rss_link)))
            .unwrap_or(false)
    }

    pub fn subscribe(&mut self, subscriber: SubscriberId, rss_link: &str, rss: &feed::Rss) -> bool {
        let feed_id = gen_hash(&rss_link);
        {
            let subscribed_feeds = self
                .subscribers
                .entry(subscriber)
                .or_insert_with(HashSet::default);
            if !subscribed_feeds.insert(feed_id) {
                return false;
            }
        }
        {
            let feed = self.feeds.entry(feed_id).or_insert_with(|| Feed {
                link: rss_link.to_owned(),
                title: rss.title.to_owned(),
                down_time: None,
                ttl: rss.ttl,
                hash_list: rss.items.iter().map(gen_item_hash).collect(),
                subscribers: HashSet::default(),
            });
            feed.subscribers.insert(subscriber);
        }
        self.save().unwrap_or_default();
        true
    }

    pub fn unsubscribe(&mut self, subscriber: SubscriberId, rss_link: &str) -> Option<Feed> {
        let feed_id = gen_hash(&rss_link);

        let clear_subscriber;
        if let Some(subscribed_feeds) = self.subscribers.get_mut(&subscriber) {
            if subscribed_feeds.remove(&feed_id) {
                clear_subscriber = subscribed_feeds.is_empty();
            } else {
                return None;
            }
        } else {
            return None;
        }
        if clear_subscriber {
            self.subscribers.remove(&subscriber);
        }

        let result;
        let clear_feed;
        if let Some(feed) = self.feeds.get_mut(&feed_id) {
            if feed.subscribers.remove(&subscriber) {
                clear_feed = feed.subscribers.is_empty();
                result = feed.clone();
            } else {
                return None;
            }
        } else {
            return None;
        };
        if clear_feed {
            self.feeds.remove(&feed_id);
        }
        self.save().unwrap_or_default();
        Some(result)
    }

    pub fn delete_subscriber(&mut self, subscriber: SubscriberId) {
        self.subscribed_feeds(subscriber)
            .map(|feeds| {
                for feed in feeds {
                    let _ = self.unsubscribe(subscriber, &feed.link);
                }
            })
            .unwrap_or_default();
    }

    pub fn update_subscriber(&mut self, from: SubscriberId, to: SubscriberId) {
        let feeds = self.subscribers.remove(&from).unwrap();
        for feed_id in &feeds {
            let feed = self.feeds.get_mut(&feed_id).unwrap();
            feed.subscribers.remove(&from);
            feed.subscribers.insert(to);
        }
        self.subscribers.insert(to, feeds);
    }

    /// Update the feed in database, return updates
    pub fn update(&mut self, rss_link: &str, new_feed: feed::Rss) -> Vec<FeedUpdate> {
        let feed_id = gen_hash(&rss_link);
        if self.feeds.get(&feed_id).is_none() {
            return Vec::new();
        }

        self.reset_down_time(rss_link);
        let feed = self.feeds.get_mut(&feed_id).unwrap();

        let mut updates = Vec::new();
        let mut new_items = Vec::new();
        let mut new_hash_list = Vec::new();
        let items_len = new_feed.items.len();
        for item in new_feed.items {
            let hash = gen_item_hash(&item);
            if !feed.hash_list.contains(&hash) {
                new_hash_list.push(hash);
                new_items.push(item);
            }
        }
        if !new_items.is_empty() {
            updates.push(FeedUpdate::Items(new_items));

            let max_size = items_len * 2;
            let mut append: Vec<u64> = feed
                .hash_list
                .iter()
                .take(max_size - new_hash_list.len())
                .cloned()
                .collect();
            new_hash_list.append(&mut append);
            feed.hash_list = new_hash_list;
        }
        if new_feed.title != feed.title {
            updates.push(FeedUpdate::Title(new_feed.title.clone()));
            feed.title = new_feed.title;
        }
        feed.ttl = new_feed.ttl;
        if !updates.is_empty() {
            self.save().unwrap_or_default();
        }
        updates
    }

    pub fn save(&self) -> Result<(), DataError> {
        let feeds_list: Vec<&Feed> = self.feeds.iter().map(|(_id, feed)| feed).collect();
        let mut file = File::create(&self.path)?;
        if let Err(e) = serde_json::to_writer(&mut file, &feeds_list) {
            if e.is_io() {
                return Err(DataError::Io(e.into()));
            } else {
                unreachable!(e);
            };
        }
        Ok(())
    }
}

pub enum FeedUpdate {
    Items(Vec<feed::Item>),
    Title(String),
}

fn gen_item_hash(item: &feed::Item) -> u64 {
    item.id.as_ref().map(|id| gen_hash(&id)).unwrap_or_else(|| {
        let title = item.title.as_ref().map(|s| s.as_str()).unwrap_or_default();
        let link = item.link.as_ref().map(|s| s.as_str()).unwrap_or_default();
        gen_hash(&format!("{}{}", title, link))
    })
}

pub type Size64 = BuildHasherDefault<Size64Hasher>;

/// A specialized hasher for u64 and i64
///
/// WARNING: Do not use it for user-controlled input
#[derive(Default)]
pub struct Size64Hasher {
    finished: bool,
    value: u64,
}

impl Hasher for Size64Hasher {
    fn finish(&self) -> u64 {
        self.value
    }

    fn write(&mut self, _bytes: &[u8]) {
        panic!("only support u64 and i64");
    }

    fn write_u64(&mut self, i: u64) {
        assert!(
            !self.finished,
            "this is a special hasher, do not write twice"
        );
        self.value = i;
        self.finished = true;
    }

    fn write_i64(&mut self, i: i64) {
        self.write_u64(i as u64);
    }
}
<|endoftext|>


<commit_before>//! Solve linear diagonal dynamics

use ndarray::*;
use ndarray_linalg::Scalar;
use super::traits::*;

/// Linear ODE with diagonalized matrix (exactly solvable)
pub struct Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    diag: Array<A, D>,
    diag_of_matrix: Array<A, D>,
    dt: A::Real,
}

impl<A, D> TimeStep for Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    type Time = A::Real;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }
    fn set_dt(&mut self, dt: Self::Time) {
        Zip::from(&mut self.diag)
            .and(&self.diag_of_matrix)
            .apply(|a, &b| { *a = b.mul_real(dt).exp(); });
    }
}

impl<A, D> ModelSize<D> for Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    fn model_size(&self) -> D::Pattern {
        self.diag.dim()
    }
}

impl<A, D> Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    pub fn new<S: Data<Elem = A>>(diag_of_matrix: &ArrayBase<S, D>, dt: A::Real) -> Self {
        let mut diag = diag_of_matrix.to_owned();
        for v in diag.iter_mut() {
            *v = v.mul_real(dt).exp();
        }
        Diagonal {
            diag: diag,
            diag_of_matrix: diag_of_matrix.to_owned(),
            dt: dt,
        }
    }
}

impl<A, Sr, D> SemiImplicitLinear<Sr, D> for Diagonal<A, D>
    where A: Scalar,
          Sr: DataMut<Elem = A>,
          D: Dimension
{
    type Scalar = A;

    fn lin<'a>(&self, mut x: &'a mut ArrayBase<Sr, D>) -> &'a mut ArrayBase<Sr, D> {
        for (val, d) in x.iter_mut().zip(self.diag.iter()) {
            *val = *val * *d;
        }
        x
    }
}
<commit_msg>Add StiffDiagonal trait<commit_after>//! Solve linear diagonal dynamics

use ndarray::*;
use ndarray_linalg::Scalar;
use super::traits::*;

pub trait StiffDiagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    fn diag(&self) -> Array<A, D>;
}

pub fn diagonal<A, D, EOM>(eom: &EOM, dt: A::Real) -> Diagonal<A, D>
    where A: Scalar,
          D: Dimension,
          EOM: StiffDiagonal<A, D>
{
    Diagonal::new(eom.diag(), dt)
}

/// Linear ODE with diagonalized matrix (exactly solvable)
pub struct Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    diag: Array<A, D>,
    diag_of_matrix: Array<A, D>,
    dt: A::Real,
}

impl<A, D> TimeStep for Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    type Time = A::Real;

    fn get_dt(&self) -> Self::Time {
        self.dt
    }
    fn set_dt(&mut self, dt: Self::Time) {
        Zip::from(&mut self.diag)
            .and(&self.diag_of_matrix)
            .apply(|a, &b| { *a = b.mul_real(dt).exp(); });
    }
}

impl<A, D> ModelSize<D> for Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    fn model_size(&self) -> D::Pattern {
        self.diag.dim()
    }
}

impl<A, D> Diagonal<A, D>
    where A: Scalar,
          D: Dimension
{
    pub fn new(diag_of_matrix: Array<A, D>, dt: A::Real) -> Self {
        let mut diag = diag_of_matrix.to_owned();
        for v in diag.iter_mut() {
            *v = v.mul_real(dt).exp();
        }
        Diagonal {
            diag: diag,
            diag_of_matrix: diag_of_matrix,
            dt: dt,
        }
    }
}

impl<A, Sr, D> SemiImplicitLinear<Sr, D> for Diagonal<A, D>
    where A: Scalar,
          Sr: DataMut<Elem = A>,
          D: Dimension
{
    type Scalar = A;

    fn lin<'a>(&self, mut x: &'a mut ArrayBase<Sr, D>) -> &'a mut ArrayBase<Sr, D> {
        for (val, d) in x.iter_mut().zip(self.diag.iter()) {
            *val = *val * *d;
        }
        x
    }
}
<|endoftext|>


<commit_before>use std::ops::Deref;

use vertex::Vertex;
use graph::TraversableToVertex;
use std::rc::Rc;
use std::sync::RwLock;

pub type EdgePointer = Rc<Box<RwLock<Edge>>>;

#[derive(Clone, Debug)]
pub struct Edge {
    // pub edge: EdgePointer,
    pub from_vertex: Vertex,
    pub to_vertex: Vertex,
    pub label: String
}

impl Edge {
    // creates a new edge
    // sets the in & out of it but doesn't touch the vertex
    pub fn new(from_vertex:Vertex, to_vertex: Vertex, label:String) -> EdgePointer {
        let edge = Rc::new(Box::new(RwLock::new(Edge{from_vertex:from_vertex,
                                      to_vertex:to_vertex,
                                      label:label})));
        edge
    }
}

impl TraversableToVertex for Edge {
	fn inV(&self) -> Vec<Vertex> {
		let mut result = Vec::new();
		let proxy = self.to_vertex.clone();
		result.push(proxy);
		result
	}
	fn outV(&self, labels: &[&str]) -> Vec<Vertex> {
		let mut result = Vec::new();
		let vertex = self.from_vertex.clone();
        result.push(vertex);
		result
	}
}
<commit_msg>removed extra box for Rc edge<commit_after>use std::ops::Deref;

use vertex::Vertex;
use graph::TraversableToVertex;
use std::rc::Rc;
use std::sync::RwLock;

pub type EdgePointer = Rc<RwLock<Edge>>;

#[derive(Clone, Debug)]
pub struct Edge {
    // pub edge: EdgePointer,
    pub from_vertex: Vertex,
    pub to_vertex: Vertex,
    pub label: String
}

impl Edge {
    // creates a new edge
    // sets the in & out of it but doesn't touch the vertex
    pub fn new(from_vertex:Vertex, to_vertex: Vertex, label:String) -> EdgePointer {
        let edge = Rc::new(RwLock::new(Edge{from_vertex:from_vertex,
                                      to_vertex:to_vertex,
                                      label:label}));
        edge
    }
}

impl TraversableToVertex for Edge {
	fn inV(&self) -> Vec<Vertex> {
		let mut result = Vec::new();
		let proxy = self.to_vertex.clone();
		result.push(proxy);
		result
	}
	fn outV(&self, labels: &[&str]) -> Vec<Vertex> {
		let mut result = Vec::new();
		let vertex = self.from_vertex.clone();
        result.push(vertex);
		result
	}
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use super::term::Term;

pub struct Context(HashMap<String, Term>);

impl Context {

    pub fn new() -> Context {
        Context(HashMap::new())
    }

    pub fn lookup(&mut self, id: String) -> Option<Term> {
        self.0.remove(&id)
    }

    pub fn bind(&mut self, id: String, term: Term) {
        self.0.insert(id, term);
    }

}

#[derive (Debug, PartialEq, Eq)]
pub enum ReferenceError { NotFound }

fn eval_shallow(ctx: &mut Context, term: Term, depth: u32) -> Result<Term, ReferenceError> {
    if depth == 0 {
        return Result::Ok(term)
    }
    match term {
        Term::Val(val) => Result::Ok(Term::Val(val)),
        Term::Abs(id, term) => Result::Ok(Term::Abs(id, term)),
        Term::Ref(id) => match ctx.lookup(id) {
            Option::None => Result::Err(ReferenceError::NotFound),
            Option::Some(val) => Result::Ok(val)
        },
        Term::App(fun, arg) => {
            match try!(eval_shallow(ctx, *fun, depth-1)) {
                Term::Abs(id, term) => {
                    ctx.bind(id, *arg);
                    Result::Ok(*term)
                },
                fun_eval => {
                    Result::Ok(Term::App(Box::new(fun_eval), arg))
                }
            }
        }
    }
}

pub fn eval(ctx: &mut Context, term: Term) -> Result<Term, ReferenceError> {
    let mut current_term = term;
     // find the head normal form of the term,
     // which is equivalent to the fixpoint of
     // the eval_shallow function
    loop {
        match current_term {
            Term::Val(val) => return Result::Ok(Term::Val(val)),
            Term::Abs(id, term) => return Result::Ok(Term::Abs(id, term)),
            term => current_term = try!(eval_shallow(ctx, term, 2))
        }
    }
}

#[test]
fn eval_val() {
    let mut c = Context::new();
    let x = eval(&mut c, Term::val_int(42));
    assert_eq!(x, Result::Ok(Term::val_int(42)));
}

#[test]
fn eval_id() {
    let mut c = Context::new();
    let x = eval(&mut c,
        Term::app(
            Term::fun(
                "x".to_string(),
                Term::id("x".to_string())
            ),
            Term::val_int(42)
        )
    );
    assert_eq!(x, Result::Ok(Term::val_int(42)));
}
<commit_msg>let context lookup return a result<commit_after>use std::collections::HashMap;

use super::term::Term;

#[derive (Debug, PartialEq, Eq)]
pub enum ReferenceError { NotFound }

pub struct Context(HashMap<String, Term>);

impl Context {

    pub fn new() -> Context {
        Context(HashMap::new())
    }

    pub fn lookup(&mut self, id: String) -> Result<Term, ReferenceError> {
        self.0.remove(&id).ok_or(ReferenceError::NotFound)
    }

    pub fn bind(&mut self, id: String, term: Term) {
        self.0.insert(id, term);
    }

}

fn eval_shallow(ctx: &mut Context, term: Term, depth: u32) -> Result<Term, ReferenceError> {
    if depth == 0 {
        return Result::Ok(term)
    }
    match term {
        Term::Val(val) => Result::Ok(Term::Val(val)),
        Term::Abs(id, term) => Result::Ok(Term::Abs(id, term)),
        Term::Ref(id) => ctx.lookup(id),
        Term::App(fun, arg) => {
            match try!(eval_shallow(ctx, *fun, depth-1)) {
                Term::Abs(id, term) => {
                    ctx.bind(id, *arg);
                    Result::Ok(*term)
                },
                fun_eval => {
                    Result::Ok(Term::App(Box::new(fun_eval), arg))
                }
            }
        }
    }
}

pub fn eval(ctx: &mut Context, term: Term) -> Result<Term, ReferenceError> {
    let mut current_term = term;
     // find the head normal form of the term,
     // which is equivalent to the fixpoint of
     // the eval_shallow function
    loop {
        match current_term {
            Term::Val(val) => return Result::Ok(Term::Val(val)),
            Term::Abs(id, term) => return Result::Ok(Term::Abs(id, term)),
            term => current_term = try!(eval_shallow(ctx, term, 2))
        }
    }
}

#[test]
fn eval_val() {
    let mut c = Context::new();
    let x = eval(&mut c, Term::val_int(42));
    assert_eq!(x, Result::Ok(Term::val_int(42)));
}

#[test]
fn eval_id() {
    let mut c = Context::new();
    let x = eval(&mut c,
        Term::app(
            Term::fun(
                "x".to_string(),
                Term::id("x".to_string())
            ),
            Term::val_int(42)
        )
    );
    assert_eq!(x, Result::Ok(Term::val_int(42)));
}
<|endoftext|>