Optimize Session::close: Avoid dup code

 - Avoid duplicate call to `close_impl` in both `process_impl` and
   `native_mux_impl`
 - Extract out code to close `TempDir` out into `Session::close`

Signed-off-by: Jiahao XU <[email protected]>

Author: NobodyXu

src/native_mux_impl/session.rs

@@ -76,17 +76,13 @@ impl Session {
         Ok(())
     }
 
-    pub(crate) async fn close(mut self) -> Result<(), Error> {
-        // This also set self.tempdir to None so that Drop::drop would do nothing.
-        if let Some(tempdir) = self.tempdir.take() {
-            self.close_impl().await?;
-
-            tempdir.close().map_err(Error::Cleanup)?;
-        } else {
-            self.close_impl().await?;
-        }
+    pub(crate) async fn close(mut self) -> Result<Option<TempDir>, Error> {
+        // Take self.tempdir so that drop would do nothing
+        let tempdir = self.tempdir.take();
 
-        Ok(())
+        self.close_impl().await?;
+
+        Ok(tempdir)
     }
 
     pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {

src/process_impl/session.rs

@@ -175,17 +175,13 @@ impl Session {
         Ok(())
     }
 
-    pub(crate) async fn close(mut self) -> Result<(), Error> {
+    pub(crate) async fn close(mut self) -> Result<Option<TempDir>, Error> {
         // Take self.tempdir so that drop would do nothing
-        if let Some(tempdir) = self.tempdir.take() {
-            self.close_impl().await?;
+        let tempdir = self.tempdir.take();
 
-            tempdir.close().map_err(Error::Cleanup)?;
-        } else {
-            self.close_impl().await?;
-        }
+        self.close_impl().await?;
 
-        Ok(())
+        Ok(tempdir)
     }
 
     pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {

src/session.rs

@@ -362,7 +362,12 @@ impl Session {
     /// This destructor terminates the ssh multiplex server
     /// regardless of how it was created.
     pub async fn close(self) -> Result<(), Error> {
-        delegate!(self.0, imp, { imp.close().await })
+        let res: Result<Option<TempDir>, Error> = delegate!(self.0, imp, { imp.close().await });
+
+        res?.map(TempDir::close)
+            .transpose()
+            .map_err(Error::Cleanup)
+            .map(|_| ())
     }
 
     /// Detach the lifetime of underlying ssh multiplex master