Reduce redundant code in disconnect() (#291)

slowtec · Oct 13, 2024 · 4823606 · 4823606
1 parent 263e32a
commit 4823606
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 25 deletions.
diff --git a/src/service/mod.rs b/src/service/mod.rs
@@ -7,6 +7,26 @@ pub(crate) mod rtu;
 #[cfg(feature = "tcp")]
 pub(crate) mod tcp;
 
+#[cfg(any(feature = "rtu", feature = "tcp"))]
+async fn disconnect<T, C>(framed: tokio_util::codec::Framed<T, C>) -> std::io::Result<()>
+where
+    T: tokio::io::AsyncWrite + Unpin,
+{
+    use tokio::io::AsyncWriteExt as _;
+
+    framed
+        .into_inner()
+        .shutdown()
+        .await
+        .or_else(|err| match err.kind() {
+            std::io::ErrorKind::NotConnected | std::io::ErrorKind::BrokenPipe => {
+                // Already disconnected.
+                Ok(())
+            }
+            _ => Err(err),
+        })
+}
+
 /// Check that `req_hdr` is the same `Header` as `rsp_hdr`.
 ///
 /// # Errors

diff --git a/src/service/rtu.rs b/src/service/rtu.rs
@@ -4,7 +4,7 @@
 use std::{fmt, io};
 
 use futures_util::{SinkExt as _, StreamExt as _};
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt as _};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::codec::Framed;
 
 use crate::{
@@ -14,7 +14,7 @@ use crate::{
     ProtocolError, Result,
 };
 
-use super::verify_response_header;
+use super::{disconnect, verify_response_header};
 
 /// Modbus RTU client
 #[derive(Debug)]
@@ -106,17 +106,7 @@ where
             // Already disconnected.
             return Ok(());
         };
-        framed
-            .into_inner()
-            .shutdown()
-            .await
-            .or_else(|err| match err.kind() {
-                io::ErrorKind::NotConnected | io::ErrorKind::BrokenPipe => {
-                    // Already disconnected.
-                    Ok(())
-                }
-                _ => Err(err),
-            })
+        disconnect(framed).await
     }
 }
 

diff --git a/src/service/tcp.rs b/src/service/tcp.rs
@@ -4,7 +4,7 @@
 use std::{fmt, io};
 
 use futures_util::{SinkExt as _, StreamExt as _};
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt as _};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::codec::Framed;
 
 use crate::{
@@ -18,6 +18,8 @@ use crate::{
     ExceptionResponse, ProtocolError, Request, Response, Result,
 };
 
+use super::disconnect;
+
 const INITIAL_TRANSACTION_ID: TransactionId = 0;
 
 #[derive(Debug)]
@@ -137,17 +139,7 @@ where
             // Already disconnected.
             return Ok(());
         };
-        framed
-            .into_inner()
-            .shutdown()
-            .await
-            .or_else(|err| match err.kind() {
-                io::ErrorKind::NotConnected | io::ErrorKind::BrokenPipe => {
-                    // Already disconnected.
-                    Ok(())
-                }
-                _ => Err(err),
-            })
+        disconnect(framed).await
     }
 }