use anyhow::Result;
use blvm_protocol::Hash;
use std::collections::{HashMap, HashSet};
use tracing::{debug, info, warn};
use super::protocol::{GetDataMessage, InventoryVector};
pub const MSG_TX: u32 = 1;
pub const MSG_BLOCK: u32 = 2;
pub const MSG_FILTERED_BLOCK: u32 = 3;
pub const MSG_CMPCT_BLOCK: u32 = 4;
pub const MSG_WITNESS_BLOCK: u32 = 0x40000002;
pub struct InventoryManager {
known_inventory: HashSet<Hash>,
pending_requests: HashMap<Hash, InventoryRequest>,
peer_inventories: HashMap<String, HashSet<Hash>>,
}
impl Default for InventoryManager {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct InventoryRequest {
pub inv_type: u32,
pub hash: Hash,
pub timestamp: u64,
pub peer: String,
}
impl InventoryManager {
pub fn new() -> Self {
Self {
known_inventory: HashSet::new(),
pending_requests: HashMap::new(),
peer_inventories: HashMap::new(),
}
}
pub fn add_inventory(&mut self, peer: &str, inventory: &[InventoryVector]) -> Result<()> {
let peer_inv = self.peer_inventories.entry(peer.to_string()).or_default();
for item in inventory {
peer_inv.insert(item.hash);
self.known_inventory.insert(item.hash);
debug!("Added inventory item {:?} from peer {}", item, peer);
}
Ok(())
}
pub fn has_inventory(&self, hash: &Hash) -> bool {
self.known_inventory.contains(hash)
}
pub fn request_data(
&mut self,
hash: Hash,
inv_type: u32,
peer: &str,
) -> Result<GetDataMessage> {
let request = InventoryRequest {
inv_type,
hash,
timestamp: crate::utils::current_timestamp(),
peer: peer.to_string(),
};
self.pending_requests.insert(hash, request.clone());
let inventory = vec![InventoryVector { inv_type, hash }];
Ok(GetDataMessage { inventory })
}
pub fn mark_fulfilled(&mut self, hash: &Hash) {
self.pending_requests.remove(hash);
debug!("Marked request for {} as fulfilled", hex::encode(hash));
}
pub fn get_pending_requests(&self) -> Vec<&InventoryRequest> {
self.pending_requests.values().collect()
}
pub fn cleanup_old_requests(&mut self, max_age_seconds: u64) {
let now = crate::utils::current_timestamp();
let old_requests: Vec<Hash> = self
.pending_requests
.iter()
.filter(|(_, request)| {
let age = now - request.timestamp;
age >= max_age_seconds
})
.map(|(hash, _)| *hash)
.collect();
for hash in old_requests {
self.pending_requests.remove(&hash);
warn!("Removed old pending request for {}", hex::encode(hash));
}
}
pub fn get_peer_inventory(&self, peer: &str) -> Option<&HashSet<Hash>> {
self.peer_inventories.get(peer)
}
pub fn remove_peer(&mut self, peer: &str) {
self.peer_inventories.remove(peer);
info!("Removed inventory for peer {}", peer);
}
pub fn inventory_count(&self) -> usize {
self.known_inventory.len()
}
pub fn pending_request_count(&self) -> usize {
self.pending_requests.len()
}
}