From 134cc4b4cd12a35afd138223c020aea825ed3593 Mon Sep 17 00:00:00 2001
From: Caleb Leinz <caleb.leinz@atlasground.com>
Date: Wed, 13 Dec 2023 08:07:42 -0800
Subject: [PATCH] v0.3.3 - Adjust shutdown on disconnect error

---
 Cargo.toml                 |  2 +-
 src/tokio/io.rs            | 20 ++++++++++++++++----
 tests/integration_tests.rs |  3 ++-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f5b9ddf..e889e3f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "stubborn-io"
-version = "0.3.2"
+version = "0.3.3"
 authors = ["David Raifaizen <david.raifaizen@protonmail.com>"]
 edition = "2021"
 description = "io traits/structs that automatically recover from potential disconnections/interruptions."
diff --git a/src/tokio/io.rs b/src/tokio/io.rs
index 7c1f5ac..d4d1473 100644
--- a/src/tokio/io.rs
+++ b/src/tokio/io.rs
@@ -95,12 +95,24 @@ enum Status<T, C> {
     FailedAndExhausted, // the way one feels after programming in dynamically typed languages
 }
 
+#[inline]
+fn poll_err<T>(
+    kind: ErrorKind,
+    reason: impl Into<Box<dyn std::error::Error + Send + Sync>>,
+) -> Poll<io::Result<T>> {
+    let io_err = io::Error::new(kind, reason);
+    Poll::Ready(Err(io_err))
+}
+
 fn exhausted_err<T>() -> Poll<io::Result<T>> {
-    let io_err = io::Error::new(
+    poll_err(
         ErrorKind::NotConnected,
         "Disconnected. Connection attempts have been exhausted.",
-    );
-    Poll::Ready(Err(io_err))
+    )
+}
+
+fn disconnected_err<T>() -> Poll<io::Result<T>> {
+    poll_err(ErrorKind::NotConnected, "Underlying I/O is disconnected.")
 }
 
 impl<T, C> Deref for StubbornIo<T, C> {
@@ -381,7 +393,7 @@ where
 
                 poll
             }
-            Status::Disconnected(_) => exhausted_err(),
+            Status::Disconnected(_) => disconnected_err(),
             Status::FailedAndExhausted => exhausted_err(),
         }
     }
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index cc81b7b..7c31127 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -23,5 +23,6 @@ async fn back_to_back_shutdown_attempts() {
     let elapsed = tokio::time::timeout(Duration::from_secs(5), connection.shutdown()).await;
 
     let result = elapsed.unwrap();
-    assert!(result.is_err());
+    let error = result.unwrap_err();
+    assert_eq!(error.kind(), std::io::ErrorKind::NotConnected);
 }