Some tests. Also, simplified set_block_status_to_processing.

rs/nns/cmc/BUILD.bazel

@@ -55,6 +55,7 @@ DEV_DEPENDENCIES = [
     "//rs/types/types_test_utils",
     "@crate_index//:candid_parser",
     "@crate_index//:futures",
+    "@crate_index//:maplit",
     "@crate_index//:serde_bytes",
 ]
 

rs/nns/cmc/Cargo.toml

@@ -45,6 +45,7 @@ yansi = "0.5.0"
 
 [dev-dependencies]
 candid_parser = { workspace = true }
+maplit = "1.0.2"
 futures = { workspace = true }
 ic-types-test-utils = { path = "../../types/types_test_utils" }
 serde_bytes = { workspace = true }

rs/nns/cmc/src/main.rs

@@ -143,7 +143,7 @@ pub enum NotificationStatus {
     NotifiedMint(NotifyMintCyclesResult),
     /// The transaction did not have a supported memo (or icrc1_memo).
     /// Therefore, we decided to send the ICP back to its source (minus fee).
-    AutomaticallyRefunded(Result<Option<BlockIndex>, NotifyError>),
+    AutomaticallyRefunded(Option<BlockIndex>),
 }
 
 /// Version of the State type.
@@ -1696,37 +1696,20 @@ fn get_u64_memo(transaction: &Transaction) -> Memo {
     Memo(u64::from_le_bytes(icrc1_memo))
 }
 
-// TODO: Use this in notify_* functions. So much repetition...
-fn set_block_status_to_processing<R>(
+/// If the block has no status in blocks_notified, set it to Processing.
+/// Otherwise, makes no changes, and returns the block's current status.
+fn set_block_status_to_processing(
     block_index: BlockIndex,
-    unless: impl Fn(&NotificationStatus) -> Option<R>,
-) -> Option<Result<R, NotifyError>> {
+) -> Result<(), NotificationStatus> {
     with_state_mut(|state| {
-        let occupied_entry = match state.blocks_notified.entry(block_index) {
-            Entry::Occupied(entry) => entry,
+        match state.blocks_notified.entry(block_index) {
+            Entry::Occupied(entry) => Err(entry.get().clone()),
 
             Entry::Vacant(entry) => {
                 entry.insert(NotificationStatus::Processing);
-                return None;
-            }
-        };
-
-        let caller_should_return = match occupied_entry.get() {
-            NotificationStatus::Processing => Err(NotifyError::Processing),
-
-            status => {
-                if let Some(ok) = unless(status) {
-                    Ok(ok)
-                } else {
-                    Err(NotifyError::InvalidTransaction(format!(
-                        "Block {} has already been processed: {:?}",
-                        block_index, status,
-                    )))
-                }
+                Ok(())
             }
-        };
-
-        Some(caller_should_return)
+        }
     })
 }
 
@@ -1860,15 +1843,37 @@ async fn issue_automatic_refund_if_memo_not_offerred(
     };
 
     // Set block's status to Processing before calling ledger.
-    let early_return_value = set_block_status_to_processing(
-        incoming_block_index,
-        |_notification_status: &NotificationStatus| None,
+    let reason_for_refund = format!(
+        "Memo ({}) in the incoming ICP transfer does not correspond to \
+         any of the operations that the Cycles Minting canister offers.",
+        memo.0,
     );
-    if let Some(early_return_value) = early_return_value {
-        return early_return_value;
+    if let Err(prior_block_status) = set_block_status_to_processing(incoming_block_index) {
+        // Do not proceed, because block is either being processed, or was
+        // finished being processed earlier.
+        use NotificationStatus::{
+            AutomaticallyRefunded, NotifiedCreateCanister, NotifiedMint, NotifiedTopUp,
+            Processing,
+        };
+        return match prior_block_status {
+            Processing => Err(NotifyError::Processing),
+
+            AutomaticallyRefunded(block_index) => Err(NotifyError::Refunded {
+                block_index,
+                reason: reason_for_refund,
+            }),
+
+            // This should not be possible, since we already verified that memo is in MEANINGFUL_MEMOS.
+            NotifiedCreateCanister(_) | NotifiedMint(_) | NotifiedTopUp(_) => {
+                Err(NotifyError::InvalidTransaction(format!(
+                    "Block has already been processed: {:?}",
+                    prior_block_status,
+                )))
+            }
+        };
     }
 
-    // Call ledger to send the ICP back.
+    // Now, it is safe to call ledger to send the ICP back, so do it.
     let refund_result = refund_icp(
         incoming_to_subaccount,
         incoming_from,
@@ -1877,22 +1882,23 @@ async fn issue_automatic_refund_if_memo_not_offerred(
     )
     .await;
     // Handle errors.
-    if let Err(err) = refund_result {
-        // Allow the user to retry.
-        clear_block_processing_status(incoming_block_index);
+    let refund_block_index = refund_result
+        .map_err(|err| {
+            // Allow the user to retry.
+            clear_block_processing_status(incoming_block_index);
 
-        return Err(err);
-    }
+            // Do not actually change the err.
+            err
+        })?;
 
     // Sending the ICP back succeeded. Therefore, update the block's status to
     // AutomaticallyRefunded.
     let old_entry_value = with_state_mut(|state| {
         state.blocks_notified.insert(
             incoming_block_index,
-            NotificationStatus::AutomaticallyRefunded(refund_result.clone()),
+            NotificationStatus::AutomaticallyRefunded(refund_block_index),
         )
     });
-
     // Log if the block's previous status somehow changed out from under us
     // while we were waiting for the ledger call to return. There is no known
     // way for this to happen (except, ofc, bugs).
@@ -1906,12 +1912,8 @@ async fn issue_automatic_refund_if_memo_not_offerred(
     }
 
     Err(NotifyError::Refunded {
-        reason: format!(
-            "Memo ({}) in the incoming ICP transfer does not correspond to \
-             any of the operations that the Cycles Minting canister offers.",
-            memo.0,
-        ),
-        block_index: refund_result.unwrap_or_default(),
+        reason: reason_for_refund,
+        block_index: refund_block_index,
     })
 }
 
@@ -2752,6 +2754,7 @@ fn get_subnet_selection(
 mod tests {
     use super::*;
     use ic_types_test_utils::ids::{subnet_test_id, user_test_id};
+    use maplit::btreemap;
     use rand::Rng;
     use serde_bytes::ByteBuf;
     use std::str::FromStr;
@@ -3637,4 +3640,56 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_set_block_status_to_processing_happy() {
+        let red_herring_block_index = 0xDEADBEEF;
+        STATE.with(|state| {
+            state.replace(Some(State {
+                blocks_notified: btreemap! {
+                    red_herring_block_index => NotificationStatus::Processing,
+                },
+                ..Default::default()
+            }))
+        });
+
+        let target_block_index = 42;
+        let result = set_block_status_to_processing(target_block_index);
+
+        assert_eq!(result, Ok(()));
+        assert_eq!(
+            with_state(|state| state.blocks_notified.clone()),
+            btreemap! {
+                // Existing data untouched.
+                red_herring_block_index => NotificationStatus::Processing,
+                // New entry.
+                target_block_index => NotificationStatus::Processing,
+            },
+        );
+    }
+
+    #[test]
+    fn test_set_block_status_to_processing_already_has_status() {
+        let target_block_index = 42;
+        let red_herring_block_index = 0xDEADBEEF;
+        let original_blocks_notified = btreemap! {
+            red_herring_block_index => NotificationStatus::Processing,
+            // Danger! Block ALREADY has status.
+            target_block_index => NotificationStatus::Processing,
+        };
+        STATE.with(|state| {
+            state.replace(Some(State {
+                blocks_notified: original_blocks_notified.clone(),
+                ..Default::default()
+            }))
+        });
+
+        let result = set_block_status_to_processing(target_block_index);
+
+        assert_eq!(result, Err(NotificationStatus::Processing));
+        assert_eq!(
+            with_state(|state| state.blocks_notified.clone()),
+            original_blocks_notified,
+        );
+    }
 }