feat: initialize Kurdistan SDK - independent fork of Polkadot SDK

pezkuwi/runtime/teyrchains/src/on_demand/benchmarking.rs

@@ -0,0 +1,115 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! On demand assigner pallet benchmarking.
+
+#![cfg(feature = "runtime-benchmarks")]
+
+use super::{Pallet, *};
+use crate::{
+	configuration::{HostConfiguration, Pallet as ConfigurationPallet},
+	paras::{Pallet as ParasPallet, ParaGenesisArgs, ParaKind, TeyrchainsCache},
+	shared::Pallet as ParasShared,
+};
+
+use alloc::vec;
+use frame_benchmarking::v2::*;
+use frame_system::RawOrigin;
+use sp_runtime::traits::Bounded;
+
+use pezkuwi_primitives::{
+	HeadData, Id as ParaId, SessionIndex, ValidationCode, ON_DEMAND_DEFAULT_QUEUE_MAX_SIZE,
+};
+
+// Constants for the benchmarking
+const SESSION_INDEX: SessionIndex = 1;
+
+// Initialize a parathread for benchmarking.
+pub fn init_parathread<T>(para_id: ParaId)
+where
+	T: Config + crate::paras::Config + crate::shared::Config,
+{
+	ParasShared::<T>::set_session_index(SESSION_INDEX);
+	let mut config = HostConfiguration::default();
+	config.scheduler_params.num_cores = 1;
+	ConfigurationPallet::<T>::force_set_active_config(config);
+	let mut teyrchains = TeyrchainsCache::new();
+	ParasPallet::<T>::initialize_para_now(
+		&mut teyrchains,
+		para_id,
+		&ParaGenesisArgs {
+			para_kind: ParaKind::Parathread,
+			genesis_head: HeadData(vec![1, 2, 3, 4]),
+			validation_code: ValidationCode(vec![1, 2, 3, 4]),
+		},
+	);
+}
+
+#[benchmarks]
+mod benchmarks {
+	/// We want to fill the queue to the maximum, so exactly one more item fits.
+	const MAX_FILL_BENCH: u32 = ON_DEMAND_DEFAULT_QUEUE_MAX_SIZE.saturating_sub(1);
+
+	use super::*;
+	#[benchmark]
+	fn place_order_keep_alive(s: Linear<1, MAX_FILL_BENCH>) {
+		// Setup
+		let caller = whitelisted_caller();
+		let para_id = ParaId::from(111u32);
+		init_parathread::<T>(para_id);
+		T::Currency::make_free_balance_be(&caller, BalanceOf::<T>::max_value());
+		Pallet::<T>::populate_queue(para_id, s);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.into()), BalanceOf::<T>::max_value(), para_id)
+	}
+
+	#[benchmark]
+	fn place_order_allow_death(s: Linear<1, MAX_FILL_BENCH>) {
+		// Setup
+		let caller = whitelisted_caller();
+		let para_id = ParaId::from(111u32);
+		init_parathread::<T>(para_id);
+		T::Currency::make_free_balance_be(&caller, BalanceOf::<T>::max_value());
+
+		Pallet::<T>::populate_queue(para_id, s);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.into()), BalanceOf::<T>::max_value(), para_id)
+	}
+
+	#[benchmark]
+	fn place_order_with_credits(s: Linear<1, MAX_FILL_BENCH>) {
+		// Setup
+		let caller: T::AccountId = whitelisted_caller();
+		let para_id = ParaId::from(111u32);
+		init_parathread::<T>(para_id);
+		Credits::<T>::insert(&caller, BalanceOf::<T>::max_value());
+
+		Pallet::<T>::populate_queue(para_id, s);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.into()), BalanceOf::<T>::max_value(), para_id)
+	}
+
+	impl_benchmark_test_suite!(
+		Pallet,
+		crate::mock::new_test_ext(
+			crate::on_demand::mock_helpers::GenesisConfigBuilder::default().build()
+		),
+		crate::mock::Test
+	);
+}

pezkuwi/runtime/teyrchains/src/on_demand/migration.rs

@@ -0,0 +1,181 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! A module that is responsible for migration of storage.
+use super::*;
+use frame_support::{
+	migrations::VersionedMigration, pallet_prelude::ValueQuery, storage_alias,
+	traits::UncheckedOnRuntimeUpgrade, weights::Weight,
+};
+
+mod v0 {
+	use super::*;
+	use alloc::collections::vec_deque::VecDeque;
+
+	#[derive(Encode, Decode, TypeInfo, Debug, PartialEq, Clone)]
+	pub(super) struct EnqueuedOrder {
+		pub para_id: ParaId,
+	}
+
+	/// Keeps track of the multiplier used to calculate the current spot price for the on demand
+	/// assigner.
+	/// NOTE: Ignoring the `OnEmpty` field for the migration.
+	#[storage_alias]
+	pub(super) type SpotTraffic<T: Config> = StorageValue<Pallet<T>, FixedU128, ValueQuery>;
+
+	/// The order storage entry. Uses a VecDeque to be able to push to the front of the
+	/// queue from the scheduler on session boundaries.
+	/// NOTE: Ignoring the `OnEmpty` field for the migration.
+	#[storage_alias]
+	pub(super) type OnDemandQueue<T: Config> =
+		StorageValue<Pallet<T>, VecDeque<EnqueuedOrder>, ValueQuery>;
+}
+
+mod v1 {
+	use super::*;
+
+	use crate::on_demand::LOG_TARGET;
+
+	/// Migration to V1
+	pub struct UncheckedMigrateToV1<T>(core::marker::PhantomData<T>);
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrateToV1<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let mut weight: Weight = Weight::zero();
+
+			// Migrate the current traffic value
+			let config = configuration::ActiveConfig::<T>::get();
+			QueueStatus::<T>::mutate(|mut queue_status| {
+				Pallet::<T>::update_spot_traffic(&config, &mut queue_status);
+
+				let v0_queue = v0::OnDemandQueue::<T>::take();
+				// Process the v0 queue into v1.
+				v0_queue.into_iter().for_each(|enqueued_order| {
+					// Readding the old orders will use the new systems.
+					Pallet::<T>::add_on_demand_order(
+						queue_status,
+						enqueued_order.para_id,
+						QueuePushDirection::Back,
+					);
+				});
+			});
+
+			// Remove the old storage.
+			v0::OnDemandQueue::<T>::kill(); // 1 write
+			v0::SpotTraffic::<T>::kill(); // 1 write
+
+			// Config read
+			weight.saturating_accrue(T::DbWeight::get().reads(1));
+			// QueueStatus read write (update_spot_traffic)
+			weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
+			// Kill x 2
+			weight.saturating_accrue(T::DbWeight::get().writes(2));
+
+			log::info!(target: LOG_TARGET, "Migrated on demand assigner storage to v1");
+			weight
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
+			let n: u32 = v0::OnDemandQueue::<T>::get().len() as u32;
+
+			log::info!(
+				target: LOG_TARGET,
+				"Number of orders waiting in the queue before: {n}",
+			);
+
+			Ok(n.encode())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(state: alloc::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+			log::info!(target: LOG_TARGET, "Running post_upgrade()");
+
+			ensure!(
+				v0::OnDemandQueue::<T>::get().is_empty(),
+				"OnDemandQueue should be empty after the migration"
+			);
+
+			let expected_len = u32::decode(&mut &state[..]).unwrap();
+			let queue_status_size = QueueStatus::<T>::get().size();
+			ensure!(
+				expected_len == queue_status_size,
+				"Number of orders should be the same before and after migration"
+			);
+
+			let n_affinity_entries: u32 =
+				AffinityEntries::<T>::iter().map(|(_index, heap)| heap.len() as u32).sum();
+			let n_para_id_affinity: u32 = ParaIdAffinity::<T>::iter()
+				.map(|(_para_id, affinity)| affinity.count as u32)
+				.sum();
+			ensure!(
+				n_para_id_affinity == n_affinity_entries,
+				"Number of affinity entries should be the same as the counts in ParaIdAffinity"
+			);
+
+			Ok(())
+		}
+	}
+}
+
+/// Migrate `V0` to `V1` of the storage format.
+pub type MigrateV0ToV1<T> = VersionedMigration<
+	0,
+	1,
+	v1::UncheckedMigrateToV1<T>,
+	Pallet<T>,
+	<T as frame_system::Config>::DbWeight,
+>;
+
+#[cfg(test)]
+mod tests {
+	use super::{v0, v1, UncheckedOnRuntimeUpgrade, Weight};
+	use crate::mock::{new_test_ext, MockGenesisConfig, OnDemand, Test};
+	use pezkuwi_primitives::Id as ParaId;
+
+	#[test]
+	fn migration_to_v1_preserves_queue_ordering() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			// Place orders for paraids 1..5
+			for i in 1..=5 {
+				v0::OnDemandQueue::<Test>::mutate(|queue| {
+					queue.push_back(v0::EnqueuedOrder { para_id: ParaId::new(i) })
+				});
+			}
+
+			// Queue has 5 orders
+			let old_queue = v0::OnDemandQueue::<Test>::get();
+			assert_eq!(old_queue.len(), 5);
+			// New queue has 0 orders
+			assert_eq!(OnDemand::get_queue_status().size(), 0);
+
+			// For tests, db weight is zero.
+			assert_eq!(
+				<v1::UncheckedMigrateToV1<Test> as UncheckedOnRuntimeUpgrade>::on_runtime_upgrade(),
+				Weight::zero()
+			);
+
+			// New queue has 5 orders
+			assert_eq!(OnDemand::get_queue_status().size(), 5);
+
+			// Compare each entry from the old queue with the entry in the new queue.
+			old_queue.iter().zip(OnDemand::get_free_entries().iter()).for_each(
+				|(old_enq, new_enq)| {
+					assert_eq!(old_enq.para_id, new_enq.para_id);
+				},
+			);
+		});
+	}
+}

pezkuwi/runtime/teyrchains/src/on_demand/mock_helpers.rs

@@ -0,0 +1,87 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Helper functions for tests, also used in runtime-benchmarks.
+
+#![cfg(test)]
+
+use super::*;
+
+use crate::{
+	mock::MockGenesisConfig,
+	paras::{ParaGenesisArgs, ParaKind},
+};
+
+use pezkuwi_primitives::{Balance, HeadData, ValidationCode};
+
+fn default_genesis_config() -> MockGenesisConfig {
+	MockGenesisConfig {
+		configuration: crate::configuration::GenesisConfig {
+			config: crate::configuration::HostConfiguration { ..Default::default() },
+		},
+		..Default::default()
+	}
+}
+
+#[derive(Debug)]
+pub struct GenesisConfigBuilder {
+	pub on_demand_cores: u32,
+	pub on_demand_base_fee: Balance,
+	pub on_demand_fee_variability: Perbill,
+	pub on_demand_max_queue_size: u32,
+	pub on_demand_target_queue_utilization: Perbill,
+	pub onboarded_on_demand_chains: Vec<ParaId>,
+}
+
+impl Default for GenesisConfigBuilder {
+	fn default() -> Self {
+		Self {
+			on_demand_cores: 10,
+			on_demand_base_fee: 10_000,
+			on_demand_fee_variability: Perbill::from_percent(1),
+			on_demand_max_queue_size: 100,
+			on_demand_target_queue_utilization: Perbill::from_percent(25),
+			onboarded_on_demand_chains: vec![],
+		}
+	}
+}
+
+impl GenesisConfigBuilder {
+	pub(super) fn build(self) -> MockGenesisConfig {
+		let mut genesis = default_genesis_config();
+		let config = &mut genesis.configuration.config;
+		config.scheduler_params.num_cores = self.on_demand_cores;
+		config.scheduler_params.on_demand_base_fee = self.on_demand_base_fee;
+		config.scheduler_params.on_demand_fee_variability = self.on_demand_fee_variability;
+		config.scheduler_params.on_demand_queue_max_size = self.on_demand_max_queue_size;
+		config.scheduler_params.on_demand_target_queue_utilization =
+			self.on_demand_target_queue_utilization;
+
+		let paras = &mut genesis.paras.paras;
+		for para_id in self.onboarded_on_demand_chains {
+			paras.push((
+				para_id,
+				ParaGenesisArgs {
+					genesis_head: HeadData::from(vec![0u8]),
+					validation_code: ValidationCode::from(vec![0u8]),
+					para_kind: ParaKind::Parathread,
+				},
+			))
+		}
+
+		genesis
+	}
+}

pezkuwi/runtime/teyrchains/src/on_demand/mod.rs

@@ -0,0 +1,813 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! The teyrchain on demand assignment module.
+//!
+//! Implements a mechanism for taking in orders for on-demand teyrchain (previously parathreads)
+//! assignments. This module is not handled by the initializer but is instead instantiated in the
+//! `construct_runtime` macro.
+//!
+//! The module currently limits parallel execution of blocks from the same `ParaId` via
+//! a core affinity mechanism. As long as there exists an affinity for a `CoreIndex` for
+//! a specific `ParaId`, orders for blockspace for that `ParaId` will only be assigned to
+//! that `CoreIndex`.
+//!
+//! NOTE: Once we have elastic scaling implemented we might want to extend this module to support
+//! ignoring core affinity up to a certain extend. This should be opt-in though as the teyrchain
+//! needs to support multiple cores in the same block. If we want to enable a single teyrchain
+//! occupying multiple cores in on-demand, we will likely add a separate order type, where the
+//! intent can be made explicit.
+
+use sp_runtime::traits::Zero;
+mod benchmarking;
+pub mod migration;
+mod mock_helpers;
+mod types;
+
+extern crate alloc;
+
+#[cfg(test)]
+mod tests;
+
+use crate::{configuration, paras, scheduler::common::Assignment};
+use alloc::collections::BinaryHeap;
+use core::mem::take;
+use frame_support::{
+	pallet_prelude::*,
+	traits::{
+		defensive_prelude::*,
+		Currency,
+		ExistenceRequirement::{self, AllowDeath, KeepAlive},
+		WithdrawReasons,
+	},
+	PalletId,
+};
+use frame_system::{pallet_prelude::*, Pallet as System};
+use pezkuwi_primitives::{CoreIndex, Id as ParaId};
+use sp_runtime::{
+	traits::{AccountIdConversion, One, SaturatedConversion},
+	FixedPointNumber, FixedPointOperand, FixedU128, Perbill, Saturating,
+};
+use types::{
+	BalanceOf, CoreAffinityCount, EnqueuedOrder, QueuePushDirection, QueueStatusType,
+	SpotTrafficCalculationErr,
+};
+
+const LOG_TARGET: &str = "runtime::teyrchains::on-demand";
+
+pub use pallet::*;
+
+pub trait WeightInfo {
+	fn place_order_allow_death(s: u32) -> Weight;
+	fn place_order_keep_alive(s: u32) -> Weight;
+	fn place_order_with_credits(s: u32) -> Weight;
+}
+
+/// A weight info that is only suitable for testing.
+pub struct TestWeightInfo;
+
+impl WeightInfo for TestWeightInfo {
+	fn place_order_allow_death(_: u32) -> Weight {
+		Weight::MAX
+	}
+
+	fn place_order_keep_alive(_: u32) -> Weight {
+		Weight::MAX
+	}
+
+	fn place_order_with_credits(_: u32) -> Weight {
+		Weight::MAX
+	}
+}
+
+/// Defines how the account wants to pay for on-demand.
+#[derive(Encode, Decode, TypeInfo, Debug, PartialEq, Clone, Eq)]
+enum PaymentType {
+	/// Use credits to purchase on-demand coretime.
+	Credits,
+	/// Use account's free balance to purchase on-demand coretime.
+	Balance,
+}
+
+#[frame_support::pallet]
+pub mod pallet {
+
+	use super::*;
+
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+
+	#[pallet::pallet]
+	#[pallet::without_storage_info]
+	#[pallet::storage_version(STORAGE_VERSION)]
+	pub struct Pallet<T>(_);
+
+	#[pallet::config]
+	pub trait Config: frame_system::Config + configuration::Config + paras::Config {
+		/// The runtime's definition of an event.
+		#[allow(deprecated)]
+		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
+
+		/// The runtime's definition of a Currency.
+		type Currency: Currency<Self::AccountId>;
+
+		/// Something that provides the weight of this pallet.
+		type WeightInfo: WeightInfo;
+
+		/// The default value for the spot traffic multiplier.
+		#[pallet::constant]
+		type TrafficDefaultValue: Get<FixedU128>;
+
+		/// The maximum number of blocks some historical revenue
+		/// information stored for.
+		#[pallet::constant]
+		type MaxHistoricalRevenue: Get<u32>;
+
+		/// Identifier for the internal revenue balance.
+		#[pallet::constant]
+		type PalletId: Get<PalletId>;
+	}
+
+	/// Creates an empty queue status for an empty queue with initial traffic value.
+	#[pallet::type_value]
+	pub(super) fn QueueStatusOnEmpty<T: Config>() -> QueueStatusType {
+		QueueStatusType { traffic: T::TrafficDefaultValue::get(), ..Default::default() }
+	}
+
+	#[pallet::type_value]
+	pub(super) fn EntriesOnEmpty<T: Config>() -> BinaryHeap<EnqueuedOrder> {
+		BinaryHeap::new()
+	}
+
+	/// Maps a `ParaId` to `CoreIndex` and keeps track of how many assignments the scheduler has in
+	/// it's lookahead. Keeping track of this affinity prevents parallel execution of the same
+	/// `ParaId` on two or more `CoreIndex`es.
+	#[pallet::storage]
+	pub(super) type ParaIdAffinity<T: Config> =
+		StorageMap<_, Twox64Concat, ParaId, CoreAffinityCount, OptionQuery>;
+
+	/// Overall status of queue (both free + affinity entries)
+	#[pallet::storage]
+	pub(super) type QueueStatus<T: Config> =
+		StorageValue<_, QueueStatusType, ValueQuery, QueueStatusOnEmpty<T>>;
+
+	/// Priority queue for all orders which don't yet (or not any more) have any core affinity.
+	#[pallet::storage]
+	pub(super) type FreeEntries<T: Config> =
+		StorageValue<_, BinaryHeap<EnqueuedOrder>, ValueQuery, EntriesOnEmpty<T>>;
+
+	/// Queue entries that are currently bound to a particular core due to core affinity.
+	#[pallet::storage]
+	pub(super) type AffinityEntries<T: Config> = StorageMap<
+		_,
+		Twox64Concat,
+		CoreIndex,
+		BinaryHeap<EnqueuedOrder>,
+		ValueQuery,
+		EntriesOnEmpty<T>,
+	>;
+
+	/// Keeps track of accumulated revenue from on demand order sales.
+	#[pallet::storage]
+	pub type Revenue<T: Config> =
+		StorageValue<_, BoundedVec<BalanceOf<T>, T::MaxHistoricalRevenue>, ValueQuery>;
+
+	/// Keeps track of credits owned by each account.
+	#[pallet::storage]
+	pub type Credits<T: Config> =
+		StorageMap<_, Blake2_128Concat, T::AccountId, BalanceOf<T>, ValueQuery>;
+
+	#[pallet::event]
+	#[pallet::generate_deposit(pub(super) fn deposit_event)]
+	pub enum Event<T: Config> {
+		/// An order was placed at some spot price amount by orderer ordered_by
+		OnDemandOrderPlaced { para_id: ParaId, spot_price: BalanceOf<T>, ordered_by: T::AccountId },
+		/// The value of the spot price has likely changed
+		SpotPriceSet { spot_price: BalanceOf<T> },
+		/// An account was given credits.
+		AccountCredited { who: T::AccountId, amount: BalanceOf<T> },
+	}
+
+	#[pallet::error]
+	pub enum Error<T> {
+		/// The order queue is full, `place_order` will not continue.
+		QueueFull,
+		/// The current spot price is higher than the max amount specified in the `place_order`
+		/// call, making it invalid.
+		SpotPriceHigherThanMaxAmount,
+		/// The account doesn't have enough credits to purchase on-demand coretime.
+		InsufficientCredits,
+	}
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn on_initialize(_now: BlockNumberFor<T>) -> Weight {
+			// Update revenue information storage.
+			Revenue::<T>::mutate(|revenue| {
+				if let Some(overdue) =
+					revenue.force_insert_keep_left(0, 0u32.into()).defensive_unwrap_or(None)
+				{
+					// We have some overdue revenue not claimed by the Coretime Chain, let's
+					// accumulate it at the oldest stored block
+					if let Some(last) = revenue.last_mut() {
+						*last = last.saturating_add(overdue);
+					}
+				}
+			});
+
+			let config = configuration::ActiveConfig::<T>::get();
+			// We need to update the spot traffic on block initialize in order to account for idle
+			// blocks.
+			QueueStatus::<T>::mutate(|queue_status| {
+				Self::update_spot_traffic(&config, queue_status);
+			});
+
+			// Reads: `Revenue`, `ActiveConfig`, `QueueStatus`
+			// Writes: `Revenue`, `QueueStatus`
+			T::DbWeight::get().reads_writes(3, 2)
+		}
+	}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {
+		/// Create a single on demand core order.
+		/// Will use the spot price for the current block and will reap the account if needed.
+		///
+		/// Parameters:
+		/// - `origin`: The sender of the call, funds will be withdrawn from this account.
+		/// - `max_amount`: The maximum balance to withdraw from the origin to place an order.
+		/// - `para_id`: A `ParaId` the origin wants to provide blockspace for.
+		///
+		/// Errors:
+		/// - `InsufficientBalance`: from the Currency implementation
+		/// - `QueueFull`
+		/// - `SpotPriceHigherThanMaxAmount`
+		///
+		/// Events:
+		/// - `OnDemandOrderPlaced`
+		#[pallet::call_index(0)]
+		#[pallet::weight(<T as Config>::WeightInfo::place_order_allow_death(QueueStatus::<T>::get().size()))]
+		#[allow(deprecated)]
+		#[deprecated(note = "This will be removed in favor of using `place_order_with_credits`")]
+		pub fn place_order_allow_death(
+			origin: OriginFor<T>,
+			max_amount: BalanceOf<T>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			Pallet::<T>::do_place_order(
+				sender,
+				max_amount,
+				para_id,
+				AllowDeath,
+				PaymentType::Balance,
+			)
+		}
+
+		/// Same as the [`place_order_allow_death`](Self::place_order_allow_death) call , but with a
+		/// check that placing the order will not reap the account.
+		///
+		/// Parameters:
+		/// - `origin`: The sender of the call, funds will be withdrawn from this account.
+		/// - `max_amount`: The maximum balance to withdraw from the origin to place an order.
+		/// - `para_id`: A `ParaId` the origin wants to provide blockspace for.
+		///
+		/// Errors:
+		/// - `InsufficientBalance`: from the Currency implementation
+		/// - `QueueFull`
+		/// - `SpotPriceHigherThanMaxAmount`
+		///
+		/// Events:
+		/// - `OnDemandOrderPlaced`
+		#[pallet::call_index(1)]
+		#[pallet::weight(<T as Config>::WeightInfo::place_order_keep_alive(QueueStatus::<T>::get().size()))]
+		#[allow(deprecated)]
+		#[deprecated(note = "This will be removed in favor of using `place_order_with_credits`")]
+		pub fn place_order_keep_alive(
+			origin: OriginFor<T>,
+			max_amount: BalanceOf<T>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			Pallet::<T>::do_place_order(
+				sender,
+				max_amount,
+				para_id,
+				KeepAlive,
+				PaymentType::Balance,
+			)
+		}
+
+		/// Create a single on demand core order with credits.
+		/// Will charge the owner's on-demand credit account the spot price for the current block.
+		///
+		/// Parameters:
+		/// - `origin`: The sender of the call, on-demand credits will be withdrawn from this
+		///   account.
+		/// - `max_amount`: The maximum number of credits to spend from the origin to place an
+		///   order.
+		/// - `para_id`: A `ParaId` the origin wants to provide blockspace for.
+		///
+		/// Errors:
+		/// - `InsufficientCredits`
+		/// - `QueueFull`
+		/// - `SpotPriceHigherThanMaxAmount`
+		///
+		/// Events:
+		/// - `OnDemandOrderPlaced`
+		#[pallet::call_index(2)]
+		#[pallet::weight(<T as Config>::WeightInfo::place_order_with_credits(QueueStatus::<T>::get().size()))]
+		pub fn place_order_with_credits(
+			origin: OriginFor<T>,
+			max_amount: BalanceOf<T>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			Pallet::<T>::do_place_order(
+				sender,
+				max_amount,
+				para_id,
+				KeepAlive,
+				PaymentType::Credits,
+			)
+		}
+	}
+}
+
+// Internal functions and interface to scheduler/wrapping assignment provider.
+impl<T: Config> Pallet<T>
+where
+	BalanceOf<T>: FixedPointOperand,
+{
+	/// Take the next queued entry that is available for a given core index.
+	///
+	/// Parameters:
+	/// - `core_index`: The core index
+	pub fn pop_assignment_for_core(core_index: CoreIndex) -> Option<Assignment> {
+		let entry: Result<EnqueuedOrder, ()> = QueueStatus::<T>::try_mutate(|queue_status| {
+			AffinityEntries::<T>::try_mutate(core_index, |affinity_entries| {
+				let free_entry = FreeEntries::<T>::try_mutate(|free_entries| {
+					let affinity_next = affinity_entries.peek();
+					let free_next = free_entries.peek();
+					let pick_free = match (affinity_next, free_next) {
+						(None, _) => true,
+						(Some(_), None) => false,
+						(Some(a), Some(f)) => f < a,
+					};
+					if pick_free {
+						let entry = free_entries.pop().ok_or(())?;
+						let (mut affinities, free): (BinaryHeap<_>, BinaryHeap<_>) =
+							take(free_entries)
+								.into_iter()
+								.partition(|e| e.para_id == entry.para_id);
+						affinity_entries.append(&mut affinities);
+						*free_entries = free;
+						Ok(entry)
+					} else {
+						Err(())
+					}
+				});
+				let entry = free_entry.or_else(|()| affinity_entries.pop().ok_or(()))?;
+				queue_status.consume_index(entry.idx);
+				Ok(entry)
+			})
+		});
+
+		let assignment = entry.map(|e| Assignment::Pool { para_id: e.para_id, core_index }).ok()?;
+
+		Pallet::<T>::increase_affinity(assignment.para_id(), core_index);
+		Some(assignment)
+	}
+
+	/// Report that an assignment was duplicated by the scheduler.
+	pub fn assignment_duplicated(para_id: ParaId, core_index: CoreIndex) {
+		Pallet::<T>::increase_affinity(para_id, core_index);
+	}
+
+	/// Report that the `para_id` & `core_index` combination was processed.
+	///
+	/// This should be called once it is clear that the assignment won't get pushed back anymore.
+	///
+	/// In other words for each `pop_assignment_for_core` a call to this function or
+	/// `push_back_assignment` must follow, but only one.
+	pub fn report_processed(para_id: ParaId, core_index: CoreIndex) {
+		Pallet::<T>::decrease_affinity_update_queue(para_id, core_index);
+	}
+
+	/// Push an assignment back to the front of the queue.
+	///
+	/// The assignment has not been processed yet. Typically used on session boundaries.
+	///
+	/// NOTE: We are not checking queue size here. So due to push backs it is possible that we
+	/// exceed the maximum queue size slightly.
+	///
+	/// Parameters:
+	/// - `para_id`: The para that did not make it.
+	/// - `core_index`: The core the para was scheduled on.
+	pub fn push_back_assignment(para_id: ParaId, core_index: CoreIndex) {
+		Pallet::<T>::decrease_affinity_update_queue(para_id, core_index);
+		QueueStatus::<T>::mutate(|queue_status| {
+			Pallet::<T>::add_on_demand_order(queue_status, para_id, QueuePushDirection::Front);
+		});
+	}
+
+	/// Adds credits to the specified account.
+	///
+	/// Parameters:
+	/// - `who`: Credit receiver.
+	/// - `amount`: The amount of new credits the account will receive.
+	pub fn credit_account(who: T::AccountId, amount: BalanceOf<T>) {
+		Credits::<T>::mutate(who.clone(), |credits| {
+			*credits = credits.saturating_add(amount);
+		});
+		Pallet::<T>::deposit_event(Event::<T>::AccountCredited { who, amount });
+	}
+
+	/// Helper function for `place_order_*` calls. Used to differentiate between placing orders
+	/// with a keep alive check or to allow the account to be reaped. The amount charged is
+	/// stored to the pallet account to be later paid out as revenue.
+	///
+	/// Parameters:
+	/// - `sender`: The sender of the call, funds will be withdrawn from this account.
+	/// - `max_amount`: The maximum balance to withdraw from the origin to place an order.
+	/// - `para_id`: A `ParaId` the origin wants to provide blockspace for.
+	/// - `existence_requirement`: Whether or not to ensure that the account will not be reaped.
+	/// - `payment_type`: Defines how the user wants to pay for on-demand.
+	///
+	/// Errors:
+	/// - `InsufficientBalance`: from the Currency implementation
+	/// - `QueueFull`
+	/// - `SpotPriceHigherThanMaxAmount`
+	///
+	/// Events:
+	/// - `OnDemandOrderPlaced`
+	fn do_place_order(
+		sender: <T as frame_system::Config>::AccountId,
+		max_amount: BalanceOf<T>,
+		para_id: ParaId,
+		existence_requirement: ExistenceRequirement,
+		payment_type: PaymentType,
+	) -> DispatchResult {
+		let config = configuration::ActiveConfig::<T>::get();
+
+		QueueStatus::<T>::mutate(|queue_status| {
+			Self::update_spot_traffic(&config, queue_status);
+			let traffic = queue_status.traffic;
+
+			// Calculate spot price
+			let spot_price: BalanceOf<T> = traffic.saturating_mul_int(
+				config.scheduler_params.on_demand_base_fee.saturated_into::<BalanceOf<T>>(),
+			);
+
+			// Is the current price higher than `max_amount`
+			ensure!(spot_price.le(&max_amount), Error::<T>::SpotPriceHigherThanMaxAmount);
+
+			ensure!(
+				queue_status.size() < config.scheduler_params.on_demand_queue_max_size,
+				Error::<T>::QueueFull
+			);
+
+			match payment_type {
+				PaymentType::Balance => {
+					// Charge the sending account the spot price. The amount will be teleported to
+					// the broker chain once it requests revenue information.
+					let amt = T::Currency::withdraw(
+						&sender,
+						spot_price,
+						WithdrawReasons::FEE,
+						existence_requirement,
+					)?;
+
+					// Consume the negative imbalance and deposit it into the pallet account. Make
+					// sure the account preserves even without the existential deposit.
+					let pot = Self::account_id();
+					if !System::<T>::account_exists(&pot) {
+						System::<T>::inc_providers(&pot);
+					}
+					T::Currency::resolve_creating(&pot, amt);
+				},
+				PaymentType::Credits => {
+					let credits = Credits::<T>::get(&sender);
+
+					// Charge the sending account the spot price in credits.
+					let new_credits_value =
+						credits.checked_sub(&spot_price).ok_or(Error::<T>::InsufficientCredits)?;
+
+					if new_credits_value.is_zero() {
+						Credits::<T>::remove(&sender);
+					} else {
+						Credits::<T>::insert(&sender, new_credits_value);
+					}
+				},
+			}
+
+			// Add the amount to the current block's (index 0) revenue information.
+			Revenue::<T>::mutate(|bounded_revenue| {
+				if let Some(current_block) = bounded_revenue.get_mut(0) {
+					*current_block = current_block.saturating_add(spot_price);
+				} else {
+					// Revenue has already been claimed in the same block, including the block
+					// itself. It shouldn't normally happen as revenue claims in the future are
+					// not allowed.
+					bounded_revenue.try_push(spot_price).defensive_ok();
+				}
+			});
+
+			Pallet::<T>::add_on_demand_order(queue_status, para_id, QueuePushDirection::Back);
+			Pallet::<T>::deposit_event(Event::<T>::OnDemandOrderPlaced {
+				para_id,
+				spot_price,
+				ordered_by: sender,
+			});
+
+			Ok(())
+		})
+	}
+
+	/// Calculate and update spot traffic.
+	fn update_spot_traffic(
+		config: &configuration::HostConfiguration<BlockNumberFor<T>>,
+		queue_status: &mut QueueStatusType,
+	) {
+		let old_traffic = queue_status.traffic;
+		match Self::calculate_spot_traffic(
+			old_traffic,
+			config.scheduler_params.on_demand_queue_max_size,
+			queue_status.size(),
+			config.scheduler_params.on_demand_target_queue_utilization,
+			config.scheduler_params.on_demand_fee_variability,
+		) {
+			Ok(new_traffic) => {
+				// Only update storage on change
+				if new_traffic != old_traffic {
+					queue_status.traffic = new_traffic;
+
+					// calculate the new spot price
+					let spot_price: BalanceOf<T> = new_traffic.saturating_mul_int(
+						config.scheduler_params.on_demand_base_fee.saturated_into::<BalanceOf<T>>(),
+					);
+
+					// emit the event for updated new price
+					Pallet::<T>::deposit_event(Event::<T>::SpotPriceSet { spot_price });
+				}
+			},
+			Err(err) => {
+				log::debug!(
+					target: LOG_TARGET,
+					"Error calculating spot traffic: {:?}", err
+				);
+			},
+		};
+	}
+
+	/// The spot price multiplier. This is based on the transaction fee calculations defined in:
+	/// https://research.web3.foundation/Polkadot/overview/token-economics#setting-transaction-fees
+	///
+	/// Parameters:
+	/// - `traffic`: The previously calculated multiplier, can never go below 1.0.
+	/// - `queue_capacity`: The max size of the order book.
+	/// - `queue_size`: How many orders are currently in the order book.
+	/// - `target_queue_utilisation`: How much of the queue_capacity should be ideally occupied,
+	///   expressed in percentages(perbill).
+	/// - `variability`: A variability factor, i.e. how quickly the spot price adjusts. This number
+	///   can be chosen by p/(k*(1-s)) where p is the desired ratio increase in spot price over k
+	///   number of blocks. s is the target_queue_utilisation. A concrete example: v =
+	///   0.05/(20*(1-0.25)) = 0.0033.
+	///
+	/// Returns:
+	/// - A `FixedU128` in the range of  `Config::TrafficDefaultValue` - `FixedU128::MAX` on
+	///   success.
+	///
+	/// Errors:
+	/// - `SpotTrafficCalculationErr::QueueCapacityIsZero`
+	/// - `SpotTrafficCalculationErr::QueueSizeLargerThanCapacity`
+	/// - `SpotTrafficCalculationErr::Division`
+	fn calculate_spot_traffic(
+		traffic: FixedU128,
+		queue_capacity: u32,
+		queue_size: u32,
+		target_queue_utilisation: Perbill,
+		variability: Perbill,
+	) -> Result<FixedU128, SpotTrafficCalculationErr> {
+		// Return early if queue has no capacity.
+		if queue_capacity == 0 {
+			return Err(SpotTrafficCalculationErr::QueueCapacityIsZero);
+		}
+
+		// Return early if queue size is greater than capacity.
+		if queue_size > queue_capacity {
+			return Err(SpotTrafficCalculationErr::QueueSizeLargerThanCapacity);
+		}
+
+		// (queue_size / queue_capacity) - target_queue_utilisation
+		let queue_util_ratio = FixedU128::from_rational(queue_size.into(), queue_capacity.into());
+		let positive = queue_util_ratio >= target_queue_utilisation.into();
+		let queue_util_diff = queue_util_ratio.max(target_queue_utilisation.into()) -
+			queue_util_ratio.min(target_queue_utilisation.into());
+
+		// variability * queue_util_diff
+		let var_times_qud = queue_util_diff.saturating_mul(variability.into());
+
+		// variability^2 * queue_util_diff^2
+		let var_times_qud_pow = var_times_qud.saturating_mul(var_times_qud);
+
+		// (variability^2 * queue_util_diff^2)/2
+		let div_by_two: FixedU128;
+		match var_times_qud_pow.const_checked_div(2.into()) {
+			Some(dbt) => div_by_two = dbt,
+			None => return Err(SpotTrafficCalculationErr::Division),
+		}
+
+		// traffic * (1 + queue_util_diff) + div_by_two
+		if positive {
+			let new_traffic = queue_util_diff
+				.saturating_add(div_by_two)
+				.saturating_add(One::one())
+				.saturating_mul(traffic);
+			Ok(new_traffic.max(<T as Config>::TrafficDefaultValue::get()))
+		} else {
+			let new_traffic = queue_util_diff.saturating_sub(div_by_two).saturating_mul(traffic);
+			Ok(new_traffic.max(<T as Config>::TrafficDefaultValue::get()))
+		}
+	}
+
+	/// Adds an order to the on demand queue.
+	///
+	/// Parameters:
+	/// - `location`: Whether to push this entry to the back or the front of the queue. Pushing an
+	///   entry to the front of the queue is only used when the scheduler wants to push back an
+	///   entry it has already popped.
+	fn add_on_demand_order(
+		queue_status: &mut QueueStatusType,
+		para_id: ParaId,
+		location: QueuePushDirection,
+	) {
+		let idx = match location {
+			QueuePushDirection::Back => queue_status.push_back(),
+			QueuePushDirection::Front => queue_status.push_front(),
+		};
+
+		let affinity = ParaIdAffinity::<T>::get(para_id);
+		let order = EnqueuedOrder::new(idx, para_id);
+		#[cfg(test)]
+		log::debug!(target: LOG_TARGET, "add_on_demand_order, order: {:?}, affinity: {:?}, direction: {:?}", order, affinity, location);
+
+		match affinity {
+			None => FreeEntries::<T>::mutate(|entries| entries.push(order)),
+			Some(affinity) =>
+				AffinityEntries::<T>::mutate(affinity.core_index, |entries| entries.push(order)),
+		}
+	}
+
+	/// Decrease core affinity for para and update queue
+	///
+	/// if affinity dropped to 0, moving entries back to `FreeEntries`.
+	fn decrease_affinity_update_queue(para_id: ParaId, core_index: CoreIndex) {
+		let affinity = Pallet::<T>::decrease_affinity(para_id, core_index);
+		#[cfg(not(test))]
+		debug_assert_ne!(
+			affinity, None,
+			"Decreased affinity for a para that has not been served on a core?"
+		);
+		if affinity != Some(0) {
+			return;
+		}
+		// No affinity more for entries on this core, free any entries:
+		//
+		// This is necessary to ensure them being served as the core might no longer exist at all.
+		AffinityEntries::<T>::mutate(core_index, |affinity_entries| {
+			FreeEntries::<T>::mutate(|free_entries| {
+				let (mut freed, affinities): (BinaryHeap<_>, BinaryHeap<_>) =
+					take(affinity_entries).into_iter().partition(|e| e.para_id == para_id);
+				free_entries.append(&mut freed);
+				*affinity_entries = affinities;
+			})
+		});
+	}
+
+	/// Decreases the affinity of a `ParaId` to a specified `CoreIndex`.
+	///
+	/// Subtracts from the count of the `CoreAffinityCount` if an entry is found and the core_index
+	/// matches. When the count reaches 0, the entry is removed.
+	/// A non-existent entry is a no-op.
+	///
+	/// Returns: The new affinity of the para on that core. `None` if there is no affinity on this
+	/// core.
+	fn decrease_affinity(para_id: ParaId, core_index: CoreIndex) -> Option<u32> {
+		ParaIdAffinity::<T>::mutate(para_id, |maybe_affinity| {
+			let affinity = maybe_affinity.as_mut()?;
+			if affinity.core_index == core_index {
+				let new_count = affinity.count.saturating_sub(1);
+				if new_count > 0 {
+					*maybe_affinity = Some(CoreAffinityCount { core_index, count: new_count });
+				} else {
+					*maybe_affinity = None;
+				}
+				return Some(new_count);
+			} else {
+				None
+			}
+		})
+	}
+
+	/// Increases the affinity of a `ParaId` to a specified `CoreIndex`.
+	/// Adds to the count of the `CoreAffinityCount` if an entry is found and the core_index
+	/// matches. A non-existent entry will be initialized with a count of 1 and uses the supplied
+	/// `CoreIndex`.
+	fn increase_affinity(para_id: ParaId, core_index: CoreIndex) {
+		ParaIdAffinity::<T>::mutate(para_id, |maybe_affinity| match maybe_affinity {
+			Some(affinity) =>
+				if affinity.core_index == core_index {
+					*maybe_affinity = Some(CoreAffinityCount {
+						core_index,
+						count: affinity.count.saturating_add(1),
+					});
+				},
+			None => {
+				*maybe_affinity = Some(CoreAffinityCount { core_index, count: 1 });
+			},
+		})
+	}
+
+	/// Collect the revenue from the `when` blockheight
+	pub fn claim_revenue_until(when: BlockNumberFor<T>) -> BalanceOf<T> {
+		let now = <frame_system::Pallet<T>>::block_number();
+		let mut amount: BalanceOf<T> = BalanceOf::<T>::zero();
+		Revenue::<T>::mutate(|revenue| {
+			while !revenue.is_empty() {
+				let index = (revenue.len() - 1) as u32;
+				if when > now.saturating_sub(index.into()) {
+					amount = amount.saturating_add(revenue.pop().defensive_unwrap_or(0u32.into()));
+				} else {
+					break;
+				}
+			}
+		});
+
+		amount
+	}
+
+	/// Account of the pallet pot, where the funds from instantaneous coretime sale are accumulated.
+	pub fn account_id() -> T::AccountId {
+		T::PalletId::get().into_account_truncating()
+	}
+
+	/// Getter for the affinity tracker.
+	#[cfg(test)]
+	fn get_affinity_map(para_id: ParaId) -> Option<CoreAffinityCount> {
+		ParaIdAffinity::<T>::get(para_id)
+	}
+
+	/// Getter for the affinity entries.
+	#[cfg(test)]
+	fn get_affinity_entries(core_index: CoreIndex) -> BinaryHeap<EnqueuedOrder> {
+		AffinityEntries::<T>::get(core_index)
+	}
+
+	/// Getter for the free entries.
+	#[cfg(test)]
+	fn get_free_entries() -> BinaryHeap<EnqueuedOrder> {
+		FreeEntries::<T>::get()
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	pub fn populate_queue(para_id: ParaId, num: u32) {
+		QueueStatus::<T>::mutate(|queue_status| {
+			for _ in 0..num {
+				Pallet::<T>::add_on_demand_order(queue_status, para_id, QueuePushDirection::Back);
+			}
+		});
+	}
+
+	#[cfg(test)]
+	fn set_queue_status(new_status: QueueStatusType) {
+		QueueStatus::<T>::set(new_status);
+	}
+
+	#[cfg(test)]
+	fn get_queue_status() -> QueueStatusType {
+		QueueStatus::<T>::get()
+	}
+
+	#[cfg(test)]
+	fn get_traffic_default_value() -> FixedU128 {
+		<T as Config>::TrafficDefaultValue::get()
+	}
+
+	#[cfg(test)]
+	fn get_revenue() -> Vec<BalanceOf<T>> {
+		Revenue::<T>::get().to_vec()
+	}
+}

pezkuwi/runtime/teyrchains/src/on_demand/tests.rs

@@ -0,0 +1,882 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+use super::*;
+
+use crate::{
+	initializer::SessionChangeNotification,
+	mock::{
+		new_test_ext, Balances, OnDemand, Paras, ParasShared, RuntimeOrigin, Scheduler, System,
+		Test,
+	},
+	on_demand::{
+		self,
+		mock_helpers::GenesisConfigBuilder,
+		types::{QueueIndex, ReverseQueueIndex},
+		Error,
+	},
+	paras::{ParaGenesisArgs, ParaKind},
+};
+use core::cmp::{Ord, Ordering};
+use frame_support::{assert_noop, assert_ok};
+use pallet_balances::Error as BalancesError;
+use pezkuwi_primitives::{BlockNumber, SessionIndex, ValidationCode, ON_DEMAND_MAX_QUEUE_MAX_SIZE};
+use sp_runtime::traits::BadOrigin;
+
+fn schedule_blank_para(id: ParaId, parakind: ParaKind) {
+	let validation_code: ValidationCode = vec![1, 2, 3].into();
+	assert_ok!(Paras::schedule_para_initialize(
+		id,
+		ParaGenesisArgs {
+			genesis_head: Vec::new().into(),
+			validation_code: validation_code.clone(),
+			para_kind: parakind,
+		}
+	));
+
+	assert_ok!(Paras::add_trusted_validation_code(RuntimeOrigin::root(), validation_code));
+}
+
+fn run_to_block(
+	to: BlockNumber,
+	new_session: impl Fn(BlockNumber) -> Option<SessionChangeNotification<BlockNumber>>,
+) {
+	while System::block_number() < to {
+		let b = System::block_number();
+
+		Scheduler::initializer_finalize();
+		Paras::initializer_finalize(b);
+
+		if let Some(notification) = new_session(b + 1) {
+			let mut notification_with_session_index = notification;
+			// We will make every session change trigger an action queue. Normally this may require
+			// 2 or more session changes.
+			if notification_with_session_index.session_index == SessionIndex::default() {
+				notification_with_session_index.session_index = ParasShared::scheduled_session();
+			}
+			Paras::initializer_on_new_session(&notification_with_session_index);
+			Scheduler::initializer_on_new_session(&notification_with_session_index);
+		}
+
+		System::on_finalize(b);
+
+		System::on_initialize(b + 1);
+		System::set_block_number(b + 1);
+
+		Paras::initializer_initialize(b + 1);
+		Scheduler::initializer_initialize(b + 1);
+
+		// Update the spot traffic and revenue on every block.
+		OnDemand::on_initialize(b + 1);
+
+		// In the real runtime this is expected to be called by the `InclusionInherent` pallet.
+		Scheduler::advance_claim_queue(&Default::default());
+	}
+}
+
+fn place_order_run_to_blocknumber(para_id: ParaId, blocknumber: Option<BlockNumber>) {
+	let alice = 100u64;
+	let amt = 10_000_000u128;
+
+	Balances::make_free_balance_be(&alice, amt);
+
+	if let Some(bn) = blocknumber {
+		run_to_block(bn, |n| if n == bn { Some(Default::default()) } else { None });
+	}
+	#[allow(deprecated)]
+	OnDemand::place_order_allow_death(RuntimeOrigin::signed(alice), amt, para_id).unwrap()
+}
+
+fn place_order_run_to_101(para_id: ParaId) {
+	place_order_run_to_blocknumber(para_id, Some(101));
+}
+
+fn place_order(para_id: ParaId) {
+	place_order_run_to_blocknumber(para_id, None);
+}
+
+#[test]
+fn spot_traffic_capacity_zero_returns_none() {
+	match OnDemand::calculate_spot_traffic(
+		FixedU128::from(u128::MAX),
+		0u32,
+		u32::MAX,
+		Perbill::from_percent(100),
+		Perbill::from_percent(1),
+	) {
+		Ok(_) => panic!("Error"),
+		Err(e) => assert_eq!(e, SpotTrafficCalculationErr::QueueCapacityIsZero),
+	};
+}
+
+#[test]
+fn spot_traffic_queue_size_larger_than_capacity_returns_none() {
+	match OnDemand::calculate_spot_traffic(
+		FixedU128::from(u128::MAX),
+		1u32,
+		2u32,
+		Perbill::from_percent(100),
+		Perbill::from_percent(1),
+	) {
+		Ok(_) => panic!("Error"),
+		Err(e) => assert_eq!(e, SpotTrafficCalculationErr::QueueSizeLargerThanCapacity),
+	}
+}
+
+#[test]
+fn spot_traffic_calculation_identity() {
+	match OnDemand::calculate_spot_traffic(
+		FixedU128::from_u32(1),
+		1000,
+		100,
+		Perbill::from_percent(10),
+		Perbill::from_percent(3),
+	) {
+		Ok(res) => {
+			assert_eq!(res, FixedU128::from_u32(1))
+		},
+		_ => (),
+	}
+}
+
+#[test]
+fn spot_traffic_calculation_u32_max() {
+	match OnDemand::calculate_spot_traffic(
+		FixedU128::from_u32(1),
+		u32::MAX,
+		u32::MAX,
+		Perbill::from_percent(100),
+		Perbill::from_percent(3),
+	) {
+		Ok(res) => {
+			assert_eq!(res, FixedU128::from_u32(1))
+		},
+		_ => panic!("Error"),
+	};
+}
+
+#[test]
+fn spot_traffic_calculation_u32_traffic_max() {
+	match OnDemand::calculate_spot_traffic(
+		FixedU128::from(u128::MAX),
+		u32::MAX,
+		u32::MAX,
+		Perbill::from_percent(1),
+		Perbill::from_percent(1),
+	) {
+		Ok(res) => assert_eq!(res, FixedU128::from(u128::MAX)),
+		_ => panic!("Error"),
+	};
+}
+
+#[test]
+fn sustained_target_increases_spot_traffic() {
+	let mut traffic = FixedU128::from_u32(1u32);
+	for _ in 0..50 {
+		traffic = OnDemand::calculate_spot_traffic(
+			traffic,
+			100,
+			12,
+			Perbill::from_percent(10),
+			Perbill::from_percent(100),
+		)
+		.unwrap()
+	}
+	assert_eq!(traffic, FixedU128::from_inner(2_718_103_312_071_174_015u128))
+}
+
+#[test]
+fn spot_traffic_can_decrease() {
+	let traffic = FixedU128::from_u32(100u32);
+	match OnDemand::calculate_spot_traffic(
+		traffic,
+		100u32,
+		0u32,
+		Perbill::from_percent(100),
+		Perbill::from_percent(100),
+	) {
+		Ok(new_traffic) => {
+			assert_eq!(new_traffic, FixedU128::from_inner(50_000_000_000_000_000_000u128))
+		},
+		_ => panic!("Error"),
+	}
+}
+
+#[test]
+fn spot_traffic_decreases_over_time() {
+	let mut traffic = FixedU128::from_u32(100u32);
+	for _ in 0..5 {
+		traffic = OnDemand::calculate_spot_traffic(
+			traffic,
+			100u32,
+			0u32,
+			Perbill::from_percent(100),
+			Perbill::from_percent(100),
+		)
+		.unwrap();
+		println!("{traffic}");
+	}
+	assert_eq!(traffic, FixedU128::from_inner(3_125_000_000_000_000_000u128))
+}
+
+#[test]
+fn spot_traffic_decreases_between_idle_blocks() {
+	// Testing spot traffic assumptions, but using the mock runtime and default on demand
+	// configuration values. Ensuring that blocks with no on demand activity (idle)
+	// decrease traffic.
+
+	let para_id = ParaId::from(111);
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		// Initialize the parathread and wait for it to be ready.
+		schedule_blank_para(para_id, ParaKind::Parathread);
+		assert!(!Paras::is_parathread(para_id));
+		run_to_block(100, |n| if n == 100 { Some(Default::default()) } else { None });
+		assert!(Paras::is_parathread(para_id));
+
+		// Set the spot traffic to a large number
+		OnDemand::set_queue_status(QueueStatusType {
+			traffic: FixedU128::from_u32(10),
+			..Default::default()
+		});
+
+		assert_eq!(OnDemand::get_queue_status().traffic, FixedU128::from_u32(10));
+
+		// Run to block 101 and ensure that the traffic decreases.
+		run_to_block(101, |n| if n == 100 { Some(Default::default()) } else { None });
+		assert!(OnDemand::get_queue_status().traffic < FixedU128::from_u32(10));
+
+		// Run to block 102 and observe that we've hit the default traffic value.
+		run_to_block(102, |n| if n == 100 { Some(Default::default()) } else { None });
+		assert_eq!(OnDemand::get_queue_status().traffic, OnDemand::get_traffic_default_value());
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn place_order_works() {
+	let alice = 1u64;
+	let amt = 10_000_000u128;
+	let para_id = ParaId::from(111);
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		// Initialize the parathread and wait for it to be ready.
+		schedule_blank_para(para_id, ParaKind::Parathread);
+
+		assert!(!Paras::is_parathread(para_id));
+
+		run_to_block(100, |n| if n == 100 { Some(Default::default()) } else { None });
+
+		assert!(Paras::is_parathread(para_id));
+
+		// Does not work unsigned
+		assert_noop!(
+			OnDemand::place_order_allow_death(RuntimeOrigin::none(), amt, para_id),
+			BadOrigin
+		);
+
+		// Does not work with max_amount lower than fee
+		let low_max_amt = 1u128;
+		assert_noop!(
+			OnDemand::place_order_allow_death(RuntimeOrigin::signed(alice), low_max_amt, para_id,),
+			Error::<Test>::SpotPriceHigherThanMaxAmount,
+		);
+
+		// Does not work with insufficient balance
+		assert_noop!(
+			OnDemand::place_order_allow_death(RuntimeOrigin::signed(alice), amt, para_id),
+			BalancesError::<Test, _>::InsufficientBalance
+		);
+
+		// Works
+		Balances::make_free_balance_be(&alice, amt);
+		run_to_block(101, |n| if n == 101 { Some(Default::default()) } else { None });
+		assert_ok!(OnDemand::place_order_allow_death(RuntimeOrigin::signed(alice), amt, para_id));
+	});
+}
+
+#[test]
+#[allow(deprecated)]
+fn place_order_keep_alive_keeps_alive() {
+	let alice = 1u64;
+	let amt = 1u128; // The same as crate::mock's EXISTENTIAL_DEPOSIT
+	let max_amt = 10_000_000u128;
+	let para_id = ParaId::from(111);
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let config = configuration::ActiveConfig::<Test>::get();
+
+		// Initialize the parathread and wait for it to be ready.
+		schedule_blank_para(para_id, ParaKind::Parathread);
+		Balances::make_free_balance_be(&alice, amt);
+
+		assert!(!Paras::is_parathread(para_id));
+		run_to_block(100, |n| if n == 100 { Some(Default::default()) } else { None });
+		assert!(Paras::is_parathread(para_id));
+
+		assert_noop!(
+			OnDemand::place_order_keep_alive(RuntimeOrigin::signed(alice), max_amt, para_id),
+			BalancesError::<Test, _>::InsufficientBalance
+		);
+
+		Balances::make_free_balance_be(&alice, max_amt);
+		assert_ok!(OnDemand::place_order_keep_alive(
+			RuntimeOrigin::signed(alice),
+			max_amt,
+			para_id
+		),);
+
+		let queue_status = QueueStatus::<Test>::get();
+		let spot_price = queue_status.traffic.saturating_mul_int(
+			config.scheduler_params.on_demand_base_fee.saturated_into::<BalanceOf<Test>>(),
+		);
+		assert_eq!(Balances::free_balance(&alice), max_amt.saturating_sub(spot_price));
+		assert_eq!(
+			FreeEntries::<Test>::get().pop(),
+			Some(EnqueuedOrder::new(QueueIndex(0), para_id))
+		);
+	});
+}
+
+#[test]
+fn place_order_with_credits() {
+	let alice = 1u64;
+	let initial_credit = 10_000_000u128;
+	let para_id = ParaId::from(111);
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let config = configuration::ActiveConfig::<Test>::get();
+
+		// Initialize the parathread and wait for it to be ready.
+		schedule_blank_para(para_id, ParaKind::Parathread);
+		OnDemand::credit_account(alice, initial_credit);
+		assert_eq!(Credits::<Test>::get(alice), initial_credit);
+
+		assert!(!Paras::is_parathread(para_id));
+		run_to_block(100, |n| if n == 100 { Some(Default::default()) } else { None });
+		assert!(Paras::is_parathread(para_id));
+
+		let queue_status = QueueStatus::<Test>::get();
+		let spot_price = queue_status.traffic.saturating_mul_int(
+			config.scheduler_params.on_demand_base_fee.saturated_into::<BalanceOf<Test>>(),
+		);
+
+		// Create an order and pay for it with credits.
+		assert_ok!(OnDemand::place_order_with_credits(
+			RuntimeOrigin::signed(alice),
+			initial_credit,
+			para_id
+		));
+		assert_eq!(Credits::<Test>::get(alice), initial_credit.saturating_sub(spot_price));
+		assert_eq!(
+			FreeEntries::<Test>::get().pop(),
+			Some(EnqueuedOrder::new(QueueIndex(0), para_id))
+		);
+
+		// Insufficient credits:
+		Credits::<Test>::insert(alice, 1u128);
+		assert_noop!(
+			OnDemand::place_order_with_credits(
+				RuntimeOrigin::signed(alice),
+				1_000_000u128,
+				para_id
+			),
+			Error::<Test>::InsufficientCredits
+		);
+	});
+}
+
+#[test]
+fn pop_assignment_for_core_works() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let para_b = ParaId::from(110);
+		schedule_blank_para(para_a, ParaKind::Parathread);
+		schedule_blank_para(para_b, ParaKind::Parathread);
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// Pop should return none with empty queue
+		assert_eq!(OnDemand::pop_assignment_for_core(CoreIndex(0)), None);
+
+		// Add enough assignments to the order queue.
+		for _ in 0..2 {
+			place_order(para_a);
+			place_order(para_b);
+		}
+
+		// Popped assignments should be for the correct paras and cores
+		assert_eq!(
+			OnDemand::pop_assignment_for_core(CoreIndex(0)).map(|a| a.para_id()),
+			Some(para_a)
+		);
+		assert_eq!(
+			OnDemand::pop_assignment_for_core(CoreIndex(1)).map(|a| a.para_id()),
+			Some(para_b)
+		);
+		assert_eq!(
+			OnDemand::pop_assignment_for_core(CoreIndex(0)).map(|a| a.para_id()),
+			Some(para_a)
+		);
+		assert_eq!(
+			OnDemand::pop_assignment_for_core(CoreIndex(1)).map(|a| a.para_id()),
+			Some(para_b)
+		);
+	});
+}
+
+#[test]
+fn push_back_assignment_works() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let para_b = ParaId::from(110);
+		schedule_blank_para(para_a, ParaKind::Parathread);
+		schedule_blank_para(para_b, ParaKind::Parathread);
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// Add enough assignments to the order queue.
+		place_order_run_to_101(para_a);
+		place_order_run_to_101(para_b);
+
+		// Pop order a
+		assert_eq!(OnDemand::pop_assignment_for_core(CoreIndex(0)).unwrap().para_id(), para_a);
+
+		// Para a should have affinity for core 0
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 1);
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().core_index, CoreIndex(0));
+
+		// Push back order a
+		OnDemand::push_back_assignment(para_a, CoreIndex(0));
+
+		// Para a should have no affinity
+		assert_eq!(OnDemand::get_affinity_map(para_a).is_none(), true);
+
+		// Queue should contain orders a, b. A in front of b.
+		assert_eq!(OnDemand::pop_assignment_for_core(CoreIndex(0)).unwrap().para_id(), para_a);
+		assert_eq!(OnDemand::pop_assignment_for_core(CoreIndex(0)).unwrap().para_id(), para_b);
+	});
+}
+
+#[test]
+fn affinity_prohibits_parallel_scheduling() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let para_b = ParaId::from(222);
+
+		schedule_blank_para(para_a, ParaKind::Parathread);
+		schedule_blank_para(para_b, ParaKind::Parathread);
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// There should be no affinity before starting.
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+		assert!(OnDemand::get_affinity_map(para_b).is_none());
+
+		// Add 2 assignments for para_a for every para_b.
+		place_order_run_to_101(para_a);
+		place_order_run_to_101(para_a);
+		place_order_run_to_101(para_b);
+
+		// Approximate having 1 core.
+		for _ in 0..3 {
+			assert!(OnDemand::pop_assignment_for_core(CoreIndex(0)).is_some());
+		}
+		assert!(OnDemand::pop_assignment_for_core(CoreIndex(0)).is_none());
+
+		// Affinity on one core is meaningless.
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 2);
+		assert_eq!(OnDemand::get_affinity_map(para_b).unwrap().count, 1);
+		assert_eq!(
+			OnDemand::get_affinity_map(para_a).unwrap().core_index,
+			OnDemand::get_affinity_map(para_b).unwrap().core_index,
+		);
+
+		// Clear affinity
+		OnDemand::report_processed(para_a, 0.into());
+		OnDemand::report_processed(para_a, 0.into());
+		OnDemand::report_processed(para_b, 0.into());
+
+		// Add 2 assignments for para_a for every para_b.
+		place_order_run_to_101(para_a);
+		place_order_run_to_101(para_a);
+		place_order_run_to_101(para_b);
+
+		// Approximate having 3 cores. CoreIndex 2 should be unable to obtain an assignment
+		for _ in 0..3 {
+			OnDemand::pop_assignment_for_core(CoreIndex(0));
+			OnDemand::pop_assignment_for_core(CoreIndex(1));
+			assert!(OnDemand::pop_assignment_for_core(CoreIndex(2)).is_none());
+		}
+
+		// Affinity should be the same as before, but on different cores.
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 2);
+		assert_eq!(OnDemand::get_affinity_map(para_b).unwrap().count, 1);
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().core_index, CoreIndex(0));
+		assert_eq!(OnDemand::get_affinity_map(para_b).unwrap().core_index, CoreIndex(1));
+
+		// Clear affinity
+		OnDemand::report_processed(para_a, CoreIndex(0));
+		OnDemand::report_processed(para_a, CoreIndex(0));
+		OnDemand::report_processed(para_b, CoreIndex(1));
+
+		// There should be no affinity after clearing.
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+		assert!(OnDemand::get_affinity_map(para_b).is_none());
+	});
+}
+
+#[test]
+fn affinity_changes_work() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let core_index = CoreIndex(0);
+		schedule_blank_para(para_a, ParaKind::Parathread);
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// There should be no affinity before starting.
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+
+		// Add enough assignments to the order queue.
+		for _ in 0..10 {
+			place_order_run_to_101(para_a);
+		}
+
+		// There should be no affinity before the scheduler pops.
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+
+		OnDemand::pop_assignment_for_core(core_index);
+
+		// Affinity count is 1 after popping.
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 1);
+
+		OnDemand::report_processed(para_a, 0.into());
+		OnDemand::pop_assignment_for_core(core_index);
+
+		// Affinity count is 1 after popping with a previous para.
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 1);
+
+		for _ in 0..3 {
+			OnDemand::pop_assignment_for_core(core_index);
+		}
+
+		// Affinity count is 4 after popping 3 times without a previous para.
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 4);
+
+		for _ in 0..5 {
+			OnDemand::report_processed(para_a, 0.into());
+			assert!(OnDemand::pop_assignment_for_core(core_index).is_some());
+		}
+
+		// Affinity count should still be 4 but queue should be empty.
+		assert!(OnDemand::pop_assignment_for_core(core_index).is_none());
+		assert_eq!(OnDemand::get_affinity_map(para_a).unwrap().count, 4);
+
+		// Pop 4 times and get to exactly 0 (None) affinity.
+		for _ in 0..4 {
+			OnDemand::report_processed(para_a, 0.into());
+			assert!(OnDemand::pop_assignment_for_core(core_index).is_none());
+		}
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+
+		// Decreasing affinity beyond 0 should still be None.
+		OnDemand::report_processed(para_a, 0.into());
+		assert!(OnDemand::pop_assignment_for_core(core_index).is_none());
+		assert!(OnDemand::get_affinity_map(para_a).is_none());
+	});
+}
+
+#[test]
+fn new_affinity_for_a_core_must_come_from_free_entries() {
+	// If affinity count for a core was zero before, and is 1 now, then the entry
+	// must have come from free_entries.
+	let teyrchains =
+		vec![ParaId::from(111), ParaId::from(222), ParaId::from(333), ParaId::from(444)];
+	let core_indices = vec![CoreIndex(0), CoreIndex(1), CoreIndex(2), CoreIndex(3)];
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		teyrchains.iter().for_each(|chain| {
+			schedule_blank_para(*chain, ParaKind::Parathread);
+		});
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// Place orders for all chains.
+		teyrchains.iter().for_each(|chain| {
+			place_order_run_to_101(*chain);
+		});
+
+		// There are 4 entries in free_entries.
+		let start_free_entries = OnDemand::get_free_entries().len();
+		assert_eq!(start_free_entries, 4);
+
+		// Pop assignments on all cores.
+		core_indices.iter().enumerate().for_each(|(n, core_index)| {
+			// There is no affinity on the core prior to popping.
+			assert!(OnDemand::get_affinity_entries(*core_index).is_empty());
+
+			// There's always an order to be popped for each core.
+			let free_entries = OnDemand::get_free_entries();
+			let next_order = free_entries.peek();
+
+			// There is no affinity on the paraid prior to popping.
+			assert!(OnDemand::get_affinity_map(next_order.unwrap().para_id).is_none());
+
+			match OnDemand::pop_assignment_for_core(*core_index) {
+				Some(assignment) => {
+					// The popped assignment came from free entries.
+					assert_eq!(start_free_entries - 1 - n, OnDemand::get_free_entries().len());
+					// The popped assignment has the same para id as the next order.
+					assert_eq!(assignment.para_id(), next_order.unwrap().para_id);
+				},
+				None => panic!("Should not happen"),
+			}
+		});
+
+		// All entries have been removed from free_entries.
+		assert!(OnDemand::get_free_entries().is_empty());
+
+		// All chains have an affinity count of 1.
+		teyrchains.iter().for_each(|chain| {
+			assert_eq!(OnDemand::get_affinity_map(*chain).unwrap().count, 1);
+		});
+	});
+}
+
+#[test]
+#[should_panic]
+fn queue_index_ordering_is_unsound_over_max_size() {
+	// NOTE: Unsoundness proof. If the number goes sufficiently over the max_queue_max_size
+	// the overflow will cause an opposite comparison to what would be expected.
+	let max_num = u32::MAX - ON_DEMAND_MAX_QUEUE_MAX_SIZE;
+	// 0 < some large number.
+	assert_eq!(QueueIndex(0).cmp(&QueueIndex(max_num + 1)), Ordering::Less);
+}
+
+#[test]
+fn queue_index_ordering_works() {
+	// The largest accepted queue size.
+	let max_num = ON_DEMAND_MAX_QUEUE_MAX_SIZE;
+
+	// 0 == 0
+	assert_eq!(QueueIndex(0).cmp(&QueueIndex(0)), Ordering::Equal);
+	// 0 < 1
+	assert_eq!(QueueIndex(0).cmp(&QueueIndex(1)), Ordering::Less);
+	// 1 > 0
+	assert_eq!(QueueIndex(1).cmp(&QueueIndex(0)), Ordering::Greater);
+	// 0 < max_num
+	assert_eq!(QueueIndex(0).cmp(&QueueIndex(max_num)), Ordering::Less);
+	// 0 > max_num + 1
+	assert_eq!(QueueIndex(0).cmp(&QueueIndex(max_num + 1)), Ordering::Less);
+
+	// Ordering within the bounds of ON_DEMAND_MAX_QUEUE_MAX_SIZE works.
+	let mut v = vec![3, 6, 2, 1, 5, 4];
+	v.sort_by_key(|&num| QueueIndex(num));
+	assert_eq!(v, vec![1, 2, 3, 4, 5, 6]);
+
+	v = vec![max_num, 4, 5, 1, 6];
+	v.sort_by_key(|&num| QueueIndex(num));
+	assert_eq!(v, vec![1, 4, 5, 6, max_num]);
+
+	// Ordering with an element outside of the bounds of the max size also works.
+	v = vec![max_num + 2, 0, 6, 2, 1, 5, 4];
+	v.sort_by_key(|&num| QueueIndex(num));
+	assert_eq!(v, vec![0, 1, 2, 4, 5, 6, max_num + 2]);
+
+	// Numbers way above the max size will overflow
+	v = vec![u32::MAX - 1, u32::MAX, 6, 2, 1, 5, 4];
+	v.sort_by_key(|&num| QueueIndex(num));
+	assert_eq!(v, vec![u32::MAX - 1, u32::MAX, 1, 2, 4, 5, 6]);
+}
+
+#[test]
+fn reverse_queue_index_does_reverse() {
+	let mut v = vec![1, 2, 3, 4, 5, 6];
+
+	// Basic reversal of a vector.
+	v.sort_by_key(|&num| ReverseQueueIndex(num));
+	assert_eq!(v, vec![6, 5, 4, 3, 2, 1]);
+
+	// Example from rust docs on `Reverse`. Should work identically.
+	v.sort_by_key(|&num| (num > 3, ReverseQueueIndex(num)));
+	assert_eq!(v, vec![3, 2, 1, 6, 5, 4]);
+
+	let mut v2 = vec![1, 2, u32::MAX];
+	v2.sort_by_key(|&num| ReverseQueueIndex(num));
+	assert_eq!(v2, vec![2, 1, u32::MAX]);
+}
+
+#[test]
+fn queue_status_size_fn_works() {
+	// Add orders to the on demand queue, and make sure that they are properly represented
+	// by the QueueStatusType::size fn.
+	let teyrchains = vec![ParaId::from(111), ParaId::from(222), ParaId::from(333)];
+	let core_indices = vec![CoreIndex(0), CoreIndex(1)];
+
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		teyrchains.iter().for_each(|chain| {
+			schedule_blank_para(*chain, ParaKind::Parathread);
+		});
+
+		assert_eq!(OnDemand::get_queue_status().size(), 0);
+
+		run_to_block(11, |n| if n == 11 { Some(Default::default()) } else { None });
+
+		// Place orders for all chains.
+		teyrchains.iter().for_each(|chain| {
+			// 2 per chain for a total of 6
+			place_order_run_to_101(*chain);
+			place_order_run_to_101(*chain);
+		});
+
+		// 6 orders in free entries
+		assert_eq!(OnDemand::get_free_entries().len(), 6);
+		// 6 orders via queue status size
+		assert_eq!(
+			OnDemand::get_free_entries().len(),
+			OnDemand::get_queue_status().size() as usize
+		);
+
+		core_indices.iter().for_each(|core_index| {
+			OnDemand::pop_assignment_for_core(*core_index);
+		});
+
+		// There should be 2 orders in the scheduler's claimqueue,
+		// 2 in assorted AffinityMaps and 2 in free.
+		// ParaId 111
+		assert_eq!(OnDemand::get_affinity_entries(core_indices[0]).len(), 1);
+		// ParaId 222
+		assert_eq!(OnDemand::get_affinity_entries(core_indices[1]).len(), 1);
+		// Free entries are from ParaId 333
+		assert_eq!(OnDemand::get_free_entries().len(), 2);
+		// For a total size of 4.
+		assert_eq!(OnDemand::get_queue_status().size(), 4)
+	});
+}
+
+#[test]
+fn revenue_information_fetching_works() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		schedule_blank_para(para_a, ParaKind::Parathread);
+		// Mock assigner sets max revenue history to 10.
+		run_to_block(10, |n| if n == 10 { Some(Default::default()) } else { None });
+		let revenue = OnDemand::claim_revenue_until(10);
+
+		// No revenue should be recorded.
+		assert_eq!(revenue, 0);
+
+		// Place one order
+		place_order_run_to_blocknumber(para_a, Some(11));
+		let revenue = OnDemand::get_revenue();
+		let amt = OnDemand::claim_revenue_until(11);
+
+		// Revenue until the current block is still zero as "until" is non-inclusive
+		assert_eq!(amt, 0);
+
+		let amt = OnDemand::claim_revenue_until(12);
+
+		// Revenue for a single order should be recorded and shouldn't have been pruned by the
+		// previous call
+		assert_eq!(amt, revenue[0]);
+
+		run_to_block(12, |n| if n == 12 { Some(Default::default()) } else { None });
+		let revenue = OnDemand::claim_revenue_until(13);
+
+		// No revenue should be recorded.
+		assert_eq!(revenue, 0);
+
+		// Place many orders
+		place_order(para_a);
+		place_order(para_a);
+
+		run_to_block(13, |n| if n == 13 { Some(Default::default()) } else { None });
+
+		place_order(para_a);
+
+		run_to_block(14, |n| if n == 14 { Some(Default::default()) } else { None });
+
+		let revenue = OnDemand::claim_revenue_until(15);
+
+		// All 3 orders should be accounted for.
+		assert_eq!(revenue, 30_000);
+
+		// Place one order
+		place_order_run_to_blocknumber(para_a, Some(16));
+
+		let revenue = OnDemand::claim_revenue_until(15);
+
+		// Order is not in range of  the revenue_until call
+		assert_eq!(revenue, 0);
+
+		run_to_block(20, |n| if n == 20 { Some(Default::default()) } else { None });
+		let revenue = OnDemand::claim_revenue_until(21);
+		assert_eq!(revenue, 10_000);
+
+		// Make sure overdue revenue is accumulated
+		for i in 21..=35 {
+			run_to_block(i, |n| if n % 10 == 0 { Some(Default::default()) } else { None });
+			place_order(para_a);
+		}
+		let revenue = OnDemand::claim_revenue_until(36);
+		assert_eq!(revenue, 150_000);
+	});
+}
+
+#[test]
+fn pot_account_is_immortal() {
+	new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
+		let para_a = ParaId::from(111);
+		let pot = OnDemand::account_id();
+		assert!(!System::account_exists(&pot));
+		schedule_blank_para(para_a, ParaKind::Parathread);
+		// Mock assigner sets max revenue history to 10.
+
+		run_to_block(10, |n| if n == 10 { Some(Default::default()) } else { None });
+		place_order_run_to_blocknumber(para_a, Some(12));
+		let purchase_revenue = Balances::free_balance(&pot);
+		assert!(purchase_revenue > 0);
+
+		run_to_block(15, |_| None);
+		let _imb = <Test as on_demand::Config>::Currency::withdraw(
+			&pot,
+			purchase_revenue,
+			WithdrawReasons::FEE,
+			ExistenceRequirement::AllowDeath,
+		);
+		assert_eq!(Balances::free_balance(&pot), 0);
+		assert!(System::account_exists(&pot));
+		assert_eq!(System::providers(&pot), 1);
+
+		// One more cycle to make sure providers are not increased on every transition from zero
+		run_to_block(20, |n| if n == 20 { Some(Default::default()) } else { None });
+		place_order_run_to_blocknumber(para_a, Some(22));
+		let purchase_revenue = Balances::free_balance(&pot);
+		assert!(purchase_revenue > 0);
+
+		run_to_block(25, |_| None);
+		let _imb = <Test as on_demand::Config>::Currency::withdraw(
+			&pot,
+			purchase_revenue,
+			WithdrawReasons::FEE,
+			ExistenceRequirement::AllowDeath,
+		);
+		assert_eq!(Balances::free_balance(&pot), 0);
+		assert!(System::account_exists(&pot));
+		assert_eq!(System::providers(&pot), 1);
+	});
+}

pezkuwi/runtime/teyrchains/src/on_demand/types.rs

@@ -0,0 +1,238 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+//! On demand module types.
+
+use super::{alloc, pallet::Config};
+use alloc::collections::BinaryHeap;
+use core::cmp::{Ord, Ordering, PartialOrd};
+use frame_support::{
+	pallet_prelude::{Decode, Encode, RuntimeDebug, TypeInfo},
+	traits::Currency,
+};
+use pezkuwi_primitives::{CoreIndex, Id as ParaId, ON_DEMAND_MAX_QUEUE_MAX_SIZE};
+use sp_runtime::FixedU128;
+
+/// Shorthand for the Balance type the runtime is using.
+pub type BalanceOf<T> =
+	<<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
+
+/// Meta data for full queue.
+///
+/// This includes elements with affinity and free entries.
+///
+/// The actual queue is implemented via multiple priority queues. One for each core, for entries
+/// which currently have a core affinity and one free queue, with entries without any affinity yet.
+///
+/// The design aims to have most queue accessess be O(1) or O(log(N)). Absolute worst case is O(N).
+/// Importantly this includes all accessess that happen in a single block. Even with 50 cores, the
+/// total complexity of all operations in the block should maintain above complexities. In
+/// particular O(N) stays O(N), it should never be O(N*cores).
+///
+/// More concrete rundown on complexity:
+///
+///  - insert: O(1) for placing an order, O(log(N)) for push backs.
+///  - pop_assignment_for_core: O(log(N)), O(N) worst case: Can only happen for one core, next core
+///  is already less work.
+///  - report_processed & push back: If affinity dropped to 0, then O(N) in the worst case. Again
+///  this divides per core.
+///
+///  Reads still exist, also improved slightly, but worst case we fetch all entries.
+#[derive(Encode, Decode, TypeInfo)]
+pub struct QueueStatusType {
+	/// Last calculated traffic value.
+	pub traffic: FixedU128,
+	/// The next index to use.
+	pub next_index: QueueIndex,
+	/// Smallest index still in use.
+	///
+	/// In case of a completely empty queue (free + affinity queues), `next_index - smallest_index
+	/// == 0`.
+	pub smallest_index: QueueIndex,
+	/// Indices that have been freed already.
+	///
+	/// But have a hole to `smallest_index`, so we can not yet bump `smallest_index`. This binary
+	/// heap is roughly bounded in the number of on demand cores:
+	///
+	/// For a single core, elements will always be processed in order. With each core added, a
+	/// level of out of order execution is added.
+	pub freed_indices: BinaryHeap<ReverseQueueIndex>,
+}
+
+impl Default for QueueStatusType {
+	fn default() -> QueueStatusType {
+		QueueStatusType {
+			traffic: FixedU128::default(),
+			next_index: QueueIndex(0),
+			smallest_index: QueueIndex(0),
+			freed_indices: BinaryHeap::new(),
+		}
+	}
+}
+
+impl QueueStatusType {
+	/// How many orders are queued in total?
+	///
+	/// This includes entries which have core affinity.
+	pub fn size(&self) -> u32 {
+		self.next_index
+			.0
+			.overflowing_sub(self.smallest_index.0)
+			.0
+			.saturating_sub(self.freed_indices.len() as u32)
+	}
+
+	/// Get current next index
+	///
+	/// to use for an element newly pushed to the back of the queue.
+	pub fn push_back(&mut self) -> QueueIndex {
+		let QueueIndex(next_index) = self.next_index;
+		self.next_index = QueueIndex(next_index.overflowing_add(1).0);
+		QueueIndex(next_index)
+	}
+
+	/// Push something to the front of the queue
+	pub fn push_front(&mut self) -> QueueIndex {
+		self.smallest_index = QueueIndex(self.smallest_index.0.overflowing_sub(1).0);
+		self.smallest_index
+	}
+
+	/// The given index is no longer part of the queue.
+	///
+	/// This updates `smallest_index` if need be.
+	pub fn consume_index(&mut self, removed_index: QueueIndex) {
+		if removed_index != self.smallest_index {
+			self.freed_indices.push(removed_index.reverse());
+			return;
+		}
+		let mut index = self.smallest_index.0.overflowing_add(1).0;
+		// Even more to advance?
+		while self.freed_indices.peek() == Some(&ReverseQueueIndex(index)) {
+			index = index.overflowing_add(1).0;
+			self.freed_indices.pop();
+		}
+		self.smallest_index = QueueIndex(index);
+	}
+}
+
+/// Type used for priority indices.
+//  NOTE: The `Ord` implementation for this type is unsound in the general case.
+//        Do not use it for anything but it's intended purpose.
+#[derive(Encode, Decode, TypeInfo, Debug, PartialEq, Clone, Eq, Copy)]
+pub struct QueueIndex(pub u32);
+
+/// QueueIndex with reverse ordering.
+///
+/// Same as `Reverse(QueueIndex)`, but with all the needed traits implemented.
+#[derive(Encode, Decode, TypeInfo, Debug, PartialEq, Clone, Eq, Copy)]
+pub struct ReverseQueueIndex(pub u32);
+
+impl QueueIndex {
+	fn reverse(self) -> ReverseQueueIndex {
+		ReverseQueueIndex(self.0)
+	}
+}
+
+impl Ord for QueueIndex {
+	fn cmp(&self, other: &Self) -> Ordering {
+		let diff = self.0.overflowing_sub(other.0).0;
+		if diff == 0 {
+			Ordering::Equal
+		} else if diff <= ON_DEMAND_MAX_QUEUE_MAX_SIZE {
+			Ordering::Greater
+		} else {
+			Ordering::Less
+		}
+	}
+}
+
+impl PartialOrd for QueueIndex {
+	fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+		Some(self.cmp(other))
+	}
+}
+
+impl Ord for ReverseQueueIndex {
+	fn cmp(&self, other: &Self) -> Ordering {
+		QueueIndex(other.0).cmp(&QueueIndex(self.0))
+	}
+}
+impl PartialOrd for ReverseQueueIndex {
+	fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+		Some(self.cmp(&other))
+	}
+}
+
+/// Internal representation of an order after it has been enqueued already.
+///
+/// This data structure is provided for a min BinaryHeap (Ord compares in reverse order with regards
+/// to its elements)
+#[derive(Encode, Decode, TypeInfo, Debug, PartialEq, Clone, Eq)]
+pub struct EnqueuedOrder {
+	pub para_id: ParaId,
+	pub idx: QueueIndex,
+}
+
+impl EnqueuedOrder {
+	pub fn new(idx: QueueIndex, para_id: ParaId) -> Self {
+		Self { idx, para_id }
+	}
+}
+
+impl PartialOrd for EnqueuedOrder {
+	fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+		match other.idx.partial_cmp(&self.idx) {
+			Some(Ordering::Equal) => other.para_id.partial_cmp(&self.para_id),
+			o => o,
+		}
+	}
+}
+
+impl Ord for EnqueuedOrder {
+	fn cmp(&self, other: &Self) -> Ordering {
+		match other.idx.cmp(&self.idx) {
+			Ordering::Equal => other.para_id.cmp(&self.para_id),
+			o => o,
+		}
+	}
+}
+
+/// Keeps track of how many assignments a scheduler currently has at a specific `CoreIndex` for a
+/// specific `ParaId`.
+#[derive(Encode, Decode, Default, Clone, Copy, TypeInfo)]
+#[cfg_attr(test, derive(PartialEq, RuntimeDebug))]
+pub struct CoreAffinityCount {
+	pub core_index: CoreIndex,
+	pub count: u32,
+}
+
+/// An indicator as to which end of the `OnDemandQueue` an assignment will be placed.
+#[cfg_attr(test, derive(RuntimeDebug))]
+pub enum QueuePushDirection {
+	Back,
+	Front,
+}
+
+/// Errors that can happen during spot traffic calculation.
+#[derive(PartialEq, RuntimeDebug)]
+pub enum SpotTrafficCalculationErr {
+	/// The order queue capacity is at 0.
+	QueueCapacityIsZero,
+	/// The queue size is larger than the queue capacity.
+	QueueSizeLargerThanCapacity,
+	/// Arithmetic error during division, either division by 0 or over/underflow.
+	Division,
+}