use std::{
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use connectrpc::{
ConnectError, RequestContext, Response, Router, ServiceRequest, ServiceResult, ServiceStream,
};
use polyc_proto::proto::polychrome::state::v1::{
AcknowledgeProjectorCursorReply, AcknowledgeProjectorCursorRequest, CompactFeedPrefixReply,
CompactFeedPrefixRequest, CreateFeedSnapshotReply, CreateFeedSnapshotRequest,
DescribeFeedStreamReply, DescribeFeedStreamRequest, FeedChunk as PbFeedChunk,
GetFeedRetentionReply, GetFeedRetentionRequest, ListProjectorsReply, ListProjectorsRequest,
PageCompleteness, RegisterProjectorReply, RegisterProjectorRequest, StateFeedService,
StateFeedServiceExt, SubscribeCommitsRequest,
};
use polyc_state::{
command::CommandMetadata,
context::CallContext,
error::{BoundKind, StateError},
feed::{
self, AcknowledgeProjectorCursor, CompactFeedPrefix, CreateSnapshot, GetFeedRetention,
JournalFeed, ListProjectors, ProjectorRegistration, RegisterProjector, SubscribeCommits,
},
id::{ConsumerId, OperationFamily},
page::{Cursor, ReadStart},
revision::JournalPosition,
stream::{StreamChunk, StreamEnd},
};
use crate::{
MAX_FEED_WIRE_MESSAGE_BYTES,
admission::{
AudienceBinding, PeerIdentity, check_audience_binding, check_call_context_version,
check_not_draining, check_transport_deadline, state_audience,
},
error::to_connect_error,
feed::wire::subscribe_request,
trace::adopt_caller_trace,
wire::{DeclaredCall, Kernel, declared_call},
};
const SUBSCRIBE_STREAM_BOUND: usize = 8;
const SHED_MARKER_ATTEMPTS: u32 = 25;
const TAIL_POLL: Duration = Duration::from_millis(20);
const MAX_TAIL_IDLE: Duration = Duration::from_secs(20);
pub const MAX_CONCURRENT_SUBSCRIPTIONS: usize = 512;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FeedStreamTuning {
max_tail_idle: Duration,
stream_bound: usize,
}
impl Default for FeedStreamTuning {
fn default() -> Self {
Self::new()
}
}
impl FeedStreamTuning {
#[must_use]
pub const fn new() -> Self {
Self {
max_tail_idle: MAX_TAIL_IDLE,
stream_bound: SUBSCRIBE_STREAM_BOUND,
}
}
#[must_use]
pub const fn with_max_tail_idle(mut self, idle: Duration) -> Self {
self.max_tail_idle = idle;
self
}
#[must_use]
pub const fn with_stream_bound(mut self, chunks: usize) -> Self {
self.stream_bound = chunks;
self
}
#[must_use]
pub const fn max_tail_idle(self) -> Duration {
self.max_tail_idle
}
#[must_use]
pub const fn stream_bound(self) -> usize {
self.stream_bound
}
}
pub struct FeedSvc {
feed: Arc<dyn JournalFeed>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
tuning: FeedStreamTuning,
}
impl FeedSvc {
#[must_use]
pub const fn new(
feed: Arc<dyn JournalFeed>,
draining: Arc<AtomicBool>,
binding: Arc<AudienceBinding>,
) -> Self {
Self {
feed,
draining,
binding,
tuning: FeedStreamTuning::new(),
}
}
#[must_use]
pub const fn with_tuning(mut self, tuning: FeedStreamTuning) -> Self {
self.tuning = tuning;
self
}
#[must_use]
pub fn register_on(self, router: Router) -> Router {
Arc::new(self).register(router)
}
fn family() -> OperationFamily {
feed::family()
}
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<polyc_proto::proto::polychrome::state::v1::CallContext>>,
method: &'static str,
) -> Result<(DeclaredCall, tracing::Span), ConnectError> {
check_not_draining(self.draining.load(Ordering::Relaxed))?;
let declared: DeclaredCall = declared_call(context).map_err(|e| to_connect_error(&e))?;
let family = Self::family();
check_call_context_version(declared.version).map_err(|e| to_connect_error(&e))?;
check_audience_binding(
&Self::peer(ctx),
&declared.audience,
&state_audience(),
&self.binding,
&family,
)
.map_err(|e| to_connect_error(&e))?;
check_transport_deadline(ctx.time_remaining(), &family)
.map_err(|e| to_connect_error(&e))?;
Ok((declared, adopt_caller_trace(ctx.headers(), method)))
}
}
type PumpedChunk = Result<PbFeedChunk, ConnectError>;
fn new_chunk_channel(
bound: usize,
) -> (
futures::channel::mpsc::Sender<PumpedChunk>,
futures::channel::mpsc::Receiver<PumpedChunk>,
) {
futures::channel::mpsc::channel(bound)
}
fn required_field(field: &str, reason: &str) -> ConnectError {
to_connect_error(&StateError::Malformed {
field: field.to_owned(),
reason: reason.to_owned(),
})
}
fn command_metadata(
command: impl Into<Option<polyc_proto::proto::polychrome::state::v1::FeedCommand>>,
) -> Result<CommandMetadata, ConnectError> {
Ok(Kernel::<CommandMetadata>::try_from(
command.into().ok_or_else(|| {
required_field("command", "a feed command carries its command identity")
})?,
)
.map_err(|e| to_connect_error(&e))?
.into_inner())
}
struct Pump {
feed: Arc<dyn JournalFeed>,
draining: Arc<AtomicBool>,
declared: DeclaredCall,
request: SubscribeCommits,
cursor: Cursor,
tuning: FeedStreamTuning,
}
fn bounded_wire_chunk(
chunk: &polyc_state::feed::FeedChunk,
previous: &Cursor,
) -> Result<(PbFeedChunk, Option<Cursor>, bool), StateError> {
let original_len = chunk.len();
let mut wire = PbFeedChunk::from(Kernel(chunk));
while usize::try_from(buffa::Message::encoded_len(&wire)).unwrap_or(usize::MAX)
> MAX_FEED_WIRE_MESSAGE_BYTES
{
if wire.records.len() <= 1 {
return Err(StateError::BoundsExceeded {
bound: BoundKind::PayloadBytes,
limit: u64::try_from(MAX_FEED_WIRE_MESSAGE_BYTES).unwrap_or(u64::MAX),
requested: u64::from(buffa::Message::encoded_len(&wire)),
});
}
wire.records.pop();
}
let shortened = wire.records.len() < original_len;
if shortened {
let delivered = wire
.records
.last()
.expect("a shortened non-empty feed chunk retains one commit")
.position;
let cursor = previous.snapshot().map_or_else(
|| Cursor::at(JournalPosition::new(delivered)),
|snapshot| Cursor::in_snapshot(snapshot.clone(), JournalPosition::new(delivered)),
);
wire.next = buffa::MessageField::some(Kernel(&cursor).into());
wire.end =
polyc_proto::proto::polychrome::state::v1::StreamEnd::from(Kernel(StreamEnd::More))
.into();
Ok((wire, Some(cursor), false))
} else {
Ok((
wire,
feed::cursor_after(chunk, Some(previous)),
chunk.is_drained(),
))
}
}
impl Pump {
fn resumed(&self) -> SubscribeCommits {
let resumed = SubscribeCommits::new(
self.request.partition().clone(),
ReadStart::Resume(self.cursor.clone()),
self.request.max_chunk_commits(),
);
match self.request.consumer() {
Some(consumer) => resumed.on_behalf_of(consumer.clone()),
None => resumed,
}
}
fn drained(&self) -> PbFeedChunk {
PbFeedChunk::from(Kernel(&feed::drained_chunk(self.cursor.clone())))
}
fn paused(&self) -> PbFeedChunk {
PbFeedChunk::from(Kernel(&StreamChunk::new(
Vec::new(),
Some(self.cursor.clone()),
StreamEnd::More,
)))
}
fn shed(&self, sender: &mut futures::channel::mpsc::Sender<PumpedChunk>) {
for _ in 0..SHED_MARKER_ATTEMPTS {
if sender.is_closed() {
return;
}
if sender.try_send(Ok(self.paused())).is_ok() {
return;
}
std::thread::sleep(TAIL_POLL);
}
}
fn run(mut self, mut sender: futures::channel::mpsc::Sender<PumpedChunk>) {
let mut idle = Duration::ZERO;
loop {
if sender.is_closed() {
return;
}
if idle >= self.tuning.max_tail_idle() {
let _ = sender.try_send(Ok(self.paused()));
return;
}
if self.draining.load(Ordering::Relaxed) {
let _ = sender.try_send(Ok(self.drained()));
return;
}
let context: CallContext = self.declared.origin_relative_context();
let chunk = match self.feed.commits(self.resumed(), &context) {
Ok(chunk) => chunk,
Err(error) => {
let _ = sender.try_send(Err(to_connect_error(&error)));
return;
}
};
if chunk.is_empty() && !chunk.is_drained() {
std::thread::sleep(TAIL_POLL);
idle = idle.saturating_add(TAIL_POLL);
continue;
}
idle = Duration::ZERO;
let (wire, advanced, drained) = match bounded_wire_chunk(&chunk, &self.cursor) {
Ok(bounded) => bounded,
Err(error) => {
let _ = sender.try_send(Err(to_connect_error(&error)));
return;
}
};
if sender.try_send(Ok(wire)).is_err() {
self.shed(&mut sender);
return;
}
if let Some(next) = advanced {
self.cursor = next;
}
if drained {
return;
}
}
}
}
#[allow(refining_impl_trait)]
impl StateFeedService for FeedSvc {
async fn create_snapshot(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, CreateFeedSnapshotRequest>,
) -> ServiceResult<CreateFeedSnapshotReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "CreateSnapshot")?;
let _entered = span.enter();
let metadata = command_metadata(message.command)?;
let snapshot = self
.feed
.create_snapshot(
CreateSnapshot::new(metadata),
&declared.origin_relative_context(),
)
.map_err(|e| to_connect_error(&e))?;
Response::ok(CreateFeedSnapshotReply {
snapshot: buffa::MessageField::some(Kernel(&snapshot).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn subscribe_commits(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, SubscribeCommitsRequest>,
) -> ServiceResult<ServiceStream<PbFeedChunk>> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "SubscribeCommits")?;
let _entered = span.enter();
let subscription = subscribe_request(
message.partition,
message.start,
message.max_chunk_commits,
message.consumer,
)
.map_err(|e| to_connect_error(&e))?;
let cursor = feed::resume_cursor(subscription.start()).map_err(|e| to_connect_error(&e))?;
let (sender, receiver) = new_chunk_channel(self.tuning.stream_bound());
let pump = Pump {
feed: Arc::clone(&self.feed),
draining: Arc::clone(&self.draining),
declared,
request: subscription,
cursor,
tuning: self.tuning,
};
tokio::task::spawn_blocking(move || pump.run(sender));
Response::stream_ok(receiver)
}
async fn register_projector(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, RegisterProjectorRequest>,
) -> ServiceResult<RegisterProjectorReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "RegisterProjector")?;
let _entered = span.enter();
let metadata = command_metadata(message.command)?;
let registration = Kernel::<ProjectorRegistration>::try_from(
message.registration.into_option().ok_or_else(|| {
required_field("registration", "a registration carries what it declares")
})?,
)
.map_err(|e| to_connect_error(&e))?
.into_inner();
let receipt = self
.feed
.register_projector(
RegisterProjector::new(metadata, registration),
&declared.origin_relative_context(),
)
.map_err(|e| to_connect_error(&e))?;
Response::ok(RegisterProjectorReply {
receipt: buffa::MessageField::some(Kernel(&receipt).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn acknowledge_projector_cursor(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, AcknowledgeProjectorCursorRequest>,
) -> ServiceResult<AcknowledgeProjectorCursorReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "AcknowledgeProjectorCursor")?;
let _entered = span.enter();
let metadata = command_metadata(message.command)?;
let cursor = Kernel::<Cursor>::from(message.cursor.into_option().ok_or_else(|| {
required_field("cursor", "an acknowledgement names how far it applied")
})?)
.into_inner();
let receipt = self
.feed
.acknowledge(
AcknowledgeProjectorCursor::new(
metadata,
ConsumerId::new(message.consumer),
cursor,
),
&declared.origin_relative_context(),
)
.map_err(|e| to_connect_error(&e))?;
Response::ok(AcknowledgeProjectorCursorReply {
receipt: buffa::MessageField::some(Kernel(&receipt).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn list_projectors(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, ListProjectorsRequest>,
) -> ServiceResult<ListProjectorsReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "ListProjectors")?;
let _entered = span.enter();
let listing = ListProjectors::new(
polyc_state::id::PartitionId::new(message.partition),
message.limit,
);
let listing = match message.consumer {
Some(consumer) => listing.for_consumer(ConsumerId::new(consumer)),
None => listing,
};
let listed = self
.feed
.projectors(listing, &declared.origin_relative_context())
.map_err(|e| to_connect_error(&e))?;
Response::ok(ListProjectorsReply {
projectors: listed
.projectors()
.iter()
.map(|status| Kernel(status).into())
.collect(),
completeness: PageCompleteness::from(Kernel(listed.completeness())).into(),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn get_feed_retention(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, GetFeedRetentionRequest>,
) -> ServiceResult<GetFeedRetentionReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "GetFeedRetention")?;
let _entered = span.enter();
let retention = self
.feed
.retention(
GetFeedRetention::new(polyc_state::id::PartitionId::new(message.partition)),
&declared.origin_relative_context(),
)
.map_err(|e| to_connect_error(&e))?;
Response::ok(GetFeedRetentionReply {
retention: buffa::MessageField::some(Kernel(&retention).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn compact_feed_prefix(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, CompactFeedPrefixRequest>,
) -> ServiceResult<CompactFeedPrefixReply> {
let message = request.to_owned_message();
let (declared, span) = self.admit(&ctx, message.context, "CompactFeedPrefix")?;
let _entered = span.enter();
let metadata = command_metadata(message.command)?;
let compaction = self
.feed
.compact(
CompactFeedPrefix::new(metadata, JournalPosition::new(message.through)),
&declared.origin_relative_context(),
)
.map_err(|e| to_connect_error(&e))?;
Response::ok(CompactFeedPrefixReply {
compaction: buffa::MessageField::some(Kernel(&compaction).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
async fn describe_feed_stream(
&self,
ctx: RequestContext,
request: ServiceRequest<'_, DescribeFeedStreamRequest>,
) -> ServiceResult<DescribeFeedStreamReply> {
let message = request.to_owned_message();
let (_declared, span) = self.admit(&ctx, message.context, "DescribeFeedStream")?;
let _entered = span.enter();
Response::ok(DescribeFeedStreamReply {
contract: buffa::MessageField::some(Kernel(self.feed.contract()).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
})
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs, clippy::unwrap_used)]
use std::time::Duration;
use futures::StreamExt as _;
use polyc_state::{
deadline::{Deadline, MonotonicInstant},
digest::ContentDigest,
feed::{CommitEnvelope, FeedRecord},
id::{CommandId, PartitionId},
journal::{JournalRecord, JournalRecordDraft, RecordKind, RecordTrust},
memory::MemoryFeed,
};
use super::*;
fn partition() -> PartitionId {
PartitionId::new("conv-1")
}
fn feed_of(commits: u64) -> Arc<MemoryFeed> {
let feed = Arc::new(MemoryFeed::new());
for n in 1..=commits {
feed.append_commit(
&partition(),
&CommandId::new(format!("commit-{n}")),
vec![JournalRecordDraft::new(
RecordKind::new("proof.commit"),
n.to_be_bytes().to_vec(),
)],
);
}
feed
}
fn pump_over(feed: Arc<dyn JournalFeed>, tuning: FeedStreamTuning) -> Pump {
Pump {
feed,
draining: Arc::new(AtomicBool::new(false)),
declared: DeclaredCall::live(state_audience(), Duration::MAX),
request: SubscribeCommits::new(
partition(),
ReadStart::Resume(Cursor::at(JournalPosition::ORIGIN)),
1,
),
cursor: Cursor::at(JournalPosition::ORIGIN),
tuning,
}
}
fn decode(chunk: PbFeedChunk) -> polyc_state::feed::FeedChunk {
Kernel::<polyc_state::feed::FeedChunk>::try_from(chunk)
.unwrap()
.into_inner()
}
fn maximally_fragmented_chunk() -> polyc_state::feed::FeedChunk {
let records = (1..=feed::MAX_CHUNK_COMMITS)
.map(|commit| {
let items = (0..polyc_state::journal::MAX_RECORDS_PER_BATCH)
.map(|index| {
JournalRecord::new(
JournalPosition::new(
u64::from(commit - 1)
* u64::from(polyc_state::journal::MAX_RECORDS_PER_BATCH)
+ u64::from(index)
+ 1,
),
RecordKind::new("k"),
RecordTrust::QuarantinedContent,
vec![b'x'; 766],
)
})
.collect::<Vec<_>>();
FeedRecord::new(
JournalPosition::new(u64::from(commit)),
CommitEnvelope::new(
PartitionId::new("p"),
CommandId::new("c"),
ContentDigest::from_bytes([7; ContentDigest::LEN]),
JournalPosition::new(u64::from(commit - 1)),
JournalPosition::new(u64::from(commit)),
u64::from(polyc_state::journal::MAX_RECORDS_PER_BATCH),
),
items,
)
})
.collect();
StreamChunk::new(
records,
Some(Cursor::at(JournalPosition::new(u64::from(
feed::MAX_CHUNK_COMMITS,
)))),
StreamEnd::More,
)
}
#[test]
fn buffa_overhead_shortens_a_feed_only_at_commit_boundaries() {
let semantic = maximally_fragmented_chunk();
assert!(feed::chunk_is_honest(&semantic, None));
let unbounded = PbFeedChunk::from(Kernel(&semantic));
assert!(
usize::try_from(buffa::Message::encoded_len(&unbounded)).unwrap()
> MAX_FEED_WIRE_MESSAGE_BYTES,
"the fixture must reproduce framing amplification beyond the old fixed headroom"
);
let (wire, advanced, drained) =
bounded_wire_chunk(&semantic, &Cursor::at(JournalPosition::ORIGIN)).unwrap();
assert!(
usize::try_from(buffa::Message::encoded_len(&wire)).unwrap()
<= MAX_FEED_WIRE_MESSAGE_BYTES
);
assert!(
!wire.records.is_empty(),
"at least one whole commit advances"
);
assert!(
wire.records.len() < semantic.len(),
"the oversized tail is deferred"
);
assert!(!drained, "a shortened chunk has more work by definition");
let delivered = wire.records.last().unwrap().position;
assert_eq!(
advanced.unwrap().position(),
JournalPosition::new(delivered)
);
let decoded = decode(wire);
assert_eq!(decoded.end(), StreamEnd::More);
assert!(feed::chunk_is_honest(&decoded, None));
assert!(decoded.resumes_without_gap(Some(&Cursor::at(JournalPosition::ORIGIN))));
}
#[test]
fn a_shed_consumer_gets_a_resume_marker_for_what_it_actually_received() {
let pump = pump_over(
feed_of(8) as Arc<dyn JournalFeed>,
FeedStreamTuning::new()
.with_stream_bound(1)
.with_max_tail_idle(Duration::from_millis(50)),
);
let (sender, mut receiver) = new_chunk_channel(1);
let reader = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(200));
futures::executor::block_on(async move {
let mut received = Vec::new();
while let Some(item) = receiver.next().await {
received.push(decode(item.unwrap()));
}
received
})
});
pump.run(sender);
let received = reader.join().unwrap();
assert!(
received.len() >= 2,
"the consumer received what fit before it fell behind: {}",
received.len()
);
let marker = received.last().unwrap();
assert!(marker.is_empty(), "the marker delivers nothing new");
assert_eq!(
marker.end(),
StreamEnd::More,
"a shed says more may follow, so the consumer resumes rather than stopping"
);
let delivered = received
.iter()
.rev()
.find(|chunk| !chunk.is_empty())
.expect("at least one chunk carried commits");
let last_delivered = feed::cursor_after(delivered, None).unwrap();
assert_eq!(
marker.next_cursor(),
Some(&last_delivered),
"a marker past a dropped chunk would manufacture the gap it exists to avoid"
);
let resumed = feed_of(8)
.commits(
SubscribeCommits::new(partition(), ReadStart::Resume(last_delivered.clone()), 8),
&CallContext::new(
Deadline::at(MonotonicInstant::from_nanos(u64::MAX)),
polyc_state::cancel::CancellationToken::new(),
),
)
.unwrap();
assert!(!feed::chunk_skips_a_commit(&resumed, Some(&last_delivered)));
assert_eq!(
resumed
.records()
.first()
.map(polyc_state::page::Positioned::position),
Some(last_delivered.position().next()),
"the resume continues immediately after what was delivered"
);
}
#[test]
fn an_idle_tail_pauses_with_the_cursor_it_reached() {
let pump = pump_over(
feed_of(2) as Arc<dyn JournalFeed>,
FeedStreamTuning::new()
.with_stream_bound(8)
.with_max_tail_idle(Duration::from_millis(50)),
);
let (sender, mut receiver) = new_chunk_channel(8);
pump.run(sender);
let received: Vec<_> = futures::executor::block_on(async move {
let mut received = Vec::new();
while let Some(item) = receiver.next().await {
received.push(decode(item.unwrap()));
}
received
});
let marker = received.last().unwrap();
assert!(marker.is_empty());
assert_eq!(marker.end(), StreamEnd::More);
assert_eq!(
marker.next_cursor().map(Cursor::position),
Some(JournalPosition::new(2)),
"the pause resumes after everything the tail delivered"
);
}
}