Repository: tokio-rs/tokio
File: tokio/src/runtime/task/list.rs (Rust)

Code before:
        notified: Notified<S>,
        id: super::Id,
    ) -> Option<Notified<S>>
    where
        S: Schedule,
    {
        unsafe {
            // safety: We just created the task, so we have exclusive access
            // to the field.
            task.header().set_owner_id(self.id);
        }
        // check close flag
        if self.closed.load(Ordering::Acquire) {
            task.shutdown();
            return None;
        }
        self.push_inner(id, task);
        Some(notified)
    }

    #[inline]
    fn push_inner(&self, task_id: super::Id, task: Task<S>) {
        let mut lock = self.lists[task_id.0 as usize & (self.grain - 1)].lock();
        lock.push_front(task);
        self.count.fetch_add(1, Ordering::Relaxed);
    }

    /// Asserts that the given task is owned by this OwnedTasks and convert it to
    /// a LocalNotified, giving the thread permission to poll this task.
    #[inline]
    pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> {
        debug_assert_eq!(task.header().get_owner_id(), Some(self.id));
        // safety: All tasks bound to this OwnedTasks are Send, so it is safe
        // to poll it on this thread no matter what thread we are on.
        LocalNotified {
            task: task.0,
            _not_send: PhantomData,
        }
    }

    /// Shuts down all tasks in the collection. This call also closes the
    /// collection, preventing new items from being added.
    pub(crate) fn close_and_shutdown_all(&self)
    where
        S: Schedule,
    {
        // The first iteration of the loop was unrolled so it can set the
        // closed bool.
        self.closed.store(true, Ordering::Release);

        for list in &self.lists {

Review comment (suggestion):
This doesn't look right to me. The `OwnedTasks` could be closed between the call to `self.closed.load(Acquire)` and the call to `self.lists[...].lock()`. If that happens, then we push something to the list, but it's never removed because the close operation has already completed.

Code after:
    }

    /// The part of `bind` that's the same for every type of future.
    unsafe fn bind_inner(&self, task: Task<S>, notified: Notified<S>) -> Option<Notified<S>>
    where
        S: Schedule,
    {
        unsafe {
            // safety: We just created the task, so we have exclusive access
            // to the field.
            task.header().set_owner_id(self.id);
        }

        let shard = self.list.lock_shard(&task);
        // Check the closed flag in the lock for ensuring all that tasks
        // will shut down after the OwnedTasks has been closed.
        if self.closed.load(Ordering::Acquire) {
            drop(shard);
            task.shutdown();
            return None;
        }
        shard.push(task);
        Some(notified)
    }

    /// Asserts that the given task is owned by this OwnedTasks and convert it to
    /// a LocalNotified, giving the thread permission to poll this task.
    #[inline]
    pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> {
        debug_assert_eq!(task.header().get_owner_id(), Some(self.id));
        // safety: All tasks bound to this OwnedTasks are Send, so it is safe
        // to poll it on this thread no matter what thread we are on.
        LocalNotified {
            task: task.0,
            _not_send: PhantomData,
        }
    }

    /// Shuts down all tasks in the collection. This call also closes the
    /// collection, preventing new items from being added.
    ///
    /// The parameter start determines which shard this method will start at.
    /// Using different values for each worker thread reduces contention.
    pub(crate) fn close_and_shutdown_all(&self, start: usize)
    where
        S: Schedule,
    {
        self.closed.store(true, Ordering::Release);
        for i in start..self.get_shard_size() + start {
            loop {
                let task = self.list.pop_back(i);