1use futures::future::BoxFuture;
2use serde::{Deserialize, Serialize};
3use std::sync::atomic::Ordering;
4use std::sync::Arc;
5use tokio::sync::broadcast;
6use tito::{
7 queue::{run_worker, EventType, QueueConfig, QueueEvent, TitoQueue},
8 types::DBUuid,
9 TiKV, TitoError, WorkerConfig,
10};
11
12#[derive(Default, Debug, Clone, Serialize, Deserialize)]
13struct UserEvent {
14 user_id: String,
15 name: String,
16 email: String,
17 action: String,
18}
19
20impl EventType for UserEvent {
21 fn event_type_name() -> &'static str {
22 "user"
23 }
24
25 fn variant_name(&self) -> &'static str {
26 "user_created"
27 }
28}
29
30#[tokio::main]
31async fn main() -> Result<(), TitoError> {
32 println!("Testing FIFO Queue\n");
33
34 let tito_db = TiKV::connect(vec!["127.0.0.1:2379"]).await?;
35
36 let queue = Arc::new(TitoQueue::new(
37 tito_db.clone(),
38 QueueConfig::with_partitions(1),
39 ));
40
41 println!("Publishing 5 events...\n");
42
43 for i in 1..=5 {
44 let user_id = DBUuid::new_v4().to_string();
45 let event = QueueEvent::new(
46 format!("user:{}", user_id),
47 UserEvent {
48 user_id,
49 name: format!("User {}", i),
50 email: format!("user{}@example.com", i),
51 action: "created".to_string(),
52 },
53 );
54
55 queue.publish(event).await?;
56 println!("Published event for User {}", i);
57 }
58
59 println!("\nStarting worker...\n");
60
61 let events_processed = Arc::new(std::sync::atomic::AtomicU32::new(0));
62 let events_processed_clone = events_processed.clone();
63
64 let (shutdown_tx, shutdown_rx) = broadcast::channel(1);
65
66 let handler = move |event: QueueEvent<UserEvent>| {
67 let counter = events_processed_clone.clone();
68 Box::pin(async move {
69 let count = counter.fetch_add(1, Ordering::SeqCst) + 1;
70 println!(
71 "[{}] {} - {} ({})",
72 count,
73 event.payload.action,
74 event.payload.name,
75 event.payload.email,
76 );
77 tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
78 Ok::<_, TitoError>(())
79 }) as BoxFuture<'static, Result<(), TitoError>>
80 };
81
82 let worker_handle = run_worker(
83 queue.clone(),
84 WorkerConfig {
85 event_type: String::from("user"),
86 consumer: String::from("example-consumer"),
87 partition_range: 0..1,
88 },
89 handler,
90 shutdown_rx,
91 )
92 .await;
93
94 tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
95
96 println!("\nShutting down...");
97 let _ = shutdown_tx.send(());
98 let _ = worker_handle.await;
99
100 let total = events_processed.load(Ordering::SeqCst);
101 println!("\nProcessed {} events", total);
102
103 Ok(())
104}