refactor(targets): unify queue/connectivity handling and coverage (#2953)

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: marshawcoco <marshawcoco@gmail.com>
This commit is contained in:
houseme
2026-05-14 12:31:23 +08:00
committed by GitHub
parent bdb98598d2
commit 81754d80b3
64 changed files with 9613 additions and 2800 deletions
+191 -422
View File
@@ -12,16 +12,14 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::{AuditEntry, AuditError, AuditRegistry, AuditResult, observability};
use hashbrown::HashMap;
use rustfs_ecstore::config::Config;
use rustfs_targets::{
StoreError, Target, TargetError,
store::{Key, Store},
target::{EntityTarget, QueuedPayload},
use crate::{
AuditEntry, AuditError, AuditRegistry, AuditResult, observability,
pipeline::{AuditPipeline, AuditRuntimeFacade, AuditRuntimeView},
};
use rustfs_ecstore::config::Config;
use rustfs_targets::{ReplayWorkerManager, Target};
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock, mpsc};
use tokio::sync::{Mutex, RwLock};
use tracing::{error, info, warn};
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -49,7 +47,7 @@ pub struct AuditSystem {
state: Arc<RwLock<AuditSystemState>>,
config: Arc<RwLock<Option<Config>>>,
/// Cancellation senders for active audit stream tasks (target_id -> cancel tx)
stream_cancellers: Arc<RwLock<HashMap<String, mpsc::Sender<()>>>>,
stream_cancellers: Arc<RwLock<ReplayWorkerManager>>,
}
impl Default for AuditSystem {
@@ -59,16 +57,68 @@ impl Default for AuditSystem {
}
impl AuditSystem {
fn pipeline(&self) -> AuditPipeline {
AuditPipeline::new(self.registry.clone())
}
fn runtime_view(&self) -> AuditRuntimeView {
AuditRuntimeView::new(self.registry.clone())
}
fn runtime_facade(&self) -> AuditRuntimeFacade {
AuditRuntimeFacade::new(self.registry.clone(), self.stream_cancellers.clone())
}
/// Creates a new audit system
pub fn new() -> Self {
Self {
registry: Arc::new(Mutex::new(AuditRegistry::new())),
state: Arc::new(RwLock::new(AuditSystemState::Stopped)),
config: Arc::new(RwLock::new(None)),
stream_cancellers: Arc::new(RwLock::new(HashMap::new())),
stream_cancellers: Arc::new(RwLock::new(ReplayWorkerManager::new())),
}
}
async fn create_targets_from_config(&self, config: &Config) -> AuditResult<Vec<Box<dyn Target<AuditEntry> + Send + Sync>>> {
let registry = self.registry.lock().await;
registry.create_audit_targets_from_config(config).await
}
async fn clear_runtime_targets(&self) -> AuditResult<()> {
{
let mut registry = self.registry.lock().await;
let mut replay_workers = self.stream_cancellers.write().await;
self.runtime_facade()
.shutdown_runtime(&mut registry, &mut replay_workers)
.await?;
}
let mut state = self.state.write().await;
*state = AuditSystemState::Stopped;
Ok(())
}
async fn commit_runtime_targets(
&self,
targets: Vec<Box<dyn Target<AuditEntry> + Send + Sync>>,
final_state: AuditSystemState,
) -> AuditResult<()> {
if targets.is_empty() {
info!("No enabled audit targets found, keeping audit system stopped");
self.clear_runtime_targets().await?;
return Ok(());
}
info!(target_count = targets.len(), "Created audit targets successfully");
let activation = self.runtime_facade().activate_targets_with_replay(targets).await;
self.runtime_facade().replace_targets(activation).await?;
let mut state = self.state.write().await;
*state = final_state;
Ok(())
}
/// Starts the audit system with the given configuration
///
/// # Arguments
@@ -103,31 +153,14 @@ impl AuditSystem {
*config_guard = Some(config.clone());
}
// Create targets from configuration
let mut registry = self.registry.lock().await;
match registry.create_audit_targets_from_config(&config).await {
match self.create_targets_from_config(&config).await {
Ok(targets) => {
if targets.is_empty() {
info!("No enabled audit targets found, keeping audit system stopped");
drop(registry);
return Ok(());
}
{
let mut state = self.state.write().await;
*state = AuditSystemState::Starting;
}
info!(target_count = targets.len(), "Created audit targets successfully");
// Initialize all targets
for target in targets {
self.init_and_register_target(target, &mut registry).await;
}
// Update state to running
let mut state = self.state.write().await;
*state = AuditSystemState::Running;
self.commit_runtime_targets(targets, AuditSystemState::Running).await?;
info!("Audit system started successfully");
Ok(())
}
@@ -207,18 +240,10 @@ impl AuditSystem {
info!("Stopping audit system");
// Stop all stream tasks first
self.stop_all_streams().await;
// Close all targets
let mut registry = self.registry.lock().await;
if let Err(e) = registry.close_all().await {
if let Err(e) = self.clear_runtime_targets().await {
error!(error = %e, "Failed to close some audit targets");
}
// Update state to stopped
let mut state = self.state.write().await;
*state = AuditSystemState::Stopped;
// Clear configuration
let mut config_guard = self.config.write().await;
*config_guard = None;
@@ -248,8 +273,6 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn dispatch(&self, entry: Arc<AuditEntry>) -> AuditResult<()> {
let start_time = std::time::Instant::now();
let state = self.state.read().await;
match *state {
@@ -262,77 +285,7 @@ impl AuditSystem {
}
}
drop(state);
// Collect cloned targets under lock, then dispatch without holding it
let targets: Vec<(String, Box<dyn Target<AuditEntry> + Send + Sync>)> = {
let registry = self.registry.lock().await;
let target_keys = registry.list_targets();
if target_keys.is_empty() {
warn!("No audit targets configured for dispatch");
return Ok(());
}
target_keys
.into_iter()
.filter_map(|key| registry.get_target(&key).map(|t| (key, t.clone_dyn())))
.collect()
};
// Dispatch to all targets concurrently (no lock held)
let mut tasks = Vec::new();
for (target_key, target) in targets {
let entity_target = EntityTarget {
object_name: entry.api.name.clone().unwrap_or_default(),
bucket_name: entry.api.bucket.clone().unwrap_or_default(),
event_name: entry.event,
data: (*entry).clone(),
};
let task = async move {
let result = target.save(Arc::new(entity_target)).await;
(target_key, result)
};
tasks.push(task);
}
// Execute all dispatch tasks
let results = futures::future::join_all(tasks).await;
let mut errors = Vec::new();
let mut success_count = 0;
for (target_key, result) in results {
match result {
Ok(_) => {
success_count += 1;
observability::record_target_success();
}
Err(e) => {
error!(target_id = %target_key, error = %e, "Failed to dispatch audit log to target");
errors.push(e);
observability::record_target_failure();
}
}
}
let dispatch_time = start_time.elapsed();
if errors.is_empty() {
observability::record_audit_success(dispatch_time);
} else {
observability::record_audit_failure(dispatch_time);
// Log errors but don't fail the entire dispatch
warn!(
error_count = errors.len(),
success_count = success_count,
"Some audit targets failed to receive log entry"
);
}
Ok(())
self.pipeline().dispatch(entry).await
}
/// Dispatches a batch of audit log entries to all active targets
@@ -343,239 +296,12 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn dispatch_batch(&self, entries: Vec<Arc<AuditEntry>>) -> AuditResult<()> {
let start_time = std::time::Instant::now();
let state = self.state.read().await;
if *state != AuditSystemState::Running {
return Err(AuditError::NotInitialized("Audit system is not running".to_string()));
}
drop(state);
// Collect targets under lock, then dispatch without holding it
let targets: Vec<(String, Box<dyn Target<AuditEntry> + Send + Sync>)> = {
let registry = self.registry.lock().await;
let target_keys = registry.list_targets();
if target_keys.is_empty() {
warn!("No audit targets configured for batch dispatch");
return Ok(());
}
target_keys
.into_iter()
.filter_map(|key| registry.get_target(&key).map(|t| (key, t.clone_dyn())))
.collect()
};
let mut tasks = Vec::new();
for (target_key, target) in targets {
let entries_clone: Vec<_> = entries.iter().map(Arc::clone).collect();
let target_key_clone = target_key.clone();
let task = async move {
let mut success_count = 0;
let mut errors = Vec::new();
for entry in entries_clone {
let entity_target = EntityTarget {
object_name: entry.api.name.clone().unwrap_or_default(),
bucket_name: entry.api.bucket.clone().unwrap_or_default(),
event_name: entry.event,
data: (*entry).clone(),
};
match target.save(Arc::new(entity_target)).await {
Ok(_) => success_count += 1,
Err(e) => errors.push(e),
}
}
(target_key_clone, success_count, errors)
};
tasks.push(task);
}
let results = futures::future::join_all(tasks).await;
let mut total_success = 0;
let mut total_errors = 0;
for (_target_id, success_count, errors) in results {
total_success += success_count;
total_errors += errors.len();
for e in errors {
error!("Batch dispatch error: {:?}", e);
}
}
let dispatch_time = start_time.elapsed();
info!(
"Batch dispatched {} entries, success: {}, errors: {}, time: {:?}",
entries.len(),
total_success,
total_errors,
dispatch_time
);
Ok(())
}
/// Stops all active audit stream tasks by sending cancellation signals.
async fn stop_all_streams(&self) {
let mut cancellers = self.stream_cancellers.write().await;
for (target_id, cancel_tx) in cancellers.drain() {
info!(target_id = %target_id, "Stopping audit stream");
let _ = cancel_tx.send(()).await;
}
}
/// Initializes a single target: runs init(), starts stream if store is present,
/// and adds it to the registry. For store-backed targets, registration and stream
/// startup proceed even if init() fails so queued entries can be drained later.
async fn init_and_register_target(
&self,
target: Box<dyn Target<AuditEntry> + Send + Sync>,
registry: &mut AuditRegistry,
) -> Option<String> {
let target_id = target.id().to_string();
let has_store = target.store().is_some();
if let Err(e) = target.init().await {
error!(target_id = %target_id, error = %e, "Failed to initialize audit target");
// Non-store targets: init failure is fatal.
if !has_store {
return None;
}
// Store-backed targets: still register and start the stream so queued
// entries can be drained when connectivity recovers.
warn!(
target_id = %target_id,
"Proceeding with store-backed audit target despite init failure"
);
}
if target.is_enabled() {
if let Some(store) = target.store() {
info!(target_id = %target_id, "Start audit stream processing for target");
let store_clone: Box<dyn Store<QueuedPayload, Error = StoreError, Key = Key> + Send> = store.boxed_clone();
let target_arc: Arc<dyn Target<AuditEntry> + Send + Sync> = Arc::from(target.clone_dyn());
let cancel_tx = self.start_audit_stream_with_batching(store_clone, target_arc);
self.stream_cancellers.write().await.insert(target_id.clone(), cancel_tx);
info!(target_id = %target_id, "Audit stream processing started");
} else {
info!(target_id = %target_id, "No store configured, skip audit stream processing");
}
} else {
info!(target_id = %target_id, "Target disabled, skip audit stream processing");
}
registry.add_target(target_id.clone(), target);
Some(target_id)
}
/// Starts the audit stream processing for a target with batching and retry logic
///
/// # Arguments
/// * `store` - The store from which to read audit entries
/// * `target` - The target to which audit entries will be sent
///
/// This function spawns a background task that continuously reads audit entries from the provided store
/// and attempts to send them to the specified target. It implements retry logic with exponential backoff
fn start_audit_stream_with_batching(
&self,
store: Box<dyn Store<QueuedPayload, Error = StoreError, Key = Key> + Send>,
target: Arc<dyn Target<AuditEntry> + Send + Sync>,
) -> mpsc::Sender<()> {
let (cancel_tx, mut cancel_rx) = mpsc::channel(1);
let state = self.state.clone();
tokio::spawn(async move {
use std::time::Duration;
use tokio::time::sleep;
info!("Starting audit stream for target: {}", target.id());
const MAX_RETRIES: usize = 5;
const BASE_RETRY_DELAY: Duration = Duration::from_secs(2);
loop {
// Check for cancellation signal
if cancel_rx.try_recv().is_ok() {
info!("Audit stream cancelled for target: {}", target.id());
break;
}
match *state.read().await {
AuditSystemState::Running | AuditSystemState::Paused | AuditSystemState::Starting => {}
_ => {
info!("Audit stream stopped for target: {}", target.id());
break;
}
}
let keys: Vec<Key> = store.list();
if keys.is_empty() {
tokio::select! {
_ = sleep(Duration::from_millis(500)) => {},
_ = cancel_rx.recv() => {
info!("Audit stream cancelled during idle for target: {}", target.id());
return;
}
}
continue;
}
for key in keys {
if cancel_rx.try_recv().is_ok() {
info!("Audit stream cancelled during processing for target: {}", target.id());
return;
}
let mut retries = 0usize;
let mut success = false;
while retries < MAX_RETRIES && !success {
match target.send_from_store(key.clone()).await {
Ok(_) => {
info!("Successfully sent audit entry, target: {}, key: {}", target.id(), key.to_string());
observability::record_target_success();
success = true;
}
Err(e) => {
match &e {
TargetError::NotConnected => {
warn!("Target {} not connected, retrying...", target.id());
}
TargetError::Timeout(_) => {
warn!("Timeout sending to target {}, retrying...", target.id());
}
TargetError::Dropped(reason) => {
warn!("Dropped queued payload for target {}: {}", target.id(), reason);
observability::record_target_failure();
break;
}
_ => {
error!("Permanent error for target {}: {}", target.id(), e);
target.record_final_failure();
observability::record_target_failure();
break;
}
}
retries += 1;
let backoff = BASE_RETRY_DELAY * (1 << retries);
sleep(backoff).await;
}
}
}
if retries >= MAX_RETRIES && !success {
warn!("Max retries exceeded for key {}, target: {}, skipping", key.to_string(), target.id());
target.record_final_failure();
observability::record_target_failure();
}
}
sleep(Duration::from_millis(100)).await;
}
});
cancel_tx
self.pipeline().dispatch_batch(entries).await
}
/// Enables a specific target
@@ -586,15 +312,7 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn enable_target(&self, target_id: &str) -> AuditResult<()> {
// This would require storing enabled/disabled state per target
// For now, just check if target exists
let registry = self.registry.lock().await;
if registry.get_target(target_id).is_some() {
info!(target_id = %target_id, "Target enabled");
Ok(())
} else {
Err(AuditError::Configuration(format!("Target not found: {target_id}"), None))
}
self.runtime_view().enable_target(target_id).await
}
/// Disables a specific target
@@ -605,15 +323,7 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn disable_target(&self, target_id: &str) -> AuditResult<()> {
// This would require storing enabled/disabled state per target
// For now, just check if target exists
let registry = self.registry.lock().await;
if registry.get_target(target_id).is_some() {
info!(target_id = %target_id, "Target disabled");
Ok(())
} else {
Err(AuditError::Configuration(format!("Target not found: {target_id}"), None))
}
self.runtime_view().disable_target(target_id).await
}
/// Removes a target from the system
@@ -624,16 +334,7 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn remove_target(&self, target_id: &str) -> AuditResult<()> {
let mut registry = self.registry.lock().await;
if let Some(target) = registry.remove_target(target_id) {
if let Err(e) = target.close().await {
error!(target_id = %target_id, error = %e, "Failed to close removed target");
}
info!(target_id = %target_id, "Target removed");
Ok(())
} else {
Err(AuditError::Configuration(format!("Target not found: {target_id}"), None))
}
self.runtime_view().remove_target(target_id).await
}
/// Updates or inserts a target
@@ -645,23 +346,7 @@ impl AuditSystem {
/// # Returns
/// * `AuditResult<()>` - Result indicating success or failure
pub async fn upsert_target(&self, target_id: String, target: Box<dyn Target<AuditEntry> + Send + Sync>) -> AuditResult<()> {
let mut registry = self.registry.lock().await;
// Initialize the target
if let Err(e) = target.init().await {
return Err(AuditError::Target(e));
}
// Remove existing target if present
if let Some(old_target) = registry.remove_target(&target_id)
&& let Err(e) = old_target.close().await
{
error!(target_id = %target_id, error = %e, "Failed to close old target during upsert");
}
registry.add_target(target_id.clone(), target);
info!(target_id = %target_id, "Target upserted");
Ok(())
self.runtime_view().upsert_target(target_id, target).await
}
/// Lists all targets
@@ -669,33 +354,27 @@ impl AuditSystem {
/// # Returns
/// * `Vec<String>` - List of target IDs
pub async fn list_targets(&self) -> Vec<String> {
let registry = self.registry.lock().await;
registry.list_targets()
self.runtime_view().list_targets().await
}
/// Returns cloned target values for read-only runtime inspection.
pub async fn get_target_values(&self) -> Vec<Box<dyn Target<AuditEntry> + Send + Sync>> {
let registry = self.registry.lock().await;
registry.list_target_values()
pub async fn get_target_values(&self) -> Vec<rustfs_targets::SharedTarget<AuditEntry>> {
self.runtime_view().get_target_values().await
}
/// Returns per-target delivery metrics for Prometheus collection.
pub async fn snapshot_target_metrics(&self) -> Vec<AuditTargetMetricSnapshot> {
let targets = self.get_target_values().await;
let mut snapshots = Vec::with_capacity(targets.len());
self.pipeline().snapshot_target_metrics().await
}
for target in targets {
let delivery = target.delivery_snapshot();
snapshots.push(AuditTargetMetricSnapshot {
failed_messages: delivery.failed_messages,
queue_length: delivery.queue_length,
target_id: target.id().to_string(),
total_messages: delivery.total_messages,
});
}
pub async fn snapshot_target_health(&self) -> Vec<rustfs_targets::RuntimeTargetHealthSnapshot> {
self.pipeline().snapshot_target_health().await
}
snapshots.sort_by(|a, b| a.target_id.cmp(&b.target_id));
snapshots
pub async fn runtime_status_snapshot(&self) -> rustfs_targets::RuntimeStatusSnapshot {
let replay_workers = self.stream_cancellers.read().await;
let registry = self.registry.lock().await;
registry.runtime_manager().status_snapshot(&replay_workers)
}
/// Gets information about a specific target
@@ -706,8 +385,7 @@ impl AuditSystem {
/// # Returns
/// * `Option<String>` - Target ID if found
pub async fn get_target(&self, target_id: &str) -> Option<String> {
let registry = self.registry.lock().await;
registry.get_target(target_id).map(|target| target.id().to_string())
self.runtime_view().get_target(target_id).await
}
/// Reloads configuration and updates targets
@@ -722,32 +400,20 @@ impl AuditSystem {
observability::record_config_reload();
// Stop all existing stream tasks first
self.stop_all_streams().await;
// Store new configuration
{
let mut config_guard = self.config.write().await;
*config_guard = Some(new_config.clone());
}
// Close all existing targets
let mut registry = self.registry.lock().await;
if let Err(e) = registry.close_all().await {
error!(error = %e, "Failed to close existing targets during reload");
}
let final_state = match self.get_state().await {
AuditSystemState::Paused => AuditSystemState::Paused,
_ => AuditSystemState::Running,
};
// Create new targets from updated configuration
match registry.create_audit_targets_from_config(&new_config).await {
match self.create_targets_from_config(&new_config).await {
Ok(targets) => {
info!(target_count = targets.len(), "Reloaded audit targets successfully");
for target in targets {
if let Some(target_id) = self.init_and_register_target(target, &mut registry).await {
info!(target_id = %target_id, "Target initialized (reload)");
}
}
self.commit_runtime_targets(targets, final_state).await?;
info!("Audit configuration reloaded successfully");
Ok(())
}
@@ -779,3 +445,106 @@ impl AuditSystem {
observability::reset_metrics().await;
}
}
#[cfg(test)]
mod tests {
use super::{AuditSystem, AuditSystemState};
use async_trait::async_trait;
use rustfs_targets::ReplayWorkerManager;
use rustfs_targets::arn::TargetID;
use rustfs_targets::store::{Key, Store};
use rustfs_targets::target::{EntityTarget, QueuedPayload, QueuedPayloadMeta};
use rustfs_targets::{StoreError, Target, TargetError};
use serde::{Serialize, de::DeserializeOwned};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::mpsc;
#[derive(Clone)]
struct TestTarget {
close_calls: Arc<AtomicUsize>,
id: TargetID,
}
impl TestTarget {
fn new(id: &str, name: &str) -> Self {
Self {
close_calls: Arc::new(AtomicUsize::new(0)),
id: TargetID::new(id.to_string(), name.to_string()),
}
}
}
#[async_trait]
impl<E> Target<E> for TestTarget
where
E: Send + Sync + 'static + Clone + Serialize + DeserializeOwned,
{
fn id(&self) -> TargetID {
self.id.clone()
}
async fn is_active(&self) -> Result<bool, TargetError> {
Ok(true)
}
async fn save(&self, _event: Arc<EntityTarget<E>>) -> Result<(), TargetError> {
Ok(())
}
async fn send_raw_from_store(&self, _key: Key, _body: Vec<u8>, _meta: QueuedPayloadMeta) -> Result<(), TargetError> {
Ok(())
}
async fn close(&self) -> Result<(), TargetError> {
self.close_calls.fetch_add(1, Ordering::SeqCst);
Ok(())
}
fn store(&self) -> Option<&(dyn Store<QueuedPayload, Error = StoreError, Key = Key> + Send + Sync)> {
None
}
fn clone_dyn(&self) -> Box<dyn Target<E> + Send + Sync> {
Box::new(self.clone())
}
fn is_enabled(&self) -> bool {
true
}
}
#[tokio::test]
async fn reload_with_empty_config_stops_existing_runtime() {
let system = AuditSystem::new();
let target = TestTarget::new("primary", "webhook");
let close_calls = Arc::clone(&target.close_calls);
{
let mut registry = system.registry.lock().await;
registry.add_target("primary:webhook".to_string(), Box::new(target));
}
{
let mut state = system.state.write().await;
*state = AuditSystemState::Running;
}
{
let mut replay_workers = system.stream_cancellers.write().await;
let (cancel_tx, _cancel_rx) = mpsc::channel(1);
replay_workers.insert("primary:webhook".to_string(), cancel_tx);
assert_eq!(replay_workers.len(), 1);
}
system
.reload_config(rustfs_ecstore::config::Config(HashMap::new()))
.await
.expect("reload with empty config should succeed");
assert_eq!(system.get_state().await, AuditSystemState::Stopped);
assert!(system.list_targets().await.is_empty());
assert_eq!(system.runtime_status_snapshot().await, ReplayWorkerManager::new().snapshot(0));
assert_eq!(close_calls.load(Ordering::SeqCst), 1);
assert_eq!(*system.config.read().await, Some(rustfs_ecstore::config::Config(HashMap::new())));
}
}