warpdrive_proxy/cache/

invalidation.rs

//! Cache invalidation via PostgreSQL LISTEN/NOTIFY
//!
//! This module provides distributed cache invalidation using PostgreSQL's
//! LISTEN/NOTIFY feature. When one WarpDrive instance invalidates a cache entry,
//! all other instances are notified via PostgreSQL and can clear their local caches.
//!
//! # Architecture
//!
//! ```text
//! Instance A                    Instance B
//! ┌────────┐                    ┌────────┐
//! │ DELETE │                    │        │
//! │  key   │                    │        │
//! └───┬────┘                    └───▲────┘
//!     │                             │
//!     │    ┌──────────────┐         │
//!     └───►│  PostgreSQL  │─────────┘
//!          │    NOTIFY    │
//!          │   "key=foo"  │
//!          └──────────────┘
//! ```
//!
//! # Example
//!
//! ```no_run
//! use warpdrive::cache::coordinator::CacheCoordinator;
//! use warpdrive::cache::invalidation::InvalidationListener;
//! use warpdrive::config::Config;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let config = Config::from_env()?;
//! let cache = CacheCoordinator::from_config(&config).await?;
//!
//! // Start listener in background (if PostgreSQL configured)
//! if let Some(db_url) = &config.database_url {
//!     InvalidationListener::spawn(
//!         db_url.clone(),
//!         config.pg_channel_cache_invalidation.clone(),
//!         cache.clone(),
//!     );
//! }
//! # Ok(())
//! # }
//! ```

use futures::StreamExt;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::task::JoinHandle;
use tracing::{error, info, warn};

use crate::cache::Cache;
use crate::cache::coordinator::CacheCoordinator;
use crate::postgres::PgListener;

/// Cache invalidation message format
///
/// Sent via PostgreSQL NOTIFY when a cache entry is invalidated.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvalidationMessage {
    /// Cache key to invalidate
    pub key: String,

    /// Source instance ID (for debugging/metrics)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source: Option<String>,

    /// Reason for invalidation (for debugging)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
}

impl InvalidationMessage {
    /// Create a new invalidation message
    pub fn new(key: String) -> Self {
        InvalidationMessage {
            key,
            source: None,
            reason: None,
        }
    }

    /// Create invalidation message with metadata
    pub fn with_metadata(key: String, source: String, reason: String) -> Self {
        InvalidationMessage {
            key,
            source: Some(source),
            reason: Some(reason),
        }
    }
}

/// Cache invalidation listener
///
/// Subscribes to PostgreSQL NOTIFY events and invalidates local cache entries.
pub struct InvalidationListener {
    _handle: JoinHandle<()>,
}

impl InvalidationListener {
    /// Spawn a cache invalidation listener in the background
    ///
    /// Creates a PostgreSQL listener and processes invalidation messages.
    /// The listener runs in a background task and will continue until the
    /// application exits or an unrecoverable error occurs.
    ///
    /// # Arguments
    ///
    /// * `database_url` - PostgreSQL connection string
    /// * `channel` - Channel name to listen on (e.g., "warpdrive:cache:invalidate")
    /// * `cache` - Cache coordinator to invalidate entries from
    ///
    /// # Panics
    ///
    /// Does not panic. Errors are logged but the listener will attempt to reconnect.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::cache::coordinator::CacheCoordinator;
    /// # use warpdrive::cache::invalidation::InvalidationListener;
    /// # async fn example(cache: CacheCoordinator) {
    /// let listener = InvalidationListener::spawn(
    ///     "postgresql://localhost/warpdrive".to_string(),
    ///     "warpdrive:cache:invalidate".to_string(),
    ///     cache,
    /// );
    /// # }
    /// ```
    pub fn spawn(
        database_url: String,
        channel: String,
        cache: CacheCoordinator,
    ) -> InvalidationListener {
        let handle =
            tokio::spawn(
                async move { Self::listen_loop(database_url, channel, Arc::new(cache)).await },
            );

        InvalidationListener { _handle: handle }
    }

    /// Main listen loop
    ///
    /// Connects to PostgreSQL, subscribes to the channel, and processes messages.
    /// Automatically reconnects on connection loss.
    async fn listen_loop(database_url: String, channel: String, cache: Arc<CacheCoordinator>) {
        info!(
            database_url = %database_url,
            channel = %channel,
            "Starting cache invalidation listener"
        );

        loop {
            match Self::try_listen(&database_url, &channel, Arc::clone(&cache)).await {
                Ok(()) => {
                    warn!("Cache invalidation listener exited cleanly, restarting...");
                }
                Err(e) => {
                    error!(
                        error = %e,
                        "Cache invalidation listener error, retrying in 5s..."
                    );
                    tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                }
            }
        }
    }

    /// Attempt to listen for invalidation messages
    ///
    /// Returns when connection is lost or an unrecoverable error occurs.
    async fn try_listen(
        database_url: &str,
        channel: &str,
        cache: Arc<CacheCoordinator>,
    ) -> anyhow::Result<()> {
        // Create listener
        let mut listener = PgListener::from_url(database_url, vec![channel.to_string()]).await?;

        info!(channel = %channel, "Cache invalidation listener connected");

        // Process messages
        let mut stream = listener.stream();
        while let Some(notification) = stream.next().await {
            if let Err(e) = Self::handle_notification(notification, Arc::clone(&cache)).await {
                error!(
                    error = %e,
                    "Failed to handle cache invalidation notification"
                );
                // Continue processing other notifications
            }
        }

        Ok(())
    }

    /// Handle a single invalidation notification
    async fn handle_notification(
        notification: crate::postgres::PgNotification,
        cache: Arc<CacheCoordinator>,
    ) -> anyhow::Result<()> {
        // Parse message
        let message: InvalidationMessage = notification.parse_payload()?;

        info!(
            key = %message.key,
            source = ?message.source,
            reason = ?message.reason,
            "Received cache invalidation notification"
        );

        // Invalidate cache entry
        cache.delete(&message.key).await?;

        // Update metrics
        crate::metrics::CACHE_INVALIDATIONS
            .with_label_values(&["pg_notify"])
            .inc();

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_invalidation_message_new() {
        let msg = InvalidationMessage::new("test_key".to_string());
        assert_eq!(msg.key, "test_key");
        assert!(msg.source.is_none());
        assert!(msg.reason.is_none());
    }

    #[test]
    fn test_invalidation_message_with_metadata() {
        let msg = InvalidationMessage::with_metadata(
            "test_key".to_string(),
            "instance-1".to_string(),
            "expired".to_string(),
        );
        assert_eq!(msg.key, "test_key");
        assert_eq!(msg.source, Some("instance-1".to_string()));
        assert_eq!(msg.reason, Some("expired".to_string()));
    }

    #[test]
    fn test_invalidation_message_serialization() {
        let msg = InvalidationMessage::new("test_key".to_string());
        let json = serde_json::to_string(&msg).unwrap();
        let parsed: InvalidationMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.key, "test_key");
    }

    #[test]
    fn test_invalidation_message_with_metadata_serialization() {
        let msg = InvalidationMessage::with_metadata(
            "test_key".to_string(),
            "instance-1".to_string(),
            "expired".to_string(),
        );
        let json = serde_json::to_string(&msg).unwrap();
        let parsed: InvalidationMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.key, "test_key");
        assert_eq!(parsed.source, Some("instance-1".to_string()));
        assert_eq!(parsed.reason, Some("expired".to_string()));
    }
}