dbeel 0.1.0

dbeel is a fast and minimal distributed thread-per-core document database.
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

use std::time::Duration;

use bincode::Options;
use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use glommio::net::TcpStream;

use crate::{
    args::Args,
    error::{Error, Result},
    messages::{NodeMetadata, ShardMessage, ShardRequest, ShardResponse},
    response_to_empty_result, response_to_result,
    utils::bincode::bincode_options,
};

#[derive(Debug, Clone)]
pub struct RemoteShardConnection {
    // The address to use to connect to the remote shard.
    pub address: String,
    connect_timeout: Duration,
    write_timeout: Duration,
    read_timeout: Duration,
}

impl RemoteShardConnection {
    #[must_use]
    pub fn from_args(address: String, args: &Args) -> Self {
        Self::new(
            address,
            Duration::from_millis(args.remote_shard_connect_timeout),
            Duration::from_millis(args.remote_shard_write_timeout),
            Duration::from_millis(args.remote_shard_read_timeout),
        )
    }

    #[must_use]
    pub fn new(
        address: String,
        connect_timeout: Duration,
        write_timeout: Duration,
        read_timeout: Duration,
    ) -> Self {
        Self {
            address,
            connect_timeout,
            write_timeout,
            read_timeout,
        }
    }

    pub async fn connect(&self) -> Result<TcpStream> {
        let stream =
            TcpStream::connect_timeout(&self.address, self.connect_timeout)
                .await?;
        stream.set_write_timeout(Some(self.write_timeout))?;
        stream.set_read_timeout(Some(self.read_timeout))?;
        Ok(stream)
    }

    pub async fn send_request(
        &self,
        request: ShardRequest,
    ) -> Result<ShardResponse> {
        let mut stream = self.connect().await?;
        send_message_to_stream(&mut stream, &ShardMessage::Request(request))
            .await?;
        let response = match get_message_from_stream(&mut stream).await? {
            ShardMessage::Response(response) => response,
            _ => return Err(Error::ResponseWrongType),
        };
        stream.close().await?;
        Ok(response)
    }

    pub async fn ping(&self) -> Result<()> {
        response_to_empty_result!(
            self.send_request(ShardRequest::Ping).await?,
            ShardResponse::Pong
        )
    }

    pub async fn get_metadata(&self) -> Result<Vec<NodeMetadata>> {
        response_to_result!(
            self.send_request(ShardRequest::GetMetadata).await?,
            ShardResponse::GetMetadata
        )
    }

    pub async fn get_collections(&self) -> Result<Vec<(String, u16)>> {
        response_to_result!(
            self.send_request(ShardRequest::GetCollections).await?,
            ShardResponse::GetCollections
        )
    }
}

pub async fn get_message_from_stream(
    stream: &mut (impl AsyncRead + Unpin),
) -> Result<ShardMessage> {
    let mut size_buf = [0; 4];
    stream.read_exact(&mut size_buf).await?;
    let size = u32::from_le_bytes(size_buf);
    let mut request_buf = vec![0; size as usize];
    stream.read_exact(&mut request_buf).await?;

    Ok(bincode_options()
        .deserialize_from::<_, ShardMessage>(&mut &request_buf[..])?)
}

pub async fn send_message_to_stream(
    stream: &mut (impl AsyncWrite + Unpin),
    message: &ShardMessage,
) -> Result<()> {
    let msg_buf = bincode_options().serialize(message)?;
    let size_buf = (msg_buf.len() as u32).to_le_bytes();

    stream.write_all(&size_buf).await?;
    stream.write_all(&msg_buf).await?;

    Ok(())
}