complete msg only for subscription

specta-rs · Jul 19, 2023 · aa997b1 · aa997b1
1 parent 52bdb3c
commit aa997b1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 23 deletions.
diff --git a/src/internal/exec/connection.rs b/src/internal/exec/connection.rs
@@ -338,10 +338,12 @@ impl<
                 }
                 StreamYield::Finished(f) => {
                     if let Some(stream) = f.take(conn.streams.as_mut()) {
-                        this.batch.as_mut().insert(exec::Response {
-                            id: stream.id(),
-                            inner: ResponseInner::Complete,
-                        });
+                        if let StreamOrFut::OwnedStream { stream } = stream {
+                            this.batch.as_mut().insert(exec::Response {
+                                id: stream.id,
+                                inner: ResponseInner::Complete,
+                            });
+                        }
 
                         PollResult::QueueSend
                     } else {
@@ -368,10 +370,12 @@ impl<
         // TODO: This can be improved by: https://github.com/jonhoo/streamunordered/pull/5
         for (token, _) in conn.steam_to_sub_id.drain() {
             if let Some(stream) = conn.streams.as_mut().take(token) {
-                this.batch.as_mut().insert(exec::Response {
-                    id: stream.id(),
-                    inner: ResponseInner::Complete,
-                });
+                if let StreamOrFut::OwnedStream { stream } = stream {
+                    this.batch.as_mut().insert(exec::Response {
+                        id: stream.id,
+                        inner: ResponseInner::Complete,
+                    });
+                }
             }
         }
         conn.steam_to_sub_id.drain().for_each(drop);

diff --git a/src/internal/exec/stream_or_fut.rs b/src/internal/exec/stream_or_fut.rs
@@ -23,20 +23,7 @@ pin_project! {
             #[pin]
             fut: ExecRequestFut,
         },
-        // TODO: Don't do it like this
-        Done {
-            id: u32,
-        },
-    }
-}
-
-impl<TCtx: 'static> StreamOrFut<TCtx> {
-    pub fn id(&self) -> u32 {
-        match self {
-            StreamOrFut::OwnedStream { stream } => stream.id,
-            StreamOrFut::ExecRequestFut { fut } => fut.id,
-            StreamOrFut::Done { id } => *id,
-        }
+        Done,
     }
 }
 
@@ -57,7 +44,7 @@ impl<TCtx: 'static> Stream for StreamOrFut<TCtx> {
                 }))
             }
             StreamOrFutProj::ExecRequestFut { fut } => fut.poll(cx).map(|v| {
-                self.set(StreamOrFut::Done { id: v.id });
+                self.set(StreamOrFut::Done);
                 Some(v)
             }),
             StreamOrFutProj::Done { .. } => Poll::Ready(None),