crates/gunnar-repair/src/grpc.rs
//! **The tonic side of the pump**, generic over a generated `pb` (gunnar-proto's,
//! gunnar-ui-server's): the twin's server streams as [`ChunkStream`], and the main's
//! `Bootstrap.Restore` client stream as a [`RestoreSink`] that can be FAILED.
//!
//! Why the poison: tonic's server reads a client's RST_STREAM(CANCEL) on a request
//! stream as a CLEAN end, so a pump that merely dropped the stream handed a fresh box
//! a short stream it finished as a restore (gunnar-ui twin-heartbeat 3a1f824; the CLI
//! measured the same). [`Frame::Poison`] fails the encoder instead: the request body
//! errors (RST_STREAM INTERNAL_ERROR) and the box stays restoring.
use std::marker::PhantomData;
use crate::pump::{ChunkStream, Fault, RestoreSink};
use crate::view::{AsRestored, Chunk, Restored};
/// gunnar's `Bootstrap.Restore`, as its generated client names it.
pub const RESTORE_PATH: &str = "/gunnar.v1.Bootstrap/Restore";
/// The `GrpcMethod` a credential signer binds the Restore credential to.
pub const RESTORE_SERVICE: &str = "gunnar.v1.Bootstrap";
pub const RESTORE_METHOD: &str = "Restore";
/// A tonic status as a [`Fault`]: the code word in snake case, the message as is.
pub fn fault(status: &tonic::Status) -> Fault {
let word = format!("{:?}", status.code());
let mut code = String::new();
for (i, ch) in word.chars().enumerate() {
if ch.is_ascii_uppercase() && i > 0 {
code.push('_');
}
code.push(ch.to_ascii_lowercase());
}
Fault {
code,
message: status.message().to_owned(),
}
}
/// How long a twin's stream may stay silent between two messages.
pub const MESSAGE_WAIT: std::time::Duration = std::time::Duration::from_secs(120);
/// **One message of a twin's stream, or a named fault after `wait` of silence**
/// (stinky #17): a twin that went silent mid-stream held the recharge, and the
/// main's restore, open for ever. The pump aborts the restore on the fault.
pub async fn within<C>(
wait: std::time::Duration,
next: impl std::future::Future<Output = Option<Result<C, Fault>>>,
) -> Option<Result<C, Fault>> {
match tokio::time::timeout(wait, next).await {
Ok(said) => said,
Err(_) => Some(Err(Fault::new(
"deadline_exceeded",
format!("nothing from the twin within {} s: it went silent mid-stream, so this stream is given up on", wait.as_secs_f32()),
))),
}
}
/// A twin's server stream (`Mirror.LogFetch` / `Mirror.SnapshotFetch`).
pub struct TonicStream<C>(pub tonic::Streaming<C>);
impl<C: Chunk> ChunkStream for TonicStream<C> {
type Chunk = C;
async fn next(&mut self) -> Option<Result<C, Fault>> {
within(MESSAGE_WAIT, async {
match self.0.message().await {
Ok(Some(c)) => Some(Ok(c)),
Ok(None) => None,
Err(status) => Some(Err(fault(&status))),
}
})
.await
}
}
/// One frame of the restore's request body: a chunk, or the poison that fails it.
pub enum Frame<C> {
Chunk(C),
Poison,
}
type Prost<C, R> = tonic_prost::ProstCodec<C, R>;
/// The generated `Bootstrap.Restore` codec, except that [`Frame::Poison`] fails the encoder.
pub struct PoisonCodec<C, R>(Prost<C, R>);
impl<C, R> Default for PoisonCodec<C, R> {
fn default() -> Self {
PoisonCodec(Prost::default())
}
}
impl<C, R> tonic::codec::Codec for PoisonCodec<C, R>
where
C: prost::Message + Send + 'static,
R: prost::Message + Default + Send + 'static,
{
type Encode = Frame<C>;
type Decode = R;
type Encoder = PoisonEncoder<C, R>;
type Decoder = <Prost<C, R> as tonic::codec::Codec>::Decoder;
fn encoder(&mut self) -> PoisonEncoder<C, R> {
PoisonEncoder(tonic::codec::Codec::encoder(&mut self.0))
}
fn decoder(&mut self) -> Self::Decoder {
tonic::codec::Codec::decoder(&mut self.0)
}
}
pub struct PoisonEncoder<C, R>(<Prost<C, R> as tonic::codec::Codec>::Encoder)
where
C: prost::Message + Send + 'static,
R: prost::Message + Default + Send + 'static;
impl<C, R> tonic::codec::Encoder for PoisonEncoder<C, R>
where
C: prost::Message + Send + 'static,
R: prost::Message + Default + Send + 'static,
{
type Item = Frame<C>;
type Error = tonic::Status;
fn encode(
&mut self,
item: Frame<C>,
dst: &mut tonic::codec::EncodeBuf<'_>,
) -> Result<(), tonic::Status> {
match item {
Frame::Chunk(chunk) => tonic::codec::Encoder::encode(&mut self.0, chunk, dst),
Frame::Poison => Err(tonic::Status::aborted(
"the twin's fetch broke; this restore stream is failed on purpose, never ended",
)),
}
}
fn buffer_settings(&self) -> tonic::codec::BufferSettings {
tonic::codec::Encoder::buffer_settings(&self.0)
}
}
/// The request body a [`restore_stream`] call is handed.
pub type RestoreBody<C> = tokio_stream::wrappers::ReceiverStream<Frame<C>>;
/// **An open `Bootstrap.Restore`**: chunks go through a channel to the call, which
/// runs as its own task.
///
/// **Dropped without [`RestoreSink::close`], it FAILS the stream** (sends the poison):
/// a pump whose future is dropped — its task aborted at shutdown, its caller gone —
/// must never end the request body cleanly, or the fresh main finishes a short stream
/// as its restore.
pub struct TonicRestore<C: Send + 'static, R> {
/// `None` once closed or aborted: nothing left for `Drop` to fail.
tx: Option<tokio::sync::mpsc::Sender<Frame<C>>>,
call: Option<tokio::task::JoinHandle<Result<tonic::Response<R>, tonic::Status>>>,
_r: PhantomData<R>,
}
impl<C: Send + 'static, R> Drop for TonicRestore<C, R> {
fn drop(&mut self) {
let Some(tx) = self.tx.take() else { return };
// The body may be full: the poison waits for room (on the runtime the call runs
// on), and the sender it holds keeps the body open until it is in.
if let Err(tokio::sync::mpsc::error::TrySendError::Full(poison)) =
tx.try_send(Frame::Poison)
{
if let Ok(rt) = tokio::runtime::Handle::try_current() {
// A32-14, stays a task: one bounded send of the poison frame, ended when it is in.
rt.spawn(async move {
let _ = tx.send(poison).await;
});
}
}
}
}
/// **Open the restore stream** with the host's own call: `call(body, codec)` runs
/// `Grpc::client_streaming(request, RESTORE_PATH, codec)` with the host's channel and
/// credential (each host signs differently; this crate never holds a key).
pub fn restore_stream<C, R, F, Fut>(call: F) -> TonicRestore<C, R>
where
C: Send + 'static,
R: Send + 'static,
F: FnOnce(RestoreBody<C>, PoisonCodec<C, R>) -> Fut,
Fut: std::future::Future<Output = Result<tonic::Response<R>, tonic::Status>> + Send + 'static,
{
let (tx, rx) = tokio::sync::mpsc::channel::<Frame<C>>(16);
// A32-14, stays a task: the one gRPC call of this restore; its handle is kept and awaited.
let call = tokio::spawn(call(
tokio_stream::wrappers::ReceiverStream::new(rx),
PoisonCodec(Prost::default()),
));
TonicRestore {
tx: Some(tx),
call: Some(call),
_r: PhantomData,
}
}
impl<C: Send + 'static, R> TonicRestore<C, R> {
async fn answer(mut self) -> Result<R, Fault> {
drop(self.tx.take());
let Some(call) = self.call.take() else {
return Err(Fault::new(
"internal",
"the restore call was already answered",
));
};
match call.await {
Ok(Ok(r)) => Ok(r.into_inner()),
Ok(Err(status)) => Err(fault(&status)),
Err(e) => Err(Fault::new("internal", format!("the restore task: {e}"))),
}
}
}
impl<C, R> RestoreSink<C> for TonicRestore<C, R>
where
C: Send + 'static,
R: AsRestored + Send + 'static,
{
async fn send(&mut self, chunk: C) -> Result<(), Fault> {
if let Some(tx) = &self.tx {
if tx.send(Frame::Chunk(chunk)).await.is_ok() {
return Ok(());
}
}
// The call ended: its status is the reason.
let Some(call) = self.call.as_mut() else {
return Err(Fault::new(
"internal",
"the restore call was already answered",
));
};
Err(match call.await {
Ok(Err(status)) => fault(&status),
Ok(Ok(_)) => Fault::new("aborted", "the main answered before the stream ended"),
Err(e) => Fault::new("internal", format!("the restore task: {e}")),
})
}
async fn close(self) -> Result<Restored, Fault> {
self.answer().await.map(|r| r.restored())
}
async fn abort(mut self, _why: &str) {
if let Some(tx) = self.tx.take() {
let _ = tx.send(Frame::Poison).await;
}
let _ = self.answer().await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
use tokio_stream::StreamExt;
struct Done;
impl AsRestored for Done {
fn restored(&self) -> Restored {
Restored::default()
}
}
/// What the main's end of the stream saw: `Some(true)` a chunk, `Some(false)` the poison.
fn call_that_records(seen: Arc<Mutex<Vec<bool>>>) -> TonicRestore<u32, Done> {
restore_stream(
move |mut body: RestoreBody<u32>, _codec: PoisonCodec<u32, Done>| async move {
while let Some(frame) = body.next().await {
seen.lock().unwrap().push(matches!(frame, Frame::Chunk(_)));
}
Ok(tonic::Response::new(Done))
},
)
}
/// A pump whose future is DROPPED mid-stream (its task aborted at shutdown, its
/// caller gone) must fail the restore, never end it: a clean end is a short
/// stream the fresh main would finish as its restore.
#[tokio::test]
async fn a_restore_sink_dropped_mid_stream_poisons_the_stream_it_never_ends_it() {
let seen = Arc::new(Mutex::new(Vec::new()));
let mut sink = call_that_records(Arc::clone(&seen));
sink.send(1).await.unwrap();
drop(sink);
for _ in 0..200 {
if seen.lock().unwrap().len() >= 2 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
assert_eq!(
*seen.lock().unwrap(),
vec![true, false],
"the chunk, then the poison"
);
}
/// ★ stinky #17 in the ONE pump (BROKEN 11:58): `next()` on a twin that went
/// silent mid-stream waited for ever, holding the recharge (and the box's restore
/// lock) open. Every tonic caller's stream goes through [`within`].
#[tokio::test]
async fn a_silent_twin_is_given_up_on_by_name_and_a_talking_one_is_not() {
let wait = std::time::Duration::from_millis(20);
let silent = tokio::time::timeout(
std::time::Duration::from_secs(5),
within::<u32>(wait, std::future::pending()),
)
.await
.expect("within must not wait for ever");
match silent {
Some(Err(f)) => assert!(
f.code == "deadline_exceeded" && f.message.contains("silent"),
"{f}"
),
other => panic!("a silent twin is a fault, not {other:?}"),
}
assert_eq!(
within(wait, async { Some(Ok::<u32, Fault>(7)) }).await,
Some(Ok(7))
);
assert_eq!(
within::<u32>(wait, async { None }).await,
None,
"an end is an end"
);
assert_eq!(
MESSAGE_WAIT,
std::time::Duration::from_secs(120),
"durability's wait (stinky #17)"
);
}
#[tokio::test]
async fn a_closed_restore_ends_the_stream_cleanly() {
let seen = Arc::new(Mutex::new(Vec::new()));
let mut sink = call_that_records(Arc::clone(&seen));
sink.send(1).await.unwrap();
sink.close().await.unwrap();
assert_eq!(*seen.lock().unwrap(), vec![true]);
}
}