Call on_eos for successful streams (#580)

Author: mladedav

tower-http/CHANGELOG.md

@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+# Unreleased
+
+## Fixed
+
+- `on_eos` is now called even for successful responses. ([#580])
+
+[#580]: https://github.com/tower-rs/tower-http/pull/580
+
 # 0.6.6
 
 ## Fixed

tower-http/src/trace/future.rs

@@ -80,7 +80,7 @@ where
                         let res = res.map(|body| ResponseBody {
                             inner: body,
                             classify_eos: None,
-                            on_eos: None,
+                            on_eos: on_eos.zip(Some(Instant::now())),
                             on_body_chunk,
                             on_failure: Some(on_failure),
                             start,

tower-http/src/trace/mod.rs

@@ -556,7 +556,7 @@ mod tests {
             .await
             .unwrap();
         assert_eq!(1, ON_BODY_CHUNK_COUNT.load(Ordering::SeqCst), "body chunk");
-        assert_eq!(0, ON_EOS.load(Ordering::SeqCst), "eos");
+        assert_eq!(1, ON_EOS.load(Ordering::SeqCst), "eos");
         assert_eq!(0, ON_FAILURE.load(Ordering::SeqCst), "failure");
     }
 
@@ -611,7 +611,7 @@ mod tests {
             .await
             .unwrap();
         assert_eq!(3, ON_BODY_CHUNK_COUNT.load(Ordering::SeqCst), "body chunk");
-        assert_eq!(0, ON_EOS.load(Ordering::SeqCst), "eos");
+        assert_eq!(1, ON_EOS.load(Ordering::SeqCst), "eos");
         assert_eq!(0, ON_FAILURE.load(Ordering::SeqCst), "failure");
     }