use super::{
NnsAuthenticatedRegistryReplaySession, NnsRegistryReplayError, NnsRegistryReplaySession,
NnsRegistryReplaySessionLimits,
};
use crate::nns::{
LiveNnsSource,
registry::{
NnsCertifiedRegistryDeltaBatchRequest, NnsCertifiedRegistryDeltaSource,
NnsRegistryHostError, fetch_nns_certified_registry_delta_batch_with_source_async,
nns_certified_registry_delta_limits,
},
};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct NnsCertifiedRegistryBootstrapRequest {
pub network: String,
pub source_endpoint: String,
pub now_unix_secs: u64,
pub limits: NnsRegistryReplaySessionLimits,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum NnsCertifiedRegistryBootstrapProbeStatus {
Complete,
CapacityReached {
field: &'static str,
maximum: u64,
required: u64,
},
}
#[derive(Debug, Eq, PartialEq)]
pub struct NnsCertifiedRegistryBootstrapProbeOutcome {
pub status: NnsCertifiedRegistryBootstrapProbeStatus,
pub session: NnsRegistryReplaySession,
}
impl NnsCertifiedRegistryBootstrapRequest {
#[must_use]
pub fn new(
network: impl Into<String>,
source_endpoint: impl Into<String>,
now_unix_secs: u64,
limits: NnsRegistryReplaySessionLimits,
) -> Self {
Self {
network: network.into(),
source_endpoint: source_endpoint.into(),
now_unix_secs,
limits,
}
}
}
pub async fn bootstrap_nns_certified_registry_async(
request: &NnsCertifiedRegistryBootstrapRequest,
) -> Result<NnsAuthenticatedRegistryReplaySession, NnsRegistryReplayError> {
let session =
bootstrap_nns_certified_registry_with_source_async(request, &LiveNnsSource).await?;
NnsAuthenticatedRegistryReplaySession::from_verified_complete(session)
}
pub async fn bootstrap_nns_certified_registry_with_source_async(
request: &NnsCertifiedRegistryBootstrapRequest,
source: &dyn NnsCertifiedRegistryDeltaSource,
) -> Result<NnsRegistryReplaySession, NnsRegistryReplayError> {
let outcome = probe_nns_certified_registry_with_source_async(request, source).await?;
match outcome.status {
NnsCertifiedRegistryBootstrapProbeStatus::Complete => Ok(outcome.session),
NnsCertifiedRegistryBootstrapProbeStatus::CapacityReached {
field,
maximum,
required,
} => Err(NnsRegistryReplayError::SessionLimitExceeded {
field,
maximum,
actual: required,
}),
}
}
pub async fn probe_nns_certified_registry_async(
request: &NnsCertifiedRegistryBootstrapRequest,
) -> Result<NnsCertifiedRegistryBootstrapProbeOutcome, NnsRegistryReplayError> {
probe_nns_certified_registry_with_source_async(request, &LiveNnsSource).await
}
pub async fn probe_nns_certified_registry_with_source_async(
request: &NnsCertifiedRegistryBootstrapRequest,
source: &dyn NnsCertifiedRegistryDeltaSource,
) -> Result<NnsCertifiedRegistryBootstrapProbeOutcome, NnsRegistryReplayError> {
crate::network::enforce_mainnet_network_with(&request.network, |network| {
NnsRegistryReplayError::InvalidBatch(NnsRegistryHostError::UnsupportedNetwork { network })
})?;
let (maximum_batch_query_calls, maximum_batch_response_bytes) = batch_reservation()?;
let mut session = NnsRegistryReplaySession::new(request.limits);
loop {
if let Err(error) = session.ensure_next_source_call_capacity(
maximum_batch_query_calls,
maximum_batch_response_bytes,
) {
return match error {
NnsRegistryReplayError::SessionLimitExceeded {
field,
maximum,
actual,
} => Ok(NnsCertifiedRegistryBootstrapProbeOutcome {
status: NnsCertifiedRegistryBootstrapProbeStatus::CapacityReached {
field,
maximum,
required: actual,
},
session,
}),
error => Err(error),
};
}
let batch_request = NnsCertifiedRegistryDeltaBatchRequest::new(
&request.network,
&request.source_endpoint,
session.state().through_version(),
request.now_unix_secs,
);
let report =
fetch_nns_certified_registry_delta_batch_with_source_async(&batch_request, source)
.await?;
session.apply_batch(&batch_request, &report)?;
if session.is_complete() {
return Ok(NnsCertifiedRegistryBootstrapProbeOutcome {
status: NnsCertifiedRegistryBootstrapProbeStatus::Complete,
session,
});
}
}
}
fn batch_reservation() -> Result<(u64, u64), NnsRegistryReplayError> {
let limits = nns_certified_registry_delta_limits();
let chunk_calls = u64::try_from(limits.max_chunk_references)
.map_err(|_| NnsRegistryReplayError::Accounting)?;
let maximum_batch_query_calls = chunk_calls
.checked_add(1)
.ok_or(NnsRegistryReplayError::Accounting)?;
let maximum_batch_response_bytes = limits
.max_response_body_bytes
.checked_add(limits.max_chunk_response_bytes)
.and_then(|bytes| u64::try_from(bytes).ok())
.ok_or(NnsRegistryReplayError::Accounting)?;
Ok((maximum_batch_query_calls, maximum_batch_response_bytes))
}