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

Code before:
        /// ```should_panic
        /// use tokio::runtime::UnhandledPanic;
        ///
        /// # #[tokio::main]
        /// # async fn main() {
        /// tokio::task::LocalSet::new()
        ///     .unhandled_panic(UnhandledPanic::ShutdownRuntime)
        ///     .run_until(async {
        ///         tokio::task::spawn_local(async { panic!("boom"); });
        ///         tokio::task::spawn_local(async {
        ///             // This task never completes
        ///         });
        ///
        ///         // Do some work, but `run_until` will panic before it completes
        /// # loop { tokio::task::yield_now().await; }
        ///     })
        ///     .await;
        /// # }
        /// ```
        ///
        /// [`JoinHandle`]: struct@crate::task::JoinHandle
        pub fn unhandled_panic(&mut self, behavior: crate::runtime::UnhandledPanic) -> &mut Self {
            // TODO: This should be set as a builder
            Arc::get_mut(&mut Rc::get_mut(&mut self.context).expect("TODO: context should not panic").shared)
                .expect("TODO: we shouldn't panic")
                .unhandled_panic = behavior;
            self
        }
    }
}

impl fmt::Debug for LocalSet {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("LocalSet").finish()
    }
}

impl Future for LocalSet {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
        // Register the waker before starting to work
        self.context.shared.waker.register_by_ref(cx.waker());

        if self.with(|| self.tick()) {
            // If `tick` returns true, we need to notify the local future again:
            // there are still tasks remaining in the run queue.
            cx.waker().wake_by_ref();
            Poll::Pending
        } else if self.context.owned.is_empty() {
            // If the scheduler has no remaining futures, we're done!

Review comment (suggestion):
```suggestion
            Rc::get_mut(&mut self.context)
                .and_then(|ctx| Arc::get_mut(&mut ctx.shared))
                .expect("Unhandled Panic behavior modified after starting LocalSet")
                .unhandled_panic = behavior;
```

Code after:
        /// tokio-rs/tokio#4516 for more details.
        ///
        /// # Examples
        ///
        /// The following demonstrates a `LocalSet` configured to shutdown on
        /// panic. The first spawned task panics and results in the `LocalSet`
        /// shutting down. The second spawned task never has a chance to
        /// execute. The call to `run_until` will panic due to the runtime being
        /// forcibly shutdown.
        ///
        /// ```should_panic
        /// use tokio::runtime::UnhandledPanic;
        ///
        /// # #[tokio::main]
        /// # async fn main() {
        /// tokio::task::LocalSet::new()
        ///     .unhandled_panic(UnhandledPanic::ShutdownRuntime)
        ///     .run_until(async {
        ///         tokio::task::spawn_local(async { panic!("boom"); });
        ///         tokio::task::spawn_local(async {
        ///             // This task never completes
        ///         });
        ///
        ///         // Do some work, but `run_until` will panic before it completes
        /// # loop { tokio::task::yield_now().await; }
        ///     })
        ///     .await;
        /// # }
        /// ```
        ///
        /// [`JoinHandle`]: struct@crate::task::JoinHandle
        pub fn unhandled_panic(&mut self, behavior: crate::runtime::UnhandledPanic) -> &mut Self {
            // TODO: This should be set as a builder
            Rc::get_mut(&mut self.context)
                .and_then(|ctx| Arc::get_mut(&mut ctx.shared))
                .expect("Unhandled Panic behavior modified after starting LocalSet")
                .unhandled_panic = behavior;
            self
        }
    }
}

impl fmt::Debug for LocalSet {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("LocalSet").finish()
    }
}

impl Future for LocalSet {
    type Output = ();