Remove all expired allocations for a client when none specified

actors/verifreg/src/lib.rs

@@ -24,8 +24,8 @@ use fil_actors_runtime::cbor::{deserialize, serialize};
 use fil_actors_runtime::runtime::builtins::Type;
 use fil_actors_runtime::runtime::{ActorCode, Policy, Runtime};
 use fil_actors_runtime::{
-    actor_error, cbor, make_map_with_root_and_bitwidth, resolve_to_actor_id, ActorDowncast,
-    ActorError, Map, STORAGE_MARKET_ACTOR_ADDR, SYSTEM_ACTOR_ADDR,
+    actor_error, cbor, make_map_with_root_and_bitwidth, parse_uint_key, resolve_to_actor_id,
+    ActorDowncast, ActorError, Map, STORAGE_MARKET_ACTOR_ADDR, SYSTEM_ACTOR_ADDR,
     UNIVERSAL_RECEIVER_HOOK_METHOD_NUM,
 };
 use fil_actors_runtime::{ActorContext, AsActorError, BatchReturnGen, DATACAP_TOKEN_ACTOR_ADDR};
@@ -334,6 +334,7 @@ impl Actor {
 
     // An allocation may be removed after its expiration epoch has passed (by anyone).
     // When removed, the DataCap tokens are transferred back to the client.
+    // If no allocations are specified, all eligible allocations are removed.
     pub fn remove_expired_allocations<BS, RT>(
         rt: &mut RT,
         params: RemoveExpiredAllocationsParams,
@@ -342,46 +343,71 @@ impl Actor {
         BS: Blockstore,
         RT: Runtime<BS>,
     {
-        // since the alloc is expired this should be safe to publically cleanup
+        // Since the allocations are expired, this is safe to be called by anyone.
         rt.validate_immediate_caller_accept_any()?;
         let client = params.client;
+        let curr_epoch = rt.curr_epoch();
         let mut ret_gen = BatchReturnGen::new(params.allocation_ids.len());
-        let mut recovered_datacap = DataCap::zero();
-        rt.transaction(|st: &mut State, rt| {
-            let mut allocs = st.load_allocs(rt.store())?;
-            for alloc_id in params.allocation_ids {
-                let maybe_alloc = allocs.get(client, alloc_id).context_code(
-                    ExitCode::USR_ILLEGAL_STATE,
-                    "HAMT lookup failure getting allocation",
-                )?;
-                let alloc = match maybe_alloc {
-                    None => {
+        // let mut recovered_datacap = DataCap::zero();
+        let recovered_datacap = rt
+            .transaction(|st: &mut State, rt| {
+                let mut allocs = st.load_allocs(rt.store())?;
+                let mut to_remove = Vec::<AllocationID>::new();
+                if params.allocation_ids.is_empty() {
+                    // Find all expired allocations for the client.
+                    allocs
+                        .for_each(params.client, |key, alloc| {
+                            if alloc.expiration <= curr_epoch {
+                                let id = parse_uint_key(key).context_code(
+                                    ExitCode::USR_ILLEGAL_STATE,
+                                    "failed to parse uint key",
+                                )?;
+                                to_remove.push(id);
+                            }
+                            Ok(())
+                        })
+                        .context_code(
+                            ExitCode::USR_ILLEGAL_STATE,
+                            "failed to iterate over allocations",
+                        )?;
+                }
+                for alloc_id in params.allocation_ids {
+                    // Check each specified allocation is expired.
+                    let maybe_alloc = allocs.get(client, alloc_id).context_code(
+                        ExitCode::USR_ILLEGAL_STATE,
+                        "HAMT lookup failure getting allocation",
+                    )?;
+                    if let Some(alloc) = maybe_alloc {
+                        if alloc.expiration <= curr_epoch {
+                            to_remove.push(alloc_id);
+                            ret_gen.add_success();
+                        } else {
+                            ret_gen.add_fail(ExitCode::USR_FORBIDDEN);
+                            info!("cannot revoke allocation {} that has not expired", alloc_id);
+                        }
+                    } else {
                         ret_gen.add_fail(ExitCode::USR_NOT_FOUND);
                         info!(
                             "claim references allocation id {} that does not belong to client {}",
                             alloc_id, client,
                         );
-                        continue;
                     }
-                    Some(a) => a,
-                };
-                if alloc.expiration > rt.curr_epoch() {
-                    ret_gen.add_fail(ExitCode::USR_FORBIDDEN);
-                    info!("cannot revoke allocation {} that has not expired", alloc_id);
-                    continue;
                 }
-                recovered_datacap += alloc.size.0;
 
-                allocs.remove(client, alloc_id).context_code(
-                    ExitCode::USR_ILLEGAL_STATE,
-                    format!("failed to remove allocation {}", alloc_id),
-                )?;
-                ret_gen.add_success();
-            }
-            st.save_allocs(&mut allocs)?;
-            Ok(())
-        })
-        .context("state transaction failed")?;
+                let mut recovered_datacap = DataCap::zero();
+                for id in to_remove {
+                    let existing = allocs.remove(client, id).context_code(
+                        ExitCode::USR_ILLEGAL_STATE,
+                        format!("failed to remove allocation {}", id),
+                    )?;
+                    // Unwrapping here as both paths to here should ensure the allocation exists.
+                    recovered_datacap += existing.unwrap().size.0;
+                }
+
+                st.save_allocs(&mut allocs)?;
+                Ok(recovered_datacap)
+            })
+            .context("state transaction failed")?;
 
         // Transfer the recovered datacap back to the client.
         transfer(rt, client, &recovered_datacap).with_context(|| {

actors/verifreg/tests/harness/mod.rs

@@ -206,19 +206,37 @@ impl Harness {
     }
 
     // TODO this should be implemented through a call to verifreg but for now it modifies state directly
-    pub fn create_alloc(&self, rt: &mut MockRuntime, alloc: &Allocation) -> Result<(), ActorError> {
+    pub fn create_alloc(
+        &self,
+        rt: &mut MockRuntime,
+        alloc: &Allocation,
+    ) -> Result<AllocationID, ActorError> {
         let mut st: State = rt.get_state();
         let mut allocs =
             MapMap::from_root(rt.store(), &st.allocations, HAMT_BIT_WIDTH, HAMT_BIT_WIDTH)
                 .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to load allocations table")?;
+        let alloc_id = st.next_allocation_id;
         assert!(allocs
-            .put_if_absent(alloc.client, st.next_allocation_id, alloc.clone())
+            .put_if_absent(alloc.client, alloc_id, alloc.clone())
             .context_code(ExitCode::USR_ILLEGAL_STATE, "faild to put")?);
         st.next_allocation_id += 1;
         st.allocations = allocs.flush().expect("failed flushing allocation table");
         rt.replace_state(&st);
+        Ok(alloc_id)
+    }
 
-        Ok(())
+    pub fn load_alloc(
+        &self,
+        rt: &mut MockRuntime,
+        client: &Address,
+        id: AllocationID,
+    ) -> Option<Allocation> {
+        let st: State = rt.get_state();
+        let mut allocs =
+            MapMap::from_root(rt.store(), &st.allocations, HAMT_BIT_WIDTH, HAMT_BIT_WIDTH)
+                .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to load allocations table")
+                .unwrap();
+        allocs.get(client.id().unwrap(), id).unwrap().cloned()
     }
 
     // Invokes the ClaimAllocations actor method

actors/verifreg/tests/verifreg_actor_test.rs

@@ -425,7 +425,6 @@ mod claims {
 
     use fil_actor_verifreg::State;
     use fil_actors_runtime::runtime::Runtime;
-    use fil_actors_runtime::BatchReturnGen;
     use harness::*;
 
     use crate::*;
@@ -438,44 +437,59 @@ mod claims {
         alloc1.expiration = 100;
         let mut alloc2 = make_alloc("2", &CLIENT, &PROVIDER, 256);
         alloc2.expiration = 200;
+        let total_size = alloc1.size.0 + alloc2.size.0;
 
-        h.create_alloc(&mut rt, &alloc1).unwrap();
-        h.create_alloc(&mut rt, &alloc2).unwrap();
+        let id1 = h.create_alloc(&mut rt, &alloc1).unwrap();
+        let id2 = h.create_alloc(&mut rt, &alloc2).unwrap();
         let state_with_allocs: State = rt.get_state();
 
         // Can't remove allocations that aren't expired
         let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![1, 2], 0).unwrap();
-        assert_eq!(
-            BatchReturnGen::new(2)
-                .add_fail(ExitCode::USR_FORBIDDEN)
-                .add_fail(ExitCode::USR_FORBIDDEN)
-                .gen(),
-            ret
-        );
+        assert_eq!(vec![ExitCode::USR_FORBIDDEN, ExitCode::USR_FORBIDDEN], ret.codes());
 
         // Remove the first alloc, which expired.
         rt.set_epoch(100);
         let ret =
             h.remove_expired_allocations(&mut rt, &CLIENT, vec![1, 2], alloc1.size.0).unwrap();
-        assert_eq!(
-            BatchReturnGen::new(2).add_success().add_fail(ExitCode::USR_FORBIDDEN).gen(),
-            ret
-        );
+        assert_eq!(vec![ExitCode::OK, ExitCode::USR_FORBIDDEN], ret.codes());
 
         // Remove the second alloc (the first is no longer found).
         rt.set_epoch(200);
         let ret =
             h.remove_expired_allocations(&mut rt, &CLIENT, vec![1, 2], alloc2.size.0).unwrap();
-        assert_eq!(
-            BatchReturnGen::new(2).add_fail(ExitCode::USR_NOT_FOUND).add_success().gen(),
-            ret
-        );
+        assert_eq!(vec![ExitCode::USR_NOT_FOUND, ExitCode::OK], ret.codes());
 
         // Reset state and show we can remove two at once.
         rt.replace_state(&state_with_allocs);
-        let total_size = alloc1.size.0 + alloc2.size.0;
         let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![1, 2], total_size).unwrap();
-        assert_eq!(BatchReturnGen::new(2).add_success().add_success().gen(), ret);
+        assert_eq!(vec![ExitCode::OK, ExitCode::OK], ret.codes());
+
+        // Reset state and show that specifying none removes only expired allocations
+        rt.set_epoch(0);
+        rt.replace_state(&state_with_allocs);
+        let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![], 0).unwrap();
+        assert_eq!(Vec::<ExitCode>::new(), ret.codes()); // Empty response when no allocs specified
+        assert!(h.load_alloc(&mut rt, &CLIENT, id1).is_some());
+        assert!(h.load_alloc(&mut rt, &CLIENT, id2).is_some());
+
+        rt.set_epoch(100);
+        let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![], alloc1.size.0).unwrap();
+        assert_eq!(Vec::<ExitCode>::new(), ret.codes());
+        assert!(h.load_alloc(&mut rt, &CLIENT, id1).is_none()); // removed
+        assert!(h.load_alloc(&mut rt, &CLIENT, id2).is_some());
+
+        rt.set_epoch(200);
+        let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![], alloc2.size.0).unwrap();
+        assert_eq!(Vec::<ExitCode>::new(), ret.codes());
+        assert!(h.load_alloc(&mut rt, &CLIENT, id1).is_none()); // removed
+        assert!(h.load_alloc(&mut rt, &CLIENT, id2).is_none()); // removed
+
+        // Reset state and show that specifying none removes *all* expired allocations
+        rt.replace_state(&state_with_allocs);
+        let ret = h.remove_expired_allocations(&mut rt, &CLIENT, vec![], total_size).unwrap();
+        assert_eq!(Vec::<ExitCode>::new(), ret.codes());
+        assert!(h.load_alloc(&mut rt, &CLIENT, id1).is_none()); // removed
+        assert!(h.load_alloc(&mut rt, &CLIENT, id2).is_none()); // removed
     }
 
     #[test]

runtime/src/util/mapmap.rs

@@ -102,6 +102,18 @@ where
         in_map.get(&inside_k.key())
     }
 
+    // Runs a function over all values for one outer key.
+    pub fn for_each<F>(&mut self, outside_k: K1, f: F) -> Result<(), Error>
+    where
+        F: FnMut(&BytesKey, &V) -> anyhow::Result<()>,
+    {
+        let (is_empty, in_map) = self.load_inner_map(outside_k)?;
+        if is_empty {
+            return Ok(());
+        }
+        in_map.for_each(f)
+    }
+
     // Puts a key value pair in the MapMap if it is not already set.  Returns true
     // if key is newly set, false if it was already set.
     pub fn put_if_absent(&mut self, outside_k: K1, inside_k: K2, value: V) -> Result<bool, Error> {