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
use std::net::SocketAddr;
use futures::executor::{self, ThreadPool};
use futures::io::AsyncReadExt;
use futures::task::SpawnExt;
use futures::{SinkExt, StreamExt};
use romio::{TcpListener, TcpStream};
use log::info;
use super::error::Result;
use agilulf_protocol::{AsyncReadBuffer, AsyncWriteBuffer};
use agilulf_protocol::{ProtocolError, Result as ProtocolResult};
use crate::storage::AsyncDatabase;
use agilulf_protocol::Command;
use std::sync::Arc;
pub struct Server {
listener: TcpListener,
database: Arc<dyn AsyncDatabase>,
}
impl Server {
pub fn new(address: &str, database: impl AsyncDatabase + 'static) -> Result<Server> {
let addr = address.parse::<SocketAddr>()?;
let listener = TcpListener::bind(&addr)?;
Ok(Server {
listener,
database: Arc::new(database),
})
}
pub async fn run_async(mut self) -> Result<()> {
let mut thread_pool = ThreadPool::new()?;
let mut incoming = self.listener.incoming();
while let Some(stream) = incoming.next().await {
let stream: TcpStream = stream.unwrap();
let database = self.database.clone();
thread_pool.spawn(async move {
match handle_stream(stream, database).await {
Ok(()) => {}
Err(err) => log::error!("Error while handling stream: {}", err),
}
})?
}
Ok(())
}
pub fn run(self) -> Result<()> {
executor::block_on(async { self.run_async().await })?;
Ok(())
}
}
async fn handle_stream(stream: TcpStream, database: Arc<dyn AsyncDatabase>) -> Result<()> {
let remote_addr = stream.peer_addr()?;
info!("Accepting stream from: {}", remote_addr);
let (reader, writer) = stream.split();
let mut command_stream = AsyncReadBuffer::new(reader).into_command_stream().fuse();
let reply_sink = AsyncWriteBuffer::new(writer).into_reply_sink();
let mut process_sink = reply_sink.with(|command: ProtocolResult<Command>| {
Box::pin(async {
match command {
Ok(command) => match command {
Command::GET(command) => {
ProtocolResult::Ok(database.get(command.key).await.into())
}
Command::PUT(command) => {
ProtocolResult::Ok(database.put(command.key, command.value).await.into())
}
Command::SCAN(command) => {
ProtocolResult::Ok(database.scan(command.start, command.end).await.into())
}
Command::DELETE(command) => {
ProtocolResult::Ok(database.delete(command.key).await.into())
}
},
Err(err) => ProtocolResult::Ok(err.into()),
}
})
});
loop {
let command = command_stream.select_next_some().await;
if let Err(err) = process_sink.send(command).await {
match &err {
ProtocolError::IOError(err) => match err.kind() {
std::io::ErrorKind::BrokenPipe => {
break;
}
_ => {}
},
_ => {}
}
log::error!("Error while sending reply {:?}", err);
break;
}
}
info!("Closing stream from: {}", remote_addr);
Ok(())
}