ruststream_kinesis/publisher.rs
1//! [`KinesisPublisher`], its [`KinesisPublish`] policy, and the crate's publish steps.
2
3// Without the `testing` feature a connection link has one variant, so a `match` on it has a
4// single arm; the match stays so that the in-process arm has its place when the feature is on.
5#![cfg_attr(
6 not(feature = "testing"),
7 allow(clippy::infallible_destructuring_match)
8)]
9
10use std::future::{Future, ready};
11use std::sync::Arc;
12
13use aws_sdk_kinesis::Client;
14use aws_sdk_kinesis::primitives::Blob;
15#[cfg(feature = "asyncapi")]
16use ruststream::asyncapi::{Binding, Bindings};
17use ruststream::runtime::{MapPublisher, PublishBuilder, PublishSink};
18use ruststream::{HeaderMap, OutgoingFor, PairError, PublishPolicy, Publisher, Take};
19#[cfg(feature = "asyncapi")]
20use serde::Serialize;
21
22#[cfg(feature = "asyncapi")]
23use crate::stream::BINDING_KEY;
24
25use crate::broker::{ConnectedKinesisBroker, CoreCell, Link};
26use crate::error::{KinesisError, sdk_err};
27use crate::message::{PARTITION_KEY_HEADER, encode_envelope};
28
29/// The per-message publish settings of this broker: what one call varies, against what the mount
30/// site fixed on the policy.
31///
32/// Every field is optional, so a call says only what it changes, and a field no step touched
33/// keeps what the policy holds. A handler body that names a setting bounds its injected slot on
34/// this type - `Out<impl Publisher<Options = KinesisPublishOptions>, Marker>` - and takes the
35/// steps themselves from [`KinesisPublishSteps`].
36///
37/// # Examples
38///
39/// ```
40/// # #[cfg(feature = "testing")]
41/// # mod demo {
42/// use std::error::Error;
43///
44/// use ruststream::testing::TestApp;
45/// use ruststream_kinesis::prelude::*;
46/// use serde::{Deserialize, Serialize};
47///
48/// #[derive(Deserialize, Serialize, Outgoing)]
49/// pub struct Order {
50/// pub id: u64,
51/// }
52///
53/// #[derive(Serialize, Outgoing)]
54/// struct Journalled {
55/// order: u64,
56/// }
57///
58/// /// The slot the journal entries leave through.
59/// #[derive(OutSlot)]
60/// #[publishes(Journalled)]
61/// struct Journal;
62///
63/// /// Every entry of one tenant lands on one shard, so they stay in order.
64/// #[subscriber(KinesisStream::new("orders"))]
65/// async fn record(
66/// order: &Order,
67/// Out(journal): Out<impl Publisher<Options = KinesisPublishOptions>, Journal>,
68/// ) -> HandlerOutcome {
69/// let sent = journal
70/// .message(&Journalled { order: order.id })
71/// .to("orders.journal")
72/// .partition_key("tenant-acme")
73/// .publish()
74/// .await;
75/// if sent.is_err() {
76/// return HandlerOutcome::retry();
77/// }
78/// HandlerOutcome::ack()
79/// }
80///
81/// /// The app `main` runs.
82/// pub fn app() -> RustStream {
83/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
84/// b.include(record).out(Journal, Publish::default()).build();
85/// })
86/// }
87///
88/// pub async fn keys_each_entry() -> Result<(), Box<dyn Error>> {
89/// let tb = TestApp::start(app()).await?;
90/// tb.broker::<KinesisBroker>()
91/// .message(&Order { id: 7 })
92/// .to("orders")
93/// .publish()
94/// .await?;
95///
96/// tb.out::<Journal>()
97/// .assert_called_once()
98/// .with_options(&KinesisPublishOptions {
99/// partition_key: Some("tenant-acme".to_owned()),
100/// });
101/// Ok(())
102/// }
103/// # }
104/// # #[cfg(feature = "testing")]
105/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
106/// # tokio::runtime::Builder::new_multi_thread()
107/// # .enable_all()
108/// # .build()?
109/// # .block_on(demo::keys_each_entry())
110/// # }
111/// # #[cfg(not(feature = "testing"))]
112/// # fn main() {}
113/// ```
114#[derive(Debug, Clone, Default, PartialEq, Eq)]
115pub struct KinesisPublishOptions {
116 /// The record's partition key, named by
117 /// [`partition_key`](KinesisPublishSteps::partition_key) on the publish builder.
118 pub partition_key: Option<String>,
119}
120
121/// Publishes records to Kinesis streams; the destination is the stream's name.
122///
123/// The record's partition key is the unit of shard routing and per-key ordering. A publish names
124/// it with [`partition_key`](KinesisPublishSteps::partition_key) on the publish builder; a record
125/// that names none takes the key the mount site named on the policy
126/// ([`KinesisPublish::partition_key`]), and without that a process-unique key spreads the records
127/// across the shards. User headers beyond the partition key travel in a small conditional envelope
128/// (Kinesis records carry only a data blob and a partition key); plain payloads stay unenveloped,
129/// so any Kinesis consumer reads them, and a record from another producer reads back as
130/// headerless. The envelope is the magic `RSK2`, a big-endian `u32` length of the header block,
131/// the block, and the payload; each header is a `u32`-prefixed name and a `u32`-prefixed value,
132/// so a value comes back byte for byte. Records in the earlier `RSK1` text envelope still read
133/// back with their headers. One publish is one `PutRecord` call.
134/// Buildable before `connect` and usable until `shutdown`; afterwards every publish reports
135/// [`KinesisError::NotConnected`].
136#[derive(Clone)]
137pub struct KinesisPublisher {
138 cell: CoreCell,
139 /// The partition key the mount site named on the policy this publisher was paired from,
140 /// used by every record that does not name one of its own.
141 partition_key: Option<Arc<str>>,
142}
143
144impl std::fmt::Debug for KinesisPublisher {
145 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146 f.debug_struct("KinesisPublisher").finish_non_exhaustive()
147 }
148}
149
150impl KinesisPublisher {
151 pub(crate) fn new(cell: CoreCell) -> Self {
152 Self {
153 cell,
154 partition_key: None,
155 }
156 }
157
158 /// The publisher a policy pairs into: the same connection, plus what the mount site named
159 /// on that policy.
160 pub(crate) fn keyed(cell: CoreCell, partition_key: Option<Arc<str>>) -> Self {
161 Self {
162 cell,
163 partition_key,
164 }
165 }
166
167 /// What this publisher speaks over, once `connect` (or the harness) has filled the cell.
168 fn link(&self) -> Result<&Link, KinesisError> {
169 self.cell.get().ok_or(KinesisError::NotConnected)
170 }
171}
172
173/// The record's partition key, in the order the four places that can name one resolve.
174///
175/// The call's step first, then the `partition-key` header a call site wrote by hand (the portable
176/// spelling, and what a message crossing from another broker carries), then the key the mount site
177/// named on the policy - the only place a reply and an injected slot can be given one - and last a
178/// process-unique key, because a record cannot go out without one.
179///
180/// One function for both transports: the in-process mode answers the key the service would, never
181/// one of its own.
182pub(crate) fn resolve_partition_key(
183 options: Option<&KinesisPublishOptions>,
184 headers: &HeaderMap,
185 from_policy: Option<&str>,
186) -> String {
187 if let Some(key) = options.and_then(|options| options.partition_key.as_deref()) {
188 return key.to_owned();
189 }
190 if let Some(named) = headers.get(PARTITION_KEY_HEADER) {
191 return String::from_utf8_lossy(named).into_owned();
192 }
193 from_policy.map_or_else(spread_key, str::to_owned)
194}
195
196fn spread_key() -> String {
197 use std::sync::atomic::{AtomicU64, Ordering};
198 static SEQ: AtomicU64 = AtomicU64::new(0);
199 format!(
200 "rs-{}-{}",
201 std::process::id(),
202 SEQ.fetch_add(1, Ordering::Relaxed)
203 )
204}
205
206impl Publisher for KinesisPublisher {
207 /// A record's data is a `Blob`, a vector the client keeps for the request, so the publisher
208 /// takes the buffer the framework wrote: a publish that carries no header hands it straight
209 /// through, and one that does spends it as the body of the envelope.
210 type Payload = Take;
211 type Error = KinesisError;
212 /// The partition key is the one setting a Kinesis record varies per message, and
213 /// [`KinesisPublishSteps`] is where a call site sets it.
214 type Options = KinesisPublishOptions;
215
216 async fn publish(
217 &self,
218 msg: OutgoingFor<'_, Take>,
219 options: Option<&Self::Options>,
220 ) -> Result<(), Self::Error> {
221 let core = match self.link()? {
222 Link::Aws(core) => core,
223 #[cfg(feature = "testing")]
224 Link::InProcess(bus) => {
225 return bus.put(msg, options, self.partition_key.as_deref());
226 }
227 };
228 core.ensure_open()?;
229 let partition_key =
230 resolve_partition_key(options, msg.headers(), self.partition_key.as_deref());
231 // The resolved key rides the record's own field, which is this transport's wire for it:
232 // a delivery reads it back into `PARTITION_KEY_HEADER`, where the cross-broker
233 // `Partitioned` capability finds it. Stamping the header here as well would cost a map
234 // copy per publish for a value the envelope drops on the next line.
235 //
236 // The stream name is the caller's string and outlives the message the record is built
237 // from.
238 let (stream, payload, headers) = msg.into_parts();
239 let data = encode_envelope(&headers, payload);
240 core.client
241 .with(Client::put_record)
242 .stream_name(stream)
243 .partition_key(partition_key)
244 .data(Blob::new(data))
245 .send()
246 .await
247 .map(|_| ())
248 .map_err(|e| KinesisError::Publish {
249 stream: stream.to_owned(),
250 source: sdk_err(&e),
251 })
252 }
253}
254
255/// The publish policy for [`KinesisPublisher`]: pure declaration, constructible anywhere,
256/// paired with the connected broker by the runtime after `connect`.
257///
258/// # Examples
259///
260/// ```
261/// # mod demo {
262/// use ruststream_kinesis::KinesisPublish;
263/// use ruststream_kinesis::prelude::*;
264/// use serde::{Deserialize, Serialize};
265///
266/// #[derive(Deserialize)]
267/// struct Order {
268/// id: u64,
269/// }
270///
271/// #[derive(Serialize, Outgoing)]
272/// struct Journalled {
273/// order: u64,
274/// }
275///
276/// /// The slot the journal entries leave through.
277/// #[derive(OutSlot)]
278/// #[publishes(Journalled)]
279/// struct Journal;
280///
281/// /// The body names no broker: the plain bound moves between brokers untouched.
282/// #[subscriber(KinesisStream::new("orders"))]
283/// async fn record(order: &Order, Out(journal): Out<impl Publisher, Journal>) -> HandlerOutcome {
284/// let sent = journal
285/// .message(&Journalled { order: order.id })
286/// .to("orders.journal")
287/// .publish()
288/// .await;
289/// if sent.is_err() {
290/// return HandlerOutcome::retry();
291/// }
292/// HandlerOutcome::ack()
293/// }
294///
295/// #[ruststream::app]
296/// fn app() -> impl App {
297/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
298/// b.include(record).out(Journal, KinesisPublish::default()).build();
299/// })
300/// }
301/// # }
302/// # fn main() {}
303/// ```
304#[derive(Debug, Clone, Default, PartialEq, Eq)]
305#[must_use]
306pub struct KinesisPublish {
307 partition_key: Option<String>,
308}
309
310impl KinesisPublish {
311 /// The partition key every record published through this policy carries unless the call
312 /// names its own.
313 ///
314 /// This is the mount site's default, and the only key a reply can be given: a replying
315 /// handler returns a value and never reaches a publish builder, so
316 /// [`partition_key`](KinesisPublishSteps::partition_key) has no call site there. A slot the
317 /// handler publishes through has one, and a step on that call wins over this. Without either
318 /// the records spread across the stream's shards under a process-unique key.
319 ///
320 /// One key means one shard, and with it one shard's throughput: name it when the records
321 /// have to stay mutually ordered, and leave it off otherwise.
322 ///
323 /// # Examples
324 ///
325 /// ```
326 /// # mod demo {
327 /// use ruststream_kinesis::KinesisPublish;
328 /// use ruststream_kinesis::prelude::*;
329 /// use serde::{Deserialize, Serialize};
330 ///
331 /// #[derive(Deserialize)]
332 /// struct Order {
333 /// id: u64,
334 /// }
335 ///
336 /// #[derive(Serialize, Outgoing)]
337 /// struct Journalled {
338 /// order: u64,
339 /// }
340 ///
341 /// /// The slot the journal entries leave through.
342 /// #[derive(OutSlot)]
343 /// #[publishes(Journalled)]
344 /// struct Journal;
345 ///
346 /// /// The body names no broker: the plain bound moves between brokers untouched.
347 /// #[subscriber(KinesisStream::new("orders"))]
348 /// async fn record(order: &Order, Out(journal): Out<impl Publisher, Journal>) -> HandlerOutcome {
349 /// let sent = journal
350 /// .message(&Journalled { order: order.id })
351 /// .to("orders.journal")
352 /// .publish()
353 /// .await;
354 /// if sent.is_err() {
355 /// return HandlerOutcome::retry();
356 /// }
357 /// HandlerOutcome::ack()
358 /// }
359 ///
360 /// #[ruststream::app]
361 /// fn app() -> impl App {
362 /// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
363 /// // Every journal entry on one shard, mutually ordered.
364 /// let journal = KinesisPublish::default().partition_key("journal-v1");
365 /// b.include(record).out(Journal, journal).build();
366 /// })
367 /// }
368 /// # }
369 /// # fn main() {}
370 /// ```
371 pub fn partition_key(mut self, key: impl Into<String>) -> Self {
372 self.partition_key = Some(key.into());
373 self
374 }
375
376 /// The key a publisher paired from this policy carries.
377 fn key(&self) -> Option<Arc<str>> {
378 self.partition_key.as_deref().map(Arc::from)
379 }
380
381 /// What the records published through this policy add to their message object in the
382 /// generated document.
383 ///
384 /// The key a single call names is not here: the document describes the mount, and a step on
385 /// the publish builder is a property of one record. A policy that names no key says nothing,
386 /// so the document carries no field the service did not fix.
387 #[cfg(feature = "asyncapi")]
388 fn message_binding(&self) -> Bindings {
389 let Some(partition_key) = self.partition_key.as_deref() else {
390 return Bindings::new();
391 };
392 let body = KinesisMessageBinding { partition_key };
393 Binding::extension(BINDING_KEY, &body)
394 .map_or_else(|_| Bindings::new(), |binding| Bindings::new().with(binding))
395 }
396}
397
398/// The stream a position publishes to, written onto that position's channel object.
399///
400/// The policy reads nothing of its own here: a policy holds this broker's settings and never a
401/// destination, so the stream is the one the mount site resolved and the runtime handed over -
402/// the reply type's `#[outgoing(name)]`, the `reply("dest")` clause, a slot entry's name. The
403/// subscription side names its stream on the same object, so both ends of a channel report the
404/// Kinesis stream behind it.
405#[cfg(feature = "asyncapi")]
406fn channel_binding(stream: &str) -> Bindings {
407 let body = KinesisPublishChannelBinding { stream };
408 Binding::extension(BINDING_KEY, &body)
409 .map_or_else(|_| Bindings::new(), |binding| Bindings::new().with(binding))
410}
411
412/// What a record published through [`KinesisPublish`] writes into the generated document.
413///
414/// The partition key is the record's own routing field, which is where the specification puts an
415/// ordering key on the brokers it does cover; the specification lists no Kinesis binding at all,
416/// so this travels as an `x-` extension.
417#[cfg(feature = "asyncapi")]
418#[derive(Debug, Serialize)]
419struct KinesisMessageBinding<'a> {
420 #[serde(rename = "partitionKey")]
421 partition_key: &'a str,
422}
423
424/// What a position publishing through [`KinesisPublish`] writes into its channel object.
425///
426/// The stream and nothing else: a publisher holds no read settings, and the document is built
427/// before anything connects, so what an existing stream is like is the service's to answer.
428#[cfg(feature = "asyncapi")]
429#[derive(Debug, Serialize)]
430struct KinesisPublishChannelBinding<'a> {
431 stream: &'a str,
432}
433
434impl crate::sealed::PolicyKey for KinesisPublish {
435 fn with_partition_key(self, key: String) -> Self {
436 self.partition_key(key)
437 }
438}
439
440impl PublishPolicy<ConnectedKinesisBroker> for KinesisPublish {
441 type Live = KinesisPublisher;
442
443 fn pair(
444 self,
445 connected: &ConnectedKinesisBroker,
446 ) -> impl Future<Output = Result<Self::Live, PairError>> {
447 ready(Ok(connected.publisher_keyed(self.key())))
448 }
449
450 #[cfg(feature = "asyncapi")]
451 fn channel_bindings(&self, channel: &str) -> Bindings {
452 channel_binding(channel)
453 }
454
455 /// The partition key is the record's own routing field, so it describes the message and not
456 /// the channel the message goes to; the destination therefore has no part in it.
457 #[cfg(feature = "asyncapi")]
458 fn message_bindings(&self, _channel: &str) -> Bindings {
459 self.message_binding()
460 }
461}
462
463/// The Kinesis publish settings, chained on a mount site after the position that named the
464/// policy: `.out_reply(policy)`, `.out_retry(policy)` or `.out(marker, policy)`.
465///
466/// The publish-side mirror of [`KinesisSubscriberExt`](crate::KinesisSubscriberExt): the
467/// framework names the position and the policy, and what this transport does with a record
468/// beyond that is the broker's own vocabulary. A call that reaches the publish builder names its
469/// own key with [`KinesisPublishSteps::partition_key`]; this settles the key for every record
470/// published from the position, including the reply, which has no call site of its own.
471///
472/// A retry copy keeps the partition key of the record it copies: the copy carries the delivery's
473/// `partition-key` header, and that header wins over the key named here, so on
474/// `.out_retry(policy)` this setting leaves the copy's key unchanged.
475///
476/// # Examples
477///
478/// ```
479/// # mod demo {
480/// use ruststream_kinesis::prelude::*;
481/// use serde::{Deserialize, Serialize};
482///
483/// #[derive(Deserialize)]
484/// struct Order {
485/// id: u64,
486/// }
487///
488/// /// The reply type names the stream it goes to; the mount site names how it is published.
489/// #[derive(Serialize, Outgoing)]
490/// #[outgoing(name = "receipts")]
491/// struct Receipt {
492/// id: u64,
493/// }
494///
495/// #[subscriber(KinesisStream::new("orders"), reply)]
496/// async fn confirm(order: &Order) -> Receipt {
497/// Receipt { id: order.id }
498/// }
499///
500/// #[ruststream::app]
501/// fn app() -> impl App {
502/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
503/// b.include(confirm)
504/// .out_reply(Publish::default())
505/// .partition_key("tenant-acme");
506/// })
507/// }
508/// # }
509/// # fn main() {}
510/// ```
511///
512/// The settings live on the policy the chain named, so a chain that named another broker's
513/// policy - or no policy at all - does not have them.
514pub trait KinesisPublishSettings: Sized {
515 /// The partition key every record published from this position carries unless the call
516 /// names its own. The mount-site spelling of [`KinesisPublish::partition_key`].
517 // Not `#[must_use]`: on the `with_broker` path a mount chain commits when it drops, so the
518 // settled chain is meant to be discarded, exactly as the framework's own steps are.
519 #[allow(clippy::return_self_not_must_use)]
520 fn partition_key(self, key: impl Into<String>) -> Self;
521}
522
523impl<T> KinesisPublishSettings for T
524where
525 T: MapPublisher<Policy: crate::sealed::PolicyKey>,
526{
527 fn partition_key(self, key: impl Into<String>) -> Self {
528 let key = key.into();
529 self.map_publisher(|policy| crate::sealed::PolicyKey::with_partition_key(policy, key))
530 }
531}
532
533/// The Kinesis publish steps, on the framework's publish builder.
534///
535/// A step sets one field of [`KinesisPublishOptions`] for the single record the call sends. It is
536/// a position on the builder rather than a value wrapping the publisher, so the record still
537/// leaves through the entry the mount site named, carrying that entry's codec and transforms and
538/// attributed to that entry's slot.
539///
540/// The bound is on the publisher's options type, so these steps appear on a builder over a
541/// Kinesis publisher and on no other broker's.
542///
543/// # Examples
544///
545/// ```
546/// # mod demo {
547/// use ruststream_kinesis::prelude::*;
548/// use serde::{Deserialize, Serialize};
549///
550/// #[derive(Deserialize)]
551/// struct Order {
552/// id: u64,
553/// }
554///
555/// #[derive(Serialize, Outgoing)]
556/// struct Journalled {
557/// order: u64,
558/// }
559///
560/// /// The slot the journal entries leave through.
561/// #[derive(OutSlot)]
562/// #[publishes(Journalled)]
563/// struct Journal;
564///
565/// /// Every entry of one tenant lands on one shard, so they stay in order.
566/// #[subscriber(KinesisStream::new("orders"))]
567/// async fn record(
568/// order: &Order,
569/// Out(journal): Out<impl Publisher<Options = KinesisPublishOptions>, Journal>,
570/// ) -> HandlerOutcome {
571/// let sent = journal
572/// .message(&Journalled { order: order.id })
573/// .to("orders.journal")
574/// .partition_key("tenant-acme")
575/// .publish()
576/// .await;
577/// if sent.is_err() {
578/// return HandlerOutcome::retry();
579/// }
580/// HandlerOutcome::ack()
581/// }
582///
583/// #[ruststream::app]
584/// fn app() -> impl App {
585/// RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
586/// b.include(record).out(Journal, Publish::default()).build();
587/// })
588/// }
589/// # }
590/// # fn main() {}
591/// ```
592pub trait KinesisPublishSteps {
593 /// Sends this one record under `key` as its partition key.
594 ///
595 /// The partition key picks the shard, and with it the order the record keeps against its
596 /// neighbours. It wins over a `partition-key` header the call wrote by hand and over the key
597 /// the mount site named on the policy.
598 #[must_use]
599 fn partition_key(self, key: impl Into<String>) -> Self;
600}
601
602impl<Sink, Body, Enc, Hdrs, Dest> KinesisPublishSteps
603 for PublishBuilder<Sink, Body, Enc, Hdrs, Dest>
604where
605 Sink: PublishSink<Options = KinesisPublishOptions>,
606{
607 fn partition_key(mut self, key: impl Into<String>) -> Self {
608 self.options_mut()
609 .get_or_insert_with(KinesisPublishOptions::default)
610 .partition_key = Some(key.into());
611 self
612 }
613}
614
615#[cfg(all(test, feature = "testing"))]
616mod tests {
617 use ruststream::runtime::PublishExt;
618 use ruststream::testing::{InProcess, TestableBroker};
619 use ruststream::{HeaderMap, Outgoing, OutgoingMessage, Serialized};
620
621 use super::*;
622 use crate::KinesisBroker;
623
624 /// Bytes the test already holds encoded: a serialized type publishes them as they are, so
625 /// these checks stay about the headers rather than about a codec.
626 #[derive(Outgoing, Serialized)]
627 struct Payload(Vec<u8>);
628
629 fn payload() -> Payload {
630 Payload(b"payload".to_vec())
631 }
632
633 async fn connected() -> ConnectedKinesisBroker {
634 KinesisBroker::new()
635 .connect_in_process()
636 .await
637 .expect("the broker connects in process")
638 }
639
640 #[tokio::test]
641 async fn the_step_carries_the_key_through_the_publish_builder() {
642 let broker = connected().await;
643 broker
644 .publisher()
645 .message(&payload())
646 .to("jobs")
647 .partition_key("tenant-acme")
648 .publish()
649 .await
650 .expect("the publish succeeds");
651
652 let published = broker.published("jobs");
653 assert_eq!(published.len(), 1);
654 assert_eq!(
655 published[0].headers().get_str(PARTITION_KEY_HEADER),
656 Some("tenant-acme")
657 );
658 }
659
660 /// The step is the call's own answer, so it beats the header spelling of the same answer.
661 #[tokio::test]
662 async fn the_step_wins_over_a_key_the_call_site_wrote_by_hand() {
663 let broker = connected().await;
664 let mut headers = HeaderMap::new();
665 headers.insert(PARTITION_KEY_HEADER, "written-by-hand");
666 broker
667 .publisher()
668 .message(&payload())
669 .to("jobs")
670 .with_headers(headers)
671 .partition_key("named-by-the-step")
672 .publish()
673 .await
674 .expect("the publish succeeds");
675
676 let published = broker.published("jobs");
677 assert_eq!(
678 published[0].headers().get_str(PARTITION_KEY_HEADER),
679 Some("named-by-the-step")
680 );
681 }
682
683 #[tokio::test]
684 async fn the_steps_key_survives_a_call_that_names_other_headers() {
685 let broker = connected().await;
686 let mut headers = HeaderMap::new();
687 headers.insert("x-tenant", "acme");
688 broker
689 .publisher()
690 .message(&payload())
691 .to("jobs")
692 .with_headers(headers)
693 .partition_key("named-by-the-step")
694 .publish()
695 .await
696 .expect("the publish succeeds");
697
698 let published = broker.published("jobs");
699 assert_eq!(
700 published[0].headers().get_str(PARTITION_KEY_HEADER),
701 Some("named-by-the-step")
702 );
703 assert_eq!(published[0].headers().get_str("x-tenant"), Some("acme"));
704 }
705
706 /// The portable spelling stays: a message that arrives carrying the header - from another
707 /// broker, or from a caller that writes headers by hand - keeps the key it carries.
708 #[tokio::test]
709 async fn a_publish_without_a_step_keeps_the_header_route() {
710 let broker = connected().await;
711 let mut headers = HeaderMap::new();
712 headers.insert(PARTITION_KEY_HEADER, "by-hand");
713 broker
714 .publisher()
715 .publish(
716 OutgoingMessage::new("jobs", b"payload").with_headers(headers),
717 None,
718 )
719 .await
720 .expect("the publish succeeds");
721
722 let published = broker.published("jobs");
723 assert_eq!(
724 published[0].headers().get_str(PARTITION_KEY_HEADER),
725 Some("by-hand")
726 );
727 }
728
729 /// A record cannot go out without a key, so the bottom of the ladder is a key that spreads
730 /// the records rather than no key at all.
731 #[tokio::test]
732 async fn a_record_no_one_named_a_key_for_still_gets_one() {
733 let broker = connected().await;
734 broker
735 .publisher()
736 .message(&payload())
737 .to("jobs")
738 .publish()
739 .await
740 .expect("the publish succeeds");
741
742 let published = broker.published("jobs");
743 let key = published[0]
744 .headers()
745 .get_str(PARTITION_KEY_HEADER)
746 .expect("the record carries a partition key");
747 assert!(
748 key.starts_with("rs-"),
749 "expected the spreading fallback key, got {key:?}",
750 );
751 }
752
753 /// The mount site's key is a default, not an override: it applies to every record that names
754 /// none, and steps aside for one that does. The policy is the production one, paired against
755 /// the broker connected in process - the same type a routes file names.
756 #[tokio::test]
757 async fn the_mount_sites_key_applies_until_a_publish_names_its_own() {
758 let broker = connected().await;
759 let publisher = KinesisPublish::default()
760 .partition_key("named-on-the-mount")
761 .pair(&broker)
762 .await
763 .expect("the policy pairs with the connected broker");
764
765 publisher
766 .message(&payload())
767 .to("jobs")
768 .publish()
769 .await
770 .expect("the publish succeeds");
771 publisher
772 .message(&payload())
773 .to("jobs")
774 .partition_key("named-on-the-call")
775 .publish()
776 .await
777 .expect("the publish succeeds");
778
779 let records = broker.published("jobs");
780 assert_eq!(
781 records[0].headers().get_str(PARTITION_KEY_HEADER),
782 Some("named-on-the-mount")
783 );
784 assert_eq!(
785 records[1].headers().get_str(PARTITION_KEY_HEADER),
786 Some("named-on-the-call")
787 );
788 }
789}