use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourceObjectRef {
pub object_id: String,
pub merkle_root_hex: String,
}
impl MultiSourceObjectRef {
#[must_use]
pub fn new(object_id: impl Into<String>, merkle_root_hex: impl Into<String>) -> Self {
Self {
object_id: object_id.into(),
merkle_root_hex: merkle_root_hex.into(),
}
}
fn validate(&self) -> Result<(), MultiSourcePlanError> {
if self.object_id.trim().is_empty() {
return Err(MultiSourcePlanError::EmptyObjectId);
}
if self.merkle_root_hex.trim().is_empty() {
return Err(MultiSourcePlanError::EmptyMerkleRoot);
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub enum MultiSourceAuth {
SymbolAuth { key_id: String },
UnauthenticatedLab,
}
impl MultiSourceAuth {
#[must_use]
pub fn mode_id(&self) -> &'static str {
match self {
Self::SymbolAuth { .. } => "symbol_auth",
Self::UnauthenticatedLab => "unauthenticated_lab",
}
}
fn is_symbol_auth(&self) -> bool {
matches!(self, Self::SymbolAuth { .. })
}
fn validate(&self) -> Result<(), MultiSourcePlanError> {
match self {
Self::SymbolAuth { key_id } if key_id.trim().is_empty() => {
Err(MultiSourcePlanError::EmptyAuthKeyId)
}
_ => Ok(()),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourcePeer {
pub peer_id: String,
pub endpoint: String,
pub priority: u32,
pub auth: MultiSourceAuth,
}
impl MultiSourcePeer {
#[must_use]
pub fn new(
peer_id: impl Into<String>,
endpoint: impl Into<String>,
priority: u32,
auth: MultiSourceAuth,
) -> Self {
Self {
peer_id: peer_id.into(),
endpoint: endpoint.into(),
priority,
auth,
}
}
fn validate(&self) -> Result<(), MultiSourcePlanError> {
if self.peer_id.trim().is_empty() {
return Err(MultiSourcePlanError::EmptyPeerId);
}
if self.endpoint.trim().is_empty() {
return Err(MultiSourcePlanError::EmptyEndpoint {
peer_id: self.peer_id.clone(),
});
}
self.auth.validate()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub enum MultiSourceSymbolBias {
SourceFirst,
RepairFirst,
Balanced,
}
impl MultiSourceSymbolBias {
#[must_use]
pub const fn bias_id(self) -> &'static str {
match self {
Self::SourceFirst => "source_first",
Self::RepairFirst => "repair_first",
Self::Balanced => "balanced",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourceSelectionConfig {
pub min_sources: usize,
pub max_sources: usize,
pub allow_unauthenticated_lab: bool,
}
impl MultiSourceSelectionConfig {
#[must_use]
pub const fn production_default() -> Self {
Self {
min_sources: 2,
max_sources: 4,
allow_unauthenticated_lab: false,
}
}
fn validate(self) -> Result<(), MultiSourcePlanError> {
if self.min_sources == 0 {
return Err(MultiSourcePlanError::ZeroMinSources);
}
if self.max_sources < self.min_sources {
return Err(MultiSourcePlanError::MaxSourcesBelowMin {
min_sources: self.min_sources,
max_sources: self.max_sources,
});
}
Ok(())
}
}
impl Default for MultiSourceSelectionConfig {
fn default() -> Self {
Self::production_default()
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourceSourcePlan {
pub peer: MultiSourcePeer,
pub symbol_bias: MultiSourceSymbolBias,
pub selection_order: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourceStopCommand {
pub peer_id: String,
pub endpoint: String,
pub object_id: String,
pub reason: MultiSourceStopReason,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub enum MultiSourceStopReason {
DecodedAndVerified,
Cancelled,
FailedClosed,
}
impl MultiSourceStopReason {
#[must_use]
pub const fn reason_id(self) -> &'static str {
match self {
Self::DecodedAndVerified => "decoded_and_verified",
Self::Cancelled => "cancelled",
Self::FailedClosed => "failed_closed",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct MultiSourceFetchPlan {
pub object: MultiSourceObjectRef,
pub selected_sources: Vec<MultiSourceSourcePlan>,
}
impl MultiSourceFetchPlan {
#[must_use]
pub fn source_count(&self) -> usize {
self.selected_sources.len()
}
#[must_use]
pub fn stop_commands(&self, reason: MultiSourceStopReason) -> Vec<MultiSourceStopCommand> {
self.selected_sources
.iter()
.map(|source| MultiSourceStopCommand {
peer_id: source.peer.peer_id.clone(),
endpoint: source.peer.endpoint.clone(),
object_id: self.object.object_id.clone(),
reason,
})
.collect()
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum MultiSourcePlanError {
#[error("multi-source object id must be non-empty")]
EmptyObjectId,
#[error("multi-source merkle root must be non-empty")]
EmptyMerkleRoot,
#[error("multi-source peer id must be non-empty")]
EmptyPeerId,
#[error("multi-source endpoint for peer {peer_id} must be non-empty")]
EmptyEndpoint {
peer_id: String,
},
#[error("multi-source symbol-auth key id must be non-empty")]
EmptyAuthKeyId,
#[error("multi-source duplicate peer id: {peer_id}")]
DuplicatePeerId {
peer_id: String,
},
#[error("multi-source duplicate endpoint: {endpoint}")]
DuplicateEndpoint {
endpoint: String,
},
#[error("multi-source peer {peer_id} is unauthenticated but lab mode is not allowed")]
UnauthenticatedPeerRejected {
peer_id: String,
},
#[error("multi-source min_sources must be greater than zero")]
ZeroMinSources,
#[error("multi-source max_sources {max_sources} is below min_sources {min_sources}")]
MaxSourcesBelowMin {
min_sources: usize,
max_sources: usize,
},
#[error("multi-source needs at least {required} eligible sources, got {available}")]
NotEnoughEligibleSources {
required: usize,
available: usize,
},
#[error("multi-source selected source count exceeds u32::MAX")]
TooManySelectedSources,
}
pub fn plan_multi_source_fetch(
object: MultiSourceObjectRef,
peers: impl IntoIterator<Item = MultiSourcePeer>,
config: MultiSourceSelectionConfig,
) -> Result<MultiSourceFetchPlan, MultiSourcePlanError> {
object.validate()?;
config.validate()?;
let mut peer_ids = BTreeSet::new();
let mut endpoints = BTreeSet::new();
let mut by_key = BTreeMap::new();
for peer in peers {
peer.validate()?;
if !peer_ids.insert(peer.peer_id.clone()) {
return Err(MultiSourcePlanError::DuplicatePeerId {
peer_id: peer.peer_id,
});
}
if !endpoints.insert(peer.endpoint.clone()) {
return Err(MultiSourcePlanError::DuplicateEndpoint {
endpoint: peer.endpoint,
});
}
if !config.allow_unauthenticated_lab && !peer.auth.is_symbol_auth() {
return Err(MultiSourcePlanError::UnauthenticatedPeerRejected {
peer_id: peer.peer_id,
});
}
by_key.insert(
(peer.priority, peer.peer_id.clone(), peer.endpoint.clone()),
peer,
);
}
if by_key.len() < config.min_sources {
return Err(MultiSourcePlanError::NotEnoughEligibleSources {
required: config.min_sources,
available: by_key.len(),
});
}
let mut selected_sources = Vec::new();
for (idx, peer) in by_key.into_values().take(config.max_sources).enumerate() {
let selection_order =
u32::try_from(idx).map_err(|_| MultiSourcePlanError::TooManySelectedSources)?;
selected_sources.push(MultiSourceSourcePlan {
peer,
symbol_bias: symbol_bias_for_order(idx),
selection_order,
});
}
Ok(MultiSourceFetchPlan {
object,
selected_sources,
})
}
fn symbol_bias_for_order(order: usize) -> MultiSourceSymbolBias {
match order {
0 => MultiSourceSymbolBias::SourceFirst,
1 => MultiSourceSymbolBias::RepairFirst,
_ => MultiSourceSymbolBias::Balanced,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn object_ref() -> MultiSourceObjectRef {
MultiSourceObjectRef::new("object-01", "abc123")
}
fn auth(key_id: &str) -> MultiSourceAuth {
MultiSourceAuth::SymbolAuth {
key_id: key_id.to_string(),
}
}
fn peer(id: &str, endpoint: &str, priority: u32) -> MultiSourcePeer {
MultiSourcePeer::new(id, endpoint, priority, auth("key-a"))
}
#[test]
fn production_default_requires_two_authenticated_sources() {
assert_eq!(
MultiSourceSelectionConfig::production_default(),
MultiSourceSelectionConfig {
min_sources: 2,
max_sources: 4,
allow_unauthenticated_lab: false,
}
);
}
#[test]
fn selection_is_priority_then_peer_then_endpoint_deterministic() {
let plan = plan_multi_source_fetch(
object_ref(),
[
peer("peer-c", "10.0.0.3:8472", 10),
peer("peer-b", "10.0.0.2:8472", 5),
peer("peer-a", "10.0.0.1:8472", 5),
],
MultiSourceSelectionConfig {
min_sources: 2,
max_sources: 2,
allow_unauthenticated_lab: false,
},
)
.unwrap();
assert_eq!(plan.source_count(), 2);
assert_eq!(plan.selected_sources[0].peer.peer_id, "peer-a");
assert_eq!(plan.selected_sources[1].peer.peer_id, "peer-b");
assert_eq!(
plan.selected_sources
.iter()
.map(|source| source.selection_order)
.collect::<Vec<_>>(),
vec![0, 1]
);
}
#[test]
fn selected_sources_get_complementary_symbol_biases() {
let plan = plan_multi_source_fetch(
object_ref(),
[
peer("peer-a", "10.0.0.1:8472", 1),
peer("peer-b", "10.0.0.2:8472", 2),
peer("peer-c", "10.0.0.3:8472", 3),
],
MultiSourceSelectionConfig {
min_sources: 2,
max_sources: 3,
allow_unauthenticated_lab: false,
},
)
.unwrap();
assert_eq!(
plan.selected_sources
.iter()
.map(|source| source.symbol_bias)
.collect::<Vec<_>>(),
vec![
MultiSourceSymbolBias::SourceFirst,
MultiSourceSymbolBias::RepairFirst,
MultiSourceSymbolBias::Balanced,
]
);
assert_eq!(
plan.selected_sources[0].symbol_bias.bias_id(),
"source_first"
);
}
#[test]
fn stop_commands_cover_every_selected_source_in_selection_order() {
let plan = plan_multi_source_fetch(
object_ref(),
[
peer("peer-b", "10.0.0.2:8472", 1),
peer("peer-a", "10.0.0.1:8472", 0),
],
MultiSourceSelectionConfig::production_default(),
)
.unwrap();
let stops = plan.stop_commands(MultiSourceStopReason::DecodedAndVerified);
assert_eq!(stops.len(), 2);
assert_eq!(stops[0].peer_id, "peer-a");
assert_eq!(stops[1].peer_id, "peer-b");
assert!(stops.iter().all(|stop| {
stop.object_id == "object-01" && stop.reason.reason_id() == "decoded_and_verified"
}));
}
#[test]
fn unauthenticated_sources_fail_closed_outside_lab_mode() {
let err = plan_multi_source_fetch(
object_ref(),
[
peer("peer-a", "10.0.0.1:8472", 0),
MultiSourcePeer::new(
"peer-b",
"10.0.0.2:8472",
1,
MultiSourceAuth::UnauthenticatedLab,
),
],
MultiSourceSelectionConfig::production_default(),
)
.unwrap_err();
assert!(matches!(
err,
MultiSourcePlanError::UnauthenticatedPeerRejected { peer_id }
if peer_id == "peer-b"
));
}
#[test]
fn lab_mode_must_still_meet_minimum_source_count() {
let plan = plan_multi_source_fetch(
object_ref(),
[
MultiSourcePeer::new(
"peer-a",
"10.0.0.1:8472",
0,
MultiSourceAuth::UnauthenticatedLab,
),
MultiSourcePeer::new(
"peer-b",
"10.0.0.2:8472",
1,
MultiSourceAuth::UnauthenticatedLab,
),
],
MultiSourceSelectionConfig {
min_sources: 2,
max_sources: 2,
allow_unauthenticated_lab: true,
},
)
.unwrap();
assert_eq!(plan.source_count(), 2);
assert_eq!(
plan.selected_sources[0].peer.auth.mode_id(),
"unauthenticated_lab"
);
}
#[test]
fn duplicate_peer_ids_and_endpoints_fail_closed() {
assert!(matches!(
plan_multi_source_fetch(
object_ref(),
[
peer("peer-a", "10.0.0.1:8472", 0),
peer("peer-a", "10.0.0.2:8472", 1),
],
MultiSourceSelectionConfig::production_default(),
),
Err(MultiSourcePlanError::DuplicatePeerId { .. })
));
assert!(matches!(
plan_multi_source_fetch(
object_ref(),
[
peer("peer-a", "10.0.0.1:8472", 0),
peer("peer-b", "10.0.0.1:8472", 1),
],
MultiSourceSelectionConfig::production_default(),
),
Err(MultiSourcePlanError::DuplicateEndpoint { .. })
));
}
#[test]
fn invalid_config_and_object_identity_fail_closed() {
assert!(matches!(
plan_multi_source_fetch(
MultiSourceObjectRef::new("", "abc"),
[peer("peer-a", "10.0.0.1:8472", 0)],
MultiSourceSelectionConfig::production_default(),
),
Err(MultiSourcePlanError::EmptyObjectId)
));
assert!(matches!(
plan_multi_source_fetch(
object_ref(),
[
peer("peer-a", "10.0.0.1:8472", 0),
peer("peer-b", "10.0.0.2:8472", 1),
],
MultiSourceSelectionConfig {
min_sources: 3,
max_sources: 2,
allow_unauthenticated_lab: false,
},
),
Err(MultiSourcePlanError::MaxSourcesBelowMin { .. })
));
}
}