use connectrpc::client::{ClientConfig, ClientTransport};
use polyc_proto::proto::polychrome::state::v1 as pb;
use std::time::Duration;
use polyc_state::{
error::StateError,
feed::{
self, AcknowledgeProjectorCursor, CompactFeedPrefix, CreateSnapshot, FeedChunk,
FeedCompaction, FeedRetention, FeedSnapshot, GetFeedRetention, ListProjectors,
ProjectorListing, ProjectorStatus, RegisterProjector, SubscribeCommits,
},
page::{Cursor, ReadStart},
receipt::Receipt,
stream::{StreamContract, StreamEnd},
};
use crate::{
MAX_FEED_WIRE_MESSAGE_BYTES,
error::{TransportFallback, from_connect_error},
feed::wire::subscription,
trace::{bounded_traced_options, streaming_traced_options},
wire::{DeclaredCall, Kernel},
};
const RESUME_BACKOFF_BASE: Duration = Duration::from_millis(20);
const RESUME_BACKOFF_CAP: Duration = Duration::from_secs(5);
const RESUME_BACKOFF_MAX_DOUBLINGS: u32 = 8;
fn jitter_seed(request: &SubscribeCommits) -> u64 {
let mut seed: u64 = 0xcbf2_9ce4_8422_2325;
let mut absorb = |bytes: &[u8]| {
for byte in bytes {
seed ^= u64::from(*byte);
seed = seed.wrapping_mul(0x0100_0000_01b3);
}
};
absorb(request.partition().as_str().as_bytes());
if let Some(consumer) = request.consumer() {
absorb(consumer.as_str().as_bytes());
}
seed | 1
}
pub struct FeedClient<T> {
inner: pb::StateFeedServiceClient<T>,
}
impl<T> FeedClient<T>
where
T: ClientTransport,
<T::ResponseBody as connectrpc::http_body::Body>::Error: std::fmt::Display,
{
pub fn new(transport: T, config: ClientConfig) -> Self {
Self {
inner: pb::StateFeedServiceClient::new(
transport,
config.with_default_max_message_size(MAX_FEED_WIRE_MESSAGE_BYTES),
),
}
}
fn fallback(attempted_bytes: usize) -> TransportFallback {
TransportFallback::new(
feed::family(),
MAX_FEED_WIRE_MESSAGE_BYTES as u64,
attempted_bytes as u64,
)
}
pub async fn create_snapshot(
&self,
declared: &DeclaredCall,
command: &CreateSnapshot,
) -> Result<FeedSnapshot, StateError> {
let request = pb::CreateFeedSnapshotRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
command: buffa::MessageField::some(Kernel(command).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.create_snapshot_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
Ok(Kernel::<FeedSnapshot>::try_from(
reply
.snapshot
.into_option()
.ok_or_else(|| missing("snapshot", "a recorded snapshot carries its binding"))?,
)?
.into_inner())
}
pub async fn register(
&self,
declared: &DeclaredCall,
command: &RegisterProjector,
) -> Result<Receipt, StateError> {
let request = pb::RegisterProjectorRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
command: buffa::MessageField::some(Kernel(command).into()),
registration: buffa::MessageField::some(Kernel(command.registration()).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.register_projector_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
receipt_of(reply.receipt)
}
pub async fn acknowledge(
&self,
declared: &DeclaredCall,
command: &AcknowledgeProjectorCursor,
) -> Result<Receipt, StateError> {
let request = pb::AcknowledgeProjectorCursorRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
command: buffa::MessageField::some(Kernel(command).into()),
consumer: command.consumer().as_str().to_owned(),
cursor: buffa::MessageField::some(pb::Cursor::from(Kernel(command.cursor()))),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.acknowledge_projector_cursor_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
receipt_of(reply.receipt)
}
pub async fn projectors(
&self,
declared: &DeclaredCall,
request: &ListProjectors,
) -> Result<ProjectorListing, StateError> {
let request = pb::ListProjectorsRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
partition: request.partition().as_str().to_owned(),
limit: request.limit(),
consumer: request
.consumer()
.map(|consumer| consumer.as_str().to_owned()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.list_projectors_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
let completeness = crate::wire::completeness("completeness", reply.completeness)?;
let projectors = reply
.projectors
.into_iter()
.map(|status| Kernel::<ProjectorStatus>::try_from(status).map(Kernel::into_inner))
.collect::<Result<Vec<_>, _>>()?;
Ok(ProjectorListing::new(projectors, completeness))
}
pub async fn retention(
&self,
declared: &DeclaredCall,
request: &GetFeedRetention,
) -> Result<FeedRetention, StateError> {
let request = pb::GetFeedRetentionRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
partition: request.partition().as_str().to_owned(),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.get_feed_retention_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
Ok(Kernel::<FeedRetention>::from(
reply
.retention
.into_option()
.ok_or_else(|| missing("retention", "a successful read carries its retention"))?,
)
.into_inner())
}
pub async fn compact(
&self,
declared: &DeclaredCall,
command: &CompactFeedPrefix,
) -> Result<FeedCompaction, StateError> {
let request = pb::CompactFeedPrefixRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
command: buffa::MessageField::some(Kernel(command).into()),
through: command.through().get(),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.compact_feed_prefix_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
Ok(Kernel::<FeedCompaction>::try_from(
reply
.compaction
.into_option()
.ok_or_else(|| missing("compaction", "a compaction reports what it did"))?,
)?
.into_inner())
}
pub async fn contract(&self, declared: &DeclaredCall) -> Result<StreamContract, StateError> {
let request = pb::DescribeFeedStreamRequest {
context: buffa::MessageField::some(Kernel(declared).into()),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let reply = self
.inner
.describe_feed_stream_with_options(request, bounded_traced_options(declared))
.await
.map_err(|e| from_connect_error(&e, &Self::fallback(attempted)))?
.into_owned();
Ok(Kernel::<StreamContract>::try_from(
reply
.contract
.into_option()
.ok_or_else(|| missing("contract", "a stream declares its contract"))?,
)?
.into_inner())
}
pub fn subscribe(
&self,
declared: DeclaredCall,
request: SubscribeCommits,
) -> Result<FeedSubscription<'_, T>, StateError> {
let cursor = feed::resume_cursor(request.start())?;
let jitter = jitter_seed(&request);
Ok(FeedSubscription {
client: &self.inner,
declared,
request,
cursor,
stream: None,
ended: false,
last_end: None,
stalls: 0,
jitter,
})
}
}
pub struct FeedSubscription<'a, T: ClientTransport> {
client: &'a pb::StateFeedServiceClient<T>,
declared: DeclaredCall,
request: SubscribeCommits,
cursor: Cursor,
stream: Option<
connectrpc::client::ServerStream<
T::ResponseBody,
pb::__buffa::view::FeedChunkView<'static>,
>,
>,
ended: bool,
last_end: Option<StreamEnd>,
stalls: u32,
jitter: u64,
}
impl<T> FeedSubscription<'_, T>
where
T: ClientTransport,
T::ResponseBody: Unpin,
<T::ResponseBody as connectrpc::http_body::Body>::Error: std::fmt::Display,
{
#[must_use]
pub const fn cursor(&self) -> &Cursor {
&self.cursor
}
#[must_use]
pub const fn is_ended(&self) -> bool {
self.ended
}
pub async fn next_chunk(&mut self) -> Result<Option<FeedChunk>, StateError> {
loop {
if self.ended {
return Ok(None);
}
if self.stream.is_none() {
if self.stalls > 0 {
tokio::time::sleep(self.backoff()).await;
}
self.dial().await?;
}
let received = match self.stream.as_mut() {
Some(stream) => stream.message::<pb::FeedChunk>().await,
None => return Ok(None),
};
match received {
Ok(Some(message)) => {
let chunk = match Kernel::<FeedChunk>::try_from(message.to_owned_message()) {
Ok(chunk) => chunk.into_inner(),
Err(typed) => return Err(self.refuse(typed)),
};
if let Some(next) = feed::cursor_after(&chunk, Some(&self.cursor)) {
self.cursor = next;
}
self.last_end = Some(chunk.end());
self.stalls = 0;
if chunk.is_drained() {
self.ended = true;
}
return Ok(Some(chunk));
}
Ok(None) => {
self.stream = None;
match self.last_end.take() {
Some(StreamEnd::Exhausted | StreamEnd::Drained) => {
self.ended = true;
return Ok(None);
}
Some(StreamEnd::More) | None => {
self.stalls = self.stalls.saturating_add(1);
}
}
}
Err(error) => {
let typed = from_connect_error(
&error,
&TransportFallback::new(
feed::family(),
MAX_FEED_WIRE_MESSAGE_BYTES as u64,
0,
),
);
return Err(self.refuse(typed));
}
}
}
}
fn refuse(&mut self, typed: StateError) -> StateError {
self.stream = None;
if !typed.is_retry_safe() {
self.ended = true;
}
typed
}
fn backoff(&mut self) -> Duration {
let exponent = (self.stalls - 1).min(RESUME_BACKOFF_MAX_DOUBLINGS);
let ceiling = RESUME_BACKOFF_BASE
.saturating_mul(1_u32 << exponent)
.min(RESUME_BACKOFF_CAP);
self.jitter ^= self.jitter << 13;
self.jitter ^= self.jitter >> 7;
self.jitter ^= self.jitter << 17;
let half = ceiling / 2;
let spread = u64::from(half.subsec_nanos())
.saturating_add(half.as_secs().saturating_mul(1_000_000_000));
let offset = if spread == 0 { 0 } else { self.jitter % spread };
half.saturating_add(Duration::from_nanos(offset))
}
async fn dial(&mut self) -> Result<(), StateError> {
let resumed = SubscribeCommits::new(
self.request.partition().clone(),
ReadStart::Resume(self.cursor.clone()),
self.request.max_chunk_commits(),
);
let resumed = match self.request.consumer() {
Some(consumer) => resumed.on_behalf_of(consumer.clone()),
None => resumed,
};
let (partition, start, max_chunk_commits, consumer) = subscription(&resumed);
let request = pb::SubscribeCommitsRequest {
context: buffa::MessageField::some(Kernel(&self.declared).into()),
partition,
start: buffa::MessageField::some(start),
max_chunk_commits,
consumer,
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let attempted = buffa::Message::encoded_len(&request) as usize;
let stream = self
.client
.subscribe_commits_with_options(request, streaming_traced_options())
.await
.map_err(|error| {
let typed = from_connect_error(
&error,
&TransportFallback::new(
feed::family(),
MAX_FEED_WIRE_MESSAGE_BYTES as u64,
attempted as u64,
),
);
if !typed.is_retry_safe() {
self.ended = true;
}
typed
})?;
self.stream = Some(stream);
Ok(())
}
}
fn receipt_of(receipt: impl Into<Option<pb::Receipt>>) -> Result<Receipt, StateError> {
Ok(Kernel::<Receipt>::try_from(
receipt
.into()
.ok_or_else(|| missing("receipt", "a committed command carries its receipt"))?,
)?
.into_inner())
}
fn missing(field: &str, reason: &str) -> StateError {
StateError::Malformed {
field: field.to_owned(),
reason: reason.to_owned(),
}
}
#[cfg(test)]
mod tests {
#![allow(
clippy::pedantic,
clippy::nursery,
missing_docs,
reason = "a test module: panics are the failure mode"
)]
use std::{
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll},
};
use bytes::Bytes;
use connectrpc::{
client::{BoxFuture, ClientBody, ClientConfig},
envelope::Envelope,
http_body::{Body, Frame},
};
use polyc_state::{
error::RetryClass, id::PartitionId, revision::JournalPosition, stream::StreamChunk,
};
use super::*;
use crate::state_audience;
struct CannedBody(std::vec::IntoIter<Bytes>);
impl Body for CannedBody {
type Data = Bytes;
type Error = std::io::Error;
fn poll_frame(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
Poll::Ready(self.0.next().map(|bytes| Ok(Frame::data(bytes))))
}
}
#[derive(Clone)]
struct CannedTransport {
dials: Arc<Mutex<std::collections::VecDeque<Vec<Bytes>>>>,
dialed: Arc<Mutex<usize>>,
content_type: &'static str,
}
impl CannedTransport {
fn serving(responses: Vec<Vec<Bytes>>) -> Self {
Self {
dials: Arc::new(Mutex::new(responses.into())),
dialed: Arc::new(Mutex::new(0)),
content_type: "application/connect+proto",
}
}
fn serving_unary(reply: Bytes) -> Self {
Self {
dials: Arc::new(Mutex::new(vec![vec![reply]].into())),
dialed: Arc::new(Mutex::new(0)),
content_type: "application/proto",
}
}
fn dials(&self) -> usize {
*self.dialed.lock().expect("the dial count")
}
}
impl ClientTransport for CannedTransport {
type ResponseBody = CannedBody;
type Error = std::io::Error;
fn send(
&self,
_request: http::Request<ClientBody>,
) -> BoxFuture<'static, Result<http::Response<Self::ResponseBody>, Self::Error>> {
*self.dialed.lock().expect("the dial count") += 1;
let frames = self
.dials
.lock()
.expect("the canned responses")
.pop_front()
.unwrap_or_default();
let content_type = self.content_type;
Box::pin(async move {
Ok(http::Response::builder()
.status(http::StatusCode::OK)
.header(http::header::CONTENT_TYPE, content_type)
.body(CannedBody(frames.into_iter()))
.expect("a canned response"))
})
}
}
fn framed(chunk: &pb::FeedChunk) -> Bytes {
Envelope::data(Bytes::from(buffa::Message::encode_to_vec(chunk))).encode()
}
fn end_of_stream() -> Bytes {
Envelope::end_stream(Bytes::from_static(b"{}")).encode()
}
fn chunk_at(position: u64, end: StreamEnd) -> pb::FeedChunk {
let chunk: FeedChunk = StreamChunk::new(
Vec::new(),
Some(Cursor::at(JournalPosition::new(position))),
end,
);
pb::FeedChunk::from(Kernel(&chunk))
}
fn declared() -> DeclaredCall {
DeclaredCall::live(state_audience(), Duration::MAX)
}
fn subscription_request() -> SubscribeCommits {
SubscribeCommits::new(
PartitionId::new("conv-feed-client"),
ReadStart::Resume(Cursor::at(JournalPosition::ORIGIN)),
8,
)
}
#[tokio::test]
async fn a_malformed_chunk_ends_the_subscription_without_moving_the_cursor() {
let mut malformed = chunk_at(7, StreamEnd::More);
malformed.end = pb::StreamEnd::STREAM_END_UNSPECIFIED.into();
let transport = CannedTransport::serving(vec![vec![
framed(&chunk_at(5, StreamEnd::More)),
framed(&malformed),
framed(&chunk_at(9, StreamEnd::Exhausted)),
end_of_stream(),
]]);
let client = FeedClient::new(
transport.clone(),
ClientConfig::new("http://feed.invalid".parse().expect("a base url")),
);
let mut subscription = client
.subscribe(declared(), subscription_request())
.expect("a resumable subscription");
let first = subscription
.next_chunk()
.await
.expect("the well-formed chunk is delivered")
.expect("a chunk, not the end of the subscription");
assert_eq!(first.end(), StreamEnd::More);
assert_eq!(
subscription.cursor(),
&Cursor::at(JournalPosition::new(5)),
"the delivered chunk moved the cursor"
);
let refused = subscription
.next_chunk()
.await
.expect_err("a chunk this build cannot read is a refusal");
assert!(
matches!(refused, StateError::Malformed { ref field, .. } if field == "end"),
"got {refused}"
);
assert_eq!(
refused.retry_class(),
RetryClass::Terminal,
"a malformed chunk is terminal, and the contract says terminal ends the subscription"
);
assert!(
subscription.is_ended(),
"a terminal refusal ends the subscription"
);
assert_eq!(
subscription.cursor(),
&Cursor::at(JournalPosition::new(5)),
"the cursor never steps over commits that were not delivered"
);
assert!(
subscription
.next_chunk()
.await
.expect("an ended subscription answers rather than failing")
.is_none(),
"an ended subscription yields nothing forever"
);
assert_eq!(
transport.dials(),
1,
"the refusal is not a stall, so nothing re-dials"
);
}
#[tokio::test]
async fn a_listing_that_does_not_say_whether_it_was_cut_is_refused() {
let reply = pb::ListProjectorsReply {
projectors: Vec::new(),
completeness: pb::PageCompleteness::PAGE_COMPLETENESS_UNSPECIFIED.into(),
__buffa_unknown_fields: buffa::UnknownFields::default(),
};
let transport =
CannedTransport::serving_unary(Bytes::from(buffa::Message::encode_to_vec(&reply)));
let client = FeedClient::new(
transport.clone(),
ClientConfig::new("http://feed.invalid".parse().expect("a base url")),
);
let refused = client
.projectors(
&declared(),
&ListProjectors::new(PartitionId::new("conv-feed-client"), 4),
)
.await
.expect_err("a listing that does not report its completeness is refused");
assert!(
matches!(&refused, StateError::Malformed { field, .. } if field == "completeness"),
"the refusal names the field it could not read: {refused:?}"
);
assert_eq!(
refused.retry_class(),
RetryClass::Terminal,
"a reply this client cannot read does not become readable on a retry"
);
assert_eq!(transport.dials(), 1, "the refusal is not a stall");
}
}