[fix] improvements to proposal handling (#690)

* [fix] improvements to proposal handling

* fix harness

* fix clippy

Author: 1xstj

dkg-gadget/src/constants.rs

@@ -0,0 +1,61 @@
+// Constants for dkg-gadget
+
+// ================= Common ======================== //
+pub const DKG_KEYGEN_PROTOCOL_NAME: &str = "/webb-tools/dkg/keygen/1";
+
+pub const DKG_SIGNING_PROTOCOL_NAME: &str = "/webb-tools/dkg/signing/1";
+
+// ================= Worker ========================== //
+pub mod worker {
+	pub const ENGINE_ID: sp_runtime::ConsensusEngineId = *b"WDKG";
+
+	pub const STORAGE_SET_RETRY_NUM: usize = 5;
+
+	pub const MAX_SUBMISSION_DELAY: u32 = 3;
+
+	pub const MAX_KEYGEN_RETRIES: usize = 5;
+
+	/// How many blocks to keep the proposal hash in out local cache.
+	pub const PROPOSAL_HASH_LIFETIME: u32 = 10;
+}
+
+// ============= Signing Manager ======================= //
+
+pub mod signing_manager {
+	// the maximum number of tasks that the work manager tries to assign
+	pub const MAX_RUNNING_TASKS: usize = 1;
+
+	// the maximum number of tasks that can be enqueued,
+	// enqueued here implies not actively running but listening for messages
+	pub const MAX_ENQUEUED_TASKS: usize = 20;
+
+	// How often to poll the jobs to check completion status
+	pub const JOB_POLL_INTERVAL_IN_MILLISECONDS: u64 = 500;
+
+	// Number of signing sets to generate for every proposal
+	pub const MAX_POTENTIAL_SIGNING_SETS_PER_PROPOSAL: u8 = 2;
+}
+
+// ============= Networking ======================= //
+
+pub mod network {
+	/// Maximum number of known messages hashes to keep for a peer.
+	pub const MAX_KNOWN_MESSAGES: usize = 4096;
+
+	/// Maximum allowed size for a DKG Signed Message notification.
+	pub const MAX_MESSAGE_SIZE: u64 = 16 * 1024 * 1024;
+
+	/// Maximum number of duplicate messages that a single peer can send us.
+	///
+	/// This is to prevent a malicious peer from spamming us with messages.
+	pub const MAX_DUPLICATED_MESSAGES_PER_PEER: usize = 8;
+}
+
+// ============= Keygen Manager ======================= //
+
+pub mod keygen_manager {
+	/// only 1 task at a time may run for keygen
+	pub const MAX_RUNNING_TASKS: usize = 1;
+	/// There should never be any job enqueueing for keygen
+	pub const MAX_ENQUEUED_TASKS: usize = 0;
+}

dkg-gadget/src/gossip_engine/network.rs

@@ -40,6 +40,7 @@
 //! engine, and it is verified then it will be added to the Engine's internal stream of DKG
 //! messages, later the DKG Gadget will read this stream and process the DKG message.
 
+pub use crate::constants::network::*;
 use crate::{debug_logger::DebugLogger, metrics::Metrics, worker::HasLatestHeader, DKGKeystore};
 use codec::{Decode, Encode};
 use dkg_primitives::types::{DKGError, SignedDKGMessage};
@@ -152,17 +153,6 @@ impl NetworkGossipEngineBuilder {
 	}
 }
 
-/// Maximum number of known messages hashes to keep for a peer.
-const MAX_KNOWN_MESSAGES: usize = 4096;
-
-/// Maximum allowed size for a DKG Signed Message notification.
-const MAX_MESSAGE_SIZE: u64 = 16 * 1024 * 1024;
-
-/// Maximum number of duplicate messages that a single peer can send us.
-///
-/// This is to prevent a malicious peer from spamming us with messages.
-const MAX_DUPLICATED_MESSAGES_PER_PEER: usize = 8;
-
 #[allow(unused)]
 mod rep {
 	use sc_peerset::ReputationChange as Rep;

dkg-gadget/src/keygen_manager/mod.rs

@@ -2,6 +2,7 @@
 
 use crate::{
 	async_protocols::{remote::AsyncProtocolRemote, KeygenPartyId, KeygenRound},
+	constants::keygen_manager::*,
 	dkg_modules::KeygenProtocolSetupParameters,
 	gossip_engine::GossipEngineIface,
 	signing_manager::work_manager::{JobMetadata, PollMethod, WorkManager},
@@ -68,11 +69,6 @@ pub enum KeygenState {
 	Failed { session_id: u64 },
 }
 
-/// only 1 task at a time may run for keygen
-const MAX_RUNNING_TASKS: usize = 1;
-/// There should never be any job enqueueing for keygen
-const MAX_ENQUEUED_TASKS: usize = 0;
-
 impl<B, BE, C, GE> KeygenManager<B, BE, C, GE>
 where
 	B: Block,

dkg-gadget/src/lib.rs

@@ -43,17 +43,16 @@ pub mod worker;
 
 pub mod async_protocols;
 pub use dkg_logging::debug_logger;
+pub mod constants;
 pub mod dkg_modules;
 pub mod gossip_messages;
 pub mod storage;
 
+pub use constants::{DKG_KEYGEN_PROTOCOL_NAME, DKG_SIGNING_PROTOCOL_NAME};
 pub use debug_logger::RoundsEventType;
 use gossip_engine::NetworkGossipEngineBuilder;
 pub use keystore::DKGKeystore;
 
-pub const DKG_KEYGEN_PROTOCOL_NAME: &str = "/webb-tools/dkg/keygen/1";
-pub const DKG_SIGNING_PROTOCOL_NAME: &str = "/webb-tools/dkg/signing/1";
-
 /// Returns the configuration value to put in
 /// [`sc_network::config::NetworkConfiguration::extra_sets`].
 pub fn dkg_peers_set_config(

dkg-gadget/src/signing_manager/mod.rs

@@ -8,6 +8,7 @@ use dkg_primitives::{
 use self::work_manager::WorkManager;
 use crate::{
 	async_protocols::KeygenPartyId,
+	constants::signing_manager::*,
 	dkg_modules::SigningProtocolSetupParameters,
 	gossip_engine::GossipEngineIface,
 	metric_inc,
@@ -21,7 +22,6 @@ use dkg_runtime_primitives::crypto::Public;
 use sp_api::HeaderT;
 use std::sync::atomic::{AtomicBool, Ordering};
 use webb_proposals::TypedChainId;
-
 /// For balancing the amount of work done by each node
 pub mod work_manager;
 
@@ -51,13 +51,6 @@ impl<B: Block, BE, C, GE> Clone for SigningManager<B, BE, C, GE> {
 	}
 }
 
-// the maximum number of tasks that the work manager tries to assign
-const MAX_RUNNING_TASKS: usize = 4;
-const MAX_ENQUEUED_TASKS: usize = 20;
-// How often to poll the jobs to check completion status
-const JOB_POLL_INTERVAL_IN_MILLISECONDS: u64 = 500;
-pub const MAX_POTENTIAL_SIGNING_SETS_PER_PROPOSAL: u8 = 2;
-
 impl<B, BE, C, GE> SigningManager<B, BE, C, GE>
 where
 	B: Block,

dkg-gadget/src/worker.rs

@@ -53,6 +53,7 @@ use dkg_runtime_primitives::{
 	GENESIS_AUTHORITY_SET_ID,
 };
 
+pub use crate::constants::worker::*;
 use crate::{
 	async_protocols::{remote::AsyncProtocolRemote, AsyncProtocolParameters},
 	dkg_modules::DKGModules,
@@ -70,17 +71,6 @@ use crate::{
 	Client,
 };
 
-pub const ENGINE_ID: sp_runtime::ConsensusEngineId = *b"WDKG";
-
-pub const STORAGE_SET_RETRY_NUM: usize = 5;
-
-pub const MAX_SUBMISSION_DELAY: u32 = 3;
-
-pub const MAX_KEYGEN_RETRIES: usize = 5;
-
-/// How many blocks to keep the proposal hash in out local cache.
-pub const PROPOSAL_HASH_LIFETIME: u32 = 10;
-
 pub type Shared<T> = Arc<RwLock<T>>;
 
 pub struct WorkerParams<B, BE, C, GE>

dkg-test-orchestrator/src/main.rs

@@ -93,7 +93,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 	);
 
 	let max_signing_sets_per_proposal =
-		dkg_gadget::signing_manager::MAX_POTENTIAL_SIGNING_SETS_PER_PROPOSAL;
+		dkg_gadget::constants::signing_manager::MAX_POTENTIAL_SIGNING_SETS_PER_PROPOSAL;
 
 	// first, spawn the orchestrator/mock-blockchain
 	let orchestrator_task = MockBlockchain::new(

pallets/dkg-proposal-handler/src/functions.rs

@@ -1,6 +1,6 @@
 use super::*;
 use dkg_runtime_primitives::handlers::decode_proposals::ProposalIdentifier;
-use sp_runtime::traits::{CheckedAdd, One};
+use sp_runtime::traits::{CheckedAdd, CheckedSub, One};
 
 impl<T: Config> Pallet<T> {
 	// *** API methods ***
@@ -272,4 +272,98 @@ impl<T: Config> Pallet<T> {
 	pub fn signed_proposals_len() -> usize {
 		SignedProposals::<T>::iter_keys().count()
 	}
+
+	pub fn on_idle_create_proposal_batches(mut remaining_weight: Weight) -> Weight {
+		// fetch all unsigned proposals
+		let unsigned_proposals: Vec<_> = UnsignedProposals::<T>::iter().collect();
+		let unsigned_proposals_len = unsigned_proposals.len() as u64;
+		remaining_weight =
+			remaining_weight.saturating_sub(T::DbWeight::get().reads(unsigned_proposals_len));
+
+		for (typed_chain_id, unsigned_proposals) in unsigned_proposals {
+			remaining_weight =
+				remaining_weight.saturating_sub(T::DbWeight::get().reads_writes(1, 3));
+
+			if remaining_weight.is_zero() {
+				break
+			}
+
+			let batch_id_res = Self::generate_next_batch_id();
+
+			if batch_id_res.is_err() {
+				log::debug!(
+					target: "runtime::dkg_proposal_handler",
+					"on_idle: Cannot generate next batch ID: {:?}",
+					batch_id_res,
+				);
+				return remaining_weight
+			}
+
+			let batch_id = batch_id_res.expect("checked above");
+
+			// create new proposal batch
+			let proposal_batch = StoredUnsignedProposalBatchOf::<T> {
+				batch_id,
+				proposals: unsigned_proposals,
+				timestamp: <frame_system::Pallet<T>>::block_number(),
+			};
+			// push the batch to unsigned proposal queue
+			UnsignedProposalQueue::<T>::insert(typed_chain_id, batch_id, proposal_batch);
+
+			// remove the batch from the unsigned proposal list
+			UnsignedProposals::<T>::remove(typed_chain_id);
+		}
+
+		remaining_weight
+	}
+
+	pub fn on_idle_remove_expired_batches(
+		now: T::BlockNumber,
+		mut remaining_weight: Weight,
+	) -> Weight {
+		use dkg_runtime_primitives::DKGPayloadKey::RefreshProposal;
+
+		// early return if we dont have enough weight to perform a read
+		if remaining_weight.is_zero() {
+			return remaining_weight
+		}
+
+		// fetch all unsigned proposals
+		let unsigned_proposals: Vec<_> = UnsignedProposalQueue::<T>::iter().collect();
+		let unsigned_proposals_len = unsigned_proposals.len() as u64;
+		remaining_weight =
+			remaining_weight.saturating_sub(T::DbWeight::get().reads(unsigned_proposals_len));
+
+		// filter out proposals to delete
+		let unsigned_proposal_past_expiry = unsigned_proposals.into_iter().filter(
+			|(_, _, StoredUnsignedProposalBatchOf::<T> { proposals, timestamp, .. })| {
+				let key = proposals.first().expect("cannot be empty").key;
+
+				// Skip expiration for keygen related proposals
+				if let RefreshProposal(_) = key {
+					return false
+				}
+
+				let time_passed = now.checked_sub(timestamp).unwrap_or_default();
+				time_passed > T::UnsignedProposalExpiry::get()
+			},
+		);
+
+		// remove unsigned proposal until we run out of weight
+		for expired_proposal in unsigned_proposal_past_expiry {
+			remaining_weight =
+				remaining_weight.saturating_sub(T::DbWeight::get().writes(One::one()));
+
+			if remaining_weight.is_zero() {
+				break
+			}
+			Self::deposit_event(Event::<T>::ProposalBatchExpired {
+				target_chain: expired_proposal.0,
+				batch_id: expired_proposal.1,
+			});
+			UnsignedProposalQueue::<T>::remove(expired_proposal.0, expired_proposal.1);
+		}
+
+		remaining_weight
+	}
 }