Skip to content

Commit

Permalink
Merge pull request #661 from input-output-hk/djo/simplify_epoch_offset
Browse files Browse the repository at this point in the history 
Simplify epoch offset computations
  • Loading branch information
@Alenar
Alenar committed Dec 19, 2022
2 parents f3fd5aa + cbfadfd commit 30c6517
Show file tree
Hide file tree
Showing 10 changed files with 138 additions and 261 deletions.
2 changes: 1 addition & 1 deletion mithril-aggregator/src/command_args.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ async fn do_first_launch_initialization_if_needed(
Epoch(0) => (Epoch(0), Epoch(1)),
epoch => (
epoch.offset_to_signer_retrieval_epoch()?,
epoch.offset_to_next_signer_retrieval_epoch()?,
epoch.offset_to_next_signer_retrieval_epoch(),
),
};

Expand Down
4 changes: 1 addition & 3 deletions mithril-aggregator/src/dependency.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,7 @@ impl DependencyManager {
let work_epoch = current_epoch
.offset_to_signer_retrieval_epoch()
.expect("epoch.offset_by SIGNER_EPOCH_RETRIEVAL_OFFSET should not fail");
let epoch_to_sign = current_epoch
.offset_to_next_signer_retrieval_epoch()
.expect("epoch.offset_by NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET should not fail");
let epoch_to_sign = current_epoch.offset_to_next_signer_retrieval_epoch();

(work_epoch, epoch_to_sign)
}
Expand Down
171 changes: 82 additions & 89 deletions mithril-aggregator/src/multi_signer.rs
View file
Original file line number Diff line number Diff line change
@@ -1,27 +1,23 @@
use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use chrono::prelude::*;
use hex::ToHex;
use slog_scope::{debug, trace, warn};
use std::{collections::HashMap, sync::Arc};
use thiserror::Error;

use mithril_common::crypto_helper::{
key_decode_hex, key_encode_hex, ProtocolAggregateVerificationKey, ProtocolAggregationError,
ProtocolClerk, ProtocolKeyRegistration, ProtocolMultiSignature, ProtocolParameters,
ProtocolPartyId, ProtocolRegistrationError, ProtocolSignerVerificationKey,
ProtocolSingleSignature, ProtocolStakeDistribution,
};
use mithril_common::entities::{self, SignerWithStake};
use mithril_common::store::{StakeStore, StakeStorer, StoreError};
use mithril_common::{
NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET, SIGNER_EPOCH_RECORDING_OFFSET,
SIGNER_EPOCH_RETRIEVAL_OFFSET,
crypto_helper::{
key_decode_hex, key_encode_hex, ProtocolAggregateVerificationKey, ProtocolAggregationError,
ProtocolClerk, ProtocolKeyRegistration, ProtocolMultiSignature, ProtocolParameters,
ProtocolPartyId, ProtocolRegistrationError, ProtocolSignerVerificationKey,
ProtocolSingleSignature, ProtocolStakeDistribution,
},
entities::{self, Epoch, SignerWithStake},
store::{StakeStore, StakeStorer, StoreError},
};

use crate::store::{SingleSignatureStorer, VerificationKeyStorer};
use crate::{
store::{SingleSignatureStorer, VerificationKeyStorer},
ProtocolParametersStore, ProtocolParametersStorer, SingleSignatureStore, VerificationKeyStore,
};

Expand Down Expand Up @@ -191,13 +187,10 @@ pub trait MultiSigner: Sync + Send {
}

/// Get signers with stake
async fn get_signers_with_stake(&self)
-> Result<Vec<entities::SignerWithStake>, ProtocolError>;
async fn get_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError>;

/// Get signers for the next epoch with their stake
async fn get_next_signers_with_stake(
&self,
) -> Result<Vec<entities::SignerWithStake>, ProtocolError>;
async fn get_next_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError>;

/// Registers a single signature
async fn register_single_signature(
Expand Down Expand Up @@ -322,18 +315,12 @@ impl MultiSignerImpl {
}
}

/// Get stake distribution with epoch offset
async fn get_stake_distribution_with_epoch_offset(
/// Get the [stake distribution][ProtocolStakeDistribution] for the given `epoch`
async fn get_stake_distribution_at_epoch(
&self,
epoch_offset: i64,
epoch: Epoch,
) -> Result<ProtocolStakeDistribution, ProtocolError> {
debug!("Get stake distribution with epoch offset"; "epoch_offset"=>epoch_offset);
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(epoch_offset)?;
debug!("Get stake distribution at epoch"; "epoch"=> #?epoch);

let stakes = self
.stake_store
Expand All @@ -343,18 +330,12 @@ impl MultiSignerImpl {
Ok(stakes.into_iter().collect::<ProtocolStakeDistribution>())
}

/// Get protocol parameters with epoch offset
async fn get_protocol_parameters_with_epoch_offset(
/// Get the [protocol parameters][ProtocolParameters] for the given `epoch`
async fn get_protocol_parameters_at_epoch(
&self,
epoch_offset: i64,
epoch: Epoch,
) -> Result<Option<ProtocolParameters>, ProtocolError> {
debug!("Get protocol parameters with epoch offset"; "epoch_offset"=>epoch_offset);
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(epoch_offset)?;
debug!("Get protocol parameters at epoch"; "epoch"=> #?epoch);

match self
.protocol_parameters_store
Expand Down Expand Up @@ -431,8 +412,13 @@ impl MultiSigner for MultiSignerImpl {
/// Get protocol parameters
async fn get_protocol_parameters(&self) -> Result<Option<ProtocolParameters>, ProtocolError> {
debug!("Get protocol parameters");
self.get_protocol_parameters_with_epoch_offset(SIGNER_EPOCH_RETRIEVAL_OFFSET)
.await
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_to_signer_retrieval_epoch()?;
self.get_protocol_parameters_at_epoch(epoch).await
}

/// Update protocol parameters
Expand All @@ -446,7 +432,7 @@ impl MultiSigner for MultiSignerImpl {
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(SIGNER_EPOCH_RECORDING_OFFSET)?;
.offset_to_recording_epoch();

self.protocol_parameters_store
.save_protocol_parameters(epoch, protocol_parameters.to_owned().into())
Expand All @@ -459,24 +445,39 @@ impl MultiSigner for MultiSignerImpl {
&self,
) -> Result<Option<ProtocolParameters>, ProtocolError> {
debug!("Get next protocol parameters");
self.get_protocol_parameters_with_epoch_offset(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)
.await
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_to_next_signer_retrieval_epoch();
self.get_protocol_parameters_at_epoch(epoch).await
}

/// Get stake distribution
async fn get_stake_distribution(&self) -> Result<ProtocolStakeDistribution, ProtocolError> {
debug!("Get stake distribution");
self.get_stake_distribution_with_epoch_offset(SIGNER_EPOCH_RETRIEVAL_OFFSET)
.await
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_to_signer_retrieval_epoch()?;
self.get_stake_distribution_at_epoch(epoch).await
}

/// Get next stake distribution
async fn get_next_stake_distribution(
&self,
) -> Result<ProtocolStakeDistribution, ProtocolError> {
debug!("Get next stake distribution");
self.get_stake_distribution_with_epoch_offset(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)
.await
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_to_next_signer_retrieval_epoch();
self.get_stake_distribution_at_epoch(epoch).await
}

/// Update stake distribution
Expand All @@ -490,7 +491,7 @@ impl MultiSigner for MultiSignerImpl {
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(SIGNER_EPOCH_RECORDING_OFFSET)?;
.offset_to_recording_epoch();
let stakes = HashMap::from_iter(stakes.iter().cloned());
self.stake_store.save_stakes(epoch, stakes).await?;

Expand Down Expand Up @@ -523,7 +524,7 @@ impl MultiSigner for MultiSignerImpl {
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
.offset_to_signer_retrieval_epoch()?;
let signers = self
.verification_key_store
.get_verification_keys(epoch)
Expand All @@ -537,16 +538,14 @@ impl MultiSigner for MultiSignerImpl {
}
}

async fn get_signers_with_stake(
&self,
) -> Result<Vec<entities::SignerWithStake>, ProtocolError> {
async fn get_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError> {
debug!("Get signers with stake");
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
.offset_to_signer_retrieval_epoch()?;
let signers = self
.verification_key_store
.get_verification_keys(epoch)
Expand All @@ -558,7 +557,7 @@ impl MultiSigner for MultiSignerImpl {
.iter()
.filter_map(|(party_id, stake)| {
signers.get(party_id).map(|signer| {
entities::SignerWithStake::new(
SignerWithStake::new(
party_id.to_owned(),
signer.verification_key.to_owned(),
signer.verification_key_signature.to_owned(),
Expand All @@ -571,16 +570,14 @@ impl MultiSigner for MultiSignerImpl {
.collect())
}

async fn get_next_signers_with_stake(
&self,
) -> Result<Vec<entities::SignerWithStake>, ProtocolError> {
async fn get_next_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError> {
debug!("Get next signers with stake");
let epoch = self
.current_beacon
.as_ref()
.ok_or_else(ProtocolError::UnavailableBeacon)?
.epoch
.offset_by(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
.offset_to_next_signer_retrieval_epoch();
let signers = self
.verification_key_store
.get_verification_keys(epoch)
Expand All @@ -592,7 +589,7 @@ impl MultiSigner for MultiSignerImpl {
.iter()
.filter_map(|(party_id, stake)| {
signers.get(party_id).map(|signer| {
entities::SignerWithStake::new(
SignerWithStake::new(
party_id.to_owned(),
signer.verification_key.to_owned(),
signer.verification_key_signature.to_owned(),
Expand Down Expand Up @@ -724,31 +721,30 @@ impl MultiSigner for MultiSignerImpl {
#[cfg(test)]
mod tests {
use super::*;
use crate::store::{SingleSignatureStore, VerificationKeyStore};
use crate::ProtocolParametersStore;

use mithril_common::crypto_helper::tests_setup::*;
use mithril_common::fake_data;
use mithril_common::store::adapter::MemoryAdapter;
use mithril_common::store::StakeStore;

use std::collections::HashMap;
use std::sync::Arc;
use crate::{
store::{SingleSignatureStore, VerificationKeyStore},
ProtocolParametersStore,
};
use mithril_common::{
crypto_helper::tests_setup::*,
fake_data,
store::{adapter::MemoryAdapter, StakeStore},
};
use std::{collections::HashMap, sync::Arc};

async fn setup_multi_signer() -> MultiSignerImpl {
let beacon = fake_data::beacon();
let verification_key_store = VerificationKeyStore::new(Box::new(
MemoryAdapter::<entities::Epoch, HashMap<entities::PartyId, entities::Signer>>::new(
None,
)
.unwrap(),
), None);
let verification_key_store = VerificationKeyStore::new(
Box::new(
MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::Signer>>::new(None)
.unwrap(),
),
None,
);
let stake_store = StakeStore::new(
Box::new(
MemoryAdapter::<entities::Epoch, HashMap<entities::PartyId, entities::Stake>>::new(
None,
)
.unwrap(),
MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::Stake>>::new(None)
.unwrap(),
),
None,
);
Expand All @@ -764,16 +760,13 @@ mod tests {
);
let protocol_parameters_store = ProtocolParametersStore::new(
Box::new(
MemoryAdapter::<entities::Epoch, entities::ProtocolParameters>::new(Some(vec![
MemoryAdapter::<Epoch, entities::ProtocolParameters>::new(Some(vec![
(
beacon.epoch.offset_to_signer_retrieval_epoch().unwrap(),
fake_data::protocol_parameters(),
),
(
beacon
.epoch
.offset_to_next_signer_retrieval_epoch()
.unwrap(),
beacon.epoch.offset_to_next_signer_retrieval_epoch(),
fake_data::protocol_parameters(),
),
]))
Expand Down Expand Up @@ -851,7 +844,7 @@ mod tests {

offset_epoch(
&mut multi_signer,
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
)
.await;

Expand Down Expand Up @@ -887,7 +880,7 @@ mod tests {

offset_epoch(
&mut multi_signer,
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
)
.await;

Expand Down Expand Up @@ -943,7 +936,7 @@ mod tests {
for (signer_with_stake, _, _) in &signers {
verification_key_store
.save_verification_key(
start_epoch.offset_to_recording_epoch().unwrap(),
start_epoch.offset_to_recording_epoch(),
signer_with_stake.to_owned().into(),
)
.await
Expand All @@ -952,7 +945,7 @@ mod tests {

offset_epoch(
&mut multi_signer,
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
)
.await;
// We have to update the current message AFTER we reached the epoch for
Expand Down
Loading

0 comments on commit 30c6517

Please sign in to comment.