Take spans on stream completion

src/lib.rs

@@ -258,7 +258,16 @@ where
 
         let body = this.body;
 
-        span.in_scope(|| dummy_span.in_scope(|| body.poll_next(cx)))
+        let res = span.in_scope(|| dummy_span.in_scope(|| body.poll_next(cx)));
+
+        match res {
+            Poll::Ready(None) => {
+                this.span.take();
+                this.dummy_span.take();
+                Poll::Ready(None)
+            }
+            otherwise => otherwise,
+        }
     }
 }