polyc-state-connect 2026.9.0

//! Server adapter for the authoritative persona-memory journal.

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use connectrpc::{ConnectError, RequestContext, Response, Router, ServiceRequest, ServiceResult};
use polyc_proto::proto::polychrome::state::v1 as pb;
use polyc_state::{
    error::{OutageReach, StateError},
    persona_memory::journal::{
        MemoryJournalPartition, MemoryReplayRange, PersonaMemoryJournalAuthority,
    },
};
use tracing::Instrument as _;

use crate::{
    admission::{
        AudienceBinding, PeerIdentity, check_audience_binding, check_call_context_version,
        check_not_draining, check_transport_deadline, state_audience,
    },
    error::to_connect_error,
    trace::adopt_caller_trace,
    wire::{DeclaredCall, Kernel, declared_call},
};

use super::wire::{
    append_from_wire, destroy_from_wire, migrate_from_wire, partition_page_to_wire, replay_to_wire,
    rewrite_from_wire,
};

/// State's authoritative persona-memory journal service.
pub struct PersonaMemoryJournalSvc {
    authority: Arc<dyn PersonaMemoryJournalAuthority>,
    draining: Arc<AtomicBool>,
    binding: Arc<AudienceBinding>,
}

impl PersonaMemoryJournalSvc {
    /// Builds the service.
    #[must_use]
    pub const fn new(
        authority: Arc<dyn PersonaMemoryJournalAuthority>,
        draining: Arc<AtomicBool>,
        binding: Arc<AudienceBinding>,
    ) -> Self {
        Self {
            authority,
            draining,
            binding,
        }
    }

    /// Registers the capability-specific routes.
    #[must_use]
    pub fn register_on(self, router: Router) -> Router {
        use pb::StatePersonaMemoryJournalServiceExt as _;
        Arc::new(self).register(router)
    }

    fn peer(ctx: &RequestContext) -> PeerIdentity {
        PeerIdentity::from_verified_leaf(
            ctx.peer_certs().and_then(<[_]>::first).map(|leaf| &**leaf),
        )
    }

    fn admit(
        &self,
        ctx: &RequestContext,
        context: impl Into<Option<pb::CallContext>>,
        method: &'static str,
    ) -> Result<(polyc_state::context::CallContext, tracing::Span), ConnectError> {
        check_not_draining(self.draining.load(Ordering::Relaxed))?;
        let declared: DeclaredCall =
            declared_call(context).map_err(|error| to_connect_error(&error))?;
        let family = polyc_state::persona_memory::journal::family();
        check_call_context_version(declared.version).map_err(|error| to_connect_error(&error))?;
        check_audience_binding(
            &Self::peer(ctx),
            &declared.audience,
            &state_audience(),
            &self.binding,
            &family,
        )
        .map_err(|error| to_connect_error(&error))?;
        check_transport_deadline(ctx.time_remaining(), &family)
            .map_err(|error| to_connect_error(&error))?;
        Ok((
            declared.origin_relative_context(),
            adopt_caller_trace(ctx.headers(), method),
        ))
    }

    fn namespace(&self, namespace: &str) -> Result<(), ConnectError> {
        if namespace == self.authority.namespace().as_str() {
            Ok(())
        } else {
            Err(to_connect_error(&StateError::Denied {
                family: polyc_state::persona_memory::journal::family(),
            }))
        }
    }

    async fn blocking<R: Send + 'static>(
        operation: impl FnOnce() -> Result<R, StateError> + Send + 'static,
    ) -> Result<R, ConnectError> {
        tokio::task::spawn_blocking(operation)
            .await
            .map_err(|_| {
                to_connect_error(&StateError::Unavailable {
                    family: polyc_state::persona_memory::journal::family(),
                    reach: OutageReach::PossiblyApplied,
                })
            })?
            .map_err(|error| to_connect_error(&error))
    }
}

#[allow(refining_impl_trait)]
impl pb::StatePersonaMemoryJournalService for PersonaMemoryJournalSvc {
    async fn append(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::AppendPersonaMemoryRecordsRequest>,
    ) -> ServiceResult<pb::AppendPersonaMemoryRecordsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "AppendPersonaMemoryRecords")?;
        let command = append_from_wire(required(message.command, "append")?)
            .map_err(|error| to_connect_error(&error))?;
        self.namespace(command.write().metadata().scope().namespace().as_str())?;
        let authority = Arc::clone(&self.authority);
        let outcome = Self::blocking(move || authority.append(command, &context))
            .instrument(span)
            .await?;
        Response::ok(pb::AppendPersonaMemoryRecordsReply {
            receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(outcome.receipt()))),
            positions: outcome.positions().to_vec(),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn rewrite(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::RewritePersonaMemoryRecordsRequest>,
    ) -> ServiceResult<pb::RewritePersonaMemoryRecordsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "RewritePersonaMemoryRecords")?;
        let command = rewrite_from_wire(required(message.command, "rewrite")?)
            .map_err(|error| to_connect_error(&error))?;
        self.namespace(command.write().metadata().scope().namespace().as_str())?;
        let authority = Arc::clone(&self.authority);
        let outcome = Self::blocking(move || authority.rewrite(command, &context))
            .instrument(span)
            .await?;
        Response::ok(pb::RewritePersonaMemoryRecordsReply {
            receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(outcome.receipt()))),
            dropped: outcome.dropped() as u64,
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn migrate(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::MigratePersonaMemoryRecordsRequest>,
    ) -> ServiceResult<pb::MigratePersonaMemoryRecordsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "MigratePersonaMemoryRecords")?;
        let command = migrate_from_wire(required(message.command, "migration")?)
            .map_err(|error| to_connect_error(&error))?;
        self.namespace(command.write().metadata().scope().namespace().as_str())?;
        let authority = Arc::clone(&self.authority);
        let outcome = Self::blocking(move || authority.migrate(command, &context))
            .instrument(span)
            .await?;
        Response::ok(pb::MigratePersonaMemoryRecordsReply {
            receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(outcome.receipt()))),
            copied: outcome.copied() as u64,
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn destroy(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::DestroyPersonaMemoryRecordsRequest>,
    ) -> ServiceResult<pb::DestroyPersonaMemoryRecordsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "DestroyPersonaMemoryRecords")?;
        let command = destroy_from_wire(required(message.command, "destruction")?)
            .map_err(|error| to_connect_error(&error))?;
        self.namespace(command.write().metadata().scope().namespace().as_str())?;
        let authority = Arc::clone(&self.authority);
        let receipt = Self::blocking(move || authority.destroy(command, &context))
            .instrument(span)
            .await?;
        Response::ok(pb::DestroyPersonaMemoryRecordsReply {
            receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(&receipt))),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn replay(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::ReplayPersonaMemoryRecordsRequest>,
    ) -> ServiceResult<pb::ReplayPersonaMemoryRecordsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "ReplayPersonaMemoryRecords")?;
        self.namespace(&message.namespace)?;
        let partition = MemoryJournalPartition::parse(message.partition)
            .map_err(|error| to_connect_error(&error))?;
        let range = MemoryReplayRange::new(message.start, message.max_records, message.max_bytes);
        let authority = Arc::clone(&self.authority);
        let page = Self::blocking(move || authority.replay(&partition, range, &context))
            .instrument(span)
            .await?;
        if page.signed_head().is_none() {
            return Err(to_connect_error(&StateError::Unavailable {
                family: polyc_state::persona_memory::journal::family(),
                reach: OutageReach::PossiblyApplied,
            }));
        }
        Response::ok(replay_to_wire(&page))
    }

    async fn list(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::ListPersonaMemoryPartitionsRequest>,
    ) -> ServiceResult<pb::ListPersonaMemoryPartitionsReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "ListPersonaMemoryPartitions")?;
        self.namespace(&message.namespace)?;
        let authority = Arc::clone(&self.authority);
        let page = Self::blocking(move || {
            authority.list(message.after.as_deref(), message.limit, &context)
        })
        .instrument(span)
        .await?;
        Response::ok(partition_page_to_wire(&page))
    }
}

// Named over `P: buffa::ProtoBox<T>` rather than `impl Into<Option<T>>`: see
// `crate::wire::required`'s doc comment.
fn required<T: Default, P: buffa::ProtoBox<T>>(
    field: buffa::MessageField<T, P>,
    name: &'static str,
) -> Result<T, ConnectError> {
    field.into_option().ok_or_else(|| {
        to_connect_error(&StateError::Malformed {
            field: name.into(),
            reason: format!("a request carries its {name}"),
        })
    })
}