Sample task poll function

Author: bacv

nomos-da/network/core/src/protocols/sampling/behaviour.rs

@@ -606,6 +606,108 @@ impl<Membership: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'sta
             }
         }
     }
+
+    fn handle_stream_response(
+        peer_id: PeerId,
+        tasks: &FuturesUnordered<SamplingStreamFuture>,
+        to_sample: &mut HashMap<PeerId, VecDeque<(SubnetworkId, BlobId)>>,
+        to_close: &mut VecDeque<SampleStream>,
+        stream_response: SampleStreamSuccess,
+        stream: SampleStream,
+    ) -> Poll<ToSwarm<<Self as NetworkBehaviour>::ToSwarm, THandlerInEvent<Self>>> {
+        match stream_response {
+            SampleStreamSuccess::Writer(sample_response) => {
+                // Handle the free stream then return the response.
+                Self::schedule_outgoing_stream_task(tasks, to_sample, to_close, stream);
+                Self::handle_sample_response(*sample_response, peer_id)
+            }
+            SampleStreamSuccess::Reader => {
+                // Writer might be hoping to send to this stream another request, wait
+                // until the writer closes the stream.
+                let (request_receiver, response_sender) =
+                    Self::schedule_incoming_stream_task(tasks, stream);
+                Poll::Ready(ToSwarm::GenerateEvent(SamplingEvent::IncomingSample {
+                    request_receiver,
+                    response_sender,
+                }))
+            }
+        }
+    }
+
+    fn handle_stream_error(
+        to_sample: &mut HashMap<PeerId, VecDeque<(SubnetworkId, BlobId)>>,
+        to_close: &mut VecDeque<SampleStream>,
+        error: SamplingError,
+        maybe_stream: Option<SampleStream>,
+    ) -> Option<Poll<ToSwarm<<Self as NetworkBehaviour>::ToSwarm, THandlerInEvent<Self>>>> {
+        match error {
+            SamplingError::Io {
+                error,
+                peer_id,
+                message,
+            } if error.kind() == std::io::ErrorKind::ConnectionReset => {
+                // Connection reset with the attached message comes from the stream that we've
+                // tried to write to - if connection reset happens during the write it's most
+                // likely because we didn't have the connection to peer or peer closed stream on
+                // it's end because it stopped waiting for messages through this stream.
+                if let Some(peer_queue) = to_sample.get_mut(&peer_id) {
+                    if let Some(m) = message {
+                        peer_queue.push_back((m.column_idx, m.blob_id));
+                    }
+                }
+                // Stream is useless if connection was reset.
+                if let Some(stream) = maybe_stream {
+                    to_close.push_back(stream);
+                }
+                None
+            }
+            SamplingError::Io { error, .. }
+                if error.kind() == std::io::ErrorKind::UnexpectedEof =>
+            {
+                // Eof is actually expected and is proper signal about remote closing the
+                // stream. Do not propagate and continue execution of behaviour poll method.
+                if let Some(stream) = maybe_stream {
+                    to_close.push_back(stream);
+                }
+                None
+            }
+            error => {
+                if let Some(stream) = maybe_stream {
+                    to_close.push_back(stream);
+                }
+                Some(Poll::Ready(ToSwarm::GenerateEvent(
+                    SamplingEvent::SamplingError { error },
+                )))
+            }
+        }
+    }
+
+    fn poll_tasks(
+        cx: &mut Context<'_>,
+        tasks: &mut FuturesUnordered<SamplingStreamFuture>,
+        to_sample: &mut HashMap<PeerId, VecDeque<(SubnetworkId, BlobId)>>,
+        to_close: &mut VecDeque<SampleStream>,
+    ) -> Option<Poll<ToSwarm<<Self as NetworkBehaviour>::ToSwarm, THandlerInEvent<Self>>>> {
+        if let Poll::Ready(Some(future_result)) = tasks.poll_next_unpin(cx) {
+            cx.waker().wake_by_ref();
+            match future_result {
+                Ok((peer_id, stream_response, stream)) => {
+                    return Some(Self::handle_stream_response(
+                        peer_id,
+                        tasks,
+                        to_sample,
+                        to_close,
+                        stream_response,
+                        stream,
+                    ));
+                }
+                Err((error, maybe_stream)) => {
+                    Self::handle_stream_error(to_sample, to_close, error, maybe_stream);
+                }
+            }
+        }
+        None
+    }
 }
 
 impl<M: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'static> NetworkBehaviour
@@ -670,10 +772,6 @@ impl<M: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'static> Netw
             .on_connection_handler_event(peer_id, connection_id, event);
     }
 
-    #[expect(
-        clippy::too_many_lines,
-        reason = "Poll method contains all the branches in small enough sections"
-    )]
     fn poll(
         &mut self,
         cx: &mut Context<'_>,
@@ -705,6 +803,7 @@ impl<M: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'static> Netw
                 control,
             );
         }
+
         // poll incoming streams
         if let Poll::Ready(Some((peer_id, stream))) = incoming_streams.poll_next_unpin(cx) {
             let sample_stream = SampleStream { stream, peer_id };
@@ -716,76 +815,14 @@ impl<M: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'static> Netw
                 response_sender,
             }));
         }
+
         // poll tasks
-        if let Poll::Ready(Some(future_result)) = tasks.poll_next_unpin(cx) {
-            cx.waker().wake_by_ref(); // Check stream_tasks until it's empty.
-            match future_result {
-                Ok((peer_id, stream_response, stream)) => {
-                    match stream_response {
-                        SampleStreamSuccess::Writer(sample_response) => {
-                            // Schedule a new task if its available or drop the stream if not.
-                            Self::schedule_outgoing_stream_task(tasks, to_sample, to_close, stream);
-                            // handle the free stream then return the success
-                            return Self::handle_sample_response(*sample_response, peer_id);
-                        }
-                        SampleStreamSuccess::Reader => {
-                            // Writer might be hoping to send to this stream
-                            // another request, wait
-                            // until the writer closes the stream.
-                            let (request_receiver, response_sender) =
-                                Self::schedule_incoming_stream_task(tasks, stream);
-                            return Poll::Ready(ToSwarm::GenerateEvent(
-                                SamplingEvent::IncomingSample {
-                                    request_receiver,
-                                    response_sender,
-                                },
-                            ));
-                        }
-                    }
-                }
-                Err((
-                    SamplingError::Io {
-                        error,
-                        peer_id,
-                        message,
-                    },
-                    maybe_stream,
-                )) if error.kind() == std::io::ErrorKind::ConnectionReset => {
-                    if let Some(peer_queue) = to_sample.get_mut(&peer_id) {
-                        if let Some(sampling::SampleRequest {
-                            column_idx,
-                            blob_id,
-                        }) = message
-                        {
-                            peer_queue.push_back((column_idx, blob_id));
-                        }
-                    }
-                    if let Some(stream) = maybe_stream {
-                        to_close.push_back(stream);
-                    }
-                }
-                Err((SamplingError::Io { error, .. }, maybe_stream))
-                    if error.kind() == std::io::ErrorKind::UnexpectedEof =>
-                {
-                    // Eof is actually expected and is proper signal about remote closing the
-                    // stream. Do not propagate and continue execution of this poll method.
-                    if let Some(stream) = maybe_stream {
-                        to_close.push_back(stream);
-                    }
-                }
-                Err((error, maybe_stream)) => {
-                    if let Some(stream) = maybe_stream {
-                        to_close.push_back(stream);
-                    }
-                    return Poll::Ready(ToSwarm::GenerateEvent(SamplingEvent::SamplingError {
-                        error,
-                    }));
-                }
-            }
+        if let Some(result) = Self::poll_tasks(cx, tasks, to_sample, to_close) {
+            return result;
         }
+
         // Deal with connection as the underlying behaviour would do
         if let Poll::Ready(ToSwarm::Dial { mut opts }) = self.stream_behaviour.poll(cx) {
-            cx.waker().wake_by_ref();
             // attach known peer address if possible
             if let Some(address) = opts
                 .get_peer_id()
@@ -796,9 +833,11 @@ impl<M: MembershipHandler<Id = PeerId, NetworkId = SubnetworkId> + 'static> Netw
                     .extend_addresses_through_behaviour()
                     .build();
                 // If we dial, some outgoing task is created, poll again.
+                cx.waker().wake_by_ref();
                 return Poll::Ready(ToSwarm::Dial { opts });
             }
         }
+
         // Discard stream, if still pending pushback to close later.
         if let Some(mut stream) = to_close.pop_front() {
             if stream.stream.close().poll_unpin(cx).is_pending() {

tests/src/topology/configs/da.rs

@@ -60,10 +60,10 @@ impl Default for DaParams {
             policy_settings: DAConnectionPolicySettings {
                 min_dispersal_peers: 1,
                 min_replication_peers: 1,
-                max_dispersal_failures: 1,
-                max_sampling_failures: 1,
-                max_replication_failures: 1,
-                malicious_threshold: 2,
+                max_dispersal_failures: 0,
+                max_sampling_failures: 0,
+                max_replication_failures: 0,
+                malicious_threshold: 0,
             },
             monitor_settings: DAConnectionMonitorSettings {
                 failure_time_window: Duration::from_secs(5),