rubin 0.3.0

In-memory key-value store with the option for persistence
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

//! Server protocol for operating a store on a network
//!
//! Creates a [`MemStore`] which operates over a network, accepting requests from clients
//! to interact with the store.
//!
//! Can be run as an asynchronus task or as a background process, usage depends on end-user wants
//! and needs.

use std::sync::Arc;

use crate::{
    errors::MessageError,
    net::parser::{parse_request, Operation},
    store::mem::MemStore,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    sync::Mutex,
};

/// Sends a formatted response to the client prefixed with the [`Operation`] tag
async fn send_response(client: &mut TcpStream, code: Operation, msg: &str) {
    let response = format!("{}::{}\n", code, msg);
    client
        .write_all(response.as_bytes())
        .await
        .expect("unable to response to client");
}

/// Reads a message from a client and converts it to a string
async fn read_from_client(client: &mut TcpStream) -> String {
    let mut buffer = vec![0; 4096];
    let n_bytes = client
        .read(&mut buffer)
        .await
        .expect("unable to read from client");

    if n_bytes == 0 {
        return String::from("");
    }

    let msg = String::from_utf8_lossy(&buffer[..n_bytes]);
    msg.trim_end().to_string()
}

/// Main handler for the server
///
/// Processes incoming requests from the client and performs the requested operation.
/// If the operation cannot be processed, an error is returned.
async fn handler(mut client: TcpStream, store: Arc<Mutex<MemStore>>) {
    let msg = read_from_client(&mut client).await;

    let message = match parse_request(&msg) {
        Ok(msg) => msg,
        Err(error) => {
            match error {
                MessageError::InvalidMessage => {
                    send_response(&mut client, Operation::Error, "invalid message").await
                }
                MessageError::InvalidFormat => {
                    send_response(&mut client, Operation::Error, "invalid request format").await
                }
            }
            return;
        }
    };

    let mut vault = store.lock().await;
    match message.op {
        Operation::StringSet => {
            let key = &message.args[0];
            let value = &message.args[1];

            let _ = vault.insert_string(key, value);
            send_response(&mut client, message.op, "OK").await;
        }
        Operation::StringGet => {
            let key = &message.args[0];

            if let Ok(value) = vault.get_string(key) {
                send_response(&mut client, message.op, &value).await;
            }
        }
        Operation::StringRemove => {
            let key = &message.args[0];

            if let Ok(value) = vault.remove_string(key) {
                send_response(&mut client, message.op, &value).await;
            }
        }
        Operation::StringClear => {
            if vault.clear_strings().is_ok() {
                send_response(&mut client, message.op, "OK").await;
            }
        }
        _ => {
            send_response(&mut client, Operation::Noop, "nothing to do").await;
        }
    }
}

/// Starts the server to accept clients
///
/// This can be run as an independent task or as part of separate binary.
///
/// The usage really depends on what the end-user wants so this offers a simple
/// function to start the server.
///
/// An example could be that this function could be started as a Daemon using a
/// crate (e.g. [Daemonize](https://crates.io/crates/daemonize))
///
/// # Usage
///
/// ```no_run
/// use rubin::net::server::start;
///
/// #[tokio::main]
/// async fn main() -> std::io::Result<()> {
///     tokio::task::spawn(start("127.0.0.1", 9876));
///
///     // Rest of the workload
///
///     Ok(())
/// }
/// ```
pub async fn start(addr: &str, port: usize) -> std::io::Result<()> {
    let store = Arc::new(Mutex::new(MemStore::new()));
    let addr = format!("{}:{}", addr, port);
    let listener = TcpListener::bind(&addr).await?;

    println!("Started Rubin server at {}", addr);
    loop {
        let (client, _) = listener.accept().await?;
        let store = Arc::clone(&store);

        let client_addr = client.peer_addr()?;
        println!("Accepted new client: {}", client_addr);

        tokio::spawn(async move {
            handler(client, store).await;
        });
    }
}