Repository: tokio-rs/tokio
File: tokio/src/runtime/time/entry.rs (Rust)

Code before:
        // The lock acquisition in clear_entry serves this purpose. All of the
        // driver manipulations happen with the lock held, so we can just take
        // the lock and be sure that this drop happens-after everything the
        // driver did so far and happens-before everything the driver does in
        // the future. While we have the lock held, we also go ahead and
        // deregister the entry if necessary.
        unsafe { self.driver().clear_entry(NonNull::from(self.inner())) };
    }

    pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant) {
        unsafe { self.as_mut().get_unchecked_mut() }.deadline = new_time;
        unsafe { self.as_mut().get_unchecked_mut() }.registered = true;

        let tick = self.driver().time_source().deadline_to_tick(new_time);

        if self.inner().extend_expiration(tick).is_ok() {
            return;
        }

        unsafe {
            self.driver()
                .reregister(&self.driver.driver().io, tick, self.inner().into());
        }
    }

    pub(crate) fn reset_without_reregister(mut self: Pin<&mut Self>, new_time: Instant) {
        unsafe { self.as_mut().get_unchecked_mut() }.deadline = new_time;
        unsafe { self.as_mut().get_unchecked_mut() }.registered = false;

        let tick = self.driver().time_source().deadline_to_tick(new_time);

        if self.inner().extend_expiration(tick).is_ok() {
            return;
        }
    }

    pub(crate) fn poll_elapsed(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), super::Error>> {
        if self.driver().is_shutdown() {
            panic!("{}", crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR);
        }

        if !self.registered {
            let deadline = self.deadline;
            self.as_mut().reset(deadline);
        }

        let this = unsafe { self.get_unchecked_mut() };


Review comment (suggestion):
This function is very similar to `reset`. Maybe `reset` should take an argument `reregister: bool`. I'm not 100% sure, but it is worth considering given the `unsafe` in the fn.

Code after:
        // The lock acquisition in clear_entry serves this purpose. All of the
        // driver manipulations happen with the lock held, so we can just take
        // the lock and be sure that this drop happens-after everything the
        // driver did so far and happens-before everything the driver does in
        // the future. While we have the lock held, we also go ahead and
        // deregister the entry if necessary.
        unsafe { self.driver().clear_entry(NonNull::from(self.inner())) };
    }

    pub(crate) fn reset(mut self: Pin<&mut Self>, new_time: Instant, reregister: bool) {
        unsafe { self.as_mut().get_unchecked_mut() }.deadline = new_time;
        unsafe { self.as_mut().get_unchecked_mut() }.registered = reregister;

        let tick = self.driver().time_source().deadline_to_tick(new_time);

        if self.inner().extend_expiration(tick).is_ok() {
            return;
        }

        if reregister {
            unsafe {
                self.driver()
                    .reregister(&self.driver.driver().io, tick, self.inner().into());
            }
        }
    }

    pub(crate) fn poll_elapsed(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), super::Error>> {
        if self.driver().is_shutdown() {
            panic!("{}", crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR);
        }

        if !self.registered {
            let deadline = self.deadline;
            self.as_mut().reset(deadline, true);
        }

        let this = unsafe { self.get_unchecked_mut() };

        this.inner().state.poll(cx.waker())
    }

    pub(crate) fn driver(&self) -> &super::Handle {
        self.driver.driver().time()
    }

    #[cfg(all(tokio_unstable, feature = "tracing"))]
    pub(crate) fn clock(&self) -> &super::Clock {