From 5f099168cb249c9102f767e131358ab063682dc4 Mon Sep 17 00:00:00 2001
From: eitanm-starkware <eitan.m@starkware.co>
Date: Wed, 24 Jul 2024 12:13:27 +0300
Subject: [PATCH] chore(network): make streamhashmap not end

---
 crates/papyrus_network/src/utils.rs | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/crates/papyrus_network/src/utils.rs b/crates/papyrus_network/src/utils.rs
index 2ac26a8c15..319ac0c44e 100644
--- a/crates/papyrus_network/src/utils.rs
+++ b/crates/papyrus_network/src/utils.rs
@@ -36,6 +36,7 @@ impl<K: Unpin + Clone + Eq + Hash, V: Stream + Unpin> StreamHashMap<K, V> {
     }
 
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        self.finished_streams.remove(&key);
         self.map.insert(key, value)
     }
 }
@@ -45,7 +46,6 @@ impl<K: Unpin + Clone + Eq + Hash, V: Stream + Unpin> Stream for StreamHashMap<K
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let unpinned_self = Pin::into_inner(self);
-        let mut finished = true;
         for (key, stream) in &mut unpinned_self.map {
             match stream.poll_next_unpin(cx) {
                 Poll::Ready(Some(value)) => {
@@ -54,15 +54,9 @@ impl<K: Unpin + Clone + Eq + Hash, V: Stream + Unpin> Stream for StreamHashMap<K
                 Poll::Ready(None) => {
                     unpinned_self.finished_streams.insert(key.clone());
                 }
-                Poll::Pending => {
-                    finished = false;
-                }
+                Poll::Pending => {}
             }
         }
-        if finished {
-            // TODO(shahak): Make StreamHashMap not end in order to accept new inserted streams.
-            return Poll::Ready(None);
-        }
         Poll::Pending
     }
 }