ruststream_kinesis/
publisher.rs1use 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#[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#[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}