use std::collections::HashMap;
#[derive(Clone, Debug, PartialEq)]
pub enum BehaviorSignal {
FastResponder,
SlowResponder,
HighBandwidth,
LowBandwidth,
FrequentChurner,
StableConnection,
DataHoarder,
GoodContributor,
}
impl BehaviorSignal {
fn discriminant_index(&self) -> usize {
match self {
BehaviorSignal::FastResponder => 0,
BehaviorSignal::SlowResponder => 1,
BehaviorSignal::HighBandwidth => 2,
BehaviorSignal::LowBandwidth => 3,
BehaviorSignal::FrequentChurner => 4,
BehaviorSignal::StableConnection => 5,
BehaviorSignal::DataHoarder => 6,
BehaviorSignal::GoodContributor => 7,
}
}
}
#[derive(Clone, Debug)]
pub struct BehaviorProfile {
pub peer_id: String,
pub signals: Vec<BehaviorSignal>,
pub observation_count: u64,
}
impl BehaviorProfile {
pub fn has_signal(&self, signal: &BehaviorSignal) -> bool {
self.signals.contains(signal)
}
pub fn signal_count(&self) -> usize {
self.signals.len()
}
}
#[derive(Clone, Debug, Default)]
pub struct ClassifierStats {
pub total_peers: usize,
pub fast_responders: usize,
pub high_bandwidth_peers: usize,
pub good_contributors: usize,
pub churners: usize,
}
#[derive(Debug, Default)]
pub struct PeerBehaviorClassifier {
pub profiles: HashMap<String, BehaviorProfile>,
}
impl PeerBehaviorClassifier {
pub fn new() -> Self {
Self {
profiles: HashMap::new(),
}
}
pub fn record_signal(&mut self, peer_id: &str, signal: BehaviorSignal) {
let profile = self
.profiles
.entry(peer_id.to_string())
.or_insert_with(|| BehaviorProfile {
peer_id: peer_id.to_string(),
signals: Vec::new(),
observation_count: 0,
});
profile.signals.push(signal);
profile.observation_count += 1;
}
pub fn classify(&self, peer_id: &str) -> Vec<BehaviorSignal> {
let profile = match self.profiles.get(peer_id) {
Some(p) => p,
None => return Vec::new(),
};
let mut counts: HashMap<usize, (usize, BehaviorSignal)> = HashMap::new();
for signal in &profile.signals {
let idx = signal.discriminant_index();
let entry = counts.entry(idx).or_insert_with(|| (0, signal.clone()));
entry.0 += 1;
}
let mut result: Vec<(usize, BehaviorSignal)> = counts
.into_iter()
.filter_map(
|(idx, (count, signal))| {
if count > 1 {
Some((idx, signal))
} else {
None
}
},
)
.collect();
result.sort_by_key(|(idx, _)| *idx);
result.into_iter().map(|(_, signal)| signal).collect()
}
pub fn peers_with_signal(&self, signal: &BehaviorSignal) -> Vec<&str> {
let mut result: Vec<&str> = self
.profiles
.keys()
.filter(|peer_id| self.classify(peer_id).contains(signal))
.map(|s| s.as_str())
.collect();
result.sort_unstable();
result
}
pub fn remove_peer(&mut self, peer_id: &str) -> bool {
self.profiles.remove(peer_id).is_some()
}
pub fn stats(&self) -> ClassifierStats {
let mut stats = ClassifierStats {
total_peers: self.profiles.len(),
..Default::default()
};
for peer_id in self.profiles.keys() {
let classified = self.classify(peer_id);
for signal in &classified {
match signal {
BehaviorSignal::FastResponder => stats.fast_responders += 1,
BehaviorSignal::HighBandwidth => stats.high_bandwidth_peers += 1,
BehaviorSignal::GoodContributor => stats.good_contributors += 1,
BehaviorSignal::FrequentChurner => stats.churners += 1,
_ => {}
}
}
}
stats
}
pub fn profile(&self, peer_id: &str) -> Option<&BehaviorProfile> {
self.profiles.get(peer_id)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_starts_empty() {
let classifier = PeerBehaviorClassifier::new();
assert!(classifier.profiles.is_empty());
}
#[test]
fn test_record_signal_creates_profile() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c.profile("peer-a").is_some());
}
#[test]
fn test_record_signal_increments_observation_count() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
c.record_signal("peer-a", BehaviorSignal::SlowResponder);
assert_eq!(
c.profile("peer-a")
.expect("test: peer-a profile should exist")
.observation_count,
2
);
}
#[test]
fn test_record_signal_appends_to_signals() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::HighBandwidth);
c.record_signal("peer-a", BehaviorSignal::HighBandwidth);
assert_eq!(
c.profile("peer-a")
.expect("test: peer-a profile should exist")
.signals
.len(),
2
);
}
#[test]
fn test_observation_count_accumulates_across_many_signals() {
let mut c = PeerBehaviorClassifier::new();
for _ in 0..10 {
c.record_signal("peer-x", BehaviorSignal::StableConnection);
}
assert_eq!(
c.profile("peer-x")
.expect("test: peer-x profile should exist")
.observation_count,
10
);
}
#[test]
fn test_classify_single_signal_not_in_result() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c.classify("peer-a").is_empty());
}
#[test]
fn test_classify_signal_appearing_twice_is_included() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c
.classify("peer-a")
.contains(&BehaviorSignal::FastResponder));
}
#[test]
fn test_classify_multiple_distinct_signals() {
let mut c = PeerBehaviorClassifier::new();
for _ in 0..2 {
c.record_signal("peer-a", BehaviorSignal::HighBandwidth);
c.record_signal("peer-a", BehaviorSignal::GoodContributor);
}
let classified = c.classify("peer-a");
assert!(classified.contains(&BehaviorSignal::HighBandwidth));
assert!(classified.contains(&BehaviorSignal::GoodContributor));
}
#[test]
fn test_classify_returns_empty_for_unknown_peer() {
let c = PeerBehaviorClassifier::new();
assert!(c.classify("ghost-peer").is_empty());
}
#[test]
fn test_classify_returns_sorted_by_declaration_order() {
let mut c = PeerBehaviorClassifier::new();
for _ in 0..2 {
c.record_signal("peer-a", BehaviorSignal::GoodContributor); c.record_signal("peer-a", BehaviorSignal::DataHoarder); c.record_signal("peer-a", BehaviorSignal::FastResponder); }
let classified = c.classify("peer-a");
assert_eq!(classified.len(), 3);
assert_eq!(classified[0], BehaviorSignal::FastResponder);
assert_eq!(classified[1], BehaviorSignal::DataHoarder);
assert_eq!(classified[2], BehaviorSignal::GoodContributor);
}
#[test]
fn test_classify_threshold_exactly_two_observations() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::LowBandwidth);
c.record_signal("peer-a", BehaviorSignal::LowBandwidth);
assert!(c.classify("peer-a").contains(&BehaviorSignal::LowBandwidth));
}
#[test]
fn test_peers_with_signal_returns_correct_peers_sorted() {
let mut c = PeerBehaviorClassifier::new();
for id in &["charlie", "alice", "bob"] {
c.record_signal(id, BehaviorSignal::FastResponder);
c.record_signal(id, BehaviorSignal::FastResponder);
}
let peers = c.peers_with_signal(&BehaviorSignal::FastResponder);
assert_eq!(peers, vec!["alice", "bob", "charlie"]);
}
#[test]
fn test_peers_with_signal_returns_empty_when_none_qualify() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c
.peers_with_signal(&BehaviorSignal::FastResponder)
.is_empty());
}
#[test]
fn test_peers_with_signal_excludes_non_qualifying_peers() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::HighBandwidth);
c.record_signal("peer-a", BehaviorSignal::HighBandwidth);
c.record_signal("peer-b", BehaviorSignal::HighBandwidth); let peers = c.peers_with_signal(&BehaviorSignal::HighBandwidth);
assert_eq!(peers, vec!["peer-a"]);
}
#[test]
fn test_remove_peer_existing_returns_true() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c.remove_peer("peer-a"));
}
#[test]
fn test_remove_peer_nonexistent_returns_false() {
let mut c = PeerBehaviorClassifier::new();
assert!(!c.remove_peer("ghost"));
}
#[test]
fn test_remove_peer_actually_removes_profile() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
c.remove_peer("peer-a");
assert!(c.profile("peer-a").is_none());
}
#[test]
fn test_profile_returns_some_for_known_peer() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::FastResponder);
assert!(c.profile("peer-a").is_some());
}
#[test]
fn test_profile_returns_none_for_unknown_peer() {
let c = PeerBehaviorClassifier::new();
assert!(c.profile("ghost").is_none());
}
#[test]
fn test_stats_total_peers() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-1", BehaviorSignal::FastResponder);
c.record_signal("peer-2", BehaviorSignal::SlowResponder);
assert_eq!(c.stats().total_peers, 2);
}
#[test]
fn test_stats_fast_responders_count() {
let mut c = PeerBehaviorClassifier::new();
for id in &["p1", "p2"] {
c.record_signal(id, BehaviorSignal::FastResponder);
c.record_signal(id, BehaviorSignal::FastResponder);
}
c.record_signal("p3", BehaviorSignal::FastResponder); assert_eq!(c.stats().fast_responders, 2);
}
#[test]
fn test_stats_high_bandwidth_peers_count() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("p1", BehaviorSignal::HighBandwidth);
c.record_signal("p1", BehaviorSignal::HighBandwidth);
assert_eq!(c.stats().high_bandwidth_peers, 1);
}
#[test]
fn test_stats_good_contributors_count() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("p1", BehaviorSignal::GoodContributor);
c.record_signal("p1", BehaviorSignal::GoodContributor);
c.record_signal("p2", BehaviorSignal::GoodContributor);
c.record_signal("p2", BehaviorSignal::GoodContributor);
assert_eq!(c.stats().good_contributors, 2);
}
#[test]
fn test_stats_churners_count() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("p1", BehaviorSignal::FrequentChurner);
c.record_signal("p1", BehaviorSignal::FrequentChurner);
assert_eq!(c.stats().churners, 1);
}
#[test]
fn test_has_signal_true() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::StableConnection);
let profile = c
.profile("peer-a")
.expect("test: peer-a profile should exist");
assert!(profile.has_signal(&BehaviorSignal::StableConnection));
}
#[test]
fn test_has_signal_false() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::StableConnection);
let profile = c
.profile("peer-a")
.expect("test: peer-a profile should exist");
assert!(!profile.has_signal(&BehaviorSignal::FastResponder));
}
#[test]
fn test_signal_count() {
let mut c = PeerBehaviorClassifier::new();
c.record_signal("peer-a", BehaviorSignal::DataHoarder);
c.record_signal("peer-a", BehaviorSignal::DataHoarder);
c.record_signal("peer-a", BehaviorSignal::GoodContributor);
let profile = c
.profile("peer-a")
.expect("test: peer-a profile should exist");
assert_eq!(profile.signal_count(), 3);
}
}