heddle_thread_api/
live_replication.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Continuous, bounded replication on caller-authenticated protobuf streams.
3//! One feed is shared by every observer/replicator of a local Thread. Incoming
4//! metadata writes never move a checkout or install unrequested source blobs.
5use std::{future::Future, sync::Arc, time::Duration};
6
7use api::v2::client::{MessageReader, MessageWriter};
8use heddle_object_model::object::ContentHash;
9use prost::Message;
10#[cfg(feature = "native")]
11use repo::thread_replication::ThreadReplica;
12use tokio::{
13    sync::{Mutex, Notify, mpsc, oneshot, watch},
14    task::{AbortHandle, JoinHandle},
15};
16use tracing::{Instrument, instrument::WithSubscriber};
17
18use crate::{
19    contract::*,
20    replication::{Frame, InputUnit, Outbound, Session, store::ReplicaStore},
21    transport,
22};
23
24#[derive(Debug, thiserror::Error)]
25pub enum Error<E: std::error::Error + 'static> {
26    #[error(transparent)]
27    Transport(#[from] transport::Error),
28    #[error("replica store: {0}")]
29    Store(#[source] E),
30    #[error(transparent)]
31    Protocol(#[from] crate::replication::Error),
32    #[error("replication worker: {0}")]
33    Worker(String),
34    #[error("replication response budget exhausted; reopen from durable state")]
35    Backpressure,
36    #[error("Thread change feed stopped")]
37    FeedClosed,
38    #[error("Thread policy changed; reconnect from durable receipts")]
39    PolicyChanged,
40}
41pub type Result<T, E> = std::result::Result<T, Error<E>>;
42
43impl<E: std::error::Error + 'static> From<crate::replication::StoreError<E>> for Error<E> {
44    fn from(error: crate::replication::StoreError<E>) -> Self {
45        match error {
46            crate::replication::StoreError::Store(error) => Self::Store(error),
47            crate::replication::StoreError::Protocol(error) => Self::Protocol(error),
48        }
49    }
50}
51
52struct AbortOnDrop(AbortHandle);
53impl Drop for AbortOnDrop {
54    fn drop(&mut self) {
55        self.0.abort();
56    }
57}
58
59/// Create once per Thread, then clone. This observes writes from other local
60/// processes using the durable generation; it does not poll once per stream.
61#[derive(Clone)]
62pub struct Feed {
63    thread: ContentHash,
64    changes: watch::Receiver<Option<i64>>,
65    _task: Option<Arc<AbortOnDrop>>,
66}
67impl Feed {
68    /// The host shares one durable-generation watcher per Thread. Subscribe
69    /// before the initial announcement; missed notifications trigger a fresh
70    /// generation check, never an assumed accepted frontier.
71    pub fn from_changes(thread: ContentHash, changes: watch::Receiver<Option<i64>>) -> Self {
72        Self {
73            thread,
74            changes,
75            _task: None,
76        }
77    }
78
79    #[cfg(feature = "native")]
80    pub async fn new(
81        replica: ThreadReplica,
82    ) -> std::result::Result<Self, crate::replication::native::Error> {
83        let thread = replica.thread_id();
84        let initial = replica.clone();
85        let generation = tokio::task::spawn_blocking(move || initial.generation())
86            .await
87            .map_err(crate::replication::native::Error::from)??;
88        let (sender, changes) = watch::channel(Some(generation));
89        let task = tokio::spawn(async move {
90            let mut interval = tokio::time::interval(Duration::from_millis(200));
91            loop {
92                interval.tick().await;
93                if sender.is_closed() {
94                    break;
95                }
96                let replica = replica.clone();
97                let result = tokio::task::spawn_blocking(move || replica.generation()).await;
98                match result {
99                    Ok(Ok(generation)) => {
100                        sender.send_if_modified(|old| {
101                            if *old == Some(generation) {
102                                false
103                            } else {
104                                *old = Some(generation);
105                                true
106                            }
107                        });
108                    }
109                    _ => {
110                        let _ = sender.send(None);
111                        break;
112                    }
113                }
114            }
115        });
116        Ok(Self {
117            thread,
118            changes,
119            _task: Some(Arc::new(AbortOnDrop(task.abort_handle()))),
120        })
121    }
122}
123
124#[derive(Clone, Copy)]
125pub enum Side {
126    Initiator,
127    Acceptor,
128}
129impl Side {
130    fn decode<E: std::error::Error + 'static>(self, bytes: &[u8]) -> Result<Frame, E> {
131        input::reservation(bytes, 64)?;
132        Ok(match self {
133            Self::Initiator => Frame::from_response(
134                ReplicateThreadResponse::decode(bytes).map_err(transport::Error::from)?,
135            )?,
136            Self::Acceptor => Frame::from_request(
137                ReplicateThreadRequest::decode(bytes).map_err(transport::Error::from)?,
138            )?,
139        })
140    }
141    fn encode(self, frame: Frame) -> Vec<u8> {
142        match self {
143            Self::Initiator => frame.request().encode_to_vec(),
144            Self::Acceptor => frame.response().encode_to_vec(),
145        }
146    }
147}
148
149#[allow(clippy::large_enum_variant)] // heddle-api's inline import-authority bundle; boxing adds a heap hop per frame
150enum Event {
151    Incoming(Frame),
152    Announce,
153    Maintain,
154}
155
156struct Queued {
157    item: Outbound,
158    immediate_receipt: bool,
159    generation: u64,
160    delivered: Option<oneshot::Sender<()>>,
161}
162async fn completion<E: std::error::Error + 'static>(
163    queue: &mpsc::Sender<Queued>,
164    item: Outbound,
165    generation: u64,
166) -> Result<(), E> {
167    let (sent, received) = oneshot::channel();
168    queue
169        .send(Queued {
170            item,
171            immediate_receipt: true,
172            generation,
173            delivered: Some(sent),
174        })
175        .await
176        .map_err(|_| Error::Backpressure)?;
177    received
178        .await
179        .map_err(|_| Error::Worker("replication sender stopped before receipt flush".into()))
180}
181fn requires_disclosure_fence<E: std::error::Error + 'static>(unit: &InputUnit) -> Result<bool, E> {
182    use heddle_object_model::object::thread_replication::{
183        ThreadOperation, ThreadOperationBody,
184        metadata::{Control, ThreadControl},
185    };
186    let InputUnit::Operation(received) = unit else {
187        return Ok(false);
188    };
189    let operation = ThreadOperation::decode(&received.original.canonical)
190        .map_err(crate::replication::Error::from)?;
191    // Canonical parents must share the exact metadata property. Only these
192    // policy properties can settle another pending policy mutation.
193    let ThreadOperationBody::Metadata(bytes) = &operation.body else {
194        return Ok(false);
195    };
196    let control = ThreadControl::decode(bytes).map_err(crate::replication::Error::from)?;
197    Ok(matches!(
198        control.control,
199        Control::Audience(_) | Control::Retention(_) | Control::Sharing(_)
200    ))
201}
202#[path = "live_replication_input.rs"]
203pub mod input;
204fn unit_acceptance(
205    unit: &InputUnit,
206) -> Option<&Arc<crypto::original_boundary_acceptance::SignedBoundaryAcceptance>> {
207    let InputUnit::Operation(received) = unit else {
208        return None;
209    };
210    received
211        .authority_admission
212        .as_ref()?
213        .boundary_acceptance
214        .as_ref()
215}
216// Canonical/signature buffers remain owned by parsed units. Shared immutable
217// acceptance bytes are counted once while any remaining unit retains the Arc.
218// The caller separately retains the entire IntoIter backing allocation.
219fn retained_unit_bytes(units: &[InputUnit]) -> usize {
220    units
221        .iter()
222        .enumerate()
223        .map(|(index, unit)| {
224            let InputUnit::Operation(received) = unit else {
225                return 0;
226            };
227            let mut bytes =
228                received.original.canonical.capacity() + received.original.signature.capacity();
229            if let Some(receipt) = &received.authority_admission {
230                bytes += receipt.canonical.capacity() + receipt.signature.capacity();
231            }
232            if let Some(acceptance) = unit_acceptance(unit)
233                && !units[..index]
234                    .iter()
235                    .filter_map(unit_acceptance)
236                    .any(|prior| Arc::ptr_eq(prior, acceptance))
237            {
238                bytes += std::mem::size_of_val(acceptance.as_ref())
239                    + 2 * std::mem::size_of::<usize>()
240                    + acceptance.canonical.capacity()
241                    + acceptance.signature.capacity();
242            }
243            bytes
244        })
245        .sum()
246}
247
248/// A permission-only recheck must not wait for an output memory reservation
249/// already retained by this stream. Work may produce one bounded output frame.
250#[derive(Clone, Copy, Debug, PartialEq, Eq)]
251pub enum Activity {
252    /// Idle clock tick: verify locally known identity, revocation and time
253    /// caveats only. Hosts must not query the store or reserve output memory.
254    /// Every actual input and disclosure still uses the fresh gates below.
255    Idle,
256    /// Shrink retained input accounting after consumed decoded allocations have
257    /// been dropped. This is an accounting callback, not authorization.
258    InputConsumed {
259        remaining_bytes: usize,
260    },
261    Check,
262    /// Advance input admission and bounded control queues. The reader already
263    /// accounts for the input; waiting for output memory here could deadlock
264    /// every receiver while it holds the memory needed by those producers.
265    Receive,
266    /// Load and encode one output frame, retaining its memory through delivery.
267    Work,
268    /// Encode the immediate receipt for this session's just-admitted input.
269    ReceiptWork,
270    /// Recheck only that immediate receipt, never prepared source disclosure.
271    ReceiptCheck,
272    /// Post-receipt peer bookkeeping for an already admitted pending input.
273    Bookkeeping,
274}
275
276/// Hosts can charge database work and output memory separately. Finishing the
277/// work releases execution slots while the returned lease covers delivery.
278/// Devices without a shared work scheduler can keep returning `()`.
279pub trait ActivityGuard: Send {
280    type Retained: Send;
281    fn finish(self, encoded_bytes: usize) -> std::result::Result<Self::Retained, transport::Error>;
282}
283impl ActivityGuard for () {
284    type Retained = ();
285    fn finish(self, _: usize) -> std::result::Result<(), transport::Error> {
286        Ok(())
287    }
288}
289
290/// Call after validating the opening, endpoint bindings, Thread, and facets.
291/// `authorize` rechecks the live host permission, including expiry/revocation.
292/// It runs before every admission and output, including queued output. Readers
293/// must enforce the negotiated frame bound before allocating message bodies.
294pub async fn run<B, R, W, G, F, A>(
295    session: Session<B>,
296    reader: R,
297    writer: W,
298    side: Side,
299    feed: &Feed,
300    authorize: G,
301) -> Result<(), B::Error>
302where
303    B: ReplicaStore,
304    R: MessageReader<Error = transport::Error>,
305    W: MessageWriter<Error = transport::Error> + 'static,
306    G: Fn(Activity) -> F + Clone + Send + Sync + 'static,
307    F: Future<Output = std::result::Result<A, transport::Error>> + Send,
308    A: ActivityGuard,
309{
310    run_with_idle_clock(
311        session,
312        reader,
313        writer,
314        side,
315        feed,
316        authorize,
317        tokio::time::interval(Duration::from_secs(1)),
318    )
319    .await
320}
321
322// Keep the idle clock injectable inside the driver so progress tests can prove
323// that queue wakeups work without a heartbeat rescuing a missed notification.
324async fn run_with_idle_clock<B, R, W, G, F, A>(
325    mut session: Session<B>,
326    mut reader: R,
327    mut writer: W,
328    side: Side,
329    feed: &Feed,
330    authorize: G,
331    mut heartbeat: tokio::time::Interval,
332) -> Result<(), B::Error>
333where
334    B: ReplicaStore,
335    R: MessageReader<Error = transport::Error>,
336    W: MessageWriter<Error = transport::Error> + 'static,
337    G: Fn(Activity) -> F + Clone + Send + Sync + 'static,
338    F: Future<Output = std::result::Result<A, transport::Error>> + Send,
339    A: ActivityGuard,
340{
341    if feed.thread != session.replica.thread_id() {
342        return Err(transport::Error::Protocol("change feed belongs to another Thread").into());
343    }
344    drop(authorize(Activity::Check).await?);
345    let mut changes = feed.changes.clone();
346    let (queue, mut outgoing) = mpsc::channel::<Queued>(256);
347    let progress = Arc::new(Notify::new());
348    let sender_progress = progress.clone();
349    let (completions, mut incoming_completions) = mpsc::channel::<Queued>(1);
350    let sender_session = session.clone();
351    let sender_authorize = authorize.clone();
352    let delivery = Arc::new(Mutex::new(()));
353    let sender_delivery = delivery.clone();
354    let (disclosures, mut disclosure_changes) = watch::channel(0u64);
355    let mut disclosure_generation = 0u64;
356    let mut sender: JoinHandle<Result<(), B::Error>> = tokio::spawn(
357        (async move {
358            let mut deferred = None;
359            let mut priority = None;
360            let mut priority_closed = false;
361            loop {
362                let next = if let Some(item) = priority.take() { Some(item) }
363                    else if let Ok(item) = incoming_completions.try_recv() { Some(item) }
364                    else if let Some(item) = deferred.take() { Some(item) }
365                    else { tokio::select! {
366                        biased;
367                        item = incoming_completions.recv(), if !priority_closed => match item { Some(item) => Some(item), None => { priority_closed = true; continue; } },
368                        item = outgoing.recv() => item,
369                    }};
370                let Some(queued) = next else { break; };
371                sender_progress.notify_one();
372                let generation = queued.generation;
373                let immediate_receipt = queued.immediate_receipt;
374                if generation != *disclosure_changes.borrow_and_update() {
375                    continue;
376                }
377                let session = sender_session.clone();
378                let gate = sender_authorize.clone();
379                // Only cancel an unacquired work reservation. Once admitted,
380                // store implementations may own non-cancellable blocking work;
381                // keep its lease until it completes, then discard stale output.
382                let activity = if immediate_receipt { gate(Activity::ReceiptWork).await? } else {
383                    tokio::select! {
384                        biased;
385                        changed = disclosure_changes.changed() => {
386                            changed.map_err(|_| Error::FeedClosed)?;
387                            continue;
388                        }
389                        received = incoming_completions.recv(), if !priority_closed => {
390                            if let Some(received) = received {
391                                deferred = Some(queued);
392                                priority = Some(received);
393                                continue;
394                            }
395                            priority_closed = true;
396                            deferred = Some(queued);
397                            continue;
398                        }
399                        activity = gate(Activity::Work) => activity?,
400                    }
401                };
402                let Queued { item, delivered, .. } = queued;
403                let prepared = async {
404                    let frame = match item {
405                        Outbound::Operation(id) => session.export_operation(id).await?,
406                        Outbound::Frame(frame) => {
407                            if let Frame::Have(have) = &frame {
408                                let allowed = session.export_facets().await?;
409                                for frontier in &have.frontiers {
410                                    if !allowed.contains(&crate::replication::native_facet(
411                                        frontier.facet,
412                                    )?) {
413                                        return Err(transport::Error::Protocol(
414                                            "sharing policy changed before disclosure",
415                                        )
416                                        .into());
417                                    }
418                                }
419                            }
420                            frame
421                        }
422                    };
423                    let encoded = side.encode(frame);
424                    let retained = activity.finish(encoded.len())?;
425                    Ok::<_, Error<B::Error>>((encoded, retained))
426                }.await;
427                if generation != *disclosure_changes.borrow_and_update() { continue; }
428                let (encoded, retained) = prepared?;
429                let delivery_guard = sender_delivery.lock().await;
430                if generation != *disclosure_changes.borrow_and_update() {
431                    continue;
432                }
433                // Permission-only checks cannot acquire work slots while this
434                // mutex is held; an input may already own the same work pool.
435                drop(gate(if immediate_receipt { Activity::ReceiptCheck } else { Activity::Check }).await?);
436                writer.send(encoded).await?;
437                drop(delivery_guard);
438                drop(retained);
439                if let Some(delivered) = delivered { let _ = delivered.send(()); }
440            }
441            writer.finish().await?;
442            Ok(())
443        })
444        .in_current_span()
445        .with_current_subscriber(),
446    );
447    // Dropping a JoinHandle detaches it. Abort explicitly so cancellation drops
448    // the transport writer too, including while its peer is applying pressure.
449    let _sender_guard = AbortOnDrop(sender.abort_handle());
450    heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
451    let mut announce = true;
452    let mut maintain = true;
453    loop {
454        let event = tokio::select! {
455            result = &mut sender => return result.map_err(worker)?,
456            result = reader.next() => match result? {
457                Some(bytes) => Event::Incoming(side.decode(&bytes)?),
458                None => break,
459            },
460            result = changes.changed() => {
461                result.map_err(|_| Error::FeedClosed)?;
462                if changes.borrow_and_update().is_none() { return Err(Error::FeedClosed); }
463                announce = true;
464                maintain = true;
465                continue
466            },
467            _ = std::future::ready(()), if announce && queue.capacity() > 128 => Event::Announce,
468            _ = std::future::ready(()), if maintain && queue.capacity() > 128 => Event::Maintain,
469            _ = progress.notified() => continue,
470            _ = heartbeat.tick() => {
471                drop(authorize(Activity::Idle).await?);
472                continue
473            }
474        };
475        let event = match event {
476            Event::Incoming(frame) => {
477                maintain = true;
478                let preflight = authorize(Activity::Receive).await?;
479                let units = session.input_units(frame)?;
480                drop(preflight);
481                let container_bytes = units.capacity() * std::mem::size_of::<InputUnit>();
482                let mut units = units.into_iter();
483                while let Some(unit) = units.next() {
484                    let immediate_receipt = matches!(&unit, InputUnit::Operation(_));
485                    let policy = requires_disclosure_fence(&unit)?;
486                    let activity = authorize(Activity::Receive).await?;
487                    let delivery_guard = if policy {
488                        let guard = delivery.lock().await;
489                        disclosure_generation = disclosure_generation
490                            .checked_add(1)
491                            .ok_or(Error::Backpressure)?;
492                        disclosures.send_replace(disclosure_generation);
493                        Some(guard)
494                    } else {
495                        None
496                    };
497                    let output = session.handle_unit(unit).await?;
498                    drop(activity);
499                    drop(delivery_guard);
500                    let remaining_bytes = if units.len() == 0 {
501                        0
502                    } else {
503                        container_bytes + retained_unit_bytes(units.as_slice())
504                    };
505                    // IntoIter retains its backing allocation until dropped,
506                    // including after its final element has been consumed.
507                    if units.len() == 0 {
508                        drop(units);
509                        units = Vec::new().into_iter();
510                    }
511                    drop(authorize(Activity::InputConsumed { remaining_bytes }).await?);
512                    for item in output {
513                        if immediate_receipt && matches!(&item, Outbound::Frame(Frame::Receipt(_)))
514                        {
515                            completion(&completions, item, disclosure_generation).await?;
516                        } else {
517                            queue
518                                .try_send(Queued {
519                                    item,
520                                    immediate_receipt: false,
521                                    generation: disclosure_generation,
522                                    delivered: None,
523                                })
524                                .map_err(|_| Error::Backpressure)?;
525                        }
526                    }
527                    if session.has_input_bookkeeping() {
528                        let bookkeeping = authorize(Activity::Bookkeeping).await?;
529                        session.finish_input_bookkeeping().await?;
530                        drop(bookkeeping);
531                    }
532                    if policy {
533                        return Err(Error::PolicyChanged);
534                    }
535                }
536                continue;
537            }
538            other => other,
539        };
540        let activity = authorize(Activity::Receive).await?;
541        let output = match event {
542            Event::Incoming(frame) => {
543                maintain = true;
544                session.handle_input(frame).await?
545            }
546            Event::Announce => {
547                let frame = session.announcement().await?;
548                announce = frame.is_some();
549                frame.into_iter().map(Outbound::Frame).collect()
550            }
551            Event::Maintain => {
552                let frame = session.control().await?;
553                maintain = frame.is_some();
554                frame.into_iter().map(Outbound::Frame).collect()
555            }
556        };
557        drop(activity);
558        for item in output {
559            queue
560                .try_send(Queued {
561                    item,
562                    immediate_receipt: false,
563                    generation: disclosure_generation,
564                    delivered: None,
565                })
566                .map_err(|_| Error::Backpressure)?;
567        }
568    }
569    drop(queue);
570    drop(completions);
571    sender.await.map_err(worker)?
572}
573
574fn worker<E: std::error::Error + 'static>(error: tokio::task::JoinError) -> Error<E> {
575    Error::Worker(error.to_string())
576}
577
578#[cfg(all(test, feature = "native"))]
579#[path = "live_replication_tests.rs"]
580mod tests;
581
582#[cfg(all(test, feature = "native"))]
583#[path = "live_replication_schedule_tests.rs"]
584mod schedule_tests;
heddle_thread_api/live_replication.rs

heddle_thread_api/
live_replication.rs
