Skip to main content

queue_fifo/
queue_fifo.rs

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}