runtime: storage: Add random distribution buckets pick.

runtime-modules/storage/src/distribution_bucket_picker.rs

@@ -0,0 +1,93 @@
+#![warn(missing_docs)]
+
+use frame_support::traits::Randomness;
+use sp_arithmetic::traits::Zero;
+use sp_runtime::SaturatedConversion;
+use sp_std::cell::RefCell;
+use sp_std::collections::btree_set::BTreeSet;
+use sp_std::marker::PhantomData;
+
+use crate::{DynamicBagType, Module, Trait};
+use sp_std::rc::Rc;
+
+// Generates distribution bucket IDs to assign to a new dynamic bag.
+pub(crate) struct DistributionBucketPicker<T> {
+    trait_marker: PhantomData<T>,
+}
+
+impl<T: Trait> DistributionBucketPicker<T> {
+    // Get random distribution buckets from distribution bucket families using the dynamic bag
+    // creation policy.
+    pub(crate) fn pick_distribution_buckets(
+        bag_type: DynamicBagType,
+    ) -> BTreeSet<T::DistributionBucketId> {
+        let creation_policy = Module::<T>::get_dynamic_bag_creation_policy(bag_type);
+
+        if creation_policy.no_distribution_buckets_required() {
+            return BTreeSet::new();
+        }
+
+        // Randomness for all bucket family.
+        // let random_seed = RefCell::new(Module::<T>::get_initial_random_seed());
+        let random_seed = Rc::new(RefCell::new(Module::<T>::get_initial_random_seed()));
+
+        creation_policy
+            .families
+            .iter()
+            .filter_map(|(family_id, bucket_num)| {
+                Module::<T>::ensure_distribution_bucket_family_exists(family_id)
+                    .ok()
+                    .map(|fam| (fam, bucket_num))
+            })
+            .map(|(family, bucket_num)| {
+                let filtered_ids = family
+                    .distribution_buckets
+                    .iter()
+                    .filter_map(|(id, bucket)| bucket.accepting_new_bags.then(|| *id))
+                    .collect::<Vec<_>>();
+
+                (filtered_ids, bucket_num)
+            })
+            .map(|(bucket_ids, bucket_num)| {
+                Self::get_random_distribution_buckets(bucket_ids, *bucket_num, random_seed.clone())
+            })
+            .flatten()
+            .collect::<BTreeSet<_>>()
+    }
+
+    // Get random bucket IDs from the ID collection.
+    pub fn get_random_distribution_buckets(
+        ids: Vec<T::DistributionBucketId>,
+        bucket_number: u32,
+        seed: Rc<RefCell<T::Hash>>, //     seed: RefCell<T::Hash>
+    ) -> BTreeSet<T::DistributionBucketId> {
+        let mut working_ids = ids.clone();
+        let mut result_ids = BTreeSet::default();
+
+        for _ in 0..bucket_number {
+            if working_ids.is_empty() {
+                break;
+            }
+
+            let current_seed = Self::advance_random_seed(seed.clone());
+
+            let upper_bound = working_ids.len() as u64 - 1;
+            let index =
+                Module::<T>::random_index(current_seed.as_ref(), upper_bound).saturated_into();
+            result_ids.insert(working_ids.remove(index));
+        }
+
+        result_ids
+    }
+
+    // Changes the internal seed value of the container and returns new random seed.
+    fn advance_random_seed(seed: Rc<RefCell<T::Hash>>) -> T::Hash {
+        // Cannot create randomness in the initial block (Substrate error).
+        if <frame_system::Module<T>>::block_number() == Zero::zero() {
+            return Module::<T>::get_initial_random_seed();
+        }
+
+        let current_seed = seed.borrow().clone();
+        seed.replace(T::Randomness::random(current_seed.as_ref()))
+    }
+}

runtime-modules/storage/src/lib.rs

@@ -109,6 +109,7 @@ mod tests;
 mod benchmarking;
 
 mod bag_manager;
+pub(crate) mod distribution_bucket_picker;
 pub(crate) mod storage_bucket_picker;
 
 use codec::{Codec, Decode, Encode};
@@ -132,6 +133,7 @@ use common::origin::ActorOriginValidator;
 use common::working_group::WorkingGroup;
 
 use bag_manager::BagManager;
+use distribution_bucket_picker::DistributionBucketPicker;
 use storage_bucket_picker::StorageBucketPicker;
 
 // TODO: constants
@@ -387,6 +389,11 @@ impl<DistributionBucketFamilyId: Ord> DynamicBagCreationPolicy<DistributionBucke
     pub(crate) fn no_storage_buckets_required(&self) -> bool {
         self.number_of_storage_buckets == 0
     }
+
+    // Verifies non-zero number of required distribution buckets.
+    pub(crate) fn no_distribution_buckets_required(&self) -> bool {
+        self.families.iter().map(|(_, num)| num).sum::<u32>() == 0
+    }
 }
 
 /// "Storage buckets per bag" value constraint type.
@@ -2296,10 +2303,14 @@ impl<T: Trait> DataObjectStorage<T> for Module<T> {
         // == MUTATION SAFE ==
         //
 
-        let storage_buckets = Self::pick_storage_buckets_for_dynamic_bag(bag_id.clone().into());
+        let bag_type: DynamicBagType = bag_id.clone().into();
+
+        let storage_buckets = Self::pick_storage_buckets_for_dynamic_bag(bag_type);
+        let distribution_buckets = Self::pick_distribution_buckets_for_dynamic_bag(bag_type);
 
         let bag = DynamicBag::<T> {
             stored_by: storage_buckets,
+            distributed_by: distribution_buckets,
             ..Default::default()
         };
 
@@ -2874,13 +2885,20 @@ impl<T: Trait> Module<T> {
         }
     }
 
-    // Selects storage bucket ID sets to assign to the storage bucket.
+    // Selects storage bucket ID sets to assign to the dynamic bag.
     pub(crate) fn pick_storage_buckets_for_dynamic_bag(
         bag_type: DynamicBagType,
     ) -> BTreeSet<T::StorageBucketId> {
         StorageBucketPicker::<T>::pick_storage_buckets(bag_type)
     }
 
+    // Selects distributed bucket ID sets to assign to the dynamic bag.
+    pub(crate) fn pick_distribution_buckets_for_dynamic_bag(
+        bag_type: DynamicBagType,
+    ) -> BTreeSet<T::DistributionBucketId> {
+        DistributionBucketPicker::<T>::pick_distribution_buckets(bag_type)
+    }
+
     // Get default dynamic bag policy by bag type.
     fn get_default_dynamic_bag_creation_policy(
         bag_type: DynamicBagType,
@@ -3008,4 +3026,27 @@ impl<T: Trait> Module<T> {
 
         Ok(())
     }
+
+    // Generate random number from zero to upper_bound (excluding).
+    pub(crate) fn random_index(seed: &[u8], upper_bound: u64) -> u64 {
+        if upper_bound == 0 {
+            return upper_bound;
+        }
+
+        let mut rand: u64 = 0;
+        for (offset, byte) in seed.iter().enumerate().take(8) {
+            rand += (*byte as u64) << offset;
+        }
+        rand % upper_bound
+    }
+
+    // Get initial random seed. It handles the error on the initial block.
+    pub(crate) fn get_initial_random_seed() -> T::Hash {
+        // Cannot create randomness in the initial block (Substrate error).
+        if <frame_system::Module<T>>::block_number() == Zero::zero() {
+            Default::default()
+        } else {
+            T::Randomness::random_seed()
+        }
+    }
 }

runtime-modules/storage/src/storage_bucket_picker.rs

@@ -30,7 +30,7 @@ impl<T: Trait> StorageBucketPicker<T> {
 
         let required_bucket_num = creation_policy.number_of_storage_buckets as usize;
 
-        // Storage IDs accumulator.
+        // Storage bucket IDs accumulator.
         let bucket_ids_cell = RefCell::new(BTreeSet::new());
 
         RandomStorageBucketIdIterator::<T>::new()
@@ -109,34 +109,18 @@ impl<T: Trait> RandomStorageBucketIdIterator<T> {
     fn random_storage_bucket_id(&self) -> T::StorageBucketId {
         let total_buckets_number = Module::<T>::next_storage_bucket_id();
 
-        let random_bucket_id: T::StorageBucketId = self
-            .random_index(total_buckets_number.saturated_into())
-            .saturated_into();
+        let random_bucket_id: T::StorageBucketId = Module::<T>::random_index(
+            self.current_seed.as_ref(),
+            total_buckets_number.saturated_into(),
+        )
+        .saturated_into();
 
         random_bucket_id
     }
 
-    // Generate random number from zero to upper_bound (excluding).
-    fn random_index(&self, upper_bound: u64) -> u64 {
-        if upper_bound == 0 {
-            return upper_bound;
-        }
-
-        let mut rand: u64 = 0;
-        for offset in 0..8 {
-            rand += (self.current_seed.as_ref()[offset] as u64) << offset;
-        }
-        rand % upper_bound
-    }
-
     // Creates new iterator.
     pub(crate) fn new() -> Self {
-        // Cannot create randomness in the initial block (Substrate error).
-        let seed = if <frame_system::Module<T>>::block_number() == Zero::zero() {
-            Default::default()
-        } else {
-            T::Randomness::random_seed()
-        };
+        let seed = Module::<T>::get_initial_random_seed();
 
         Self {
             current_iteration: 0,

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

@@ -52,8 +52,8 @@ impl balances::Trait for Test {
 parameter_types! {
     pub const MaxStorageBucketNumber: u64 = 1000;
     pub const MaxNumberOfDataObjectsPerBag: u64 = 4;
-    pub const MaxDistributionBucketFamilyNumber: u64 = 1;
-    pub const MaxDistributionBucketNumberPerFamily: u64 = 1;
+    pub const MaxDistributionBucketFamilyNumber: u64 = 4;
+    pub const MaxDistributionBucketNumberPerFamily: u64 = 10;
     pub const DataObjectDeletionPrize: u64 = 10;
     pub const StorageModuleId: ModuleId = ModuleId(*b"mstorage"); // module storage
     pub const BlacklistSizeLimit: u64 = 1;

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

@@ -10,7 +10,7 @@ use frame_system::RawOrigin;
 use sp_runtime::SaturatedConversion;
 use sp_std::collections::btree_map::BTreeMap;
 use sp_std::collections::btree_set::BTreeSet;
-use sp_std::iter::FromIterator;
+use sp_std::iter::{repeat, FromIterator};
 
 use common::working_group::WorkingGroup;
 
@@ -23,10 +23,11 @@ use crate::{
 use mocks::{
     build_test_externalities, Balances, DataObjectDeletionPrize,
     DefaultChannelDynamicBagNumberOfStorageBuckets, DefaultMemberDynamicBagNumberOfStorageBuckets,
-    InitialStorageBucketsNumberForDynamicBag, MaxNumberOfDataObjectsPerBag,
-    MaxRandomIterationNumber, MaxStorageBucketNumber, Storage, Test, ANOTHER_STORAGE_PROVIDER_ID,
-    DEFAULT_MEMBER_ACCOUNT_ID, DEFAULT_MEMBER_ID, DEFAULT_STORAGE_PROVIDER_ACCOUNT_ID,
-    DEFAULT_STORAGE_PROVIDER_ID, DISTRIBUTION_WG_LEADER_ACCOUNT_ID, STORAGE_WG_LEADER_ACCOUNT_ID,
+    InitialStorageBucketsNumberForDynamicBag, MaxDistributionBucketFamilyNumber,
+    MaxDistributionBucketNumberPerFamily, MaxNumberOfDataObjectsPerBag, MaxRandomIterationNumber,
+    MaxStorageBucketNumber, Storage, Test, ANOTHER_STORAGE_PROVIDER_ID, DEFAULT_MEMBER_ACCOUNT_ID,
+    DEFAULT_MEMBER_ID, DEFAULT_STORAGE_PROVIDER_ACCOUNT_ID, DEFAULT_STORAGE_PROVIDER_ID,
+    DISTRIBUTION_WG_LEADER_ACCOUNT_ID, STORAGE_WG_LEADER_ACCOUNT_ID,
 };
 
 use fixtures::*;
@@ -3513,9 +3514,11 @@ fn create_distribution_bucket_family_fails_with_non_signed_origin() {
 #[test]
 fn create_distribution_bucket_family_fails_with_exceeding_family_number_limit() {
     build_test_externalities().execute_with(|| {
-        CreateDistributionBucketFamilyFixture::default()
-            .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
-            .call_and_assert(Ok(()));
+        for _ in 0..MaxDistributionBucketFamilyNumber::get() {
+            CreateDistributionBucketFamilyFixture::default()
+                .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
+                .call_and_assert(Ok(()));
+        }
 
         CreateDistributionBucketFamilyFixture::default()
             .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
@@ -3620,15 +3623,9 @@ fn create_distribution_bucket_fails_with_non_existing_family() {
 #[test]
 fn create_distribution_bucket_fails_with_exceeding_max_bucket_number() {
     build_test_externalities().execute_with(|| {
-        let family_id = CreateDistributionBucketFamilyFixture::default()
-            .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
-            .call_and_assert(Ok(()))
-            .unwrap();
-
-        CreateDistributionBucketFixture::default()
-            .with_family_id(family_id)
-            .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
-            .call_and_assert(Ok(()));
+        let (family_id, _) = create_distribution_bucket_family_with_buckets(
+            MaxDistributionBucketNumberPerFamily::get(),
+        );
 
         CreateDistributionBucketFixture::default()
             .with_family_id(family_id)
@@ -4131,3 +4128,80 @@ fn update_families_in_dynamic_bag_creation_policy_fails_with_invalid_family_id()
             ));
     });
 }
+
+fn create_distribution_bucket_family_with_buckets(bucket_number: u64) -> (u64, Vec<u64>) {
+    let family_id = CreateDistributionBucketFamilyFixture::default()
+        .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
+        .call_and_assert(Ok(()))
+        .unwrap();
+
+    let bucket_ids = repeat(family_id)
+        .take(bucket_number as usize)
+        .map(|fam_id| {
+            CreateDistributionBucketFixture::default()
+                .with_family_id(fam_id)
+                .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
+                .with_accept_new_bags(true)
+                .call_and_assert(Ok(()))
+                .unwrap()
+        })
+        .collect::<Vec<_>>();
+
+    (family_id, bucket_ids)
+}
+
+#[test]
+fn distribution_bucket_family_pick_during_dynamic_bag_creation_succeeded() {
+    build_test_externalities().execute_with(|| {
+        // Enable randomness (disabled at the initial block).
+        let starting_block = 6;
+        run_to_block(starting_block);
+
+        let dynamic_bag_type = DynamicBagType::Channel;
+        let new_bucket_number = 5;
+
+        let (family_id1, bucket_ids1) = create_distribution_bucket_family_with_buckets(
+            MaxDistributionBucketNumberPerFamily::get(),
+        );
+        let (family_id2, bucket_ids2) = create_distribution_bucket_family_with_buckets(
+            MaxDistributionBucketNumberPerFamily::get(),
+        );
+        let (family_id3, _) = create_distribution_bucket_family_with_buckets(
+            MaxDistributionBucketNumberPerFamily::get(),
+        );
+        let (family_id4, _) = create_distribution_bucket_family_with_buckets(0);
+
+        let families = BTreeMap::from_iter(vec![
+            (family_id1, new_bucket_number),
+            (family_id2, new_bucket_number),
+            (family_id3, 0),
+            (family_id4, new_bucket_number),
+        ]);
+
+        UpdateFamiliesInDynamicBagCreationPolicyFixture::default()
+            .with_origin(RawOrigin::Signed(DISTRIBUTION_WG_LEADER_ACCOUNT_ID))
+            .with_families(families)
+            .with_dynamic_bag_type(dynamic_bag_type)
+            .call_and_assert(Ok(()));
+
+        let picked_bucket_ids =
+            Storage::pick_distribution_buckets_for_dynamic_bag(dynamic_bag_type);
+
+        assert_eq!(picked_bucket_ids.len(), (new_bucket_number * 2) as usize); // buckets from two families
+
+        let total_ids1 = BTreeSet::from_iter(
+            bucket_ids1
+                .iter()
+                .cloned()
+                .chain(bucket_ids2.iter().cloned()),
+        );
+        let total_ids2 = BTreeSet::from_iter(
+            total_ids1
+                .iter()
+                .cloned()
+                .chain(picked_bucket_ids.iter().cloned()),
+        );
+
+        assert_eq!(total_ids1, total_ids2); // picked IDS are from total ID set.
+    });
+}