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::{
command::CommandScope,
id::{AggregateId, NamespaceId, PartitionId},
versioned::{
CreateVersionedDirectorySnapshot, GetVersionedSourceHead, ListVersionedDirectorySnapshot,
ReleaseVersionedDirectorySnapshot, VersionedDirectorySnapshotId, VersionedRead,
},
};
use crate::{
admission::{
AudienceBinding, PeerIdentity, check_audience_binding, check_call_context_version,
check_not_draining, check_transport_deadline, slot_wait, state_audience,
},
blocking,
error::to_connect_error,
trace::adopt_caller_trace,
wire::{DeclaredCall, declared_call},
};
pub struct VersionedSourceSvc {
authority: Arc<dyn VersionedRead>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
}
impl VersionedSourceSvc {
#[must_use]
pub const fn new(
authority: Arc<dyn VersionedRead>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
) -> Self {
Self {
authority,
draining,
binding,
}
}
#[must_use]
pub fn register_on(self, router: Router) -> Router {
use pb::StateVersionedSourceServiceExt as _;
Arc::new(self).register(router)
}
fn admit(
&self,
ctx: &RequestContext,
value: impl Into<Option<pb::CallContext>>,
method: &'static str,
) -> Result<
(
polyc_state::context::CallContext,
tracing::Span,
std::time::Duration,
),
ConnectError,
> {
check_not_draining(self.draining.load(Ordering::Relaxed))?;
let declared: DeclaredCall =
declared_call(value).map_err(|error| to_connect_error(&error))?;
let family = polyc_state::versioned::family();
check_call_context_version(declared.version).map_err(|error| to_connect_error(&error))?;
let peer = PeerIdentity::from_verified_leaf(
ctx.peer_certs().and_then(<[_]>::first).map(|leaf| &**leaf),
);
check_audience_binding(
&peer,
&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),
slot_wait(ctx.time_remaining(), &declared),
))
}
}
#[allow(refining_impl_trait)]
impl pb::StateVersionedSourceService for VersionedSourceSvc {
async fn get_source_head(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::GetVersionedSourceHeadRequest>,
) -> ServiceResult<pb::GetVersionedSourceHeadReply> {
let pb::GetVersionedSourceHeadRequest {
context,
aggregate,
partition,
namespace,
__buffa_unknown_fields: _,
} = request.to_owned_message();
let (context, span, wait) = self.admit(&ctx, context, "GetVersionedSourceHead")?;
let scope = CommandScope::new(
AggregateId::new(aggregate.clone()),
PartitionId::new(partition.clone()),
NamespaceId::new(namespace.clone()),
);
let authority = Arc::clone(&self.authority);
let head = blocking::read(span, polyc_state::versioned::family(), wait, move || {
authority.source_head(GetVersionedSourceHead::new(scope), &context)
})
.await?;
Response::ok(pb::GetVersionedSourceHeadReply {
incarnation: head.incarnation().as_bytes().to_vec(),
head: buffa::MessageField::some(pb::JournalHead {
position: head.head().position().get(),
__buffa_unknown_fields: buffa::UnknownFields::default(),
}),
aggregate,
partition,
namespace,
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn create_directory_snapshot(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::CreateVersionedDirectorySnapshotRequest>,
) -> ServiceResult<pb::CreateVersionedDirectorySnapshotReply> {
let pb::CreateVersionedDirectorySnapshotRequest {
context,
family,
__buffa_unknown_fields: _,
} = request.to_owned_message();
let (context, span, wait) =
self.admit(&ctx, context, "CreateVersionedDirectorySnapshot")?;
let requested = CreateVersionedDirectorySnapshot::new(family.clone());
let authority = Arc::clone(&self.authority);
let snapshot = blocking::read(span, polyc_state::versioned::family(), wait, move || {
authority.create_directory_snapshot(requested, &context)
})
.await?;
Response::ok(pb::CreateVersionedDirectorySnapshotReply {
snapshot: snapshot.id().as_str().to_owned(),
namespace_count: snapshot.namespace_count(),
incarnation: snapshot.lineage().as_bytes().to_vec(),
family,
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn directory_page(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::ListVersionedDirectorySnapshotRequest>,
) -> ServiceResult<pb::ListVersionedDirectorySnapshotReply> {
let pb::ListVersionedDirectorySnapshotRequest {
context,
snapshot,
after,
limit,
__buffa_unknown_fields: _,
} = request.to_owned_message();
let (context, span, wait) = self.admit(&ctx, context, "ListVersionedDirectorySnapshot")?;
let mut requested =
ListVersionedDirectorySnapshot::new(VersionedDirectorySnapshotId::new(snapshot), limit);
if let Some(after) = after {
requested = requested.after(after);
}
let authority = Arc::clone(&self.authority);
let page = blocking::read(span, polyc_state::versioned::family(), wait, move || {
authority.directory_page(requested, &context)
})
.await?;
let complete = page.completeness() == polyc_state::page::PageCompleteness::Complete;
Response::ok(pb::ListVersionedDirectorySnapshotReply {
snapshot: page.snapshot().as_str().to_owned(),
namespaces: page.namespaces().to_vec(),
complete,
incarnation: page.lineage().as_bytes().to_vec(),
next_after: page.next_after().map(str::to_owned),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn release_directory_snapshot(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::ReleaseVersionedDirectorySnapshotRequest>,
) -> ServiceResult<pb::ReleaseVersionedDirectorySnapshotReply> {
let pb::ReleaseVersionedDirectorySnapshotRequest {
context,
snapshot,
__buffa_unknown_fields: _,
} = request.to_owned_message();
let (context, span, wait) =
self.admit(&ctx, context, "ReleaseVersionedDirectorySnapshot")?;
let requested =
ReleaseVersionedDirectorySnapshot::new(VersionedDirectorySnapshotId::new(snapshot));
let authority = Arc::clone(&self.authority);
blocking::read(span, polyc_state::versioned::family(), wait, move || {
authority.release_directory_snapshot(requested, &context)
})
.await?;
Response::ok(pb::ReleaseVersionedDirectorySnapshotReply {
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}