swactor/crates/data-plane/src/control.rs

//! Reusable namespace service lifecycle and public file-registration control.

use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use swactor::actor::ActorAddress;
use swactor::runtime::Runtime;
use swactor_engine::EngineHandle;

use crate::blob::FileRegistration;
use crate::blob_transfer::BlobTransferSender;
use crate::host::HostRouteRegistrar;
use crate::namespace::{
    DataDirectoryActor, DirectoryClient, NamespaceError, OperationId, RetirementRetry,
};
use crate::namespace_store::SourceRecovery;
use crate::path::DataPath;
use crate::source::{BlobSourceIn, BlobSourcePublisher, FileBlobSourceActor};

/// How often the directory re-sends retirements that no source has
/// acknowledged yet. Retire frames are at-most-once; this retry bounds how
/// long a lost frame can strand a published source and its binding.
const RETIREMENT_RETRY_PERIOD: Duration = Duration::from_millis(250);

pub struct DataNamespaceService {
    directory: ActorAddress,
    control: DataPlaneControl,
    authority_epoch: u64,
}

impl DataNamespaceService {
    pub fn recover(
        runtime: Runtime,
        engine: EngineHandle,
        store_path: impl AsRef<Path>,
        source_sender: Arc<dyn BlobTransferSender>,
        source_publisher: Arc<dyn BlobSourcePublisher>,
        routes: Option<Arc<dyn HostRouteRegistrar>>,
    ) -> Result<Self, NamespaceError> {
        let recovery_runtime = runtime.clone();
        let recovery_sender = Arc::clone(&source_sender);
        let recovery_publisher = Arc::clone(&source_publisher);
        let retire_retry = RetirementRetry::new(
            engine,
            runtime.create_sender(),
            RETIREMENT_RETRY_PERIOD,
            routes,
        );
        let directory = DataDirectoryActor::recover(
            store_path,
            Some(retire_retry),
            move |recovery, expected_length| match recovery {
                SourceRecovery::File { path } => {
                    let source = FileBlobSourceActor::recover(
                        recovery_runtime.clone(),
                        Arc::clone(&recovery_sender),
                        path,
                        expected_length,
                    )?;
                    let source = recovery_runtime.spawn(source).map_err(|error| {
                        NamespaceError::SourceRecovery(format!(
                            "spawn recovered file source {}: {error}",
                            path.display()
                        ))
                    })?;
                    let node = match recovery_publisher.publish_source(source) {
                        Ok(node) => node,
                        Err(error) => {
                            let _ = recovery_runtime
                                .send_to(source, BlobSourceIn::Retire { reply_to: None });
                            return Err(NamespaceError::SourceRecovery(error));
                        }
                    };
                    Ok((source, node))
                }
                SourceRecovery::Actor { actor, node, .. } => Ok((*actor, *node)),
            },
        )?;
        let authority_epoch = directory.authority_epoch();
        let directory = runtime.spawn(directory).map_err(|error| {
            NamespaceError::SourceRecovery(format!("spawn data directory: {error}"))
        })?;
        let control = DataPlaneControl {
            runtime: runtime.clone(),
            directory: DirectoryClient::new(runtime, directory),
            source_sender,
            source_publisher,
        };
        Ok(Self {
            directory,
            authority_epoch,
            control,
        })
    }

    pub fn directory(&self) -> ActorAddress {
        self.directory
    }

    pub fn control(&self) -> DataPlaneControl {
        self.control.clone()
    }

    pub fn authority_epoch(&self) -> u64 {
        self.authority_epoch
    }
}

#[derive(Clone)]
pub struct DataPlaneControl {
    runtime: Runtime,
    directory: DirectoryClient,
    source_sender: Arc<dyn BlobTransferSender>,
    source_publisher: Arc<dyn BlobSourcePublisher>,
}

pub enum RetainedNamespaceResources {
    Absent,
    QuiescentStream { revision: u64 },
    Blob(crate::namespace::BlobBinding),
}

pub enum NamespaceCleanupStatus {
    Absent,
    ActiveStream,
    Removed {
        quiescent_stream_revision: Option<u64>,
    },
}

impl DataPlaneControl {
    pub async fn retained_blob_resources(
        &self,
        path: DataPath,
    ) -> Result<RetainedNamespaceResources, NamespaceError> {
        let node = match self.directory.lookup(path.clone()).await {
            Ok(node) => node,
            Err(NamespaceError::PathNotFound(_)) => return Ok(RetainedNamespaceResources::Absent),
            Err(error) => return Err(error),
        };
        if node.kind == crate::namespace::EntryKind::Stream {
            return if node.active {
                Err(NamespaceError::Protocol(
                    "retained stream is still active".to_owned(),
                ))
            } else {
                Ok(RetainedNamespaceResources::QuiescentStream {
                    revision: node.revision,
                })
            };
        }
        let binding = self.directory.resolve(path).await?;
        if binding.revision != node.revision {
            return Err(NamespaceError::Protocol(
                "retained blob changed during ownership lookup".to_owned(),
            ));
        }
        Ok(RetainedNamespaceResources::Blob(binding))
    }

    pub async fn ensure(
        &self,
        path: DataPath,
        registration: FileRegistration,
    ) -> Result<(), NamespaceError> {
        match self.directory.resolve(path.clone()).await {
            Ok(_) => Ok(()),
            Err(NamespaceError::PathNotFound(_)) | Err(NamespaceError::SourceRecovery(_)) => {
                self.register(path, registration).await
            }
            Err(error) => Err(error),
        }
    }

    pub async fn register(
        &self,
        path: DataPath,
        registration: FileRegistration,
    ) -> Result<(), NamespaceError> {
        let source = FileBlobSourceActor::open(
            self.runtime.clone(),
            Arc::clone(&self.source_sender),
            registration.path(),
        )?;
        let length = source.length();
        let recovery = source.recovery();
        let source = self
            .runtime
            .spawn(source)
            .map_err(|error| NamespaceError::SourceRecovery(error.to_string()))?;
        let mut cleanup = PendingSourceRegistration {
            runtime: self.runtime.clone(),
            source,
            armed: true,
        };
        let source_node = self
            .source_publisher
            .publish_source(source)
            .map_err(NamespaceError::SourceRecovery)?;
        self.directory
            .register(
                path,
                source,
                source_node,
                length,
                recovery,
                random_operation_id(),
            )
            .await?;
        cleanup.armed = false;
        Ok(())
    }

    pub async fn unregister(&self, path: DataPath) -> Result<(), NamespaceError> {
        self.directory
            .unregister(path, random_operation_id())
            .await?;
        Ok(())
    }

    /// Remove one namespace entry only after any stream writer has quiesced,
    /// then prove the committed directory state no longer contains the path.
    pub async fn try_unregister_quiescent(
        &self,
        path: DataPath,
    ) -> Result<NamespaceCleanupStatus, NamespaceError> {
        let node = match self.directory.lookup(path.clone()).await {
            Ok(node) => node,
            Err(NamespaceError::PathNotFound(_)) => return Ok(NamespaceCleanupStatus::Absent),
            Err(error) => return Err(error),
        };
        if node.kind == crate::namespace::EntryKind::Stream && node.active {
            return Ok(NamespaceCleanupStatus::ActiveStream);
        }
        let quiescent_stream_revision =
            (node.kind == crate::namespace::EntryKind::Stream).then_some(node.revision);
        self.directory
            .unregister(path.clone(), random_operation_id())
            .await?;
        match self.directory.lookup(path).await {
            Err(NamespaceError::PathNotFound(_)) => Ok(NamespaceCleanupStatus::Removed {
                quiescent_stream_revision,
            }),
            Ok(_) => Err(NamespaceError::Protocol(
                "namespace entry remained after unregister acknowledgement".to_owned(),
            )),
            Err(error) => Err(error),
        }
    }

    pub async fn rename(
        &self,
        source: DataPath,
        destination: DataPath,
        replace: bool,
    ) -> Result<(), NamespaceError> {
        self.directory
            .rename(source, destination, replace, random_operation_id())
            .await?;
        Ok(())
    }
}

struct PendingSourceRegistration {
    runtime: Runtime,
    source: ActorAddress,
    armed: bool,
}

impl Drop for PendingSourceRegistration {
    fn drop(&mut self) {
        if self.armed {
            let _ = self
                .runtime
                .send_to(self.source, BlobSourceIn::Retire { reply_to: None });
        }
    }
}

fn random_operation_id() -> OperationId {
    let actor = ActorAddress::new_random();
    let mut bytes = [0_u8; 16];
    bytes.copy_from_slice(&actor.0[..16]);
    OperationId::from_u128(u128::from_be_bytes(bytes))
}