starpc 0.52.0

Streaming protobuf RPC framework
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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! RpcStreamWriter - PacketWriter implementation for RpcStream.

use async_trait::async_trait;
use prost::Message;

use crate::error::Result;
use crate::proto::Packet;
use crate::rpc::PacketWriter;

use super::rpcstream::RpcStream;
use super::RpcStreamPacket;

/// RpcStreamWriter wraps an RpcStream and implements PacketWriter.
///
/// This allows using an RpcStream as the transport for nested RPC calls,
/// converting Packet messages to RpcStreamPacket::Data messages.
pub struct RpcStreamWriter<S> {
    /// The underlying RpcStream.
    pub(crate) inner: S,
}

impl<S> RpcStreamWriter<S> {
    /// Creates a new RpcStreamWriter.
    pub fn new(stream: S) -> Self {
        Self { inner: stream }
    }

    /// Returns a reference to the inner stream.
    pub fn inner(&self) -> &S {
        &self.inner
    }
}

#[async_trait]
impl<S: RpcStream + Send + Sync> PacketWriter for RpcStreamWriter<S> {
    async fn write_packet(&self, packet: Packet) -> Result<()> {
        let data = packet.encode_to_vec();
        let rpc_packet = RpcStreamPacket::new_data(data);
        self.inner.send_packet(&rpc_packet).await
    }

    async fn close(&self) -> Result<()> {
        self.inner.close_send().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::rpcstream::rpc_stream_packet;
    use crate::stream::{Context, Stream};
    use bytes::Bytes;
    use std::collections::VecDeque;
    use std::sync::atomic::{AtomicBool, Ordering};
    use tokio::sync::Mutex;

    struct MockRpcStream {
        ctx: Context,
        packets: Mutex<VecDeque<RpcStreamPacket>>,
        closed: AtomicBool,
    }

    impl MockRpcStream {
        fn new() -> Self {
            Self {
                ctx: Context::new(),
                packets: Mutex::new(VecDeque::new()),
                closed: AtomicBool::new(false),
            }
        }

        async fn get_packets(&self) -> Vec<RpcStreamPacket> {
            self.packets.lock().await.iter().cloned().collect()
        }
    }

    #[async_trait]
    impl Stream for MockRpcStream {
        fn context(&self) -> &Context {
            &self.ctx
        }

        async fn send_bytes(&self, _data: Bytes) -> Result<()> {
            Ok(())
        }

        async fn recv_bytes(&self) -> Result<Bytes> {
            Err(crate::Error::StreamClosed)
        }

        async fn close_send(&self) -> Result<()> {
            self.closed.store(true, Ordering::SeqCst);
            Ok(())
        }

        async fn close(&self) -> Result<()> {
            self.closed.store(true, Ordering::SeqCst);
            Ok(())
        }
    }

    #[async_trait]
    impl RpcStream for MockRpcStream {
        async fn send_packet(&self, packet: &RpcStreamPacket) -> Result<()> {
            self.packets.lock().await.push_back(packet.clone());
            Ok(())
        }

        async fn recv_packet(&self) -> Result<RpcStreamPacket> {
            Err(crate::Error::StreamClosed)
        }
    }

    #[tokio::test]
    async fn test_rpc_stream_writer_write_packet() {
        use std::sync::Arc;
        let stream = Arc::new(MockRpcStream::new());
        let writer = RpcStreamWriter::new(stream.clone());

        let packet = crate::packet::new_call_start(
            "test.Service".to_string(),
            "TestMethod".to_string(),
            Some(Bytes::from(vec![1, 2, 3])),
        );

        writer.write_packet(packet).await.unwrap();

        let packets = stream.get_packets().await;
        assert_eq!(packets.len(), 1);

        match &packets[0].body {
            Some(rpc_stream_packet::Body::Data(data)) => {
                // Verify we can decode the inner packet
                let inner = Packet::decode(&data[..]).unwrap();
                assert!(inner.body.is_some());
            }
            _ => panic!("Expected Data packet"),
        }
    }

    #[tokio::test]
    async fn test_rpc_stream_writer_close() {
        use std::sync::Arc;
        let stream = Arc::new(MockRpcStream::new());
        let writer = RpcStreamWriter::new(stream.clone());

        writer.close().await.unwrap();
        assert!(stream.closed.load(Ordering::SeqCst));
    }
}