procref/

service.rs

//! SharedService - High-level API for managing shared service lifecycles.

use crate::types::ServiceInfoFile;
use crate::{Error, PlatformRefCounter, RefCounter, Result, ServiceInfo};
use parking_lot::RwLock;
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;

/// Type alias for async callbacks.
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;

/// Callback for first acquire (service startup).
pub type OnFirstAcquire = Box<dyn Fn() -> BoxFuture<'static, Result<ServiceInfo>> + Send + Sync>;

/// Callback for last release (service shutdown).
pub type OnLastRelease =
    Box<dyn Fn(ServiceInfo) -> BoxFuture<'static, Result<()>> + Send + Sync>;

/// Callback for health check.
pub type OnHealthCheck = Box<dyn Fn(&ServiceInfo) -> BoxFuture<'static, bool> + Send + Sync>;

/// Callback for recovery.
pub type OnRecover =
    Box<dyn Fn(ServiceInfo) -> BoxFuture<'static, Result<ServiceInfo>> + Send + Sync>;

/// A handle to a shared service.
///
/// Dropping this handle releases the reference.
/// If it's the last reference, the service may be shut down.
pub struct ServiceHandle {
    service: Arc<SharedServiceInner>,
    info: ServiceInfo,
}

impl ServiceHandle {
    /// Get information about the service.
    pub fn info(&self) -> &ServiceInfo {
        &self.info
    }

    /// Get the service port.
    pub fn port(&self) -> u16 {
        self.info.port()
    }

    /// Get the service PID.
    pub fn pid(&self) -> u32 {
        self.info.pid()
    }
}

impl Drop for ServiceHandle {
    fn drop(&mut self) {
        // Release is handled by the service
        // We can't do async in drop, so we spawn a task or use sync release
        let _ = self.service.release_sync();
    }
}

/// Inner state for SharedService.
struct SharedServiceInner {
    name: String,
    ref_counter: PlatformRefCounter,
    info_path: PathBuf,
    current_info: RwLock<Option<ServiceInfo>>,
    on_first_acquire: Option<OnFirstAcquire>,
    on_last_release: Option<OnLastRelease>,
    on_health_check: Option<OnHealthCheck>,
    on_recover: Option<OnRecover>,
}

impl SharedServiceInner {
    /// Synchronous release (for Drop).
    fn release_sync(&self) -> Result<()> {
        let count = self.ref_counter.release()?;

        if count == 0 {
            // We're the last client
            if let Some(ref _callback) = self.on_last_release {
                if let Some(info) = self.current_info.read().clone() {
                    // We can't call async in sync context easily
                    // For now, just log. Proper async drop needs runtime support.
                    tracing::info!(
                        "Last client released, service {} should be stopped",
                        self.name
                    );

                    // Try to stop the process directly
                    crate::process::stop(info.pid(), 5000);

                    // Clean up info file
                    let _ = std::fs::remove_file(&self.info_path);
                }
            }
        }

        Ok(())
    }
}

/// Shared service manager.
///
/// Manages a service that is shared across multiple processes.
/// Uses kernel-level reference counting to track clients.
pub struct SharedService {
    inner: Arc<SharedServiceInner>,
}

impl SharedService {
    /// Create a new builder for SharedService.
    pub fn builder(name: &str) -> SharedServiceBuilder {
        SharedServiceBuilder::new(name)
    }

    /// Acquire a reference to the service.
    ///
    /// If this is the first client, `on_first_acquire` is called to start the service.
    /// Otherwise, the existing service info is returned after health check.
    pub async fn acquire(&self) -> Result<ServiceHandle> {
        // Step 1: Increment reference count
        let count = self.inner.ref_counter.acquire()?;
        tracing::debug!("Acquired reference, count={}", count);

        // Step 2: Check if we're the first client
        if count == 1 {
            // Try to get startup lock
            if self.inner.ref_counter.try_lock()? {
                // We have the lock, start the service
                let info = self.start_service().await?;
                self.inner.ref_counter.unlock()?;
                return Ok(ServiceHandle {
                    service: Arc::clone(&self.inner),
                    info,
                });
            }
            // Someone else has the lock, wait for them to finish
            self.wait_for_service().await?;
        }

        // Step 3: Get existing service info
        let info = self.get_or_recover_service().await?;

        Ok(ServiceHandle {
            service: Arc::clone(&self.inner),
            info,
        })
    }

    /// Start the service (called when we're the first client).
    async fn start_service(&self) -> Result<ServiceInfo> {
        tracing::info!("Starting service {}", self.inner.name);

        let info = if let Some(ref callback) = self.inner.on_first_acquire {
            callback().await?
        } else {
            return Err(Error::ServiceStart(
                "No on_first_acquire callback registered".to_string(),
            ));
        };

        // Save service info to file
        self.save_info(&info)?;
        *self.inner.current_info.write() = Some(info.clone());

        Ok(info)
    }

    /// Wait for another process to start the service.
    async fn wait_for_service(&self) -> Result<()> {
        let start = std::time::Instant::now();
        let timeout = std::time::Duration::from_secs(30);

        while start.elapsed() < timeout {
            if self.inner.info_path.exists() {
                if let Ok(info) = self.load_info() {
                    if info.is_alive() {
                        return Ok(());
                    }
                }
            }
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        }

        Err(Error::ServiceStart(
            "Timeout waiting for service to start".to_string(),
        ))
    }

    /// Get existing service info, or recover if unhealthy.
    async fn get_or_recover_service(&self) -> Result<ServiceInfo> {
        // Try to load from file
        let info = match self.load_info() {
            Ok(info) => info,
            Err(_) => {
                // No info file, try to recover
                return self.recover_service(None).await;
            }
        };

        // Check if healthy
        let is_healthy = if let Some(ref check) = self.inner.on_health_check {
            check(&info).await
        } else {
            // Default health check: process is alive
            info.is_alive()
        };

        if is_healthy {
            *self.inner.current_info.write() = Some(info.clone());
            Ok(info)
        } else {
            // Try to recover
            self.recover_service(Some(info)).await
        }
    }

    /// Recover an unhealthy service.
    async fn recover_service(&self, old_info: Option<ServiceInfo>) -> Result<ServiceInfo> {
        tracing::warn!("Service {} needs recovery", self.inner.name);

        // Try to get startup lock for recovery
        if !self.inner.ref_counter.try_lock()? {
            // Someone else is recovering, wait
            return self.wait_for_service().await.and_then(|_| {
                self.load_info()
            });
        }

        let info = if let Some(ref callback) = self.inner.on_recover {
            if let Some(old) = old_info {
                callback(old).await?
            } else if let Some(ref start) = self.inner.on_first_acquire {
                start().await?
            } else {
                self.inner.ref_counter.unlock()?;
                return Err(Error::ServiceRecovery(
                    "No recovery or startup callback".to_string(),
                ));
            }
        } else if let Some(ref start) = self.inner.on_first_acquire {
            // Fall back to restart
            if let Some(old) = old_info {
                crate::process::stop(old.pid(), 5000);
            }
            start().await?
        } else {
            self.inner.ref_counter.unlock()?;
            return Err(Error::ServiceRecovery(
                "No recovery or startup callback".to_string(),
            ));
        };

        self.save_info(&info)?;
        *self.inner.current_info.write() = Some(info.clone());
        self.inner.ref_counter.unlock()?;

        Ok(info)
    }

    /// Save service info to file.
    fn save_info(&self, info: &ServiceInfo) -> Result<()> {
        if let Some(parent) = self.inner.info_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let file_info = ServiceInfoFile::from(info);
        let content = serde_json::to_string_pretty(&file_info)
            .map_err(|e| Error::ServiceInfo(format!("Serialization failed: {}", e)))?;

        std::fs::write(&self.inner.info_path, content)?;
        Ok(())
    }

    /// Load service info from file.
    fn load_info(&self) -> Result<ServiceInfo> {
        let content = std::fs::read_to_string(&self.inner.info_path)?;
        let file_info: ServiceInfoFile = serde_json::from_str(&content)
            .map_err(|e| Error::ServiceInfo(format!("Deserialization failed: {}", e)))?;
        Ok(ServiceInfo::from(file_info))
    }

    /// Get current reference count.
    pub fn count(&self) -> Result<u32> {
        self.inner.ref_counter.count()
    }

    /// Get service name.
    pub fn name(&self) -> &str {
        &self.inner.name
    }
}

/// Builder for SharedService.
pub struct SharedServiceBuilder {
    name: String,
    base_dir: Option<PathBuf>,
    on_first_acquire: Option<OnFirstAcquire>,
    on_last_release: Option<OnLastRelease>,
    on_health_check: Option<OnHealthCheck>,
    on_recover: Option<OnRecover>,
}

impl SharedServiceBuilder {
    /// Create a new builder.
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            base_dir: None,
            on_first_acquire: None,
            on_last_release: None,
            on_health_check: None,
            on_recover: None,
        }
    }

    /// Set base directory for service info files.
    pub fn base_dir(mut self, dir: impl Into<PathBuf>) -> Self {
        self.base_dir = Some(dir.into());
        self
    }

    /// Set callback for first client (service startup).
    ///
    /// This is called when the first client acquires a reference.
    /// It should start the service and return ServiceInfo.
    pub fn on_first_acquire<F, Fut>(mut self, f: F) -> Self
    where
        F: Fn() -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ServiceInfo>> + Send + 'static,
    {
        self.on_first_acquire = Some(Box::new(move || Box::pin(f())));
        self
    }

    /// Set callback for last client (service shutdown).
    ///
    /// This is called when the last client releases their reference.
    /// It should stop the service.
    pub fn on_last_release<F, Fut>(mut self, f: F) -> Self
    where
        F: Fn(ServiceInfo) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        self.on_last_release = Some(Box::new(move |info| Box::pin(f(info))));
        self
    }

    /// Set callback for health check.
    ///
    /// This is called to verify the service is healthy.
    /// Return true if healthy, false if recovery is needed.
    pub fn on_health_check<F, Fut>(mut self, f: F) -> Self
    where
        F: Fn(&ServiceInfo) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = bool> + Send + 'static,
    {
        let f = Arc::new(f);
        self.on_health_check = Some(Box::new(move |info| {
            let info = info.clone();
            let f = Arc::clone(&f);
            Box::pin(async move { f(&info).await })
        }));
        self
    }

    /// Set callback for recovery.
    ///
    /// This is called when health check fails.
    /// It should recover the service and return new ServiceInfo.
    pub fn on_recover<F, Fut>(mut self, f: F) -> Self
    where
        F: Fn(ServiceInfo) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<ServiceInfo>> + Send + 'static,
    {
        self.on_recover = Some(Box::new(move |info| Box::pin(f(info))));
        self
    }

    /// Build the SharedService.
    pub fn build(self) -> Result<SharedService> {
        let base_dir = self.base_dir.unwrap_or_else(|| {
            dirs::home_dir()
                .unwrap_or_else(|| PathBuf::from("."))
                .join(".procref")
        });

        std::fs::create_dir_all(&base_dir)?;

        let info_path = base_dir.join(format!("{}.json", self.name));
        let ref_counter = PlatformRefCounter::new(&self.name)?;

        let inner = SharedServiceInner {
            name: self.name,
            ref_counter,
            info_path,
            current_info: RwLock::new(None),
            on_first_acquire: self.on_first_acquire,
            on_last_release: self.on_last_release,
            on_health_check: self.on_health_check,
            on_recover: self.on_recover,
        };

        Ok(SharedService {
            inner: Arc::new(inner),
        })
    }
}