From 28ecc3864641170f9344c41d3b983ec358710054 Mon Sep 17 00:00:00 2001
From: Robert Habermeier <rphmeier@gmail.com>
Date: Tue, 15 Jun 2021 17:09:03 +0100
Subject: [PATCH] Drop guard for detecting stale approvals (#3251)

* Drop guard for detecting stale approvals

* address nits with different API
---
 polkadot/node/core/approval-voting/src/lib.rs | 36 +++++++++++++++----
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/polkadot/node/core/approval-voting/src/lib.rs b/polkadot/node/core/approval-voting/src/lib.rs
index e7e00a85f0..597bf378cc 100644
--- a/polkadot/node/core/approval-voting/src/lib.rs
+++ b/polkadot/node/core/approval-voting/src/lib.rs
@@ -1936,6 +1936,29 @@ async fn launch_approval(
 	let (a_tx, a_rx) = oneshot::channel();
 	let (code_tx, code_rx) = oneshot::channel();
 
+	// The background future returned by this function may
+	// be dropped before completing. This guard is used to ensure that the approval
+	// work is correctly counted as stale even if so.
+	struct StaleGuard(Option<Metrics>);
+
+	impl StaleGuard {
+		fn take(mut self) -> Metrics {
+			self.0.take().expect("
+				consumed after take; so this cannot be called twice; \
+				nothing in this function reaches into the struct to avoid this API; \
+				qed
+			")
+		}
+	}
+
+	impl Drop for StaleGuard {
+		fn drop(&mut self) {
+			if let Some(metrics) = self.0.as_ref() {
+				metrics.on_approval_stale();
+			}
+		}
+	}
+
 	let candidate_hash = candidate.hash();
 
 	tracing::trace!(
@@ -1964,6 +1987,7 @@ async fn launch_approval(
 	).await;
 
 	let candidate = candidate.clone();
+	let metrics_guard = StaleGuard(Some(metrics));
 	let background = async move {
 		// Force the move of the timer into the background task.
 		let _timer = timer;
@@ -1982,7 +2006,7 @@ async fn launch_approval(
 					(candidate_hash, candidate.descriptor.para_id),
 				);
 				// do nothing. we'll just be a no-show and that'll cause others to rise up.
-				metrics.on_approval_unavailable();
+				metrics_guard.take().on_approval_unavailable();
 				return;
 			}
 			Ok(Err(RecoveryError::Invalid)) => {
@@ -1994,7 +2018,7 @@ async fn launch_approval(
 
 				// TODO: dispute. Either the merkle trie is bad or the erasure root is.
 				// https://github.com/paritytech/polkadot/issues/2176
-				metrics.on_approval_invalid();
+				metrics_guard.take().on_approval_invalid();
 				return;
 			}
 		};
@@ -2013,7 +2037,7 @@ async fn launch_approval(
 
 				// No dispute necessary, as this indicates that the chain is not behaving
 				// according to expectations.
-				metrics.on_approval_unavailable();
+				metrics_guard.take().on_approval_unavailable();
 				return;
 			}
 		};
@@ -2042,6 +2066,7 @@ async fn launch_approval(
 					"Candidate Valid",
 				);
 
+				let _ = metrics_guard.take();
 				let _ = background_tx.send(BackgroundRequest::ApprovalVote(ApprovalVoteRequest {
 					validator_index,
 					block_hash,
@@ -2059,7 +2084,7 @@ async fn launch_approval(
 
 				// TODO: issue dispute, but not for timeouts.
 				// https://github.com/paritytech/polkadot/issues/2176
-				metrics.on_approval_invalid();
+				metrics_guard.take().on_approval_invalid();
 			}
 			Ok(Err(e)) => {
 				tracing::error!(
@@ -2067,8 +2092,7 @@ async fn launch_approval(
 					err = ?e,
 					"Failed to validate candidate due to internal error",
 				);
-
-				metrics.on_approval_error();
+				metrics_guard.take().on_approval_error();
 				return
 			}
 		}