From b0f82837683c9fde081df275060c77744a88cf54 Mon Sep 17 00:00:00 2001
From: Stiopa Koltsov <stepan.koltsov@gmail.com>
Date: Thu, 25 Feb 2021 00:45:12 +0000
Subject: [PATCH] Allow oneshot::Receiver::close after successful try_recv

Before this commit `close` after successful `try_recv` panics.

My use case is this: on drop, I call `close` to prevent pushing a
message to the queue, and then fetch the message if any and process
it.

But if message is already processed, `close` panics.  And there is
no API to know if message was already fetched or not (except for
writing a wrapped which would track that info, which would be an
overkill).

But generally `close` operation should be safe to be called any
time.
---
 tokio/src/sync/oneshot.rs   |  8 ++++++--
 tokio/tests/sync_oneshot.rs | 10 ++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/tokio/src/sync/oneshot.rs b/tokio/src/sync/oneshot.rs
index 20d39dc1a1e..4954da67ec2 100644
--- a/tokio/src/sync/oneshot.rs
+++ b/tokio/src/sync/oneshot.rs
@@ -443,6 +443,9 @@ impl<T> Receiver<T> {
     /// This function is useful to perform a graceful shutdown and ensure that a
     /// value will not be sent into the channel and never received.
     ///
+    /// `close` is no-op if a message is already received or the channel
+    /// is already closed.
+    ///
     /// [`Sender`]: Sender
     /// [`try_recv`]: Receiver::try_recv
     ///
@@ -490,8 +493,9 @@ impl<T> Receiver<T> {
     /// }
     /// ```
     pub fn close(&mut self) {
-        let inner = self.inner.as_ref().unwrap();
-        inner.close();
+        if let Some(inner) = self.inner.as_ref() {
+            inner.close();
+        }
     }
 
     /// Attempts to receive a value.
diff --git a/tokio/tests/sync_oneshot.rs b/tokio/tests/sync_oneshot.rs
index 195c2553145..b1ea46a46ef 100644
--- a/tokio/tests/sync_oneshot.rs
+++ b/tokio/tests/sync_oneshot.rs
@@ -180,6 +180,16 @@ fn close_try_recv_poll() {
     let _ = rx.poll();
 }
 
+#[test]
+fn close_after_recv() {
+    let (tx, mut rx) = oneshot::channel::<i32>();
+
+    tx.send(17).unwrap();
+
+    assert_eq!(17, rx.try_recv().unwrap());
+    rx.close();
+}
+
 #[test]
 fn drops_tasks() {
     let (mut tx, mut rx) = oneshot::channel::<i32>();