diff --git a/tokio-sync/src/lock.rs b/tokio-sync/src/lock.rs index 3b8ce36f9..10f6a112d 100644 --- a/tokio-sync/src/lock.rs +++ b/tokio-sync/src/lock.rs @@ -44,7 +44,9 @@ use crate::semaphore; use std::cell::UnsafeCell; use std::fmt; +use std::future::Future; use std::ops::{Deref, DerefMut}; +use std::pin::Pin; use std::sync::Arc; use std::task::Poll::Ready; use std::task::{Context, Poll}; @@ -71,6 +73,12 @@ pub struct Lock { #[derive(Debug)] pub struct LockGuard(Lock); +/// A future that resolves to a `LockGuard`. +#[derive(Debug)] +pub struct LockFuture<'a, T> { + lock: &'a mut Lock, +} + // As long as T: Send, it's fine to send and share Lock between threads. // If T was not Send, sending and sharing a Lock would be bad, since you can access T through // Lock. @@ -120,6 +128,11 @@ impl Lock { }; Ready(LockGuard(acquired)) } + + /// A future that resolves on acquiring the lock and returns the `LockGuard`. + pub fn lock(&mut self) -> LockFuture<'_, T> { + LockFuture { lock: self } + } } impl Drop for LockGuard { @@ -179,3 +192,12 @@ impl fmt::Display for LockGuard { fmt::Display::fmt(&**self, f) } } + +impl Future for LockFuture<'_, T> { + type Output = LockGuard; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let me = &mut *self; + Pin::new(&mut *me.lock).poll_lock(cx) + } +} diff --git a/tokio-sync/tests/lock.rs b/tokio-sync/tests/lock.rs index 33cebef53..bce874d42 100644 --- a/tokio-sync/tests/lock.rs +++ b/tokio-sync/tests/lock.rs @@ -1,5 +1,7 @@ #![deny(warnings, rust_2018_idioms)] +use pin_utils::pin_mut; +use std::future::Future; use tokio_sync::lock::Lock; use tokio_test::task::MockTask; use tokio_test::{assert_pending, assert_ready}; @@ -48,3 +50,29 @@ fn readiness() { assert!(t2.is_woken()); assert_ready!(t2.enter(|cx| l.poll_lock(cx))); } + +#[test] +fn lock() { + let mut lock = Lock::new(false); + + let mut lock2 = lock.clone(); + std::thread::spawn(move || { + let l = lock2.lock(); + pin_mut!(l); + + let mut task = MockTask::new(); + let mut g = task.enter(|cx| assert_ready!(l.poll(cx))); + std::thread::sleep(std::time::Duration::from_millis(500)); + *g = true; + drop(g); + }); + + std::thread::sleep(std::time::Duration::from_millis(50)); + let mut task = MockTask::new(); + let l = lock.lock(); + pin_mut!(l); + + task.enter(|cx| assert_pending!(l.poll(cx))); + std::thread::sleep(std::time::Duration::from_millis(500)); + assert!(task.is_woken()); +}