Simply IdOrSpan

src/lib.rs

@@ -139,26 +139,36 @@ where
 }
 
 enum IdOrSpan {
-    Empty,
-    Id(Id),
+    Id(Option<Id>),
     Span(Span),
 }
 
 impl IdOrSpan {
     fn from_id(id: Option<Id>) -> Self {
-        id.map(IdOrSpan::Id).unwrap_or(IdOrSpan::Empty)
+        IdOrSpan::Id(id)
     }
 
     fn take(&mut self) -> Self {
-        std::mem::replace(self, IdOrSpan::Empty)
+        std::mem::replace(self, IdOrSpan::Id(None))
     }
 
-    fn as_span(&self) -> Option<&Span> {
+    fn span(&self) -> Option<&Span> {
         match self {
             IdOrSpan::Span(ref span) => Some(span),
             _ => None,
         }
     }
+
+    fn as_span(&mut self, f: impl Fn(Option<Id>) -> Span) -> &Span {
+        let span = match self.take() {
+            Self::Id(id) => f(id),
+            Self::Span(span) => span,
+        };
+
+        *self = Self::Span(span);
+
+        self.span().expect("Span should always exist")
+    }
 }
 
 pin_project_lite::pin_project! {
@@ -258,28 +268,14 @@ where
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.as_mut().project();
 
-        let span = this.span.take();
+        let span = this.span.as_span(|id| {
+            let span = tracing::info_span!("HTTP Client Response Body");
+            span.follows_from(id);
+            span
+        });
         let body = this.body;
 
-        let span = match span {
-            IdOrSpan::Empty => {
-                return body.poll_next(cx);
-            }
-            IdOrSpan::Id(id) => {
-                let span = tracing::info_span!("HTTP Client Response Body");
-                span.follows_from(id);
-                span
-            }
-            IdOrSpan::Span(span) => span,
-        };
-
-        *this.span = IdOrSpan::Span(span);
-
-        if let Some(span) = this.span.as_span() {
-            span.in_scope(|| body.poll_next(cx))
-        } else {
-            unreachable!("The span should always exist by this point")
-        }
+        span.in_scope(|| body.poll_next(cx))
     }
 }