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
149enum Event {
150    Incoming(Frame),
151    Announce,
152    Maintain,
153}
154
155struct Queued {
156    item: Outbound,
157    immediate_receipt: bool,
158    generation: u64,
159    delivered: Option<oneshot::Sender<()>>,
160}
161async fn completion<E: std::error::Error + 'static>(
162    queue: &mpsc::Sender<Queued>,
163    item: Outbound,
164    generation: u64,
165) -> Result<(), E> {
166    let (sent, received) = oneshot::channel();
167    queue
168        .send(Queued {
169            item,
170            immediate_receipt: true,
171            generation,
172            delivered: Some(sent),
173        })
174        .await
175        .map_err(|_| Error::Backpressure)?;
176    received
177        .await
178        .map_err(|_| Error::Worker("replication sender stopped before receipt flush".into()))
179}
180fn requires_disclosure_fence<E: std::error::Error + 'static>(unit: &InputUnit) -> Result<bool, E> {
181    use heddle_object_model::object::thread_replication::{
182        ThreadOperation, ThreadOperationBody,
183        metadata::{Control, ThreadControl},
184    };
185    let InputUnit::Operation(received) = unit else {
186        return Ok(false);
187    };
188    let operation = ThreadOperation::decode(&received.original.canonical)
189        .map_err(crate::replication::Error::from)?;
190    // Canonical parents must share the exact metadata property. Only these
191    // policy properties can settle another pending policy mutation.
192    let ThreadOperationBody::Metadata(bytes) = &operation.body else {
193        return Ok(false);
194    };
195    let control = ThreadControl::decode(bytes).map_err(crate::replication::Error::from)?;
196    Ok(matches!(
197        control.control,
198        Control::Audience(_) | Control::Retention(_) | Control::Sharing(_)
199    ))
200}
201#[path = "live_replication_input.rs"]
202pub mod input;
203fn unit_acceptance(
204    unit: &InputUnit,
205) -> Option<&Arc<crypto::original_boundary_acceptance::SignedBoundaryAcceptance>> {
206    let InputUnit::Operation(received) = unit else {
207        return None;
208    };
209    received
210        .authority_admission
211        .as_ref()?
212        .boundary_acceptance
213        .as_ref()
214}
215// Canonical/signature buffers remain owned by parsed units. Shared immutable
216// acceptance bytes are counted once while any remaining unit retains the Arc.
217// The caller separately retains the entire IntoIter backing allocation.
218fn retained_unit_bytes(units: &[InputUnit]) -> usize {
219    units
220        .iter()
221        .enumerate()
222        .map(|(index, unit)| {
223            let InputUnit::Operation(received) = unit else {
224                return 0;
225            };
226            let mut bytes =
227                received.original.canonical.capacity() + received.original.signature.capacity();
228            if let Some(receipt) = &received.authority_admission {
229                bytes += receipt.canonical.capacity() + receipt.signature.capacity();
230            }
231            if let Some(acceptance) = unit_acceptance(unit)
232                && !units[..index]
233                    .iter()
234                    .filter_map(unit_acceptance)
235                    .any(|prior| Arc::ptr_eq(prior, acceptance))
236            {
237                bytes += std::mem::size_of_val(acceptance.as_ref())
238                    + 2 * std::mem::size_of::<usize>()
239                    + acceptance.canonical.capacity()
240                    + acceptance.signature.capacity();
241            }
242            bytes
243        })
244        .sum()
245}
246
247/// A permission-only recheck must not wait for an output memory reservation
248/// already retained by this stream. Work may produce one bounded output frame.
249#[derive(Clone, Copy, Debug, PartialEq, Eq)]
250pub enum Activity {
251    /// Idle clock tick: verify locally known identity, revocation and time
252    /// caveats only. Hosts must not query the store or reserve output memory.
253    /// Every actual input and disclosure still uses the fresh gates below.
254    Idle,
255    /// Shrink retained input accounting after consumed decoded allocations have
256    /// been dropped. This is an accounting callback, not authorization.
257    InputConsumed {
258        remaining_bytes: usize,
259    },
260    Check,
261    /// Advance input admission and bounded control queues. The reader already
262    /// accounts for the input; waiting for output memory here could deadlock
263    /// every receiver while it holds the memory needed by those producers.
264    Receive,
265    /// Load and encode one output frame, retaining its memory through delivery.
266    Work,
267    /// Encode the immediate receipt for this session's just-admitted input.
268    ReceiptWork,
269    /// Recheck only that immediate receipt, never prepared source disclosure.
270    ReceiptCheck,
271    /// Post-receipt peer bookkeeping for an already admitted pending input.
272    Bookkeeping,
273}
274
275/// Hosts can charge database work and output memory separately. Finishing the
276/// work releases execution slots while the returned lease covers delivery.
277/// Devices without a shared work scheduler can keep returning `()`.
278pub trait ActivityGuard: Send {
279    type Retained: Send;
280    fn finish(self, encoded_bytes: usize) -> std::result::Result<Self::Retained, transport::Error>;
281}
282impl ActivityGuard for () {
283    type Retained = ();
284    fn finish(self, _: usize) -> std::result::Result<(), transport::Error> {
285        Ok(())
286    }
287}
288
289/// Call after validating the opening, endpoint bindings, Thread, and facets.
290/// `authorize` rechecks the live host permission, including expiry/revocation.
291/// It runs before every admission and output, including queued output. Readers
292/// must enforce the negotiated frame bound before allocating message bodies.
293pub async fn run<B, R, W, G, F, A>(
294    session: Session<B>,
295    reader: R,
296    writer: W,
297    side: Side,
298    feed: &Feed,
299    authorize: G,
300) -> Result<(), B::Error>
301where
302    B: ReplicaStore,
303    R: MessageReader<Error = transport::Error>,
304    W: MessageWriter<Error = transport::Error> + 'static,
305    G: Fn(Activity) -> F + Clone + Send + Sync + 'static,
306    F: Future<Output = std::result::Result<A, transport::Error>> + Send,
307    A: ActivityGuard,
308{
309    run_with_idle_clock(
310        session,
311        reader,
312        writer,
313        side,
314        feed,
315        authorize,
316        tokio::time::interval(Duration::from_secs(1)),
317    )
318    .await
319}
320
321// Keep the idle clock injectable inside the driver so progress tests can prove
322// that queue wakeups work without a heartbeat rescuing a missed notification.
323async fn run_with_idle_clock<B, R, W, G, F, A>(
324    mut session: Session<B>,
325    mut reader: R,
326    mut writer: W,
327    side: Side,
328    feed: &Feed,
329    authorize: G,
330    mut heartbeat: tokio::time::Interval,
331) -> Result<(), B::Error>
332where
333    B: ReplicaStore,
334    R: MessageReader<Error = transport::Error>,
335    W: MessageWriter<Error = transport::Error> + 'static,
336    G: Fn(Activity) -> F + Clone + Send + Sync + 'static,
337    F: Future<Output = std::result::Result<A, transport::Error>> + Send,
338    A: ActivityGuard,
339{
340    if feed.thread != session.replica.thread_id() {
341        return Err(transport::Error::Protocol("change feed belongs to another Thread").into());
342    }
343    drop(authorize(Activity::Check).await?);
344    let mut changes = feed.changes.clone();
345    let (queue, mut outgoing) = mpsc::channel::<Queued>(256);
346    let progress = Arc::new(Notify::new());
347    let sender_progress = progress.clone();
348    let (completions, mut incoming_completions) = mpsc::channel::<Queued>(1);
349    let sender_session = session.clone();
350    let sender_authorize = authorize.clone();
351    let delivery = Arc::new(Mutex::new(()));
352    let sender_delivery = delivery.clone();
353    let (disclosures, mut disclosure_changes) = watch::channel(0u64);
354    let mut disclosure_generation = 0u64;
355    let mut sender: JoinHandle<Result<(), B::Error>> = tokio::spawn(
356        (async move {
357            let mut deferred = None;
358            let mut priority = None;
359            let mut priority_closed = false;
360            loop {
361                let next = if let Some(item) = priority.take() { Some(item) }
362                    else if let Ok(item) = incoming_completions.try_recv() { Some(item) }
363                    else if let Some(item) = deferred.take() { Some(item) }
364                    else { tokio::select! {
365                        biased;
366                        item = incoming_completions.recv(), if !priority_closed => match item { Some(item) => Some(item), None => { priority_closed = true; continue; } },
367                        item = outgoing.recv() => item,
368                    }};
369                let Some(queued) = next else { break; };
370                sender_progress.notify_one();
371                let generation = queued.generation;
372                let immediate_receipt = queued.immediate_receipt;
373                if generation != *disclosure_changes.borrow_and_update() {
374                    continue;
375                }
376                let session = sender_session.clone();
377                let gate = sender_authorize.clone();
378                // Only cancel an unacquired work reservation. Once admitted,
379                // store implementations may own non-cancellable blocking work;
380                // keep its lease until it completes, then discard stale output.
381                let activity = if immediate_receipt { gate(Activity::ReceiptWork).await? } else {
382                    tokio::select! {
383                        biased;
384                        changed = disclosure_changes.changed() => {
385                            changed.map_err(|_| Error::FeedClosed)?;
386                            continue;
387                        }
388                        received = incoming_completions.recv(), if !priority_closed => {
389                            if let Some(received) = received {
390                                deferred = Some(queued);
391                                priority = Some(received);
392                                continue;
393                            }
394                            priority_closed = true;
395                            deferred = Some(queued);
396                            continue;
397                        }
398                        activity = gate(Activity::Work) => activity?,
399                    }
400                };
401                let Queued { item, delivered, .. } = queued;
402                let prepared = async {
403                    let frame = match item {
404                        Outbound::Operation(id) => session.export_operation(id).await?,
405                        Outbound::Frame(frame) => {
406                            if let Frame::Have(have) = &frame {
407                                let allowed = session.export_facets().await?;
408                                for frontier in &have.frontiers {
409                                    if !allowed.contains(&crate::replication::native_facet(
410                                        frontier.facet,
411                                    )?) {
412                                        return Err(transport::Error::Protocol(
413                                            "sharing policy changed before disclosure",
414                                        )
415                                        .into());
416                                    }
417                                }
418                            }
419                            frame
420                        }
421                    };
422                    let encoded = side.encode(frame);
423                    let retained = activity.finish(encoded.len())?;
424                    Ok::<_, Error<B::Error>>((encoded, retained))
425                }.await;
426                if generation != *disclosure_changes.borrow_and_update() { continue; }
427                let (encoded, retained) = prepared?;
428                let delivery_guard = sender_delivery.lock().await;
429                if generation != *disclosure_changes.borrow_and_update() {
430                    continue;
431                }
432                // Permission-only checks cannot acquire work slots while this
433                // mutex is held; an input may already own the same work pool.
434                drop(gate(if immediate_receipt { Activity::ReceiptCheck } else { Activity::Check }).await?);
435                writer.send(encoded).await?;
436                drop(delivery_guard);
437                drop(retained);
438                if let Some(delivered) = delivered { let _ = delivered.send(()); }
439            }
440            writer.finish().await?;
441            Ok(())
442        })
443        .in_current_span()
444        .with_current_subscriber(),
445    );
446    // Dropping a JoinHandle detaches it. Abort explicitly so cancellation drops
447    // the transport writer too, including while its peer is applying pressure.
448    let _sender_guard = AbortOnDrop(sender.abort_handle());
449    heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
450    let mut announce = true;
451    let mut maintain = true;
452    loop {
453        let event = tokio::select! {
454            result = &mut sender => return result.map_err(worker)?,
455            result = reader.next() => match result? {
456                Some(bytes) => Event::Incoming(side.decode(&bytes)?),
457                None => break,
458            },
459            result = changes.changed() => {
460                result.map_err(|_| Error::FeedClosed)?;
461                if changes.borrow_and_update().is_none() { return Err(Error::FeedClosed); }
462                announce = true;
463                maintain = true;
464                continue
465            },
466            _ = std::future::ready(()), if announce && queue.capacity() > 128 => Event::Announce,
467            _ = std::future::ready(()), if maintain && queue.capacity() > 128 => Event::Maintain,
468            _ = progress.notified() => continue,
469            _ = heartbeat.tick() => {
470                drop(authorize(Activity::Idle).await?);
471                continue
472            }
473        };
474        let event = match event {
475            Event::Incoming(frame) => {
476                maintain = true;
477                let preflight = authorize(Activity::Receive).await?;
478                let units = session.input_units(frame)?;
479                drop(preflight);
480                let container_bytes = units.capacity() * std::mem::size_of::<InputUnit>();
481                let mut units = units.into_iter();
482                while let Some(unit) = units.next() {
483                    let immediate_receipt = matches!(&unit, InputUnit::Operation(_));
484                    let policy = requires_disclosure_fence(&unit)?;
485                    let activity = authorize(Activity::Receive).await?;
486                    let delivery_guard = if policy {
487                        let guard = delivery.lock().await;
488                        disclosure_generation = disclosure_generation
489                            .checked_add(1)
490                            .ok_or(Error::Backpressure)?;
491                        disclosures.send_replace(disclosure_generation);
492                        Some(guard)
493                    } else {
494                        None
495                    };
496                    let output = session.handle_unit(unit).await?;
497                    drop(activity);
498                    drop(delivery_guard);
499                    let remaining_bytes = if units.len() == 0 {
500                        0
501                    } else {
502                        container_bytes + retained_unit_bytes(units.as_slice())
503                    };
504                    // IntoIter retains its backing allocation until dropped,
505                    // including after its final element has been consumed.
506                    if units.len() == 0 {
507                        drop(units);
508                        units = Vec::new().into_iter();
509                    }
510                    drop(authorize(Activity::InputConsumed { remaining_bytes }).await?);
511                    for item in output {
512                        if immediate_receipt && matches!(&item, Outbound::Frame(Frame::Receipt(_)))
513                        {
514                            completion(&completions, item, disclosure_generation).await?;
515                        } else {
516                            queue
517                                .try_send(Queued {
518                                    item,
519                                    immediate_receipt: false,
520                                    generation: disclosure_generation,
521                                    delivered: None,
522                                })
523                                .map_err(|_| Error::Backpressure)?;
524                        }
525                    }
526                    if session.has_input_bookkeeping() {
527                        let bookkeeping = authorize(Activity::Bookkeeping).await?;
528                        session.finish_input_bookkeeping().await?;
529                        drop(bookkeeping);
530                    }
531                    if policy {
532                        return Err(Error::PolicyChanged);
533                    }
534                }
535                continue;
536            }
537            other => other,
538        };
539        let activity = authorize(Activity::Receive).await?;
540        let output = match event {
541            Event::Incoming(frame) => {
542                maintain = true;
543                session.handle_input(frame).await?
544            }
545            Event::Announce => {
546                let frame = session.announcement().await?;
547                announce = frame.is_some();
548                frame.into_iter().map(Outbound::Frame).collect()
549            }
550            Event::Maintain => {
551                let frame = session.control().await?;
552                maintain = frame.is_some();
553                frame.into_iter().map(Outbound::Frame).collect()
554            }
555        };
556        drop(activity);
557        for item in output {
558            queue
559                .try_send(Queued {
560                    item,
561                    immediate_receipt: false,
562                    generation: disclosure_generation,
563                    delivered: None,
564                })
565                .map_err(|_| Error::Backpressure)?;
566        }
567    }
568    drop(queue);
569    drop(completions);
570    sender.await.map_err(worker)?
571}
572
573fn worker<E: std::error::Error + 'static>(error: tokio::task::JoinError) -> Error<E> {
574    Error::Worker(error.to_string())
575}
576
577#[cfg(all(test, feature = "native"))]
578#[path = "live_replication_tests.rs"]
579mod tests;
580
581#[cfg(all(test, feature = "native"))]
582#[path = "live_replication_schedule_tests.rs"]
583mod schedule_tests;
heddle_thread_api/live_replication.rs

heddle_thread_api/
live_replication.rs
