diff --git a/tests/global.rs b/tests/global.rs index c3047e344..cb83c145a 100644 --- a/tests/global.rs +++ b/tests/global.rs @@ -5,6 +5,8 @@ extern crate env_logger; use std::{io, thread}; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::Relaxed; use futures::prelude::*; use tokio::net::{TcpStream, TcpListener}; @@ -18,7 +20,7 @@ macro_rules! t { } #[test] -fn hammer() { +fn hammer_old() { let _ = env_logger::init(); let threads = (0..10).map(|_| { @@ -73,48 +75,62 @@ fn hammer_split() { use tokio_io::io; const N: usize = 100; + const ITER: usize = 100; let _ = env_logger::init(); - let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap())); - let addr = t!(srv.local_addr()); + for _ in 0..ITER { + let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap())); + let addr = t!(srv.local_addr()); - let mut rt = Runtime::new().unwrap(); + let cnt = Arc::new(AtomicUsize::new(0)); - fn split(socket: TcpStream) { - let socket = Arc::new(socket); - let rd = Rd(socket.clone()); - let wr = Wr(socket); + let mut rt = Runtime::new().unwrap(); - let rd = io::read(rd, vec![0; 1]) - .map(|_| ()) - .map_err(|e| panic!("read error = {:?}", e)); + fn split(socket: TcpStream, cnt: Arc) { + let socket = Arc::new(socket); + let rd = Rd(socket.clone()); + let wr = Wr(socket); - let wr = io::write_all(wr, b"1") - .map(|_| ()) - .map_err(|e| panic!("write error = {:?}", e)); + let cnt2 = cnt.clone(); - tokio::spawn(rd); - tokio::spawn(wr); - } + let rd = io::read(rd, vec![0; 1]) + .map(move |_| { + cnt2.fetch_add(1, Relaxed); + }) + .map_err(|e| panic!("read error = {:?}", e)); - rt.spawn({ - srv.incoming() - .map_err(|e| panic!("accept error = {:?}", e)) - .take(N as u64) - .for_each(|socket| { - split(socket); - Ok(()) - }) - }); + let wr = io::write_all(wr, b"1") + .map(move |_| { + cnt.fetch_add(1, Relaxed); + }) + .map_err(move |e| panic!("write error = {:?}", e)); + + tokio::spawn(rd); + tokio::spawn(wr); + } - for _ in 0..N { rt.spawn({ - TcpStream::connect(&addr) - .map_err(|e| panic!("connect error = {:?}", e)) - .map(|socket| split(socket)) + let cnt = cnt.clone(); + srv.incoming() + .map_err(|e| panic!("accept error = {:?}", e)) + .take(N as u64) + .for_each(move |socket| { + split(socket, cnt.clone()); + Ok(()) + }) }); - } - rt.shutdown_on_idle().wait().unwrap(); + for _ in 0..N { + rt.spawn({ + let cnt = cnt.clone(); + TcpStream::connect(&addr) + .map_err(move |e| panic!("connect error = {:?}", e)) + .map(move |socket| split(socket, cnt)) + }); + } + + rt.shutdown_on_idle().wait().unwrap(); + assert_eq!(N * 4, cnt.load(Relaxed)); + } } diff --git a/tokio-reactor/src/atomic_task.rs b/tokio-reactor/src/atomic_task.rs index 6a4788e6c..9b996988b 100644 --- a/tokio-reactor/src/atomic_task.rs +++ b/tokio-reactor/src/atomic_task.rs @@ -1,9 +1,11 @@ +#![allow(dead_code)] + +use super::Task; + use std::fmt; use std::cell::UnsafeCell; use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering::{Acquire, Release}; - -use Task; +use std::sync::atomic::Ordering::{Acquire, Release, AcqRel}; /// A synchronization primitive for task notification. /// @@ -24,37 +26,115 @@ use Task; /// `AtomicTask` does not provide any memory ordering guarantees, as such the /// user should use caution and use other synchronization primitives to guard /// the result of the underlying computation. -pub struct AtomicTask { +pub(crate) struct AtomicTask { state: AtomicUsize, task: UnsafeCell>, } -/// Initial state, the `AtomicTask` is currently not being used. -/// -/// The value `2` is picked specifically because it between the write lock & -/// read lock values. Since the read lock is represented by an incrementing -/// counter, this enables an atomic fetch_sub operation to be used for releasing -/// a lock. -const WAITING: usize = 2; +// `AtomicTask` is a multi-consumer, single-producer transfer cell. The cell +// stores a `Task` value produced by calls to `register` and many threads can +// race to take the task (to notify it) by calling `notify. +// +// If a new `Task` instance is produced by calling `register` before an existing +// one is consumed, then the existing one is overwritten. +// +// While `AtomicTask` is single-producer, the implementation ensures memory +// safety. In the event of concurrent calls to `register`, there will be a +// single winner whose task will get stored in the cell. The losers will not +// have their tasks notified. As such, callers should ensure to add +// synchronization to calls to `register`. +// +// The implementation uses a single `AtomicUsize` value to coordinate access to +// the `Task` cell. There are two bits that are operated on independently. These +// are represented by `REGISTERING` and `NOTIFYING`. +// +// The `REGISTERING` bit is set when a producer enters the critical section. The +// `NOTIFYING` 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 task cell by transitioning the +// state from `WAITING` to `REGISTERING` or `NOTIFYING`, depending on the +// operation the thread wishes to perform. When this transition is made, it is +// guaranteed that no other thread will access the task 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 task cell. +// +// If the lock is obtained, then the thread sets the task cell to the task +// 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 `notify` will observe the task. +// +// If the transition fails, then there was a concurrent call to `notify` that +// was unable to access the task cell (due to the registering thread holding the +// lock). To handle this, the registering thread removes the task it just set +// from the cell and calls `notify` on it. This call to notify represents the +// attempt to notify by the other thread (that set the `NOTIFYING` bit). The +// state is then transitioned from `REGISTERING | NOTIFYING` back to `WAITING`. +// This transition must succeed because, at this point, the state cannot be +// transitioned by another thread. +// +// # Notifying +// +// On a call to `notify`, an attempt to transition the state from `WAITING` to +// `NOTIFYING` is made. On success, the caller obtains a lock on the task cell. +// +// If the lock is obtained, then the thread takes ownership of the current value +// in teh task cell, and calls `notify` 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 `NOTIFYING` 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 +// `NOTIFYING` critical section. Either way, no action must be taken. +// +// If the current thread is the only concurrent call to `notify` and another +// thread is in the `register` critical section, when the other thread **exits** +// the `register` critical section, it will observe the `NOTIFYING` bit and +// handle the notify itself. +// +// If another thread is in the `notify` critical section, then it will handle +// notifying the task. +// +// # A potential race (is safely handled). +// +// Imagine the following situation: +// +// * Thread A obtains the `notify` lock and notifies a task. +// +// * Before thread A releases the `notify` lock, the notified task is scheduled. +// +// * Thread B attempts to notify the task. In theory this should result in the +// task being notified, but it cannot because thread A still holds the notify +// lock. +// +// This case is handled by requiring users of `AtomicTask` to call `register` +// **before** attempting to observe the application state change that resulted +// in the task being notified. The notifiers also change the application state +// before calling notify. +// +// Because of this, the task will do one of two things. +// +// 1) Observe the application state change that Thread B is notifying on. In +// this case, it is OK for Thread B's notification to be lost. +// +// 2) Call register before attempting to observe the application state. Since +// Thread A still holds the `notify` lock, the call to `register` will result +// in the task notifying itself and get scheduled again. -/// The `register` function has determined that the task is no longer current. -/// This implies that `AtomicTask::register` is being called from a different -/// task than is represented by the currently stored task. The write lock is -/// obtained to update the task cell. -const LOCKED_WRITE: usize = 0; +/// Idle state +const WAITING: usize = 0; -/// At least one call to `notify` happened concurrently to `register` updating -/// the task cell. This state is detected when `register` exits the mutation -/// code and signals to `register` that it is responsible for notifying its own -/// task. -const LOCKED_WRITE_NOTIFIED: usize = 1; +/// A new task value is being registered with the `AtomicTask` cell. +const REGISTERING: usize = 0b01; - -/// The `notify` function has locked access to the task cell for notification. -/// -/// The constant is left here mostly for documentation reasons. -#[allow(dead_code)] -const LOCKED_READ: usize = 3; +/// The task currently registered with the `AtomicTask` cell is being notified. +const NOTIFYING: usize = 0b10; impl AtomicTask { /// Create an `AtomicTask` initialized with the given `Task` @@ -69,7 +149,7 @@ impl AtomicTask { } } - /// Registers the task to be notified on calls to `notify`. + /// Registers the provided task to be notified on calls to `notify`. /// /// The new task will take place of any previous tasks that were registered /// by previous calls to `register`. Any calls to `notify` that happen after @@ -84,36 +164,75 @@ impl AtomicTask { /// 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. - pub(crate) fn register(&self, task: Task) { - match self.state.compare_and_swap(WAITING, LOCKED_WRITE, Acquire) { + pub fn register_task(&self, task: Task) { + match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) { WAITING => { unsafe { - // Locked acquired, update the task cell - *self.task.get() = Some(task); + // Locked acquired, update the waker cell + *self.task.get() = Some(task.clone()); - // Release the lock. If the state transitioned to - // `LOCKED_NOTIFIED`, this means that an notify has been - // signaled, so notify the task. - if LOCKED_WRITE_NOTIFIED == self.state.swap(WAITING, Release) { - (*self.task.get()).as_ref().unwrap().notify(); + // Release the lock. If the state transitioned to include + // the `NOTIFYING` bit, this means that a notify has been + // called concurrently, so we have to remove the task and + // notify it.` + // + // Start by assuming that the state is `REGISTERING` as this + // is what we jut set it to. + let mut curr = REGISTERING; + + // If a task has to be notified, it will be set here. + let mut notify: Option = None; + + loop { + let res = self.state.compare_exchange( + curr, WAITING, AcqRel, Acquire); + + match res { + Ok(_) => { + // The atomic exchange was successful, now + // notify the task (if set) and return. + if let Some(task) = notify { + task.notify(); + } + + return; + } + Err(actual) => { + // This branch can only be reached if a + // concurrent thread called `notify`. In this + // case, `actual` **must** be `REGISTERING | + // `NOTIFYING`. + debug_assert_eq!(actual, REGISTERING | NOTIFYING); + + // Take the task to notify once the atomic operation has + // completed. + notify = (*self.task.get()).take(); + + // Update `curr` for the next iteration of the + // loop + curr = actual; + } + } } } } - LOCKED_WRITE | LOCKED_WRITE_NOTIFIED => { - // A thread is concurrently calling `register`. This shouldn't - // happen as it doesn't really make much sense, but it isn't - // unsafe per se. Since two threads are concurrently trying to - // update the task, it's undefined which one "wins" (no ordering - // guarantees), so we can just do nothing. + NOTIFYING => { + // Currently in the process of notifying the task, i.e., + // `notify` is currently being called on the old task handle. + // So, we call notify on the new task handle + task.notify(); } state => { - debug_assert!(state != LOCKED_WRITE, "unexpected state LOCKED_WRITE"); - debug_assert!(state != LOCKED_WRITE_NOTIFIED, "unexpected state LOCKED_WRITE_NOTIFIED"); - - // Currently in a read locked state, this implies that `notify` - // is currently being called on the old task handle. So, we call - // notify on the new task handle - task.notify(); + // 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. + // + // We just want to maintain memory safety. It is ok to drop the + // call to `register`. + debug_assert!( + state == REGISTERING || + state == REGISTERING | NOTIFYING); } } } @@ -122,49 +241,33 @@ impl AtomicTask { /// /// If `register` has not been called yet, then this does nothing. pub fn notify(&self) { - let mut curr = WAITING; + // 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 `AtomicTask` is associated with. + match self.state.fetch_or(NOTIFYING, AcqRel) { + WAITING => { + // The notifying lock has been acquired. + let task = unsafe { (*self.task.get()).take() }; - loop { - if curr == LOCKED_WRITE { - // Transition the state to LOCKED_NOTIFIED - let actual = self.state.compare_and_swap(LOCKED_WRITE, LOCKED_WRITE_NOTIFIED, Release); + // Release the lock + self.state.fetch_and(!NOTIFYING, Release); - if curr == actual { - // Success, return - return; + if let Some(task) = task { + task.notify(); } - - // update current state variable and try again - curr = actual; - - } else if curr == LOCKED_WRITE_NOTIFIED { - // Currently in `LOCKED_WRITE_NOTIFIED` state, nothing else to do. - return; - - } else { - // Currently in a LOCKED_READ state, so attempt to increment the - // lock count. - let actual = self.state.compare_and_swap(curr, curr + 1, Acquire); - - // Locked acquired - if actual == curr { - // Notify the task - unsafe { - if let Some(ref task) = (*self.task.get()).take() { - task.notify(); - } - } - - // Release the lock - self.state.fetch_sub(1, Release); - - // Done - return; - } - - // update current state variable and try again - curr = actual; - + } + state => { + // There is a concurrent thread currently updating the + // associated task. + // + // Nothing more to do as the `NOTIFYING` bit has been set. It + // doesn't matter if there are concurrent registering threads or + // not. + // + debug_assert!( + state == REGISTERING || + state == REGISTERING | NOTIFYING || + state == NOTIFYING); } } } diff --git a/tokio-reactor/src/background.rs b/tokio-reactor/src/background.rs index 03f057cbe..88f78a8bc 100644 --- a/tokio-reactor/src/background.rs +++ b/tokio-reactor/src/background.rs @@ -137,7 +137,7 @@ impl Future for Shutdown { fn poll(&mut self) -> Poll<(), ()> { let task = Task::Futures1(task::current()); - self.inner.shared.shutdown_task.register(task); + self.inner.shared.shutdown_task.register_task(task); if !self.inner.is_shutdown() { return Ok(Async::NotReady); diff --git a/tokio-reactor/src/lib.rs b/tokio-reactor/src/lib.rs index a0ce1788c..c46b8c6b2 100644 --- a/tokio-reactor/src/lib.rs +++ b/tokio-reactor/src/lib.rs @@ -42,8 +42,8 @@ extern crate tokio_io; #[cfg(feature = "unstable-futures")] extern crate futures2; -pub(crate) mod background; mod atomic_task; +pub(crate) mod background; mod poll_evented; mod registration; @@ -120,6 +120,9 @@ struct Inner { /// The underlying system event queue. io: mio::Poll, + /// ABA guard counter + next_aba_guard: AtomicUsize, + /// Dispatch slabs for I/O and futures events io_dispatch: RwLock>, @@ -128,6 +131,7 @@ struct Inner { } struct ScheduledIo { + aba_guard: usize, readiness: AtomicUsize, reader: AtomicTask, writer: AtomicTask, @@ -145,11 +149,11 @@ static HANDLE_FALLBACK: AtomicUsize = ATOMIC_USIZE_INIT; /// Tracks the reactor for the current execution context. thread_local!(static CURRENT_REACTOR: RefCell> = RefCell::new(None)); -const TOKEN_WAKEUP: mio::Token = mio::Token(0); -const TOKEN_START: usize = 1; +const TOKEN_SHIFT: usize = 22; // Kind of arbitrary, but this reserves some token space for later usage. -const MAX_SOURCES: usize = usize::MAX >> 4; +const MAX_SOURCES: usize = (1 << TOKEN_SHIFT) - 1; +const TOKEN_WAKEUP: mio::Token = mio::Token(MAX_SOURCES); fn _assert_kinds() { fn _assert() {} @@ -221,6 +225,7 @@ impl Reactor { _wakeup_registration: wakeup_pair.0, inner: Arc::new(Inner { io: io, + next_aba_guard: AtomicUsize::new(0), io_dispatch: RwLock::new(Slab::with_capacity(1)), wakeup: wakeup_pair.1, }), @@ -358,10 +363,16 @@ impl Reactor { } fn dispatch(&self, token: mio::Token, ready: mio::Ready) { - let token = usize::from(token) - TOKEN_START; + let aba_guard = token.0 & !MAX_SOURCES; + let token = token.0 & MAX_SOURCES; + let io_dispatch = self.inner.io_dispatch.read().unwrap(); if let Some(io) = io_dispatch.get(token) { + if aba_guard != io.aba_guard { + return; + } + io.readiness.fetch_or(ready.as_usize(), Relaxed); if ready.is_writable() || platform::is_hup(&ready) { @@ -545,6 +556,9 @@ impl Inner { fn add_source(&self, source: &Evented) -> io::Result { + // Get an ABA guard value + let aba_guard = self.next_aba_guard.fetch_add(1 << TOKEN_SHIFT, Relaxed); + let mut io_dispatch = self.io_dispatch.write().unwrap(); if io_dispatch.len() == MAX_SOURCES { @@ -554,13 +568,14 @@ impl Inner { // Acquire a write lock let key = io_dispatch.insert(ScheduledIo { + aba_guard, readiness: AtomicUsize::new(0), reader: AtomicTask::new(), writer: AtomicTask::new(), }); try!(self.io.register(source, - mio::Token(TOKEN_START + key), + mio::Token(aba_guard | key), mio::Ready::all(), mio::PollOpt::edge())); @@ -588,7 +603,7 @@ impl Inner { Direction::Write => (&sched.writer, mio::Ready::writable()), }; - task.register(t); + task.register_task(t); if sched.readiness.load(SeqCst) & ready.as_usize() != 0 { task.notify(); diff --git a/tokio-reactor/src/poll_evented.rs b/tokio-reactor/src/poll_evented.rs index c9f5b77d2..58a65038c 100644 --- a/tokio-reactor/src/poll_evented.rs +++ b/tokio-reactor/src/poll_evented.rs @@ -647,9 +647,9 @@ impl fmt::Debug for PollEvented { impl Drop for PollEvented { fn drop(&mut self) { - if let Some(io) = self.io.as_ref() { + if let Some(io) = self.io.take() { // Ignore errors - let _ = self.inner.registration.deregister(io); + let _ = self.inner.registration.deregister(&io); } } } diff --git a/tokio-reactor/src/registration.rs b/tokio-reactor/src/registration.rs index 84199ea23..840406339 100644 --- a/tokio-reactor/src/registration.rs +++ b/tokio-reactor/src/registration.rs @@ -6,7 +6,7 @@ use mio::{self, Evented}; #[cfg(feature = "unstable-futures")] use futures2; -use std::{io, mem, usize}; +use std::{io, ptr, usize}; use std::cell::UnsafeCell; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::SeqCst; @@ -68,7 +68,7 @@ struct Inner { struct Node { direction: Direction, task: Task, - next: Option>, + next: *mut Node, } /// Initial state. The handle is not set and the registration is idle. @@ -197,40 +197,35 @@ impl Registration { // are pending readiness notifications. let actual = self.state.swap(READY, SeqCst); - // Consume the stack of nodes. - let ptr = actual & !LIFECYCLE_MASK; + // Consume the stack of nodes - if ptr != 0 { - let mut read = false; - let mut write = false; - let mut curr = unsafe { Box::from_raw(ptr as *mut Node) }; + let mut read = false; + let mut write = false; + let mut ptr = (actual & !LIFECYCLE_MASK) as *mut Node; - let inner = unsafe { (*self.inner.get()).as_ref().unwrap() }; + let inner = unsafe { (*self.inner.get()).as_ref().unwrap() }; - loop { - let node = *curr; - let Node { - direction, - task, - next, - } = node; + while !ptr.is_null() { + let node = unsafe { Box::from_raw(ptr) }; + let node = *node; + let Node { + direction, + task, + next, + } = node; - let flag = match direction { - Direction::Read => &mut read, - Direction::Write => &mut write, - }; + let flag = match direction { + Direction::Read => &mut read, + Direction::Write => &mut write, + }; - if !*flag { - *flag = true; + if !*flag { + *flag = true; - inner.register(direction, task); - } - - match next { - Some(next) => curr = next, - None => break, - } + inner.register(direction, task); } + + ptr = next; } return res.map(|_| true); @@ -388,41 +383,35 @@ impl Registration { let inner = unsafe { (*self.inner.get()).as_ref().unwrap() }; return inner.poll_ready(direction, notify, task); } - _ => { + LOCKED => { if !notify { // Skip the notification tracking junk. return Ok(None); } - let ptr = state & !LIFECYCLE_MASK; + let next_ptr = (state & !LIFECYCLE_MASK) as *mut Node; + + let task = task(); // Get the node let mut n = node.take().unwrap_or_else(|| { Box::new(Node { direction, - task: task(), - next: None, + task: task, + next: ptr::null_mut(), }) }); - n.next = if ptr == 0 { - None - } else { - // Great care must be taken of the CAS fails - Some(unsafe { Box::from_raw(ptr as *mut Node) }) - }; + n.next = next_ptr; - let ptr = Box::into_raw(n); - let next = ptr as usize | (state & LIFECYCLE_MASK); + let node_ptr = Box::into_raw(n); + let next = node_ptr as usize | (state & LIFECYCLE_MASK); let actual = self.state.compare_and_swap(state, next, SeqCst); if actual != state { // Back out of the node boxing - let mut n = unsafe { Box::from_raw(ptr) }; - - // We don't really own this - mem::forget(n.next.take()); + let n = unsafe { Box::from_raw(node_ptr) }; // Save this for next loop node = Some(n); @@ -433,6 +422,7 @@ impl Registration { return Ok(None); } + _ => unreachable!(), } } } @@ -532,10 +522,11 @@ impl Inner { sched.readiness.fetch_and(!mask_no_hup, SeqCst)); if ready.is_empty() && notify { + let task = task(); // Update the task info match direction { - Direction::Read => sched.reader.register(task()), - Direction::Write => sched.writer.register(task()), + Direction::Read => sched.reader.register_task(task), + Direction::Write => sched.writer.register_task(task), } // Try again diff --git a/tokio-threadpool/src/lib.rs b/tokio-threadpool/src/lib.rs index d809f27fc..e4e829bef 100644 --- a/tokio-threadpool/src/lib.rs +++ b/tokio-threadpool/src/lib.rs @@ -936,9 +936,10 @@ impl Future for Shutdown { type Error = (); fn poll(&mut self) -> Poll<(), ()> { + use futures::task; trace!("Shutdown::poll"); - self.inner().shutdown_task.task1.register(); + self.inner().shutdown_task.task1.register_task(task::current()); if 0 != self.inner().num_workers.load(Acquire) { return Ok(Async::NotReady); diff --git a/tokio-threadpool/src/task.rs b/tokio-threadpool/src/task.rs index 59deca07b..6da38ec6a 100644 --- a/tokio-threadpool/src/task.rs +++ b/tokio-threadpool/src/task.rs @@ -222,13 +222,13 @@ impl Task { let actual = self.inner().state.compare_and_swap( Idle.into(), Scheduled.into(), - Relaxed).into(); + AcqRel).into(); match actual { Idle => return true, Running => { let actual = self.inner().state.compare_and_swap( - Running.into(), Notified.into(), Relaxed).into(); + Running.into(), Notified.into(), AcqRel).into(); match actual { Idle => continue, diff --git a/tokio-threadpool/tests/threadpool.rs b/tokio-threadpool/tests/threadpool.rs index 5edc4028b..b881fc45e 100644 --- a/tokio-threadpool/tests/threadpool.rs +++ b/tokio-threadpool/tests/threadpool.rs @@ -186,7 +186,7 @@ fn force_shutdown_drops_futures() { let a = num_inc.clone(); let b = num_dec.clone(); - let mut pool = Builder::new() + let pool = Builder::new() .around_worker(move |w, _| { a.fetch_add(1, Relaxed); w.run(); @@ -548,3 +548,103 @@ fn panic_in_task() { await_shutdown(pool.shutdown_on_idle()); } + +#[test] +#[cfg(not(feature = "unstable-futures"))] +fn hammer() { + use futures::future; + use futures::sync::{oneshot, mpsc}; + + const N: usize = 1000; + const ITER: usize = 20; + + struct Counted { + cnt: Arc, + inner: T, + } + + impl Future for Counted { + type Item = T::Item; + type Error = T::Error; + + fn poll(&mut self) -> Poll { + self.inner.poll() + } + } + + impl Drop for Counted { + fn drop(&mut self) { + self.cnt.fetch_add(1, Relaxed); + } + } + + for i in 0.. ITER { + println!("~~~ ITER {} ~~~", i); + + let pool = Builder::new() + // .pool_size(30) + .build(); + + let cnt = Arc::new(AtomicUsize::new(0)); + + let (listen_tx, listen_rx) = mpsc::unbounded::>>(); + let mut listen_tx = listen_tx.wait(); + + pool.spawn({ + let c1 = cnt.clone(); + let c2 = cnt.clone(); + let pool = pool.sender().clone(); + let task = listen_rx + .map_err(|e| panic!("accept error = {:?}", e)) + .for_each(move |tx| { + let task = future::lazy(|| { + let (tx2, rx2) = oneshot::channel(); + + tx.send(tx2).unwrap(); + rx2 + }) + .map_err(|e| panic!("e={:?}", e)) + .and_then(|_| { + Ok(()) + }); + + pool.spawn(Counted { + inner: task, + cnt: c1.clone(), + }).unwrap(); + + Ok(()) + }); + + Counted { + inner: task, + cnt: c2, + } + }); + + for _ in 0..N { + let cnt = cnt.clone(); + let (tx, rx) = oneshot::channel(); + listen_tx.send(tx).unwrap(); + + pool.spawn({ + let task = rx + .map_err(|e| panic!("rx err={:?}", e)) + .and_then(|tx| { + tx.send(()).unwrap(); + Ok(()) + }); + + Counted { + inner: task, + cnt, + } + }); + } + + drop(listen_tx); + + pool.shutdown_on_idle().wait().unwrap(); + assert_eq!(N * 2 + 1, cnt.load(Relaxed)); + } +}