Remove shard partial offline state

pallets/shards/src/lib.rs

@@ -14,7 +14,7 @@ pub mod pallet {
 	use frame_support::pallet_prelude::{ValueQuery, *};
 	use frame_system::pallet_prelude::*;
 	use schnorr_evm::VerifyingKey;
-	use sp_runtime::Saturating;
+	use sp_runtime::{traits::Zero, Saturating};
 	use sp_std::vec;
 	use sp_std::vec::Vec;
 	use time_primitives::{
@@ -107,6 +107,10 @@ pub mod pallet {
 		OptionQuery,
 	>;
 
+	#[pallet::storage]
+	pub type ShardMembersOnline<T: Config> =
+		StorageMap<_, Blake2_128Concat, ShardId, u16, ValueQuery>;
+
 	#[pallet::event]
 	#[pallet::generate_deposit(pub(super) fn deposit_event)]
 	pub enum Event<T: Config> {
@@ -297,14 +301,49 @@ pub mod pallet {
 		pub fn get_shard_commitment(shard_id: ShardId) -> Option<Vec<TssPublicKey>> {
 			ShardCommitment::<T>::get(shard_id)
 		}
+		pub fn do_member_online(shard_id: ShardId, status: ShardStatus) {
+			let new_status = match status {
+				ShardStatus::Created | ShardStatus::Committed => ShardStatus::Offline,
+				_ => status,
+			};
+			ShardState::<T>::insert(shard_id, new_status);
+		}
+		pub fn do_member_offline(shard_id: ShardId, status: ShardStatus, max: u16) {
+			let new_status = match status {
+				// if a member goes offline before the group key is submitted,
+				// then the shard will never go online
+				ShardStatus::Created | ShardStatus::Committed => ShardStatus::Offline,
+				ShardStatus::Online => {
+					if max.is_zero() {
+						ShardStatus::Offline
+					} else {
+						let shard_members_offline = (Self::get_shard_members(shard_id).len()
+							as u16)
+							.saturating_sub(ShardMembersOnline::<T>::get(shard_id));
+						if shard_members_offline > max {
+							ShardStatus::Offline
+						} else {
+							ShardStatus::Online
+						}
+					}
+				},
+				_ => status,
+			};
+			if matches!(new_status, ShardStatus::Offline) && !matches!(status, ShardStatus::Offline)
+			{
+				Self::remove_shard_offline(shard_id);
+			} else if !matches!(new_status, ShardStatus::Offline) {
+				ShardState::<T>::insert(shard_id, new_status);
+			}
+		}
 	}
 
 	impl<T: Config> MemberEvents for Pallet<T> {
 		fn member_online(id: &AccountId, _network: NetworkId) {
 			let Some(shard_id) = MemberShard::<T>::get(id) else { return };
 			let Some(old_status) = ShardState::<T>::get(shard_id) else { return };
-			let new_status = old_status.online_member();
-			ShardState::<T>::insert(shard_id, new_status);
+			ShardMembersOnline::<T>::mutate(shard_id, |x| *x = x.saturating_plus_one());
+			Self::do_member_online(shard_id, old_status);
 		}
 
 		fn member_offline(id: &AccountId, _: NetworkId) -> Weight {
@@ -317,29 +356,20 @@ pub mod pallet {
 			let Some(shard_threshold) = ShardThreshold::<T>::get(shard_id) else {
 				return T::DbWeight::get().reads(3);
 			};
+			ShardMembersOnline::<T>::mutate(shard_id, |x| *x = x.saturating_less_one());
 			let total_members = Self::get_shard_members(shard_id).len();
 			let max_members_offline = total_members.saturating_sub(shard_threshold.into());
 			let Ok(max_members_offline) = max_members_offline.try_into() else {
 				return T::DbWeight::get().reads(4);
 			};
-			let new_status = old_status.offline_member(max_members_offline);
-			if matches!(new_status, ShardStatus::Offline)
-				&& !matches!(old_status, ShardStatus::Offline)
-			{
-				Self::remove_shard_offline(shard_id);
-			} else if !matches!(new_status, ShardStatus::Offline) {
-				ShardState::<T>::insert(shard_id, new_status);
-			}
+			Self::do_member_offline(shard_id, old_status, max_members_offline);
 			T::WeightInfo::member_offline()
 		}
 	}
 
 	impl<T: Config> ShardsInterface for Pallet<T> {
 		fn is_shard_online(shard_id: ShardId) -> bool {
-			matches!(
-				ShardState::<T>::get(shard_id),
-				Some(ShardStatus::Online | ShardStatus::PartialOffline(_))
-			)
+			matches!(ShardState::<T>::get(shard_id), Some(ShardStatus::Online))
 		}
 
 		fn is_shard_member(member: &AccountId) -> bool {

pallets/shards/src/tests.rs

@@ -122,16 +122,6 @@ fn dkg_times_out() {
 	});
 }
 
-#[test]
-fn member_offline_sets_online_shard_partially_offline() {
-	let shard = shard();
-	new_test_ext().execute_with(|| {
-		create_shard(0, &shard, 1);
-		Shards::member_offline(&shard[0].account_id, ETHEREUM);
-		assert_eq!(ShardState::<Test>::get(0), Some(ShardStatus::PartialOffline(1)));
-	});
-}
-
 #[test]
 fn member_offline_above_threshold_sets_online_shard_offline() {
 	let shard = shard();
@@ -143,12 +133,12 @@ fn member_offline_above_threshold_sets_online_shard_offline() {
 }
 
 #[test]
-fn member_online_sets_partially_offline_shard_back_online() {
+fn member_offline_below_thresholds_keeps_online_shard_online() {
 	let shard = shard();
 	new_test_ext().execute_with(|| {
 		create_shard(0, &shard, 1);
 		Shards::member_offline(&shard[0].account_id, ETHEREUM);
-		assert_eq!(ShardState::<Test>::get(0), Some(ShardStatus::PartialOffline(1)));
+		assert_eq!(ShardState::<Test>::get(0), Some(ShardStatus::Online));
 		Shards::member_online(&shard[0].account_id, ETHEREUM);
 		assert_eq!(ShardState::<Test>::get(0), Some(ShardStatus::Online));
 	});

primitives/src/shard.rs

@@ -8,7 +8,6 @@ use scale_info::prelude::string::String;
 use scale_info::TypeInfo;
 #[cfg(feature = "std")]
 use serde::{Deserialize, Serialize};
-use sp_runtime::traits::{Saturating, Zero};
 use sp_std::vec::Vec;
 pub type TssPublicKey = [u8; 33];
 pub type TssSignature = [u8; 64];
@@ -55,7 +54,6 @@ pub enum ShardStatus {
 	Created,
 	Committed,
 	Online,
-	PartialOffline(u16),
 	Offline,
 }
 
@@ -65,47 +63,6 @@ impl Default for ShardStatus {
 	}
 }
 
-impl ShardStatus {
-	pub fn online_member(&self) -> Self {
-		match self {
-			ShardStatus::PartialOffline(count) => {
-				let new_count = count.saturating_less_one();
-				if new_count.is_zero() {
-					ShardStatus::Online
-				} else {
-					ShardStatus::PartialOffline(new_count)
-				}
-			},
-			ShardStatus::Created | ShardStatus::Committed => ShardStatus::Offline,
-			_ => *self,
-		}
-	}
-
-	pub fn offline_member(&self, max: u16) -> Self {
-		match self {
-			ShardStatus::PartialOffline(count) => {
-				let new_count = count.saturating_plus_one();
-				if new_count > max {
-					ShardStatus::Offline
-				} else {
-					ShardStatus::PartialOffline(new_count)
-				}
-			},
-			// if a member goes offline before the group key is submitted,
-			// then the shard will never go online
-			ShardStatus::Created | ShardStatus::Committed => ShardStatus::Offline,
-			ShardStatus::Online => {
-				if max.is_zero() {
-					ShardStatus::Offline
-				} else {
-					ShardStatus::PartialOffline(1)
-				}
-			},
-			_ => *self,
-		}
-	}
-}
-
 #[cfg(feature = "std")]
 pub struct TssSigningRequest {
 	pub request_id: TssId,