Skip to content

Commit

Permalink
Allow oneshot::Receiver::close after successful try_recv
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
stepancheg committed Feb 25, 2021
1 parent c9d2a36 commit b0f8283
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
8 changes: 6 additions & 2 deletions tokio/src/sync/oneshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
///
Expand Down Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions tokio/tests/sync_oneshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>();
Expand Down

0 comments on commit b0f8283

Please sign in to comment.