amq-rpc 0.1.0

RabbitMQ RPC library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use futures::StreamExt;

use crate::connection::AmqpConnection;
use crate::error::RpcResult;
use crate::types::{Command, CommandResponse, CommandHandler};

pub struct RpcServer {
    connection: AmqpConnection,
    queue: String,
    handlers: Arc<RwLock<HashMap<String, Box<dyn CommandHandler>>>>,
}

impl RpcServer {
    pub fn new(connection: AmqpConnection, queue: String) -> Self {
        Self {
            connection,
            queue,
            handlers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn register_handler<H>(&self, command: &str, handler: H)
    where
        H: CommandHandler + 'static,
    {
        self.handlers.write().await.insert(command.to_string(), Box::new(handler));
    }

    pub async fn start(&self) -> RpcResult<()> {
        self.connection.connect().await?;
        self.connection.declare_queue(&self.queue, false).await?;

        let consumer = self.connection.consume(&self.queue, "rpc_server").await?;
        let handlers = self.handlers.clone();
        let connection = self.connection.clone();

        tokio::spawn(async move {
            let mut consumer = consumer;
            while let Some(delivery) = consumer.next().await {
                if let Ok(delivery) = delivery {
                    let correlation_id = delivery.properties.correlation_id()
                        .as_ref()
                        .map(|s| s.as_str().to_string())
                        .unwrap_or_default();

                    let reply_to = delivery.properties.reply_to()
                        .as_ref()
                        .map(|s| s.as_str().to_string());

                    if let Some(reply_queue) = reply_to {
                        let response = match serde_json::from_slice::<Command>(&delivery.data) {
                            Ok(cmd) => {
                                let handlers_guard = handlers.read().await;
                                match handlers_guard.get(&cmd.command) {
                                    Some(handler) => {
                                        match handler.handle(&cmd.command, cmd.args).await {
                                            Ok(result) => CommandResponse {
                                                success: true,
                                                result: Some(result),
                                                error: None,
                                            },
                                            Err(e) => CommandResponse {
                                                success: false,
                                                result: None,
                                                error: Some(e.to_string()),
                                            },
                                        }
                                    }
                                    None => CommandResponse {
                                        success: false,
                                        result: None,
                                        error: Some(format!("Unknown command: {}", cmd.command)),
                                    },
                                }
                            }
                            Err(e) => CommandResponse {
                                success: false,
                                result: None,
                                error: Some(format!("Failed to parse command: {}", e)),
                            },
                        };

                        if let Ok(response_payload) = serde_json::to_vec(&response) {
                            let _ = connection.publish(
                                &reply_queue,
                                &response_payload,
                                &correlation_id,
                                None,
                                None,
                            ).await;
                        }
                    }

                    let _ = connection.ack(delivery.delivery_tag).await;
                }
            }
        });

        Ok(())
    }

    pub async fn close(&self) -> RpcResult<()> {
        self.connection.close().await
    }
}