use std::future::Future;
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::{
context::CallContext,
error::StateError,
id::{CommandId, NamespaceId},
tasks::{ContextId, EdgeId, TaskDirectory, TaskId, TaskRead, TaskWrite},
versioned::{VersionedRead, VersionedTransact},
};
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, Kernel, declared_call},
};
use super::wire::{command_from_wire, index_to_wire, record_to_wire};
pub trait TaskAuthority: TaskRead + TaskWrite {
fn namespace(&self) -> &NamespaceId;
}
impl<V> TaskAuthority for TaskDirectory<V>
where
V: VersionedRead + VersionedTransact,
{
fn namespace(&self) -> &NamespaceId {
self.namespace()
}
}
pub struct TaskSvc {
authority: Arc<dyn TaskAuthority>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
}
impl TaskSvc {
#[must_use]
pub const fn new(
authority: Arc<dyn TaskAuthority>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
) -> Self {
Self {
authority,
draining,
binding,
}
}
#[must_use]
pub fn register_on(self, router: Router) -> Router {
use pb::StateTaskServiceExt 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<(CallContext, tracing::Span, std::time::Duration), 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::tasks::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),
slot_wait(ctx.time_remaining(), &declared),
))
}
fn namespace(&self, value: &str) -> Result<(), ConnectError> {
if value == self.authority.namespace().as_str() {
Ok(())
} else {
Err(to_connect_error(&StateError::Denied {
family: polyc_state::tasks::family(),
}))
}
}
}
#[allow(refining_impl_trait)]
impl pb::StateTaskService for TaskSvc {
fn transact(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::TransactStateTaskRequest>,
) -> impl Future<Output = ServiceResult<pb::TransactStateTaskReply>> {
let message = request.to_owned_message();
async move {
let (context, span, wait) = self.admit(&ctx, message.context, "TransactStateTask")?;
let metadata = message
.metadata
.into_option()
.ok_or_else(|| malformed_connect("metadata", "a task command carries metadata"))?;
self.namespace(&metadata.namespace)?;
let operation = message.operation.into_option().ok_or_else(|| {
malformed_connect("operation", "a task command carries one operation")
})?;
let command =
command_from_wire(metadata, operation).map_err(|error| to_connect_error(&error))?;
let authority = Arc::clone(&self.authority);
let receipt = blocking::mutation(span, polyc_state::tasks::family(), wait, move || {
authority.transact(command, &context)
})
.await?;
Response::ok(pb::TransactStateTaskReply {
receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(&receipt))),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
fn get_task(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::GetStateTaskRequest>,
) -> impl Future<Output = ServiceResult<pb::GetStateTaskReply>> {
let message = request.to_owned_message();
async move {
let (context, span, wait) = self.admit(&ctx, message.context, "GetStateTask")?;
self.namespace(&message.namespace)?;
let owner = EdgeId::new(message.owner_edge_id);
let task_id = TaskId::new(message.task_id);
let authority = Arc::clone(&self.authority);
let fact = blocking::read(span, polyc_state::tasks::family(), wait, move || {
authority.task(&owner, &task_id, &context)
})
.await?;
let snapshot_revision = fact.snapshot_revision().get();
let entry_revision = fact
.entry_revision()
.map(polyc_state::revision::Revision::get);
Response::ok(pb::GetStateTaskReply {
task: fact
.value()
.as_ref()
.map_or_else(buffa::MessageField::none, |record| {
buffa::MessageField::some(record_to_wire(record))
}),
snapshot_revision,
entry_revision,
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
fn get_context_index(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::GetStateTaskContextIndexRequest>,
) -> impl Future<Output = ServiceResult<pb::GetStateTaskContextIndexReply>> {
let message = request.to_owned_message();
async move {
let (context, span, wait) =
self.admit(&ctx, message.context, "GetStateTaskContextIndex")?;
self.namespace(&message.namespace)?;
let owner = EdgeId::new(message.owner_edge_id);
let context_id = ContextId::new(message.context_id);
let authority = Arc::clone(&self.authority);
let fact = blocking::read(span, polyc_state::tasks::family(), wait, move || {
authority.context_index(&owner, &context_id, &context)
})
.await?;
Response::ok(pb::GetStateTaskContextIndexReply {
index: buffa::MessageField::some(index_to_wire(fact.value())),
snapshot_revision: fact.snapshot_revision().get(),
entry_revision: fact
.entry_revision()
.map(polyc_state::revision::Revision::get),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
fn list_tasks(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::ListStateTasksRequest>,
) -> impl Future<Output = ServiceResult<pb::ListStateTasksReply>> {
let message = request.to_owned_message();
async move {
let (context, span, wait) = self.admit(&ctx, message.context, "ListStateTasks")?;
self.namespace(&message.namespace)?;
let owner = EdgeId::new(message.owner_edge_id);
let context_id = ContextId::new(message.context_id);
let after = message.after.map(TaskId::new);
let page_size = usize::try_from(message.page_size).unwrap_or(usize::MAX);
let authority = Arc::clone(&self.authority);
let fact = blocking::read(span, polyc_state::tasks::family(), wait, move || {
authority.list_by_context(&owner, &context_id, after.as_ref(), page_size, &context)
})
.await?;
let snapshot_revision = fact.snapshot_revision().get();
let entry_revision = fact
.entry_revision()
.map(polyc_state::revision::Revision::get);
let page = fact.into_value();
Response::ok(pb::ListStateTasksReply {
tasks: page.tasks().iter().map(record_to_wire).collect(),
next_after: page.next_after().map(|id| id.as_str().to_owned()),
snapshot_revision,
entry_revision,
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
fn get_receipt(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, pb::GetStateTaskReceiptRequest>,
) -> impl Future<Output = ServiceResult<pb::GetStateTaskReceiptReply>> {
let message = request.to_owned_message();
async move {
let (_context, span, wait) =
self.admit(&ctx, message.context, "GetStateTaskReceipt")?;
self.namespace(&message.namespace)?;
let namespace = NamespaceId::new(message.namespace);
let command_id = CommandId::new(message.command_id);
let authority = Arc::clone(&self.authority);
let receipt = blocking::read(span, polyc_state::tasks::family(), wait, move || {
authority.committed_receipt(&namespace, &command_id)
})
.await?;
Response::ok(pb::GetStateTaskReceiptReply {
receipt: receipt
.as_ref()
.map_or_else(buffa::MessageField::none, |value| {
buffa::MessageField::some(pb::Receipt::from(Kernel(value)))
}),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
}
fn malformed_connect(field: &str, reason: &str) -> ConnectError {
to_connect_error(&StateError::Malformed {
field: field.into(),
reason: reason.into(),
})
}