runtime: storage: Fix storage bucket-to-bag assignment.

runtime-modules/storage/src/lib.rs

@@ -824,6 +824,26 @@ pub struct StorageBucketRecord<WorkerId, AccountId> {
 
     /// Defines limits for a bucket.
     pub voucher: Voucher,
+
+    /// Number of assigned bags.
+    pub assigned_bags: u64,
+}
+
+impl<WorkerId, AccountId> StorageBucketRecord<WorkerId, AccountId> {
+    // Increment the assigned bags number.
+    fn register_bag_assignment(&mut self) {
+        self.assigned_bags = self.assigned_bags.saturating_add(1);
+    }
+
+    // Decrement the assigned bags number.
+    fn unregister_bag_assignment(&mut self) {
+        self.assigned_bags = self.assigned_bags.saturating_sub(1);
+    }
+
+    // Checks the bag assignment number. Returns true if it equals zero.
+    fn no_bags_assigned(&self) -> bool {
+        self.assigned_bags == 0
+    }
 }
 
 // Helper-struct for the data object uploading.
@@ -1537,6 +1557,9 @@ decl_module! {
                 Error::<T>::CannotDeleteNonEmptyStorageBucket
             );
 
+            // Check that no assigned bags left.
+            ensure!(bucket.no_bags_assigned(), Error::<T>::StorageBucketIsBoundToBag);
+
             //
             // == MUTATION SAFE ==
             //
@@ -1716,6 +1739,7 @@ decl_module! {
                 operator_status,
                 accepting_new_bags,
                 voucher,
+                assigned_bags: 0,
             };
 
             let storage_bucket_id = Self::next_storage_bucket_id();
@@ -1773,6 +1797,9 @@ decl_module! {
                 );
             }
 
+            // Update bag counters.
+            Self::change_bag_assignments_for_storage_buckets(&add_buckets, &remove_buckets);
+
             Bags::<T>::mutate(&bag_id, |bag| {
                 bag.update_storage_buckets(&mut add_buckets.clone(), &remove_buckets);
             });
@@ -2223,7 +2250,10 @@ decl_module! {
                 bag.update_distribution_buckets(&mut add_buckets_ids.clone(), &remove_buckets_ids);
             });
 
-            Self::change_bag_assignments(&add_buckets_ids, &remove_buckets_ids);
+            Self::change_bag_assignments_for_distribution_buckets(
+                &add_buckets_ids,
+                &remove_buckets_ids
+            );
 
             Self::deposit_event(
                 RawEvent::DistributionBucketsUpdatedForBag(
@@ -2657,7 +2687,15 @@ impl<T: Trait> DataObjectStorage<T> for Module<T> {
 
         <Bags<T>>::remove(&bag_id);
 
-        Self::change_bag_assignments(&BTreeSet::new(), &deleted_dynamic_bag.distributed_by);
+        Self::change_bag_assignments_for_distribution_buckets(
+            &BTreeSet::new(),
+            &deleted_dynamic_bag.distributed_by,
+        );
+
+        Self::change_bag_assignments_for_storage_buckets(
+            &BTreeSet::new(),
+            &deleted_dynamic_bag.stored_by,
+        );
 
         Self::deposit_event(RawEvent::DynamicBagDeleted(
             deletion_prize_account_id,
@@ -2708,6 +2746,7 @@ impl<T: Trait> DataObjectStorage<T> for Module<T> {
         //
         // == MUTATION SAFE ==
         //
+
         Self::create_dynamic_bag_inner(
             &dynamic_bag_id,
             &deletion_prize,
@@ -2757,10 +2796,6 @@ impl<T: Trait> Module<T> {
         storage_buckets: &BTreeSet<T::StorageBucketId>,
         distribution_buckets: &BTreeSet<DistributionBucketId<T>>,
     ) -> DispatchResult {
-        //
-        // = MUTATION SAFE =
-        //
-
         if let Some(deletion_prize) = deletion_prize.clone() {
             <StorageTreasury<T>>::deposit(&deletion_prize.account_id, deletion_prize.prize)?;
         }
@@ -2776,7 +2811,12 @@ impl<T: Trait> Module<T> {
 
         <Bags<T>>::insert(&bag_id, bag);
 
-        Self::change_bag_assignments(&distribution_buckets, &BTreeSet::new());
+        Self::change_bag_assignments_for_distribution_buckets(
+            &distribution_buckets,
+            &BTreeSet::new(),
+        );
+
+        Self::change_bag_assignments_for_storage_buckets(&storage_buckets, &BTreeSet::new());
 
         Self::deposit_event(RawEvent::DynamicBagCreated(
             dynamic_bag_id.clone(),
@@ -3832,7 +3872,7 @@ impl<T: Trait> Module<T> {
     }
 
     // Add and/or remove distribution buckets assignments to bags.
-    fn change_bag_assignments(
+    fn change_bag_assignments_for_distribution_buckets(
         add_buckets: &BTreeSet<DistributionBucketId<T>>,
         remove_buckets: &BTreeSet<DistributionBucketId<T>>,
     ) {
@@ -3867,6 +3907,28 @@ impl<T: Trait> Module<T> {
         }
     }
 
+    // Add and/or remove storage buckets assignments to bags.
+    fn change_bag_assignments_for_storage_buckets(
+        add_buckets: &BTreeSet<T::StorageBucketId>,
+        remove_buckets: &BTreeSet<T::StorageBucketId>,
+    ) {
+        for bucket_id in add_buckets.iter() {
+            if StorageBucketById::<T>::contains_key(bucket_id) {
+                StorageBucketById::<T>::mutate(bucket_id, |bucket| {
+                    bucket.register_bag_assignment();
+                })
+            }
+        }
+
+        for bucket_id in remove_buckets.iter() {
+            if StorageBucketById::<T>::contains_key(bucket_id) {
+                StorageBucketById::<T>::mutate(bucket_id, |bucket| {
+                    bucket.unregister_bag_assignment();
+                })
+            }
+        }
+    }
+
     // Checks distribution buckets for bag assignment number. Returns true only if all 'assigned_bags' are
     // zero.
     fn no_bags_assigned(family_id: &T::DistributionBucketFamilyId) -> bool {

runtime-modules/storage/src/tests/mod.rs

@@ -329,6 +329,49 @@ fn update_storage_buckets_for_bags_succeeded() {
     });
 }
 
+#[test]
+fn update_storage_buckets_for_bags_succeeded_with_additioonal_checks_on_adding_and_removing() {
+    build_test_externalities().execute_with(|| {
+        set_default_update_storage_buckets_per_bag_limit();
+
+        let static_bag_id = StaticBagId::Council;
+        let bag_id: BagId<Test> = static_bag_id.into();
+
+        let bucket_id = CreateStorageBucketFixture::default()
+            .with_origin(RawOrigin::Signed(STORAGE_WG_LEADER_ACCOUNT_ID))
+            .call_and_assert(Ok(()))
+            .unwrap();
+
+        let add_buckets_ids = BTreeSet::from_iter(vec![bucket_id]);
+
+        UpdateStorageBucketForBagsFixture::default()
+            .with_origin(RawOrigin::Signed(STORAGE_WG_LEADER_ACCOUNT_ID))
+            .with_bag_id(bag_id.clone())
+            .with_add_bucket_ids(add_buckets_ids.clone())
+            .call_and_assert(Ok(()));
+
+        // Add check
+        let bag = Storage::bag(&bag_id);
+        assert_eq!(bag.stored_by, add_buckets_ids);
+
+        let bucket = Storage::storage_bucket_by_id(&bucket_id);
+        assert_eq!(bucket.assigned_bags, 1);
+
+        // ******
+        UpdateStorageBucketForBagsFixture::default()
+            .with_origin(RawOrigin::Signed(STORAGE_WG_LEADER_ACCOUNT_ID))
+            .with_bag_id(bag_id.clone())
+            .with_remove_bucket_ids(add_buckets_ids.clone())
+            .call_and_assert(Ok(()));
+
+        let bag = Storage::bag(&bag_id);
+        assert_eq!(bag.stored_by.len(), 0);
+
+        let bucket = Storage::storage_bucket_by_id(&bucket_id);
+        assert_eq!(bucket.assigned_bags, 0);
+    });
+}
+
 #[test]
 fn update_storage_buckets_for_bags_fails_with_non_existing_dynamic_bag() {
     build_test_externalities().execute_with(|| {
@@ -2690,6 +2733,49 @@ fn delete_dynamic_bags_succeeded_with_assigned_distribution_buckets() {
     });
 }
 
+#[test]
+fn delete_dynamic_bags_succeeded_with_assigned_storage_buckets() {
+    build_test_externalities().execute_with(|| {
+        let initial_balance = 1000;
+        let deletion_prize_value = 77;
+        increase_account_balance(&DEFAULT_MEMBER_ACCOUNT_ID, initial_balance);
+
+        let storage_buckets_number = DefaultMemberDynamicBagNumberOfStorageBuckets::get();
+        create_storage_buckets(storage_buckets_number);
+
+        let dynamic_bag_id = DynamicBagId::<Test>::Member(DEFAULT_MEMBER_ID);
+        CreateDynamicBagFixture::default()
+            .with_bag_id(dynamic_bag_id.clone())
+            .with_deletion_prize(DynamicBagDeletionPrize::<Test> {
+                account_id: DEFAULT_MEMBER_ACCOUNT_ID,
+                prize: deletion_prize_value,
+            })
+            .call_and_assert(Ok(()));
+
+        let bag = Storage::dynamic_bag(&dynamic_bag_id);
+
+        assert_eq!(bag.stored_by.len(), storage_buckets_number as usize);
+
+        let stored_by_bag = bag.stored_by.clone();
+        for bucket_id in &stored_by_bag {
+            let bucket = Storage::storage_bucket_by_id(bucket_id);
+
+            assert_eq!(bucket.assigned_bags, 1);
+        }
+
+        DeleteDynamicBagFixture::default()
+            .with_bag_id(dynamic_bag_id.clone())
+            .with_deletion_account_id(DEFAULT_MEMBER_ACCOUNT_ID)
+            .call_and_assert(Ok(()));
+
+        for bucket_id in &stored_by_bag {
+            let bucket = Storage::storage_bucket_by_id(bucket_id);
+
+            assert_eq!(bucket.assigned_bags, 0);
+        }
+    });
+}
+
 #[test]
 fn delete_dynamic_bags_fails_with_non_existent_dynamic_bag() {
     build_test_externalities().execute_with(|| {
@@ -2801,6 +2887,20 @@ fn delete_storage_bucket_fails_with_non_empty_bucket() {
     });
 }
 
+#[test]
+fn delete_storage_bucket_fails_with_assigned_bag() {
+    build_test_externalities().execute_with(|| {
+        let bag_id = BagId::<Test>::Static(StaticBagId::Council);
+
+        let bucket_id = create_default_storage_bucket_and_assign_to_bag(bag_id.clone());
+
+        DeleteStorageBucketFixture::default()
+            .with_origin(RawOrigin::Signed(STORAGE_WG_LEADER_ACCOUNT_ID))
+            .with_storage_bucket_id(bucket_id)
+            .call_and_assert(Err(Error::<Test>::StorageBucketIsBoundToBag.into()));
+    });
+}
+
 #[test]
 fn remove_storage_bucket_operator_succeeded() {
     build_test_externalities().execute_with(|| {