ruststream-gcp-pubsub 0.7.0

//! [`PubSubSubscriber`]: a stream of deliveries backed by a pump task.
//!
//! The client's `MessageStream` exposes an inherent `next()` that is not documented
//! cancel-safe, so the crate owns a pump task per subscription: it drives `next`, converts
//! deliveries, and forwards them into a bounded channel the subscriber polls. Settlement needs
//! no round trip through the pump - the client's `Handler` travels inside each message and is
//! consumed by `ack`/`nack` directly.
//!
//! Batches are assembled on top of that channel. The streaming pull hands over one delivery at a
//! time, so the subscriber wraps itself in the core's [`BufferedSubscriber`]: the batch size
//! arrives per subscription from the mount site's `batch(n)`, and the deadline that closes a
//! partial batch is this crate's, tunable with
//! [`GooglePubSub::batch_wait`](crate::GooglePubSub::batch_wait).

use std::num::NonZeroUsize;
#[cfg(feature = "testing")]
use std::time::Duration;

use futures::Stream;

use google_cloud_pubsub::subscriber::{MessageStream, ShutdownToken};
use ruststream::{BatchSubscriber, BufferedSubscriber, Subscriber};
use tokio::runtime::Handle;
use tokio::sync::mpsc;

use crate::broker::Core;
use crate::error::{PubSubError, box_err};
#[cfg(feature = "testing")]
use crate::in_process::Consumer;
use crate::message::PubSubMessage;
use crate::runtime_slot::RuntimeSlot;
use crate::subscription::{DeliveryLimits, GooglePubSub};

/// How many converted deliveries may sit between the pump and the consumer. Real prefetch is
/// the client's own flow control (`max_outstanding`); this only decouples the two loops.
const CHANNEL_CAPACITY: usize = 16;

/// A subscription to one Pub/Sub subscription; yields [`PubSubMessage`]s one at a time, or in
/// batches of at most the size the mount site named.
///
/// Dropping the subscriber signals the client's shutdown token, which drains the stream and
/// stops the pump task.
pub struct PubSubSubscriber {
    subscription: String,
    // Named for what it is rather than what it yields: `buffer.batches(size)` reads, and
    // `batches.batches(size)` would not.
    buffer: BufferedSubscriber<Deliveries>,
}

impl std::fmt::Debug for PubSubSubscriber {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PubSubSubscriber")
            .field("subscription", &self.subscription)
            .finish_non_exhaustive()
    }
}

impl PubSubSubscriber {
    /// The full resource name of the subscription this stream consumes from.
    #[must_use]
    pub fn subscription(&self) -> &str {
        &self.subscription
    }

    /// Opens the stream synchronously (the client connects lazily) and spawns the pump on
    /// `runtime`, the one the broker connected on.
    pub(crate) fn open(core: &Core, descriptor: &GooglePubSub, runtime: &Handle) -> Self {
        let name = core.subscription_name(descriptor.subscription());
        let mut builder = core.subscriber.subscribe(name.clone());
        if let Some(messages) = descriptor.max_outstanding_value() {
            builder = builder.set_max_outstanding_messages(messages);
        }
        if let Some(extension) = descriptor.ack_extension_value() {
            builder = builder.set_max_lease_extension(extension);
        }
        // The budget a held delivery spends: the client stops extending past it, so the value the
        // descriptor reports and the value the client honours have to be the same one.
        let limits = descriptor.limits();
        builder = builder.set_max_lease(limits.max_lease);
        // The client starts the subscription's lease loop on the runtime `build` runs in, and
        // the subscription may be opened from a runtime that stops right after (a dedicated
        // thread): the loop belongs beside the pump, on the runtime the broker connected on.
        let stream = {
            let _connect_runtime = runtime.enter();
            builder.build()
        };
        let shutdown = stream.shutdown_token();

        let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY);
        runtime.spawn(pump(stream, tx, name.clone(), limits, core.runtime_slot()));

        Self {
            subscription: name,
            buffer: BufferedSubscriber::new(Deliveries::Pull(Pull {
                rx,
                shutdown,
                runtime: runtime.clone(),
            }))
            .max_wait(descriptor.batch_wait_value()),
        }
    }

    /// A subscription of the in-process transport, batched over the same buffer and deadline.
    #[cfg(feature = "testing")]
    pub(crate) fn in_process(
        subscription: String,
        consumer: Consumer,
        batch_wait: Duration,
    ) -> Self {
        Self {
            subscription,
            buffer: BufferedSubscriber::new(Deliveries::InProcess(consumer)).max_wait(batch_wait),
        }
    }
}

impl Subscriber for PubSubSubscriber {
    type Message = PubSubMessage;
    type Error = PubSubError;

    fn stream(&mut self) -> impl Stream<Item = Result<PubSubMessage, PubSubError>> + Send + '_ {
        self.buffer.stream()
    }
}

/// Batches come off the same channel the single-message stream reads, closed by the size the
/// registration named or by the descriptor's [`batch_wait`](GooglePubSub::batch_wait),
/// whichever comes first. Nothing at the mount site says the batches are assembled here rather
/// than on the wire.
impl BatchSubscriber for PubSubSubscriber {
    type Batch = Vec<PubSubMessage>;

    fn batches(
        &mut self,
        size: NonZeroUsize,
    ) -> impl Stream<Item = Result<Self::Batch, PubSubError>> + Send + '_ {
        self.buffer.batches(size)
    }
}

/// The wire side of a subscription, read one delivery at a time: the streaming pull, or, under
/// the `testing` feature, a consumer of the in-process transport. Private, because a service
/// reaches it through [`PubSubSubscriber`], which owns the buffer that turns these deliveries into
/// batches.
///
/// Without the feature there is one variant, so the type is the pull itself and the `match` in
/// [`stream`](Subscriber::stream) is irrefutable.
enum Deliveries {
    Pull(Pull),
    #[cfg(feature = "testing")]
    InProcess(Consumer),
}

// The zero-cost promise of the in-process mode: a subscription built without it holds exactly the
// streaming pull it reads.
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Deliveries>() == size_of::<Pull>());

/// The streaming pull: the pump's output channel, the token that stops the client's stream, and
/// the runtime the pump runs on.
struct Pull {
    rx: mpsc::Receiver<Result<PubSubMessage, PubSubError>>,
    shutdown: ShutdownToken,
    runtime: Handle,
}

impl Drop for Pull {
    fn drop(&mut self) {
        // `shutdown` is async and destructors are sync; the spawned signal drains the stream,
        // which ends the pump task. It runs beside the pump, on the runtime the broker connected
        // on, whichever thread drops the subscriber. Best effort by design: if that runtime is
        // already gone, the pump died with it.
        let token = self.shutdown.clone();
        self.runtime.spawn(async move {
            token.shutdown().await;
        });
    }
}

impl Subscriber for Deliveries {
    type Message = PubSubMessage;
    type Error = PubSubError;

    fn stream(&mut self) -> impl Stream<Item = Result<PubSubMessage, PubSubError>> + Send + '_ {
        // Poll the channel in place rather than wrapping it in an owning stream, so `stream`
        // can be called again after the returned stream is dropped (the runtime and the
        // conformance helpers re-enter it per call).
        futures::stream::poll_fn(move |cx| match self {
            Self::Pull(pull) => pull.rx.poll_recv(cx),
            #[cfg(feature = "testing")]
            Self::InProcess(consumer) => consumer.poll_next(cx),
        })
    }
}

async fn pump(
    mut stream: MessageStream,
    out: mpsc::Sender<Result<PubSubMessage, PubSubError>>,
    subscription: String,
    limits: DeliveryLimits,
    runtime: RuntimeSlot,
) {
    while let Some(item) = stream.next().await {
        match item {
            Ok((message, handler)) => {
                if out
                    .send(Ok(PubSubMessage::new(message, handler, limits, runtime)))
                    .await
                    .is_err()
                {
                    // Subscriber dropped; its Drop has already signalled shutdown.
                    break;
                }
            }
            Err(err) => {
                // The client absorbs transient failures internally; an error here is
                // permanent and closes the stream.
                let _ = out
                    .send(Err(PubSubError::Receive {
                        subscription: subscription.clone(),
                        source: box_err(err),
                    }))
                    .await;
                break;
            }
        }
    }
}