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

Code before:
    /// 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);
        }
        // check close flag
        if self.closed.load(Ordering::Acquire) {
            task.shutdown();
            return None;
        }
        self.list.push(task);

        // double check close flag for ensuring all tasks will shutdown after OwnedTasks has been closed,
        // it should be completed quickly.
        if self.closed.load(Ordering::Acquire) {
            for i in 0..self.get_shard_size() {
                if let Some(task) = self.list.pop_back(i) {
                    task.shutdown();
                }
            }
        }
        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 should be random among different worker threads
    /// to reduce lock conflicts during shutdown.
    /// Initiate shutting down the shard indexed by the start, and reset to 0
    /// once the shar_size is reached, continuing until start - 1, it works like a ring.
    pub(crate) fn close_and_shutdown_all(&self, start: usize)
    where

Review comment (suggestion):
You can avoid this by defining a method like this:
```
impl<L: ShardedListItem> ShardedList<L, L::Target> {
    pub fn lock_shard(&self, val: &L::Handle) -> ShardGuard<'_, L, L::Target> {
        let id = unsafe { L::get_shared_id(L::as_raw(&val)) };
        ShardGuard {
            lock: self.shard_inner(id),
            count: &self.count,
            id,
        }
    }
}

struct ShardGuard<'a, L, T> {
    lock: MutexGuard<'a, LinkedList<L, T>>
    count: &'a AtomicUsize,
    id: usize,
}

impl<L: ShardedListItem> ShardGuard<L, L::Target> {
    /// Push a value to this shard.
    pub fn push_shard(self, val: L::Handle) {
        let id = unsafe { L::get_shared_id(L::as_raw(&val)) };
        assert_eq!(id, self.id);
        self.lock.push_front(val);
        self.count.fetch_add(1, Ordering::Relaxed);
    }
}
```
This will let you check `self.closed` while the mutex is locked.

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);