Skip to content

Commit

Permalink
adding, stream method to consumer delivery struct and adding examples
Browse files Browse the repository at this point in the history
  • Loading branch information
DanielePalaia committed Oct 20, 2024
1 parent ffa07fb commit 989be13
Show file tree
Hide file tree
Showing 5 changed files with 162 additions and 7 deletions.
74 changes: 74 additions & 0 deletions examples/receive_super_stream.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
use futures::StreamExt;
use rabbitmq_stream_client::error::StreamCreateError;
use rabbitmq_stream_client::types::{
ByteCapacity, OffsetSpecification, ResponseCode, SuperStreamConsumer,
};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::task;
use tokio::time::sleep;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
use rabbitmq_stream_client::Environment;
let environment = Environment::builder().build().await?;
let message_count = 10;
let super_stream = "hello-rust-stream";

let create_response = environment
.stream_creator()
.max_length(ByteCapacity::GB(5))
.create_super_stream(super_stream, 3, None)
.await;

if let Err(e) = create_response {
if let StreamCreateError::Create { stream, status } = e {
match status {
// we can ignore this error because the stream already exists
ResponseCode::StreamAlreadyExists => {}
err => {
println!("Error creating stream: {:?} {:?}", stream, err);
}
}
}
}

let mut super_stream_consumer: SuperStreamConsumer = environment
.super_stream_consumer()
.offset(OffsetSpecification::First)
.build(super_stream)
.await
.unwrap();

let received_messages = Arc::new(AtomicU32::new(0));

for mut consumer in super_stream_consumer.get_consumers().await.into_iter() {
let received_messages_outer = received_messages.clone();

task::spawn(async move {
let mut inner_received_messages = received_messages_outer.clone();
while let Some(delivery) = consumer.next().await {
let d = delivery.unwrap();
println!(
"Got message: {:#?} from stream: {} with offset: {}",
d.message()
.data()
.map(|data| String::from_utf8(data.to_vec()).unwrap()),
d.stream(),
d.offset(),
);
let value = inner_received_messages.fetch_add(1, Ordering::Relaxed);
if value == message_count {
let handle = consumer.handle();
_ = handle.close().await;
break;
}
}
});
}

sleep(Duration::from_millis(20000)).await;

Ok(())
}
78 changes: 78 additions & 0 deletions examples/send_super_stream.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
use rabbitmq_stream_client::error::StreamCreateError;
use rabbitmq_stream_client::types::{
ByteCapacity, HashRoutingMurmurStrategy, Message, ResponseCode, RoutingStrategy,
};
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::Notify;

fn hash_strategy_value_extractor(message: &Message) -> String {
String::from_utf8(Vec::from(message.data().unwrap())).expect("Found invalid UTF-8")
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
use rabbitmq_stream_client::Environment;
let environment = Environment::builder().build().await?;
let message_count = 100;
let stream = "hello-rust-stream";
let confirmed_messages = Arc::new(AtomicU32::new(0));
let notify_on_send = Arc::new(Notify::new());

let _ = environment
.stream_creator()
.max_length(ByteCapacity::GB(5))
.create_super_stream(stream, 3, None)
.await;

let create_response = environment
.stream_creator()
.max_length(ByteCapacity::GB(5))
.create_super_stream(stream, 3, None)
.await;

if let Err(e) = create_response {
if let StreamCreateError::Create { stream, status } = e {
match status {
// we can ignore this error because the stream already exists
ResponseCode::StreamAlreadyExists => {}
err => {
println!("Error creating stream: {:?} {:?}", stream, err);
}
}
}
}

let mut super_stream_producer = environment
.super_stream_producer(RoutingStrategy::HashRoutingStrategy(
HashRoutingMurmurStrategy {
routing_extractor: &hash_strategy_value_extractor,
},
))
.build(stream)
.await
.unwrap();

for i in 0..message_count {
println!("sending message {}", i);
let counter = confirmed_messages.clone();
let notifier = notify_on_send.clone();
let msg = Message::builder().body(format!("message{}", i)).build();
super_stream_producer
.send(msg, move |_| {
let inner_counter = counter.clone();
let inner_notifier = notifier.clone();
async move {
if inner_counter.fetch_add(1, Ordering::Relaxed) == message_count - 1 {
inner_notifier.notify_one();
}
}
})
.await
.unwrap();
}

notify_on_send.notified().await;
let _ = super_stream_producer.close().await;
Ok(())
}
7 changes: 7 additions & 0 deletions src/consumer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -330,6 +330,7 @@ impl MessageHandler for ConsumerMessageHandler {
.0
.sender
.send(Ok(Delivery {
stream: self.0.stream.clone(),
subscription_id: self.0.subscription_id,
message,
offset,
Expand Down Expand Up @@ -358,6 +359,7 @@ impl MessageHandler for ConsumerMessageHandler {
/// Envelope from incoming message
#[derive(Debug, Clone)]
pub struct Delivery {
stream: String,
subscription_id: u8,
message: Message,
offset: u64,
Expand All @@ -369,6 +371,11 @@ impl Delivery {
self.subscription_id
}

/// Get a reference to the delivery's stream name.
pub fn stream(&self) -> &String {
&self.stream
}

/// Get a reference to the delivery's message.
pub fn message(&self) -> &Message {
&self.message
Expand Down
8 changes: 2 additions & 6 deletions src/superstream_consumer.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use crate::superstream::DefaultSuperStreamMetadata;
use crate::{error::ConsumerCreateError, Client, Consumer, Environment};
use crate::{error::ConsumerCreateError, Consumer, Environment};
use rabbitmq_stream_protocol::commands::subscribe::OffsetSpecification;
use std::sync::Arc;
//type FilterPredicate = Option<Arc<dyn Fn(&Message) -> bool + Send + Sync>>;
Expand All @@ -11,8 +11,6 @@ pub struct SuperStreamConsumer {
}

struct SuperStreamConsumerInternal {
client: Client,
super_stream: String,
offset_specification: OffsetSpecification,
consumers: Vec<Consumer>,
}
Expand Down Expand Up @@ -54,8 +52,6 @@ impl SuperStreamConsumerBuilder {
}

let super_stream_consumer_internal = Arc::new(SuperStreamConsumerInternal {
super_stream: super_stream.to_string(),
client: client.clone(),
offset_specification: self.offset_specification.clone(),
consumers,
});
Expand All @@ -65,7 +61,7 @@ impl SuperStreamConsumerBuilder {
})
}

pub async fn offset(mut self, offset_specification: OffsetSpecification) -> Self {
pub fn offset(mut self, offset_specification: OffsetSpecification) -> Self {
self.offset_specification = offset_specification;
self
}
Expand Down
2 changes: 1 addition & 1 deletion tests/integration/consumer_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ async fn super_stream_consumer_test() {

task::spawn(async move {
let mut inner_received_messages = received_messages_outer.clone();
while let d = consumer.next().await.unwrap() {
while let _ = consumer.next().await.unwrap() {
let value = inner_received_messages.fetch_add(1, Ordering::Relaxed);
if value == message_count {
let handle = consumer.handle();
Expand Down

0 comments on commit 989be13

Please sign in to comment.