From 106223fb8550f6bfa734f72ae7155270ff30fc9c Mon Sep 17 00:00:00 2001 From: Matthew Hounslow Date: Wed, 24 Dec 2025 00:21:05 -0700 Subject: [PATCH] Add module unregistration to buswatch-sdk (#20) - Add Instrumentor::unregister(name) method to remove modules - Add GlobalState::unregister_module(name) for cleanup - Returns true if module was found and removed, false otherwise - Comprehensive tests covering: - Unregister existing module - Unregister non-existent module - Verify metrics are removed after unregister - Verify re-registering works after unregister - Old handles still work but don't appear in snapshots - Unregistering one module doesn't affect others - Multiple unregister calls are safe - Update README.md with unregistration documentation - Update CHANGELOG.md This enables clean lifecycle management for temporary or dynamic modules. --- CHANGELOG.md | 11 ++ buswatch-sdk/README.md | 34 +++++ buswatch-sdk/src/instrumentor.rs | 214 +++++++++++++++++++++++++++++++ buswatch-sdk/src/state.rs | 135 +++++++++++++++++++ 4 files changed, 394 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 56894c0..2160aa6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **buswatch-tui**: CSV export format (#32) + - Press `E` (uppercase) to export current view to CSV format + - `--export-csv` CLI flag to export snapshot and exit + - CSV includes module name, topic, type (Read/Write), count, backlog, pending duration, rate, and health status + - Proper CSV escaping for special characters (commas, quotes, newlines) +- **buswatch-sdk**: Module unregistration support (#20) + - `Instrumentor::unregister(name)` method to remove modules from internal state + - `GlobalState::unregister_module(name)` for module cleanup + - Returns `true` if module was found and removed, `false` otherwise + - Enables clean lifecycle management for temporary or dynamic modules + - Supports re-registration with fresh state after unregister - **buswatch-sdk**: Prometheus exposition format export (`prometheus` feature) - HTTP server serving metrics at configurable endpoint - All metrics include `module` and `topic` labels diff --git a/buswatch-sdk/README.md b/buswatch-sdk/README.md index 1853c1b..2d4a8f5 100644 --- a/buswatch-sdk/README.md +++ b/buswatch-sdk/README.md @@ -110,6 +110,40 @@ This starts an HTTP server that Prometheus can scrape. Metrics available: Health check endpoints (`/health`, `/healthz`) are also available for Kubernetes probes. +## Module Lifecycle + +### Registering Modules + +```rust +let handle = instrumentor.register("my-service"); +``` + +If a module with the same name already exists, `register()` returns a handle to the existing module. + +### Unregistering Modules + +When a module is no longer needed (e.g., a temporary worker, a completed task), you can unregister it: + +```rust +// Unregister the module +let removed = instrumentor.unregister("my-service"); + +if removed { + println!("Module was successfully unregistered"); +} else { + println!("Module didn't exist"); +} +``` + +Unregistering a module: +- Removes it from future snapshots +- Clears all associated metrics +- Returns `true` if the module existed, `false` otherwise +- Can be safely called multiple times +- Allows re-registration with fresh state + +**Note:** Existing `ModuleHandle` instances will continue to work after unregistration, but their metrics won't appear in snapshots unless the module is re-registered. + ## Recording Metrics ### Basic Counting diff --git a/buswatch-sdk/src/instrumentor.rs b/buswatch-sdk/src/instrumentor.rs index c9ebc34..4d09f28 100644 --- a/buswatch-sdk/src/instrumentor.rs +++ b/buswatch-sdk/src/instrumentor.rs @@ -79,6 +79,37 @@ impl Instrumentor { } } + /// Unregister a module and remove it from internal state. + /// + /// This removes the module and all its associated metrics from future snapshots. + /// Returns `true` if the module was found and removed, `false` if it didn't exist. + /// + /// Note: Any existing `ModuleHandle` instances for this module will continue to work + /// and accumulate metrics, but those metrics will not appear in snapshots unless + /// the module is re-registered. + /// + /// # Arguments + /// + /// * `name` - The module name to unregister + /// + /// # Example + /// + /// ```rust + /// use buswatch_sdk::Instrumentor; + /// + /// let instrumentor = Instrumentor::new(); + /// let handle = instrumentor.register("temp-worker"); + /// + /// // ... use the handle ... + /// + /// // When done, unregister the module + /// let removed = instrumentor.unregister("temp-worker"); + /// assert!(removed); + /// ``` + pub fn unregister(&self, name: &str) -> bool { + self.state.unregister_module(name) + } + /// Collect a snapshot of all current metrics. /// /// This is useful if you want to emit snapshots manually rather than @@ -380,4 +411,187 @@ mod tests { Some(50) ); // 950 - 900 } + + #[test] + fn unregister_existing_module() { + let instrumentor = Instrumentor::new(); + + let handle = instrumentor.register("service-a"); + handle.record_read("topic", 100); + + // Verify module exists + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 1); + assert!(snapshot.modules.contains_key("service-a")); + + // Unregister the module + let removed = instrumentor.unregister("service-a"); + assert!(removed, "Should return true when module exists"); + + // Verify module is gone + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 0); + assert!(!snapshot.modules.contains_key("service-a")); + } + + #[test] + fn unregister_nonexistent_module() { + let instrumentor = Instrumentor::new(); + + let _handle = instrumentor.register("service-a"); + + // Try to unregister a module that doesn't exist + let removed = instrumentor.unregister("service-b"); + assert!(!removed, "Should return false when module doesn't exist"); + + // Original module should still be there + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 1); + assert!(snapshot.modules.contains_key("service-a")); + } + + #[test] + fn unregister_removes_metrics_from_snapshot() { + let instrumentor = Instrumentor::new(); + + let producer = instrumentor.register("producer"); + let consumer = instrumentor.register("consumer"); + + producer.record_write("events", 100); + consumer.record_read("events", 50); + + // Both modules present + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 2); + assert_eq!( + snapshot.modules.get("producer").unwrap().writes.get("events").unwrap().count, + 100 + ); + + // Unregister producer + instrumentor.unregister("producer"); + + // Only consumer remains + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 1); + assert!(snapshot.modules.contains_key("consumer")); + assert!(!snapshot.modules.contains_key("producer")); + } + + #[test] + fn can_reregister_after_unregister() { + let instrumentor = Instrumentor::new(); + + // Register and record metrics + let handle1 = instrumentor.register("worker"); + handle1.record_read("jobs", 50); + + let snapshot = instrumentor.collect(); + assert_eq!( + snapshot.modules.get("worker").unwrap().reads.get("jobs").unwrap().count, + 50 + ); + + // Unregister + instrumentor.unregister("worker"); + + // Re-register - should get fresh state + let handle2 = instrumentor.register("worker"); + let snapshot = instrumentor.collect(); + + // New module should start with empty metrics + let metrics = snapshot.modules.get("worker").unwrap(); + assert_eq!(metrics.reads.len(), 0, "Re-registered module should have no topics"); + + // Add new metrics + handle2.record_read("jobs", 25); + + let snapshot = instrumentor.collect(); + assert_eq!( + snapshot.modules.get("worker").unwrap().reads.get("jobs").unwrap().count, + 25 + ); + } + + #[test] + fn old_handle_still_works_after_unregister() { + let instrumentor = Instrumentor::new(); + + let handle = instrumentor.register("service"); + handle.record_write("events", 100); + + // Unregister the module + instrumentor.unregister("service"); + + // Handle can still record metrics (to its internal state) + handle.record_write("events", 50); + + // But they won't appear in snapshots + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 0); + + // Re-registering creates a new module with fresh state + let new_handle = instrumentor.register("service"); + let snapshot = instrumentor.collect(); + + // The new module starts fresh (old handle's metrics are not visible) + let metrics = snapshot.modules.get("service").unwrap(); + assert_eq!(metrics.writes.len(), 0); + + // But the old handle's internal state is still at 150 + // (this is implementation detail - old handle keeps its Arc to old state) + new_handle.record_write("events", 10); + let snapshot = instrumentor.collect(); + assert_eq!( + snapshot.modules.get("service").unwrap().writes.get("events").unwrap().count, + 10 + ); + } + + #[test] + fn unregister_one_module_does_not_affect_others() { + let instrumentor = Instrumentor::new(); + + let service_a = instrumentor.register("service-a"); + let service_b = instrumentor.register("service-b"); + let service_c = instrumentor.register("service-c"); + + service_a.record_read("topic", 10); + service_b.record_read("topic", 20); + service_c.record_read("topic", 30); + + // Unregister service-b + instrumentor.unregister("service-b"); + + // service-a and service-c should remain + let snapshot = instrumentor.collect(); + assert_eq!(snapshot.modules.len(), 2); + assert!(snapshot.modules.contains_key("service-a")); + assert!(!snapshot.modules.contains_key("service-b")); + assert!(snapshot.modules.contains_key("service-c")); + + assert_eq!( + snapshot.modules.get("service-a").unwrap().reads.get("topic").unwrap().count, + 10 + ); + assert_eq!( + snapshot.modules.get("service-c").unwrap().reads.get("topic").unwrap().count, + 30 + ); + } + + #[test] + fn unregister_multiple_times_is_safe() { + let instrumentor = Instrumentor::new(); + + instrumentor.register("service"); + + // First unregister succeeds + assert!(instrumentor.unregister("service")); + + // Subsequent unregisters return false + assert!(!instrumentor.unregister("service")); + assert!(!instrumentor.unregister("service")); + } + } diff --git a/buswatch-sdk/src/state.rs b/buswatch-sdk/src/state.rs index 4ae55ea..b7a01af 100644 --- a/buswatch-sdk/src/state.rs +++ b/buswatch-sdk/src/state.rs @@ -220,6 +220,17 @@ impl GlobalState { .clone() } + /// Unregister a module and remove it from internal state. + /// + /// Returns `true` if the module was found and removed, `false` if it didn't exist. + /// + /// Note: This does not affect global topic write counters, as those are shared + /// across all modules and used for backlog computation. + pub fn unregister_module(&self, name: &str) -> bool { + let mut modules = self.modules.write(); + modules.remove(name).is_some() + } + /// Collect all modules into a Snapshot. pub fn collect(&self) -> Snapshot { let modules = self.modules.read(); @@ -589,4 +600,128 @@ mod tests { let rate = compute_rate(None, 100, now); assert!(rate.is_none()); } + + #[test] + fn unregister_module_removes_module() { + let global = GlobalState::default(); + + let module1 = global.register_module("service-a"); + let module2 = global.register_module("service-b"); + + module1 + .get_or_create_read("topic") + .count + .fetch_add(10, Ordering::Relaxed); + module2 + .get_or_create_read("topic") + .count + .fetch_add(20, Ordering::Relaxed); + + // Verify both modules exist + let snapshot = global.collect(); + assert_eq!(snapshot.modules.len(), 2); + + // Unregister service-a + let removed = global.unregister_module("service-a"); + assert!(removed, "Should return true when module is removed"); + + // Verify only service-b remains + let snapshot = global.collect(); + assert_eq!(snapshot.modules.len(), 1); + assert!(snapshot.modules.contains_key("service-b")); + assert!(!snapshot.modules.contains_key("service-a")); + } + + #[test] + fn unregister_nonexistent_module_returns_false() { + let global = GlobalState::default(); + + let _module = global.register_module("service-a"); + + // Try to unregister a module that doesn't exist + let removed = global.unregister_module("service-b"); + assert!(!removed, "Should return false when module doesn't exist"); + + // service-a should still be there + let snapshot = global.collect(); + assert_eq!(snapshot.modules.len(), 1); + assert!(snapshot.modules.contains_key("service-a")); + } + + #[test] + fn unregister_module_does_not_affect_global_write_counters() { + let global = GlobalState::default(); + + let producer = global.register_module("producer"); + producer + .get_or_create_write("events") + .count + .fetch_add(100, Ordering::Relaxed); + global + .get_topic_write_counter("events") + .fetch_add(100, Ordering::Relaxed); + + // Verify global counter is set + let counter = global.get_topic_write_counter("events"); + assert_eq!(counter.load(Ordering::Relaxed), 100); + + // Unregister the producer + global.unregister_module("producer"); + + // Global counter should still be 100 + let counter = global.get_topic_write_counter("events"); + assert_eq!(counter.load(Ordering::Relaxed), 100); + } + + #[test] + fn can_reregister_after_unregister() { + let global = GlobalState::default(); + + // Register and record some metrics + let module1 = global.register_module("service"); + module1 + .get_or_create_read("topic") + .count + .fetch_add(50, Ordering::Relaxed); + + let snapshot = global.collect(); + assert_eq!(snapshot.modules.get("service").unwrap().reads.get("topic").unwrap().count, 50); + + // Unregister + global.unregister_module("service"); + + // Re-register and verify metrics start fresh + let module2 = global.register_module("service"); + let snapshot = global.collect(); + + // New module should have fresh state (count = 0) + let metrics = snapshot.modules.get("service").unwrap(); + assert_eq!(metrics.reads.len(), 0, "Re-registered module should start with no topics"); + + // Add some new metrics + module2 + .get_or_create_read("topic") + .count + .fetch_add(10, Ordering::Relaxed); + + let snapshot = global.collect(); + assert_eq!(snapshot.modules.get("service").unwrap().reads.get("topic").unwrap().count, 10); + } + + #[test] + fn unregister_module_multiple_times_is_safe() { + let global = GlobalState::default(); + + global.register_module("service"); + + // First unregister should succeed + assert!(global.unregister_module("service")); + + // Second unregister should return false + assert!(!global.unregister_module("service")); + + // Third time for good measure + assert!(!global.unregister_module("service")); + } + }