Split payment and cleanup implementations

webmaster128 · webmaster128 · commit a4a6fbfc7245 · 2023-03-14T17:41:21.000+01:00
diff --git a/contracts/burner/schema/burner.json b/contracts/burner/schema/burner.json
@@ -20,6 +20,7 @@
     ],
     "properties": {
       "payout": {
+        "description": "The address we send all remaining balance to",
         "type": "string"
       }
     },
diff --git a/contracts/burner/schema/raw/migrate.json b/contracts/burner/schema/raw/migrate.json
@@ -7,6 +7,7 @@
   ],
   "properties": {
     "payout": {
+      "description": "The address we send all remaining balance to",
       "type": "string"
     }
   },
diff --git a/contracts/burner/src/contract.rs b/contracts/burner/src/contract.rs
@@ -2,7 +2,7 @@ use cosmwasm_std::{
     entry_point, BankMsg, DepsMut, Env, MessageInfo, Order, Response, StdError, StdResult,
 };
 
-use crate::msg::{InstantiateMsg, MigrateMsg};
+use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg};
 
 #[entry_point]
 pub fn instantiate(
@@ -18,31 +18,58 @@ pub fn instantiate(
 
 #[entry_point]
 pub fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> StdResult<Response> {
-    // delete all state
-    let keys: Vec<_> = deps
-        .storage
-        .range(None, None, Order::Ascending)
-        .map(|(k, _)| k)
-        .collect();
-    let count = keys.len();
-    for k in keys {
-        deps.storage.remove(&k);
-    }
-
     // get balance and send all to recipient
     let balance = deps.querier.query_all_balances(env.contract.address)?;
     let send = BankMsg::Send {
         to_address: msg.payout.clone(),
         amount: balance,
     };
+    Ok(Response::new()
+        .add_message(send)
+        .add_attribute("action", "burn")
+        .add_attribute("payout", msg.payout))
+}
+
+#[entry_point]
+pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> StdResult<Response> {
+    match msg {
+        ExecuteMsg::Cleanup { limit } => execute_cleanup(deps, env, info, limit),
+    }
+}
+
+pub fn execute_cleanup(
+    deps: DepsMut,
+    _env: Env,
+    _info: MessageInfo,
+    limit: Option<u32>,
+) -> StdResult<Response> {
+    // the number of elements we can still take (decreasing over time)
+    let mut limit = limit.unwrap_or(u32::MAX) as usize;
 
-    let data_msg = format!("burnt {} keys", count).into_bytes();
+    let mut deleted = 0;
+    const PER_SCAN: usize = 20;
+    loop {
+        let take_this_scan = std::cmp::min(PER_SCAN, limit);
+        let keys: Vec<_> = deps
+            .storage
+            .range(None, None, Order::Ascending)
+            .take(take_this_scan)
+            .map(|(k, _)| k)
+            .collect();
+        let deleted_this_scan = keys.len();
+        for k in keys {
+            deps.storage.remove(&k);
+        }
+        deleted += deleted_this_scan;
+        limit -= deleted_this_scan;
+        if limit == 0 || deleted_this_scan < take_this_scan {
+            break;
+        }
+    }
 
     Ok(Response::new()
-        .add_message(send)
         .add_attribute("action", "burn")
-        .add_attribute("payout", msg.payout)
-        .set_data(data_msg))
+        .add_attribute("deleted_entries", deleted.to_string()))
 }
 
 #[cfg(test)]
@@ -51,7 +78,18 @@ mod tests {
     use cosmwasm_std::testing::{
         mock_dependencies, mock_dependencies_with_balance, mock_env, mock_info,
     };
-    use cosmwasm_std::{coins, StdError, Storage, SubMsg};
+    use cosmwasm_std::{coins, Attribute, StdError, Storage, SubMsg};
+
+    /// Gets the value of the first attribute with the given key
+    fn first_attr(data: impl AsRef<[Attribute]>, search_key: &str) -> Option<String> {
+        data.as_ref().iter().find_map(|a| {
+            if a.key == search_key {
+                Some(a.value.clone())
+            } else {
+                None
+            }
+        })
+    }
 
     #[test]
     fn instantiate_fails() {
@@ -70,16 +108,9 @@ mod tests {
     }
 
     #[test]
-    fn migrate_cleans_up_data() {
+    fn migrate_sends_funds() {
         let mut deps = mock_dependencies_with_balance(&coins(123456, "gold"));
 
-        // store some sample data
-        deps.storage.set(b"foo", b"bar");
-        deps.storage.set(b"key2", b"data2");
-        deps.storage.set(b"key3", b"cool stuff");
-        let cnt = deps.storage.range(None, None, Order::Ascending).count();
-        assert_eq!(3, cnt);
-
         // change the verifier via migrate
         let payout = String::from("someone else");
         let msg = MigrateMsg {
@@ -96,9 +127,48 @@ mod tests {
                 amount: coins(123456, "gold"),
             })
         );
+    }
+
+    #[test]
+    fn execute_cleans_up_data() {
+        let mut deps = mock_dependencies_with_balance(&coins(123456, "gold"));
+
+        // store some sample data
+        deps.storage.set(b"foo", b"bar");
+        deps.storage.set(b"key2", b"data2");
+        deps.storage.set(b"key3", b"cool stuff");
+        let cnt = deps.storage.range(None, None, Order::Ascending).count();
+        assert_eq!(cnt, 3);
+
+        // change the verifier via migrate
+        let payout = String::from("someone else");
+        let msg = MigrateMsg { payout };
+        let _res = migrate(deps.as_mut(), mock_env(), msg).unwrap();
+
+        let res = execute(
+            deps.as_mut(),
+            mock_env(),
+            mock_info("anon", &[]),
+            ExecuteMsg::Cleanup { limit: Some(2) },
+        )
+        .unwrap();
+        assert_eq!(first_attr(res.attributes, "deleted_entries").unwrap(), "2");
+
+        // One item should be left
+        let cnt = deps.storage.range(None, None, Order::Ascending).count();
+        assert_eq!(cnt, 1);
+
+        let res = execute(
+            deps.as_mut(),
+            mock_env(),
+            mock_info("anon", &[]),
+            ExecuteMsg::Cleanup { limit: Some(2) },
+        )
+        .unwrap();
+        assert_eq!(first_attr(res.attributes, "deleted_entries").unwrap(), "1");
 
-        // check there is no data in storage
+        // Now all are gone
         let cnt = deps.storage.range(None, None, Order::Ascending).count();
-        assert_eq!(0, cnt);
+        assert_eq!(cnt, 0);
     }
 }
diff --git a/contracts/burner/src/msg.rs b/contracts/burner/src/msg.rs
@@ -2,9 +2,22 @@ use cosmwasm_schema::cw_serde;
 
 #[cw_serde]
 pub struct MigrateMsg {
+    /// The address we send all remaining balance to
     pub payout: String,
 }
 
 /// A placeholder where we don't take any input
 #[cw_serde]
 pub struct InstantiateMsg {}
+
+#[cw_serde]
+pub enum ExecuteMsg {
+    /// Cleans up the given number of state elements.
+    /// Call this multiple times to increamentally clean up state.
+    Cleanup {
+        /// The number of state elements to delete.
+        ///
+        /// Set this to None for unlimited cleanup (if your state is small or you are feeling YOLO)
+        limit: Option<u32>,
+    },
+}
diff --git a/contracts/burner/tests/integration.rs b/contracts/burner/tests/integration.rs
@@ -17,17 +17,28 @@
 //!      });
 //! 4. Anywhere you see query(&deps, ...) you must replace it with query(&mut deps, ...)
 
-use cosmwasm_std::{coins, BankMsg, ContractResult, Order, Response, SubMsg};
-use cosmwasm_vm::testing::{instantiate, migrate, mock_env, mock_info, mock_instance};
+use cosmwasm_std::{coins, Attribute, BankMsg, ContractResult, Order, Response, SubMsg};
+use cosmwasm_vm::testing::{execute, instantiate, migrate, mock_env, mock_info, mock_instance};
 
-use burner::msg::{InstantiateMsg, MigrateMsg};
+use burner::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg};
 use cosmwasm_vm::Storage;
 
 // This line will test the output of cargo wasm
 static WASM: &[u8] = include_bytes!("../target/wasm32-unknown-unknown/release/burner.wasm");
 // You can uncomment this line instead to test productionified build from rust-optimizer
 // static WASM: &[u8] = include_bytes!("../contract.wasm");
 
+/// Gets the value of the first attribute with the given key
+fn first_attr(data: impl AsRef<[Attribute]>, search_key: &str) -> Option<String> {
+    data.as_ref().iter().find_map(|a| {
+        if a.key == search_key {
+            Some(a.value.clone())
+        } else {
+            None
+        }
+    })
+}
+
 #[test]
 fn instantiate_fails() {
     let mut deps = mock_instance(WASM, &[]);
@@ -44,21 +55,9 @@ fn instantiate_fails() {
 }
 
 #[test]
-fn migrate_cleans_up_data() {
+fn migrate_sends_funds() {
     let mut deps = mock_instance(WASM, &coins(123456, "gold"));
 
-    // store some sample data
-    deps.with_storage(|storage| {
-        storage.set(b"foo", b"bar").0.unwrap();
-        storage.set(b"key2", b"data2").0.unwrap();
-        storage.set(b"key3", b"cool stuff").0.unwrap();
-        let iter_id = storage.scan(None, None, Order::Ascending).0.unwrap();
-        let cnt = storage.all(iter_id).0.unwrap().len();
-        assert_eq!(3, cnt);
-        Ok(())
-    })
-    .unwrap();
-
     // change the verifier via migrate
     let payout = String::from("someone else");
     let msg = MigrateMsg {
@@ -75,12 +74,61 @@ fn migrate_cleans_up_data() {
             amount: coins(123456, "gold"),
         }),
     );
+}
+
+#[test]
+fn execute_cleans_up_data() {
+    let mut deps = mock_instance(WASM, &coins(123456, "gold"));
+
+    // store some sample data
+    deps.with_storage(|storage| {
+        storage.set(b"foo", b"bar").0.unwrap();
+        storage.set(b"key2", b"data2").0.unwrap();
+        storage.set(b"key3", b"cool stuff").0.unwrap();
+        let iter_id = storage.scan(None, None, Order::Ascending).0.unwrap();
+        let cnt = storage.all(iter_id).0.unwrap().len();
+        assert_eq!(cnt, 3);
+        Ok(())
+    })
+    .unwrap();
+
+    // change the verifier via migrate
+    let payout = String::from("someone else");
+    let msg = MigrateMsg { payout };
+    let _res: Response = migrate(&mut deps, mock_env(), msg).unwrap();
+
+    let res: Response = execute(
+        &mut deps,
+        mock_env(),
+        mock_info("anon", &[]),
+        ExecuteMsg::Cleanup { limit: Some(2) },
+    )
+    .unwrap();
+    assert_eq!(first_attr(res.attributes, "deleted_entries").unwrap(), "2");
+
+    // One item should be left
+    deps.with_storage(|storage| {
+        let iter_id = storage.scan(None, None, Order::Ascending).0.unwrap();
+        let cnt = storage.all(iter_id).0.unwrap().len();
+        assert_eq!(cnt, 1);
+        Ok(())
+    })
+    .unwrap();
+
+    let res: Response = execute(
+        &mut deps,
+        mock_env(),
+        mock_info("anon", &[]),
+        ExecuteMsg::Cleanup { limit: Some(2) },
+    )
+    .unwrap();
+    assert_eq!(first_attr(res.attributes, "deleted_entries").unwrap(), "1");
 
     // check there is no data in storage
     deps.with_storage(|storage| {
         let iter_id = storage.scan(None, None, Order::Ascending).0.unwrap();
         let cnt = storage.all(iter_id).0.unwrap().len();
-        assert_eq!(0, cnt);
+        assert_eq!(cnt, 0);
         Ok(())
     })
     .unwrap();

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@`
`20`	`20`	`],`
`21`	`21`	`"properties": {`
`22`	`22`	`"payout": {`
	`23`	`+ "description": "The address we send all remaining balance to",`
`23`	`24`	`"type": "string"`
`24`	`25`	`}`
`25`	`26`	`},`
Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@`
`7`	`7`	`],`
`8`	`8`	`"properties": {`
`9`	`9`	`"payout": {`
	`10`	`+ "description": "The address we send all remaining balance to",`
`10`	`11`	`"type": "string"`
`11`	`12`	`}`
`12`	`13`	`},`