Skip to main content

ruststream_kinesis/
publisher.rs

1//! [`KinesisPublisher`] and its [`KinesisPublish`] policy.
2
3use aws_sdk_kinesis::primitives::Blob;
4use ruststream::{OutgoingMessage, PairError, PublishPolicy, Publisher};
5
6use crate::broker::{ConnectedKinesisBroker, Core, CoreCell};
7use crate::error::{KinesisError, sdk_err};
8use crate::message::{PARTITION_KEY_HEADER, encode_envelope};
9
10/// Publishes records to Kinesis streams (the destination is the stream name or ARN).
11///
12/// The `partition-key` header becomes the record's partition key - the unit of shard routing
13/// and per-key ordering; without one, a process-unique key spreads records across shards.
14/// User headers beyond the partition key travel in a small conditional envelope (Kinesis
15/// records carry only a data blob and a partition key); plain payloads stay unenveloped.
16/// Buildable before `connect` and usable until `shutdown`; afterwards every publish reports
17/// [`KinesisError::NotConnected`].
18#[derive(Clone)]
19pub struct KinesisPublisher {
20    cell: CoreCell,
21}
22
23impl std::fmt::Debug for KinesisPublisher {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("KinesisPublisher").finish_non_exhaustive()
26    }
27}
28
29impl KinesisPublisher {
30    pub(crate) fn new(cell: CoreCell) -> Self {
31        Self { cell }
32    }
33
34    fn core(&self) -> Result<&Core, KinesisError> {
35        let core = self.cell.get().ok_or(KinesisError::NotConnected)?;
36        core.ensure_open()?;
37        Ok(core)
38    }
39}
40
41fn spread_key() -> String {
42    use std::sync::atomic::{AtomicU64, Ordering};
43    static SEQ: AtomicU64 = AtomicU64::new(0);
44    format!(
45        "rs-{}-{}",
46        std::process::id(),
47        SEQ.fetch_add(1, Ordering::Relaxed)
48    )
49}
50
51impl Publisher for KinesisPublisher {
52    type Error = KinesisError;
53
54    async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
55        let core = self.core()?;
56        let partition_key = msg
57            .headers()
58            .get(PARTITION_KEY_HEADER)
59            .map_or_else(spread_key, |key| String::from_utf8_lossy(key).into_owned());
60        let data = encode_envelope(msg.headers(), msg.payload());
61        core.client
62            .put_record()
63            .stream_name(msg.name())
64            .partition_key(partition_key)
65            .data(Blob::new(data))
66            .send()
67            .await
68            .map(|_| ())
69            .map_err(|e| KinesisError::Publish {
70                stream: msg.name().to_owned(),
71                source: sdk_err(&e),
72            })
73    }
74}
75
76/// The publish policy for [`KinesisPublisher`]: pure declaration, constructible anywhere,
77/// paired with the connected broker by the runtime after `connect`.
78///
79/// # Examples
80///
81/// ```
82/// use ruststream_kinesis::KinesisPublish;
83///
84/// let policy = KinesisPublish::default();
85/// # let _ = policy;
86/// ```
87#[derive(Debug, Clone, Copy, Default)]
88#[must_use]
89pub struct KinesisPublish;
90
91impl PublishPolicy<ConnectedKinesisBroker> for KinesisPublish {
92    type Live = KinesisPublisher;
93
94    async fn pair(self, connected: &ConnectedKinesisBroker) -> Result<Self::Live, PairError> {
95        Ok(connected.publisher())
96    }
97}