本エントリーではKey Value Storeをtokioで作るうえで学んだことを書いていきます。
RustのLT会 Shinjuku.rs #13で話させていただいた内容です。
発表時のスライド
作ったもの
TCP(TLS)でremoteに接続してKey/ValueをCRUDするだけのserverです。
❯ export KVSD_HOST=kvsd.ymgyt.io ❯ kvsd set hello rust OK ❯ kvsd get hello rust ❯ kvsd set hello 'rust!!!' OK ❯ kvsd delete hello OK old value: rust!!!
TCP Server
ClientからのTCP接続ごとにtaskをtokio::spawn()して専用のHandlerを生成します。
Key/Valueを保存するFileごとにもtaskを生成しておき、共有resourceの処理でlock等が必要ないようにします。
(図の四角がtokio::spawn()したtaskを表しています)
Max Connection
tokioを利用することで、clientごとにthreadを生成する必要がなくなったのですが無制限にtaskを生成するわけにもいかないので最大connection数を制御したいです。
最大connectionに達したときはTcpListener.accept()をよびださないように実装します。
struct SemaphoreListener { inner: TcpListener, max_connections: Arc<Semaphore>, } impl SemaphoreListener { fn new(listener: TcpListener, max_connections: u32) -> Self { Self { inner: listener, max_connections: Arc::new(Semaphore::new(max_connections as usize)), } } async fn accept(&mut self) -> std::io::Result<(TcpStream, std::net::SocketAddr)> { self.max_connections.acquire().await.forget(); self.inner.accept().await } }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L494
tokio::sync::SemaphoreでTcpListenerをwrapしてaccept()時にacquire()をよぶことで最大connectionに達したときはblockするようにします。
所有権の関係でforget()を呼んでいるので、Clientとの接続が終了したときのHandlerのdrop処理で帳尻をあわせる必要があります。
struct Handler { // ... max_connections: Arc<Semaphore>, } impl Drop for Handler { fn drop(&mut self) { self.max_connections.add_permits(1); } }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L488
Graceful shutdown
終了処理時(SIGINT等)にすべてのClientとの接続が切れるまで待機する必要があります。
Graceful shutdownに必要な処理はそれぞれ以下のtokioのAPIを利用することで実現できました。
- Signal handling =>
tokio::signal::ctrl_c() - Shutdownの通知 =>
tokio::sync::broadcast::channel() - Shutdownの待機 =>
tokio::sync::mpsc::channel()
impl Server { pub(crate) async fn run( mut self, listener: TcpListener, shutdown: impl Future, ) -> Result<()> { tokio::select! { result = self.serve(listener) => { if let Err(err) = result { error!(cause = %err, "Failed to accept"); } } _ = shutdown => { info!("Shutdown signal received"); } } info!("Notify shutdown to all handlers"); self.graceful_shutdown.shutdown().await; info!("Shutdown successfully completed"); Ok(()) } } struct GracefulShutdown { notify_shutdown: broadcast::Sender<ShutdownSignal>, shutdown_complete_tx: mpsc::Sender<ShutdownCompleteSignal>, shutdown_complete_rx: mpsc::Receiver<ShutdownCompleteSignal>, } impl GracefulShutdown { fn new() -> Self { let (notify_shutdown, _) = broadcast::channel(1); let (shutdown_complete_tx, shutdown_complete_rx) = mpsc::channel(1); Self { notify_shutdown, shutdown_complete_tx, shutdown_complete_rx, } } // Notify handlers of the shutdown and wait for it to be completed. async fn shutdown(mut self) { // Notify shutdown to all handler. drop(self.notify_shutdown); // Drop final Sender so the Receiver below can complete. drop(self.shutdown_complete_tx); // Wait for all handler to finish. let _ = self.shutdown_complete_rx.recv().await; } }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L204
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L159
tokio::select!を利用することで、task内で２つのfutureを同時にawaitすることができるようになります。
impl Futureとしておくことでtest時にはtokio::sync::Notifyを利用することでできます。
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/tests/integration_test.rs#L38
Stream read/write
TCP上で自分で定義したデータ構造での通信を目指します。
MessageはClient-Server間でやりとりする単位です(Authenticate, Set, Success,...)
FrameはMessageの構成要素です(String, Bytes,Null,...)
実体はredis serialization protocolの劣化版みたいなものです。
TcpStream/TlsStreamをwrapした構造体を定義して、client/serverにprotocol実装を提供します。
use bytes::{Buf, BytesMut}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufWriter}; pub struct Connection<T = TcpStream> { stream: BufWriter<T>, // The buffer for reading frames. buffer: BytesMut, }
bufferはtcpからreadしてMessageを構成するために利用します。
write
client/serverから通信したいMessageをうけとるとそれをFramesに分解してそれぞれserializationしていきます。
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufWriter}; impl<T> Connection<T> where T: AsyncWrite + AsyncRead + Unpin, { pub(crate) async fn write_message(&mut self, message: impl Into<MessageFrames>) -> Result<()> { let frames = message.into(); self.stream.write_u8(frameprefix::MESSAGE_FRAMES).await?; self.write_decimal(frames.len()).await?; for frame in frames { self.write_frame(frame).await? } self.stream.flush().await?; Ok(()) } async fn write_frame(&mut self, frame: Frame) -> Result<()> { match frame { Frame::MessageType(mt) => { self.stream.write_u8(frameprefix::MESSAGE_TYPE).await?; self.stream.write_u8(mt.into()).await?; } Frame::String(val) => { self.stream.write_u8(frameprefix::STRING).await?; self.stream.write_all(val.as_bytes()).await?; self.stream.write_all(DELIMITER).await?; } Frame::Bytes(val) => { self.stream.write_u8(frameprefix::BYTES).await?; self.write_decimal(val.len() as u64).await?; self.stream.write_all(val.as_ref()).await?; self.stream.write_all(DELIMITER).await?; } Frame::Time(val) => { self.stream.write_u8(frameprefix::TIME).await?; self.stream.write_all(val.to_rfc3339().as_bytes()).await?; self.stream.write_all(DELIMITER).await?; } Frame::Null => { self.stream.write_u8(frameprefix::NULL).await?; } } Ok(()) } }
AsyncWriteExtにはendianを考慮したwrite methodが定義されているので便利です。
read
impl<T> Connection<T> where T: AsyncWrite + AsyncRead + Unpin, { pub(crate) async fn read_message_with_timeout( &mut self, duration: Duration, ) -> Result<Option<Message>> { match tokio::time::timeout(duration, self.read_message()).await { Ok(read_result) => read_result, Err(elapsed) => Err(Error::from(elapsed)), } } pub(crate) async fn read_message(&mut self) -> Result<Option<Message>> { match self.read_message_frames().await? { Some(message_frames) => Ok(Some(Message::from_frames(message_frames)?)), None => Ok(None), } } async fn read_message_frames(&mut self) -> Result<Option<MessageFrames>> { loop { if let Some(message_frames) = self.parse_message_frames()? { return Ok(Some(message_frames)); } if 0 == self.stream.read_buf(&mut self.buffer).await? { return if self.buffer.is_empty() { Ok(None) } else { Err(ErrorKind::ConnectionResetByPeer.into()) }; } } } fn parse_message_frames(&mut self) -> Result<Option<MessageFrames>> { use FrameError::Incomplete; let mut buf = Cursor::new(&self.buffer[..]); match MessageFrames::check_parse(&mut buf) { Ok(_) => { let len = buf.position() as usize; buf.set_position(0); let message_frames = MessageFrames::parse(&mut buf)?; self.buffer.advance(len); Ok(Some(message_frames)) } Err(Incomplete) => Ok(None), Err(e) => Err(e.into()), } } }
futureにtimeoutを設定するにはtokio::time::timeoutを利用します。
tokio::time::timeout()でwrapするとtimeoutつきのfutureを作れ、timeoutするとtokio::time::Elapsedを返してくれます。
tcpからreadしたときに意味ある単位でreadできているとは限らないのでbufferに読み込んだタイミングで一度parseを試みます。
parse時の実装をsimpleにするためにstd::io::Cursorでwrapしておくことで、parse側では必要な単位でbufferを読み込めます。
bufferにMessageを構成するに十分なFrameがない場合はloopでtcpのreadを繰り返します。
bufferに十分なFramesがある場合は、呼び出し元にMessageを返して、bufferをclearします。
TcpStreamのtestにはtokio::io::duplex()が便利
TcpStreamをwrapしたような型のtestではtokio::io::duplex()が便利でした。
#[test] fn message_frames() { tokio_test::block_on(async move { let (client, server) = tokio::io::duplex(1024); let mut client_conn = Connection::new(client, None); let mut server_conn = Connection::new(server, None); let messages: Vec<Message> = vec![ Message::Authenticate(Authenticate::new("user", "pass")), Message::Ping(Ping::new().record_client_time()), // ... ]; let messages_clone = messages.clone(); let write_handle = tokio::spawn(async move { for message in messages { client_conn.write_message(message).await.unwrap(); } }); let read_handle = tokio::spawn(async move { for want in messages_clone { let got = server_conn.read_message().await.unwrap().unwrap(); assert_eq!(want, got); } }); write_handle.await.unwrap(); read_handle.await.unwrap(); }) }
このような感じでtcp listenerを用意せずにメモリ上でtestが完結できました。
async codeのtestではruntimeを起動しておく必要があるので、tokio_testを利用しました。
TLS
TCPをTLSで保護するにはtokio_rustlsを利用できます。
server
use tokio_rustls::rustls; use tokio_rustls::server::TlsStream; use tokio_rustls::TlsAcceptor; let mut tls_config = rustls::ServerConfig::new(rustls::NoClientAuth::new()); let certs: Vec<rustls::Certificate> = self.config.load_certs(); let mut keys: Vec<rustls::PrivateKey> = self.config.load_keys(); tls_config.set_single_cert(certs, keys.remove(0)).unwrap(); let tls_acceptor = TlsAcceptor::from(Arc::new(tls_config)); let tls_stream: TlsStream<TcpStream> = acceptor.accept(stream).await?;
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L251
設定を生成して、鍵と証明書をloadして、TcpStreamをwrapします。
client
use tokio_rustls::client::TlsStream; use tokio_rustls::{rustls, webpki, TlsConnector}; // tokio-rustls = { version = "0.20.0", features = ["dangerous_configuration"] } let mut tls_config = rustls::ClientConfig::new(); tls_config .dangerous() .set_certificate_verifier(Arc::new(DangerousServerCertVerifier::new())); let connector = TlsConnector::from(Arc::new(tls_config)); let domain = webpki::DNSNameRef::try_from_ascii_str(host) .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid host"))?; let stream = tokio::net::TcpStream::connect(addr).await?; let tls_stream: TlsStream<TcpStream> = connector.connect(domain, stream).await?;
clientも同様に設定を生成してTcpStreamをwrapしてやります。
local環境ではオレオレ証明書を使いたかったのでdangerous_configuration featureを有効にして証明書検証処理を自分で定義することができます。
struct DangerousServerCertVerifier {} impl DangerousServerCertVerifier { fn new() -> Self { Self {} } } impl rustls::ServerCertVerifier for DangerousServerCertVerifier { fn verify_server_cert( &self, _roots: &rustls::RootCertStore, _presented_certs: &[rustls::Certificate], _dns_name: webpki::DNSNameRef<'_>, _oscp_response: &[u8], ) -> Result<rustls::ServerCertVerified, rustls::TLSError> { Ok(ServerCertVerified::assertion()) } }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/client/tcp.rs#L168
Shared stateをMutexからchannelで管理
Shared state(file)を排他的に利用したい場合にMutex等でlockするのではなく専用のtaskを生成してchannel経由で処理を依頼します。
処理の依頼をchannelに書き込むのはいいのですが、結果をどう受け取るのかが問題になってきます。
ここではtokio::sync::oneshotを利用しました。
use tokio::sync::oneshot; pub(crate) struct Work<Req, Res> { // ... pub(crate) request: Req, // Wrap with option so that response can be sent via mut reference. pub(crate) response_sender: Option<oneshot::Sender<Result<Res>>>, }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/core/uow.rs#L27
処理を依頼する型のfieldにoneshot::Sender<T>を定義しておき呼び出し元で保持しているReceiverで結果をうけとります。
serverがclientからMessageをうけとって、Key/Valueを取得する処理は以下のようになりました。
Message::Get(get) => { let (work: Uow, rx: Receiver<Result<_>>) = UnitOfWork::new_get(get); self.sender.send(work).await?; match rx.await? { Ok(Some(value)) => { connection.write_message(Success::with_value(value)).await? } Ok(None) => connection.write_message(Success::new()).await?, _ => unreachable!(), } }
https://github.com/ymgyt/kvsd/blob/6b40c6a3edad0f416631f7ae66674fb5d7922cd7/src/server/tcp.rs#L444
Fileへのappend onlyな保存
ようやくKey/Valueの話に。
永続化するKey/Valueはfileにappend onlyで保存していきます。メモリ上ではKeyとoffsetを保持しておきます。
Key/Valueをfileから読み込む処理はいかのようになりました。
use tokio::sync::mpsc::Receiver; use tokio::io::{AsyncRead, AsyncSeek, AsyncSeekExt, AsyncWrite, SeekFrom}; pub(crate) struct Table<File = fs::File> { file: File, index: Index, // HashMap<String,usize>, receiver: Receiver<UnitOfWork>, } impl<File> Table<File> where File: AsyncWrite + AsyncRead + AsyncSeek + Unpin, { async fn lookup_entry(&mut self, key: &Key) -> Result<Option<Entry>> { let offset = match self.index.lookup_offset(key){ Some(offset) => offset, None => return Ok(None), }; let current = self.file.seek(SeekFrom::Current(0)).await?; self.file.seek(SeekFrom::Start(offset as u64)).await?; let (_, entry) = Entry::decode_from(&mut self.file).await?; self.file.seek(SeekFrom::Start(current)).await?; Ok(Some(entry)) } }
tokio::io::AsyncSeekがあるのでoffsetにseekする処理も同期と同じように書けました。
まとめ
- はじめてtokioを利用しようと思ったときmini redisの実装が非常に参考になりました。コメントもたくさんあり親切でした。
- tokioを利用してremoteと通信してデータを保存するまでを自分で作れて楽しかったです
- tokioのAPIを利用してGoに近い形でconcurrentな処理書けそうと思いました。
tokio::select!(select)tokio::spawn()(goroutine)sync::{mpsc,oneshot,broadcast}(chan)
この記事を書いている日にtokio v1.0.0がreleaseされました。
最低でも5年はメンテしていく旨も発表されておりtokioを利用しはじめるにはちょうどいい時期なのではないでしょうか。