Skip to content

Commit

Permalink
draft structure
Browse files Browse the repository at this point in the history
  • Loading branch information
@ParthDesai
ParthDesai committed Sep 25, 2024
1 parent 98e467c commit b655d6b
Show file tree
Hide file tree
Showing 8 changed files with 415 additions and 258 deletions.
66 changes: 14 additions & 52 deletions pallets/collator-assignment/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@ use {
},
sp_std::{collections::btree_set::BTreeSet, fmt::Debug, prelude::*, vec},
tp_traits::{
CollatorAssignmentHook, CollatorAssignmentTip, GetContainerChainAuthor,
GetHostConfiguration, GetSessionContainerChains, ParaId, RemoveInvulnerables,
RemoveParaIdsWithNoCredits, ShouldRotateAllCollators, Slot,
CollatorAssignmentTip, GetContainerChainAuthor, GetHostConfiguration,
GetSessionContainerChains, ParaId, RemoveInvulnerables, RemoveParaIdsWithNoCredits,
ShouldRotateAllCollators, Slot,
},
};
pub use {dp_collator_assignment::AssignedCollators, pallet::*};
Expand Down Expand Up @@ -105,8 +105,10 @@ pub mod pallet {
type ShouldRotateAllCollators: ShouldRotateAllCollators<Self::SessionIndex>;
type GetRandomnessForNextBlock: GetRandomnessForNextBlock<BlockNumberFor<Self>>;
type RemoveInvulnerables: RemoveInvulnerables<Self::AccountId>;
type RemoveParaIdsWithNoCredits: RemoveParaIdsWithNoCredits;
type CollatorAssignmentHook: CollatorAssignmentHook<BalanceOf<Self>>;
type RemoveParaIdsWithNoCredits: RemoveParaIdsWithNoCredits<
BalanceOf<Self>,
Self::AccountId,
>;
type Currency: Currency<Self::AccountId>;
type CollatorAssignmentTip: CollatorAssignmentTip<BalanceOf<Self>>;
type ForceEmptyOrchestrator: Get<bool>;
Expand Down Expand Up @@ -309,13 +311,13 @@ pub mod pallet {
old_assigned.container_chains.keys().cloned().collect();

// Remove the containerChains that do not have enough credits for block production
T::RemoveParaIdsWithNoCredits::remove_para_ids_with_no_credits(
T::RemoveParaIdsWithNoCredits::pre_assignment_remove_para_ids_with_no_credits(
&mut container_chain_ids,
&old_assigned_para_ids,
);
// TODO: parathreads should be treated a bit differently, they don't need to have the same amount of credits
// as parathreads because they will not be producing blocks on every slot.
T::RemoveParaIdsWithNoCredits::remove_para_ids_with_no_credits(
T::RemoveParaIdsWithNoCredits::pre_assignment_remove_para_ids_with_no_credits(
&mut parathreads,
&old_assigned_para_ids,
);
Expand Down Expand Up @@ -467,51 +469,11 @@ pub mod pallet {

// TODO: this probably is asking for a refactor
// only apply the onCollatorAssignedHook if sufficient collators
for para_id in &container_chain_ids {
if !new_assigned
.container_chains
.get(para_id)
.unwrap_or(&vec![])
.is_empty()
{
if let Err(e) = T::CollatorAssignmentHook::on_collators_assigned(
*para_id,
maybe_tip.as_ref(),
false,
) {
// On error remove para from assignment
log::warn!(
"CollatorAssignmentHook error! Removing para {} from assignment: {:?}",
u32::from(*para_id),
e
);
new_assigned.container_chains.remove(para_id);
}
}
}

for para_id in &parathreads {
if !new_assigned
.container_chains
.get(para_id)
.unwrap_or(&vec![])
.is_empty()
{
if let Err(e) = T::CollatorAssignmentHook::on_collators_assigned(
*para_id,
maybe_tip.as_ref(),
true,
) {
// On error remove para from assignment
log::warn!(
"CollatorAssignmentHook error! Removing para {} from assignment: {:?}",
u32::from(*para_id),
e
);
new_assigned.container_chains.remove(para_id);
}
}
}
T::RemoveParaIdsWithNoCredits::post_assignment_remove_para_ids_with_no_credits(
&old_assigned_para_ids,
&mut new_assigned.container_chains,
&maybe_tip,
);

Self::store_collator_fullness(
&new_assigned,
Expand Down
37 changes: 14 additions & 23 deletions pallets/collator-assignment/src/mock.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ use {
},
sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet},
tp_traits::{
CollatorAssignmentHook, CollatorAssignmentTip, ParaId, ParathreadParams,
RemoveInvulnerables, RemoveParaIdsWithNoCredits, SessionContainerChains,
CollatorAssignmentTip, ParaId, ParathreadParams, RemoveInvulnerables,
RemoveParaIdsWithNoCredits, SessionContainerChains,
},
tracing_subscriber::{layer::SubscriberExt, FmtSubscriber},
};
Expand Down Expand Up @@ -310,23 +310,6 @@ impl CollatorAssignmentTip<u32> for MockCollatorAssignmentTip {
}
}
}
pub struct MockCollatorAssignmentHook;

impl CollatorAssignmentHook<u32> for MockCollatorAssignmentHook {
fn on_collators_assigned(
para_id: ParaId,
_maybe_tip: Option<&u32>,
_is_parathread: bool,
) -> Result<Weight, sp_runtime::DispatchError> {
// Only fail for para 1001
if MockData::mock().assignment_hook_errors && para_id == 1001.into() {
// The error doesn't matter
Err(sp_runtime::DispatchError::Unavailable)
} else {
Ok(Weight::default())
}
}
}

pub struct GetCoreAllocationConfigurationImpl;

Expand All @@ -353,7 +336,6 @@ impl pallet_collator_assignment::Config for Test {
type GetRandomnessForNextBlock = MockGetRandomnessForNextBlock;
type RemoveInvulnerables = RemoveAccountIdsAbove100;
type RemoveParaIdsWithNoCredits = RemoveParaIdsAbove5000;
type CollatorAssignmentHook = MockCollatorAssignmentHook;
type CollatorAssignmentTip = MockCollatorAssignmentTip;
type ForceEmptyOrchestrator = ConstBool<false>;
type Currency = ();
Expand Down Expand Up @@ -416,14 +398,23 @@ impl RemoveInvulnerables<u64> for RemoveAccountIdsAbove100 {
/// Any ParaId >= 5000 will be considered to not have enough credits
pub struct RemoveParaIdsAbove5000;

impl RemoveParaIdsWithNoCredits for RemoveParaIdsAbove5000 {
fn remove_para_ids_with_no_credits(
impl<AC> RemoveParaIdsWithNoCredits<u32, AC> for RemoveParaIdsAbove5000 {
fn pre_assignment_remove_para_ids_with_no_credits(
para_ids: &mut Vec<ParaId>,
_currently_assigned: &BTreeSet<ParaId>,
_old_assigned: &BTreeSet<ParaId>,
) {
para_ids.retain(|para_id| *para_id <= ParaId::from(5000));
}

fn post_assignment_remove_para_ids_with_no_credits(
_current_assigned: &BTreeSet<ParaId>,
new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
_maybe_tip: &Option<u32>,
) -> Weight {
new_assigned.retain(|para_id, _| *para_id <= ParaId::from(5000));
Weight::zero()
}

#[cfg(feature = "runtime-benchmarks")]
fn make_valid_para_ids(_para_ids: &[ParaId]) {}
}
Expand Down
58 changes: 0 additions & 58 deletions pallets/collator-assignment/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1344,64 +1344,6 @@ fn assign_collators_prioritizing_tip() {
});
}

#[test]
fn on_collators_assigned_hook_failure_removes_para_from_assignment() {
new_test_ext().execute_with(|| {
run_to_block(1);

MockData::mutate(|m| {
m.collators_per_container = 2;
m.collators_per_parathread = 2;
m.min_orchestrator_chain_collators = 5;
m.max_orchestrator_chain_collators = 5;

m.collators = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
m.container_chains = vec![1001, 1002, 1003, 1004];
m.assignment_hook_errors = false;
});
run_to_block(11);

assert_eq!(
assigned_collators(),
BTreeMap::from_iter(vec![
(1, 1000),
(2, 1000),
(3, 1000),
(4, 1000),
(5, 1000),
(6, 1001),
(7, 1001),
(8, 1002),
(9, 1002),
(10, 1003),
(11, 1003),
]),
);

// Para 1001 will fail on_assignment_hook
MockData::mutate(|m| {
m.assignment_hook_errors = true;
});

run_to_block(21);

assert_eq!(
assigned_collators(),
BTreeMap::from_iter(vec![
(1, 1000),
(2, 1000),
(3, 1000),
(4, 1000),
(5, 1000),
(8, 1002),
(9, 1002),
(10, 1003),
(11, 1003),
]),
);
});
}

#[test]
fn assign_collators_truncates_before_shuffling() {
// Check that if there are more collators than needed, we only assign the first collators
Expand Down
21 changes: 21 additions & 0 deletions pallets/services-payment/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -422,6 +422,27 @@ pub mod pallet {
});
}

pub fn charge_tip(para_id: &ParaId, tip: &BalanceOf<T>) -> Result<(), DispatchError> {
// Only charge the tip to the paras that had a max tip set
// (aka were willing to tip for being assigned a collator)
if MaxTip::<T>::get(para_id).is_some() {
let tip_imbalance = T::Currency::withdraw(
&Self::parachain_tank(*para_id),
*tip,
WithdrawReasons::TIP,
ExistenceRequirement::KeepAlive,
)?;

Self::deposit_event(Event::<T>::CollatorAssignmentTipCollected {
para_id: *para_id,
payer: Self::parachain_tank(*para_id),
tip: *tip,
});
T::OnChargeForCollatorAssignmentTip::on_unbalanced(tip_imbalance);
}
Ok(())
}

pub fn free_block_production_credits(para_id: ParaId) -> Option<BlockNumberFor<T>> {
BlockProductionCredits::<T>::get(para_id)
}
Expand Down
25 changes: 19 additions & 6 deletions primitives/traits/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ use {
traits::{CheckedAdd, CheckedMul},
ArithmeticError, DispatchResult, Perbill,
},
sp_std::{collections::btree_set::BTreeSet, vec::Vec},
sp_std::{collections::btree_map::BTreeMap, collections::btree_set::BTreeSet, vec::Vec},
};

// Separate import as rustfmt wrongly change it to `sp_std::vec::self`, which is the module instead
Expand Down Expand Up @@ -271,26 +271,39 @@ impl<AccountId: Clone> RemoveInvulnerables<AccountId> for () {

/// Helper trait for pallet_collator_assignment to be able to not assign collators to container chains with no credits
/// in pallet_services_payment
pub trait RemoveParaIdsWithNoCredits {
pub trait RemoveParaIdsWithNoCredits<B, AC> {
/// Remove para ids with not enough credits. The resulting order will affect priority: the first para id in the list
/// will be the first one to get collators.
fn remove_para_ids_with_no_credits(
fn pre_assignment_remove_para_ids_with_no_credits(
para_ids: &mut Vec<ParaId>,
currently_assigned: &BTreeSet<ParaId>,
old_assigned: &BTreeSet<ParaId>,
);
fn post_assignment_remove_para_ids_with_no_credits(
current_assigned: &BTreeSet<ParaId>,
new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
maybe_tip: &Option<B>,
) -> Weight;

/// Make those para ids valid by giving them enough credits, for benchmarking.
#[cfg(feature = "runtime-benchmarks")]
fn make_valid_para_ids(para_ids: &[ParaId]);
}

impl RemoveParaIdsWithNoCredits for () {
fn remove_para_ids_with_no_credits(
impl<B, AC> RemoveParaIdsWithNoCredits<B, AC> for () {
fn pre_assignment_remove_para_ids_with_no_credits(
_para_ids: &mut Vec<ParaId>,
_currently_assigned: &BTreeSet<ParaId>,
) {
}

fn post_assignment_remove_para_ids_with_no_credits(
_current_assigned: &BTreeSet<ParaId>,
_new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
_maybe_tip: &Option<B>,
) -> Weight {
Default::default()
}

#[cfg(feature = "runtime-benchmarks")]
fn make_valid_para_ids(_para_ids: &[ParaId]) {}
}
Expand Down
Loading

0 comments on commit b655d6b

Please sign in to comment.