use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
use tokio::sync::broadcast;
use serde::{Deserialize, Serialize};
use super::ResourceKind;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ResourceVersion(pub u64);
impl ResourceVersion {
pub fn new(version: u64) -> Self {
Self(version)
}
pub fn value(&self) -> u64 {
self.0
}
}
impl std::fmt::Display for ResourceVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum WatchEvent {
Added {
kind: ResourceKind,
key: String,
value: serde_json::Value,
version: u64,
},
Modified {
kind: ResourceKind,
key: String,
value: serde_json::Value,
version: u64,
},
Deleted {
kind: ResourceKind,
key: String,
value: serde_json::Value,
version: u64,
},
Bookmark {
kind: ResourceKind,
version: u64,
},
Error {
message: String,
code: u32,
},
}
impl WatchEvent {
pub fn event_type(&self) -> &str {
match self {
WatchEvent::Added { .. } => "ADDED",
WatchEvent::Modified { .. } => "MODIFIED",
WatchEvent::Deleted { .. } => "DELETED",
WatchEvent::Bookmark { .. } => "BOOKMARK",
WatchEvent::Error { .. } => "ERROR",
}
}
pub fn version(&self) -> Option<u64> {
match self {
WatchEvent::Added { version, .. } => Some(*version),
WatchEvent::Modified { version, .. } => Some(*version),
WatchEvent::Deleted { version, .. } => Some(*version),
WatchEvent::Bookmark { version, .. } => Some(*version),
WatchEvent::Error { .. } => None,
}
}
}
pub struct WatchStream {
rx: broadcast::Receiver<WatchEvent>,
kind: ResourceKind,
min_version: Option<u64>,
}
impl WatchStream {
pub fn new(rx: broadcast::Receiver<WatchEvent>, kind: ResourceKind) -> Self {
Self {
rx,
kind,
min_version: None,
}
}
pub fn from_version(mut self, version: u64) -> Self {
self.min_version = Some(version);
self
}
pub async fn recv(&mut self) -> Option<WatchEvent> {
loop {
match self.rx.recv().await {
Ok(event) => {
let matches_kind = match &event {
WatchEvent::Added { kind, .. } => kind == &self.kind,
WatchEvent::Modified { kind, .. } => kind == &self.kind,
WatchEvent::Deleted { kind, .. } => kind == &self.kind,
WatchEvent::Bookmark { kind, .. } => kind == &self.kind,
WatchEvent::Error { .. } => true,
};
if !matches_kind {
continue;
}
if let Some(min_version) = self.min_version {
if let Some(version) = event.version() {
if version <= min_version {
continue;
}
}
}
return Some(event);
}
Err(broadcast::error::RecvError::Lagged(n)) => {
return Some(WatchEvent::Error {
message: format!("Watch lagged by {} events", n),
code: 410, });
}
Err(broadcast::error::RecvError::Closed) => {
return None;
}
}
}
}
}
pub struct WatchRegistry {
tx: broadcast::Sender<WatchEvent>,
watchers: RwLock<HashMap<ResourceKind, usize>>,
}
impl WatchRegistry {
pub fn new() -> Self {
let (tx, _) = broadcast::channel(1024);
Self {
tx,
watchers: RwLock::new(HashMap::new()),
}
}
pub fn subscribe(&self, kind: ResourceKind) -> WatchStream {
let rx = self.tx.subscribe();
let mut watchers = self.watchers.write();
*watchers.entry(kind.clone()).or_insert(0) += 1;
WatchStream::new(rx, kind)
}
pub fn notify(&self, event: WatchEvent) {
let _ = self.tx.send(event);
}
pub fn watcher_count(&self, kind: &ResourceKind) -> usize {
self.watchers.read().get(kind).copied().unwrap_or(0)
}
pub fn send_bookmarks(&self, version: u64) {
let kinds: Vec<_> = self.watchers.read().keys().cloned().collect();
for kind in kinds {
self.notify(WatchEvent::Bookmark {
kind,
version,
});
}
}
}
impl Default for WatchRegistry {
fn default() -> Self {
Self::new()
}
}
pub struct WatchCache<T> {
items: RwLock<HashMap<String, (T, u64)>>,
version: RwLock<u64>,
registry: Arc<WatchRegistry>,
kind: ResourceKind,
}
impl<T: Clone + Send + Sync + 'static> WatchCache<T> {
pub fn new(kind: ResourceKind, registry: Arc<WatchRegistry>) -> Self {
Self {
items: RwLock::new(HashMap::new()),
version: RwLock::new(0),
registry,
kind,
}
}
pub fn set(&self, key: impl Into<String>, value: T, version: u64) {
let key = key.into();
let mut items = self.items.write();
let is_new = !items.contains_key(&key);
items.insert(key.clone(), (value, version));
*self.version.write() = version;
}
pub fn get(&self, key: &str) -> Option<T> {
self.items.read().get(key).map(|(v, _)| v.clone())
}
pub fn remove(&self, key: &str) -> Option<T> {
self.items.write().remove(key).map(|(v, _)| v)
}
pub fn list(&self) -> Vec<T> {
self.items.read().values().map(|(v, _)| v.clone()).collect()
}
pub fn version(&self) -> u64 {
*self.version.read()
}
pub fn watch(&self) -> WatchStream {
self.registry.subscribe(self.kind.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_watch_event_type() {
let event = WatchEvent::Added {
kind: ResourceKind::Workload,
key: "test".to_string(),
value: serde_json::json!({}),
version: 1,
};
assert_eq!(event.event_type(), "ADDED");
assert_eq!(event.version(), Some(1));
}
#[test]
fn test_watch_registry() {
let registry = WatchRegistry::new();
let _stream = registry.subscribe(ResourceKind::Workload);
assert_eq!(registry.watcher_count(&ResourceKind::Workload), 1);
}
#[tokio::test]
async fn test_watch_stream() {
let registry = Arc::new(WatchRegistry::new());
let mut stream = registry.subscribe(ResourceKind::Workload);
let registry_clone = registry.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
registry_clone.notify(WatchEvent::Added {
kind: ResourceKind::Workload,
key: "test".to_string(),
value: serde_json::json!({"name": "test"}),
version: 1,
});
});
let event = tokio::time::timeout(
std::time::Duration::from_millis(100),
stream.recv()
).await.unwrap();
assert!(event.is_some());
assert_eq!(event.unwrap().event_type(), "ADDED");
}
}