sync: Add LockFuture for Lock (#1184)

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<T> {
 #[derive(Debug)]
 pub struct LockGuard<T>(Lock<T>);
 
+/// A future that resolves to a `LockGuard`.
+#[derive(Debug)]
+pub struct LockFuture<'a, T> {
+    lock: &'a mut Lock<T>,
+}
+
 // As long as T: Send, it's fine to send and share Lock<T> between threads.
 // If T was not Send, sending and sharing a Lock<T> would be bad, since you can access T through
 // Lock<T>.
@@ -120,6 +128,11 @@ impl<T> Lock<T> {
         };
         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<T> Drop for LockGuard<T> {
@@ -179,3 +192,12 @@ impl<T: fmt::Display> fmt::Display for LockGuard<T> {
         fmt::Display::fmt(&**self, f)
     }
 }
+
+impl<T> Future for LockFuture<'_, T> {
+    type Output = LockGuard<T>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let me = &mut *self;
+        Pin::new(&mut *me.lock).poll_lock(cx)
+    }
+}

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());
+}