use std::sync::Arc;
use futures::StreamExt;
use tracing::{info, warn};
use crate::nats::HasJetStream;
use crate::{BoxFuture, FlatbedWorkerError};
pub trait KvWorker: Send + Sync + 'static {
type Value: Send + 'static;
type Context: HasJetStream + Send + Sync + 'static;
type ParseError: std::fmt::Display + Send + 'static;
const NAME: &'static str;
const DESCRIPTION: Option<&'static str> = None;
const RESTART: ::core::option::Option<crate::RestartPolicy> = ::core::option::Option::None;
const BUCKET: &'static str;
fn on_put(&self, ctx: Arc<Self::Context>, key: String, value: Self::Value) -> BoxFuture<()>;
fn on_delete(&self, ctx: Arc<Self::Context>, key: String) -> BoxFuture<()>;
fn parse_value(bytes: &[u8]) -> Result<Self::Value, Self::ParseError>;
fn drain(&self, _ctx: Arc<Self::Context>) -> BoxFuture<Result<(), FlatbedWorkerError>> {
Box::pin(async { Ok(()) })
}
}
pub fn run_kv_worker_drain<W, C>(
ctx: Arc<dyn std::any::Any + Send + Sync>,
) -> BoxFuture<Result<(), FlatbedWorkerError>>
where
W: KvWorker<Context = C> + Default,
C: Send + Sync + 'static,
{
Box::pin(async move {
let ctx: Arc<C> = ctx
.downcast::<C>()
.unwrap_or_else(|_| panic!("kv_worker '{}' context type mismatch", W::NAME));
W::default().drain(ctx).await
})
}
pub async fn run_kv_worker<W, C>(
ctx: Arc<dyn std::any::Any + Send + Sync>,
) -> Result<(), FlatbedWorkerError>
where
W: KvWorker<Context = C> + Default,
C: HasJetStream + Send + Sync + 'static,
{
let ctx: Arc<C> = ctx
.downcast::<C>()
.unwrap_or_else(|_| panic!("kv_worker '{}' context type mismatch", W::NAME));
let worker = W::default();
loop {
let store_result = {
let jetstream = HasJetStream::jetstream(&*ctx);
jetstream.get_key_value(W::BUCKET).await
};
let store = match store_result {
Ok(s) => s,
Err(e) => {
warn!(
worker = W::NAME,
bucket = W::BUCKET,
error = %e,
"KV bucket lookup failed; retrying in 5s",
);
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
continue;
}
};
watch_cycle::<W, C>(&worker, Arc::clone(&ctx), store).await;
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
}
const WATCH_FAILURE_BACKOFF: std::time::Duration = std::time::Duration::from_secs(5);
async fn watch_cycle<W, C>(worker: &W, ctx: Arc<C>, store: async_nats::jetstream::kv::Store)
where
W: KvWorker<Context = C>,
C: HasJetStream + Send + Sync + 'static,
{
let mut watch = match store.watch_all().await {
Ok(w) => w,
Err(e) => {
warn!(
worker = W::NAME,
bucket = W::BUCKET,
error = %e,
"KV watch_all failed; will reopen",
);
tokio::time::sleep(WATCH_FAILURE_BACKOFF).await;
return;
}
};
info!(worker = W::NAME, bucket = W::BUCKET, "subscribed");
while let Some(entry_result) = watch.next().await {
let entry = match entry_result {
Ok(e) => e,
Err(e) => {
warn!(
worker = W::NAME,
bucket = W::BUCKET,
error = %e,
"KV entry stream error; reopening",
);
tokio::time::sleep(WATCH_FAILURE_BACKOFF).await;
return;
}
};
let event = map_operation(entry.operation, entry.key, &entry.value);
dispatch_event::<W, C>(worker, Arc::clone(&ctx), event).await;
}
info!(
worker = W::NAME,
bucket = W::BUCKET,
"KV watch stream ended; reopening"
);
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum KvEvent {
Put { key: String, value: Vec<u8> },
Delete { key: String },
}
fn map_operation(op: async_nats::jetstream::kv::Operation, key: String, value: &[u8]) -> KvEvent {
match op {
async_nats::jetstream::kv::Operation::Put => KvEvent::Put {
key,
value: value.to_vec(),
},
async_nats::jetstream::kv::Operation::Delete
| async_nats::jetstream::kv::Operation::Purge => KvEvent::Delete { key },
}
}
async fn dispatch_event<W, C>(worker: &W, ctx: Arc<C>, event: KvEvent)
where
W: KvWorker<Context = C>,
C: HasJetStream + Send + Sync + 'static,
{
match event {
KvEvent::Put { key, value } => match W::parse_value(&value) {
Ok(decoded) => {
worker.on_put(ctx, key, decoded).await;
}
Err(e) => {
warn!(
worker = W::NAME,
bucket = W::BUCKET,
key = %key,
error = %e,
"KV value decode failed; skipping entry",
);
}
},
KvEvent::Delete { key } => {
worker.on_delete(ctx, key).await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug, PartialEq)]
struct TestValue {
n: u32,
}
struct TestCtx;
impl HasJetStream for TestCtx {
fn jetstream(&self) -> &async_nats::jetstream::Context {
panic!("stub jetstream: not expected to be called in parse_value tests");
}
}
struct TestKvWorker;
impl Default for TestKvWorker {
fn default() -> Self {
Self
}
}
impl KvWorker for TestKvWorker {
type Value = TestValue;
type Context = TestCtx;
type ParseError = String;
const NAME: &'static str = "test-kv-worker";
const BUCKET: &'static str = "TEST_BUCKET";
fn on_put(
&self,
_ctx: Arc<Self::Context>,
_key: String,
_value: Self::Value,
) -> BoxFuture<()> {
Box::pin(async {})
}
fn on_delete(&self, _ctx: Arc<Self::Context>, _key: String) -> BoxFuture<()> {
Box::pin(async {})
}
fn parse_value(bytes: &[u8]) -> Result<Self::Value, Self::ParseError> {
if bytes.len() != 4 {
return Err(format!("expected 4 bytes, got {}", bytes.len()));
}
let n = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
Ok(TestValue { n })
}
}
#[test]
fn parse_value_valid_bytes() {
let bytes = 42u32.to_le_bytes();
let result = TestKvWorker::parse_value(&bytes);
assert_eq!(result, Ok(TestValue { n: 42 }));
}
#[test]
fn parse_value_invalid_bytes_returns_error() {
let result = TestKvWorker::parse_value(&[1, 2]);
assert!(result.is_err());
assert!(result.unwrap_err().contains("expected 4 bytes"));
}
#[test]
fn parse_value_empty_bytes_returns_error() {
let result = TestKvWorker::parse_value(&[]);
assert!(result.is_err());
}
#[test]
fn kv_worker_description_default_is_none() {
assert!(<TestKvWorker as KvWorker>::DESCRIPTION.is_none());
}
#[test]
fn kv_worker_constants_match() {
assert_eq!(<TestKvWorker as KvWorker>::NAME, "test-kv-worker");
assert_eq!(<TestKvWorker as KvWorker>::BUCKET, "TEST_BUCKET");
}
#[derive(Debug, PartialEq)]
enum CustomKvParseError {
WrongLen { got: usize, expected: usize },
}
impl std::fmt::Display for CustomKvParseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::WrongLen { got, expected } => {
write!(f, "wrong length: got {got}, expected {expected}")
}
}
}
}
struct TypedErrorKvWorker;
impl KvWorker for TypedErrorKvWorker {
type Value = TestValue;
type Context = TestCtx;
type ParseError = CustomKvParseError;
const NAME: &'static str = "typed-error-kv-worker";
const BUCKET: &'static str = "TYPED_BUCKET";
fn on_put(
&self,
_ctx: Arc<Self::Context>,
_key: String,
_value: Self::Value,
) -> BoxFuture<()> {
Box::pin(async {})
}
fn on_delete(&self, _ctx: Arc<Self::Context>, _key: String) -> BoxFuture<()> {
Box::pin(async {})
}
fn parse_value(bytes: &[u8]) -> Result<Self::Value, Self::ParseError> {
if bytes.len() != 4 {
return Err(CustomKvParseError::WrongLen {
got: bytes.len(),
expected: 4,
});
}
let n = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
Ok(TestValue { n })
}
}
#[test]
fn typed_kv_parse_error_round_trips_structured_variant() {
let err = TypedErrorKvWorker::parse_value(&[1, 2]).unwrap_err();
assert_eq!(
err,
CustomKvParseError::WrongLen {
got: 2,
expected: 4
}
);
}
#[test]
fn typed_kv_parse_error_display_matches_expected_template() {
let err = CustomKvParseError::WrongLen {
got: 2,
expected: 4,
};
assert_eq!(err.to_string(), "wrong length: got 2, expected 4");
}
use std::sync::Mutex;
#[derive(Default)]
struct RecorderWorker {
puts: Arc<Mutex<Vec<(String, TestValue)>>>,
deletes: Arc<Mutex<Vec<String>>>,
}
impl KvWorker for RecorderWorker {
type Value = TestValue;
type Context = TestCtx;
type ParseError = String;
const NAME: &'static str = "recorder-worker";
const BUCKET: &'static str = "RECORDER_BUCKET";
fn on_put(
&self,
_ctx: Arc<Self::Context>,
key: String,
value: Self::Value,
) -> BoxFuture<()> {
let puts = Arc::clone(&self.puts);
Box::pin(async move {
puts.lock().unwrap().push((key, value));
})
}
fn on_delete(&self, _ctx: Arc<Self::Context>, key: String) -> BoxFuture<()> {
let deletes = Arc::clone(&self.deletes);
Box::pin(async move {
deletes.lock().unwrap().push(key);
})
}
fn parse_value(bytes: &[u8]) -> Result<Self::Value, Self::ParseError> {
<TestKvWorker as KvWorker>::parse_value(bytes)
}
}
#[tokio::test]
async fn dispatch_event_routes_put_to_on_put() {
let worker = RecorderWorker::default();
let ctx = Arc::new(TestCtx);
dispatch_event::<RecorderWorker, TestCtx>(
&worker,
ctx,
KvEvent::Put {
key: "tenant-a/worker".into(),
value: 42u32.to_le_bytes().to_vec(),
},
)
.await;
let puts = worker.puts.lock().unwrap();
assert_eq!(puts.len(), 1);
assert_eq!(puts[0].0, "tenant-a/worker");
assert_eq!(puts[0].1, TestValue { n: 42 });
assert!(worker.deletes.lock().unwrap().is_empty());
}
#[tokio::test]
async fn dispatch_event_routes_delete_to_on_delete() {
let worker = RecorderWorker::default();
let ctx = Arc::new(TestCtx);
dispatch_event::<RecorderWorker, TestCtx>(
&worker,
ctx,
KvEvent::Delete {
key: "tenant-a/worker".into(),
},
)
.await;
let deletes = worker.deletes.lock().unwrap();
assert_eq!(deletes.len(), 1);
assert_eq!(deletes[0], "tenant-a/worker");
assert!(worker.puts.lock().unwrap().is_empty());
}
#[tokio::test]
async fn dispatch_event_skips_put_when_decode_fails() {
let worker = RecorderWorker::default();
let ctx = Arc::new(TestCtx);
dispatch_event::<RecorderWorker, TestCtx>(
&worker,
ctx,
KvEvent::Put {
key: "tenant-a/worker".into(),
value: vec![0xFF, 0xFE], },
)
.await;
assert!(
worker.puts.lock().unwrap().is_empty(),
"on_put must not fire when parse_value fails",
);
assert!(
worker.deletes.lock().unwrap().is_empty(),
"on_delete must not fire either — bad bytes are silently skipped",
);
}
#[test]
fn map_operation_routes_put_to_put_event() {
let event = map_operation(
async_nats::jetstream::kv::Operation::Put,
"tenant-a/worker".to_string(),
&42u32.to_le_bytes(),
);
assert_eq!(
event,
KvEvent::Put {
key: "tenant-a/worker".to_string(),
value: 42u32.to_le_bytes().to_vec(),
}
);
}
#[test]
fn map_operation_routes_delete_to_delete_event() {
let event = map_operation(
async_nats::jetstream::kv::Operation::Delete,
"tenant-a/worker".to_string(),
b"",
);
assert_eq!(
event,
KvEvent::Delete {
key: "tenant-a/worker".to_string(),
}
);
}
#[test]
fn map_operation_routes_purge_to_delete_event() {
let event = map_operation(
async_nats::jetstream::kv::Operation::Purge,
"tenant-a/worker".to_string(),
b"",
);
assert_eq!(
event,
KvEvent::Delete {
key: "tenant-a/worker".to_string(),
}
);
}
crate::register_kv_worker!(TestKvWorker, TestCtx);
#[test]
fn register_kv_worker_macro_makes_worker_discoverable() {
let workers = crate::get_workers();
let found = workers
.iter()
.find(|w| w.name == <TestKvWorker as KvWorker>::NAME);
assert!(
found.is_some(),
"TestKvWorker should be discoverable via inventory"
);
let info = found.unwrap();
assert_eq!(info.name, "test-kv-worker");
assert_eq!(info.description, <TestKvWorker as KvWorker>::DESCRIPTION);
}
}