Document correctness of select! and Poll::Pending

tower-batch/src/future.rs

@@ -49,6 +49,13 @@ where
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut this = self.project();
 
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        // This loop ensures that the task is scheduled as required, because it
+        // only returns Pending when another future returns Pending.
         loop {
             match this.state.as_mut().project() {
                 ResponseStateProj::Failed(e) => {

tower-batch/src/service.rs

@@ -134,12 +134,19 @@ where
             return Poll::Ready(Err(self.get_worker_error()));
         }
 
-        // Then, poll to acquire a semaphore permit. If we acquire a permit,
-        // then there's enough buffer capacity to send a new request. Otherwise,
-        // we need to wait for capacity.
-
-        // In tokio 0.3.7, `acquire_owned` panics if its semaphore returns an error,
-        // so we don't need to handle errors until we upgrade to tokio 1.0.
+        // CORRECTNESS
+        //
+        // Poll to acquire a semaphore permit. If we acquire a permit, then
+        // there's enough buffer capacity to send a new request. Otherwise, we
+        // need to wait for capacity.
+        //
+        // In tokio 0.3.7, `acquire_owned` panics if its semaphore returns an
+        // error, so we don't need to handle errors until we upgrade to
+        // tokio 1.0.
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`. If it returns Pending, the semaphore also schedules
+        // the task for wakeup when the next permit is available.
         ready!(self.semaphore.poll_acquire(cx));
 
         Poll::Ready(Ok(()))

tower-batch/src/worker.rs

@@ -143,6 +143,7 @@ where
                 },
                 Some(mut sleep) => {
                     // Wait on either a new message or the batch timer.
+                    // If both are ready, select! chooses one of them at random.
                     tokio::select! {
                         maybe_msg = self.rx.recv() => match maybe_msg {
                             Some(msg) => {

tower-fallback/src/future.rs

@@ -74,6 +74,13 @@ where
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let mut this = self.project();
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        // This loop ensures that the task is scheduled as required, because it
+        // only returns Pending when a future or service returns Pending.
         loop {
             match this.state.as_mut().project() {
                 ResponseStateProj::PollResponse1 { fut, .. } => match ready!(fut.poll(cx)) {

zebra-consensus/src/chain.rs

@@ -71,7 +71,13 @@ where
         Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        // Correctness:
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        // If either verifier is unready, this task is scheduled for wakeup when it becomes
+        // ready.
         //
         // We acquire checkpoint readiness before block readiness, to avoid an unlikely
         // hang during the checkpoint to block verifier transition. If the checkpoint and

zebra-network/src/peer/client.rs

@@ -113,7 +113,13 @@ impl Stream for ClientRequestReceiver {
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match self.inner.poll_next_unpin(cx) {
             Poll::Ready(client_request) => Poll::Ready(client_request.map(Into::into)),
-            // `inner.poll_next_unpin` parks the task for this future
+            // CORRECTNESS
+            //
+            // The current task must be scheduled for wakeup every time we
+            // return `Poll::Pending`.
+            //
+            // inner.poll_next_unpin` schedules this task for wakeup when
+            // there are new items available in the inner stream.
             Poll::Pending => Poll::Pending,
         }
     }
@@ -198,6 +204,19 @@ impl Service<Request> for Client {
         Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        //`ready!` returns `Poll::Pending` when `server_tx` is unready, and
+        // schedules this task for wakeup.
+        //
+        // Since `shutdown_tx` is used for oneshot communication to the heartbeat
+        // task, it will never be `Pending`.
+        //
+        // TODO: should the Client exit if the heartbeat task exits and drops
+        // `shutdown_tx`?
         if ready!(self.server_tx.poll_ready(cx)).is_err() {
             Poll::Ready(Err(self
                 .error_slot

zebra-network/src/peer_set/set.rs

@@ -464,6 +464,25 @@ where
             if self.preselected_p2c_index.is_none() {
                 trace!("no ready services, sending demand signal");
                 let _ = self.demand_signal.try_send(());
+                // CORRECTNESS
+                //
+                // The current task must be scheduled for wakeup every time we
+                // return `Poll::Pending`.
+                //
+                // As long as there are unready or new peers, this task will run,
+                // because:
+                // - `poll_discover` schedules this task for wakeup when new
+                //   peers arrive.
+                // - if there are unready peers, `poll_unready` schedules this
+                //   task for wakeup when peer services become ready.
+                // - if the preselected peer is not ready, `service.poll_ready`
+                //   schedules this task for wakeup when that service becomes
+                //   ready.
+                //
+                // To avoid peers blocking on a full background error channel:
+                // - if no background tasks have exited since the last poll,
+                //   `poll_background_errors` schedules this task for wakeup when
+                //   the next task exits.
                 return Poll::Pending;
             }
         }

zebra-network/src/peer_set/unready_service.rs

@@ -40,6 +40,15 @@ impl<K, S: Service<Req>, Req> Future for UnreadyService<K, S, Req> {
             return Poll::Ready(Err((key, Error::Canceled)));
         }
 
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        //`ready!` returns `Poll::Pending` when the service is unready, and
+        // schedules this task for wakeup.
+        //
+        // `cancel.poll` also schedules this task for wakeup if it is canceled.
         let res = ready!(this
             .service
             .as_mut()

zebrad/src/components/inbound/downloads.rs

@@ -108,6 +108,15 @@ where
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
         let this = self.project();
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        // If no download and verify tasks have exited since the last poll, this
+        // task is scheduled for wakeup when the next task becomes ready.
+        //
+        // TODO:
         // This would be cleaner with poll_map #63514, but that's nightly only.
         if let Some(join_result) = ready!(this.pending.poll_next(cx)) {
             match join_result.expect("block download and verify tasks must not panic") {
@@ -223,6 +232,8 @@ where
         .in_current_span();
 
         let task = tokio::spawn(async move {
+            // TODO: if the verifier and cancel are both ready, which should we
+            //       prefer? (Currently, select! chooses one at random.)
             tokio::select! {
                 _ = &mut cancel_rx => {
                     tracing::trace!("task cancelled prior to completion");

zebrad/src/components/sync/downloads.rs

@@ -72,6 +72,15 @@ where
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
         let this = self.project();
+        // CORRECTNESS
+        //
+        // The current task must be scheduled for wakeup every time we return
+        // `Poll::Pending`.
+        //
+        // If no download and verify tasks have exited since the last poll, this
+        // task is scheduled for wakeup when the next task becomes ready.
+        //
+        // TODO:
         // This would be cleaner with poll_map #63514, but that's nightly only.
         if let Some(join_result) = ready!(this.pending.poll_next(cx)) {
             match join_result.expect("block download and verify tasks must not panic") {
@@ -149,6 +158,8 @@ where
         let mut verifier = self.verifier.clone();
         let task = tokio::spawn(
             async move {
+                // TODO: if the verifier and cancel are both ready, which should
+                //       we prefer? (Currently, select! chooses one at random.)
                 let rsp = tokio::select! {
                     _ = &mut cancel_rx => {
                         tracing::trace!("task cancelled prior to download completion");
@@ -169,6 +180,8 @@ where
                 metrics::counter!("sync.downloaded.block.count", 1);
 
                 let rsp = verifier.ready_and().await?.call(block);
+                // TODO: if the verifier and cancel are both ready, which should
+                //       we prefer? (Currently, select! chooses one at random.)
                 let verification = tokio::select! {
                     _ = &mut cancel_rx => {
                         tracing::trace!("task cancelled prior to verification");

zebrad/src/components/tokio.rs

@@ -46,6 +46,8 @@ pub(crate) trait RuntimeRun {
 impl RuntimeRun for Runtime {
     fn run(&mut self, fut: impl Future<Output = Result<(), Report>>) {
         let result = self.block_on(async move {
+            // If the run task and shutdown are both ready, select! chooses
+            // one of them at random.
             tokio::select! {
                 result = fut => result,
                 _ = shutdown() => Ok(()),
@@ -68,6 +70,7 @@ mod imp {
     use tracing::info;
 
     pub(super) async fn shutdown() {
+        // If both signals are received, select! chooses one of them at random.
         tokio::select! {
             // SIGINT  - Terminal interrupt signal. Typically generated by shells in response to Ctrl-C.
             () = sig(SignalKind::interrupt(), "SIGINT") => {}