├── .gitignore ├── test_data ├── test.jpg ├── test.pdf ├── test.png └── test.html ├── src ├── lib.rs ├── bin │ └── segment.rs ├── config.rs ├── server.rs ├── db.rs ├── command │ ├── mod.rs │ └── test.rs ├── connection.rs └── frame.rs ├── Cargo.toml ├── .github └── workflows │ └── build.yml ├── segment.conf ├── LICENSE ├── docs └── protocol.v1.md └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | /Cargo.lock 3 | .DS_Store -------------------------------------------------------------------------------- /test_data/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thetinygoat/segment/HEAD/test_data/test.jpg -------------------------------------------------------------------------------- /test_data/test.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thetinygoat/segment/HEAD/test_data/test.pdf -------------------------------------------------------------------------------- /test_data/test.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thetinygoat/segment/HEAD/test_data/test.png -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | mod command; 2 | pub mod config; 3 | mod connection; 4 | mod db; 5 | mod frame; 6 | pub mod server; 7 | -------------------------------------------------------------------------------- /test_data/test.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | Document 8 | 9 | 10 | This is a test html doc for segment binary safety test 11 | 12 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "segment" 3 | version = "0.0.1" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | tokio = { version = "1", features = ["rt-multi-thread", "time", "macros", "sync", "signal", "net", "io-util"] } 10 | anyhow = "1.0.66" 11 | thiserror = "1.0.37" 12 | clap = { version = "4.0.18", features = ["derive"] } 13 | tokio-util = "0.7.4" 14 | crossbeam = "0.8.2" 15 | bytes = "1.2.1" 16 | tracing = "0.1.37" 17 | tracing-subscriber = "0.3.16" 18 | atoi = "2.0.0" 19 | parking_lot = "0.12.1" 20 | tokio-test = "0.4.2" 21 | async-recursion = "1.0.0" 22 | sysinfo = "0.26.8" -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | push: 5 | branches: [ "dev" ] 6 | pull_request: 7 | branches: [ "dev" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | build-linux: 14 | 15 | runs-on: ubuntu-latest 16 | 17 | steps: 18 | - uses: actions/checkout@v3 19 | - name: Build 20 | run: cargo build --verbose 21 | - name: Run tests 22 | run: cargo test --verbose 23 | 24 | build-macos: 25 | 26 | runs-on: macos-12 27 | 28 | steps: 29 | - uses: actions/checkout@v3 30 | - name: Build 31 | run: cargo build --verbose 32 | - name: Run tests 33 | run: cargo test --verbose 34 | -------------------------------------------------------------------------------- /segment.conf: -------------------------------------------------------------------------------- 1 | port=1698 2 | 3 | # specifies the max memory that is availabe to the server. Once the server reaches 4 | # this memory limit the server will start evicting keys according to the max memory 5 | # policy configured for that keyspace. Only two units of memory are supported which are *mb* and *gb*. 6 | # If you want the server to not have any memory limit set this as 0 (0mb or 0gb). 7 | # Examples: 8 | # max_memory=200gb 9 | # max_memory=100mb 10 | # max_memory=0mb 11 | max_memory=512mb 12 | 13 | # connection buffer size is the size of the connection buffer in *bytes* which is used to read 14 | # data from the socket. You can tune acording to size of data that you expect. A larger buffer size will 15 | # use more memory. Only change this if you know what you are doing 16 | connection_buffer_size=4096 17 | 18 | # bind tells the segment server which interface to listen on 19 | bind=127.0.0.1 20 | -------------------------------------------------------------------------------- /src/bin/segment.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use clap::Parser; 3 | use segment::config::ServerConfig; 4 | use segment::server; 5 | use tokio::net::TcpListener; 6 | use tracing::Level; 7 | 8 | #[derive(Debug, Parser)] 9 | struct Args { 10 | /// path to segment config file 11 | #[arg(long, default_value = "segment.conf")] 12 | config: String, 13 | 14 | /// start the server in debug mode 15 | #[arg(long)] 16 | debug: bool, 17 | } 18 | 19 | #[tokio::main] 20 | async fn main() -> Result<()> { 21 | let args = Args::parse(); 22 | let mut log_level = Level::INFO; 23 | if args.debug { 24 | log_level = Level::DEBUG; 25 | } 26 | let subscriber = tracing_subscriber::fmt().with_max_level(log_level).finish(); 27 | tracing::subscriber::set_global_default(subscriber)?; 28 | let cfg = ServerConfig::load_from_disk(&args.config)?; 29 | let ln = TcpListener::bind(format!("{}:{}", cfg.bind(), cfg.port())).await?; 30 | server::start(ln, cfg).await?; 31 | Ok(()) 32 | } 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Sachin Saini 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /src/config.rs: -------------------------------------------------------------------------------- 1 | use std::fs::File; 2 | use std::io::{self, BufRead, BufReader}; 3 | use std::net::{AddrParseError, IpAddr, Ipv4Addr}; 4 | use std::num::ParseIntError; 5 | use std::str::FromStr; 6 | use thiserror::Error; 7 | 8 | const PORT_LABEL: &str = "port"; 9 | const MAX_MEMORY_LABEL: &str = "max_memory"; 10 | const CONNECTION_BUFFER_SIZE_LABEL: &str = "connection_buffer_size"; 11 | const BIND_LABEL: &str = "bind"; 12 | 13 | #[derive(Debug)] 14 | pub struct ServerConfig { 15 | port: u16, 16 | max_memory: u64, 17 | connection_buffer_size: usize, 18 | bind: IpAddr, 19 | } 20 | 21 | #[derive(Debug, Error)] 22 | pub enum ServerConfigError { 23 | #[error(transparent)] 24 | FileRead(#[from] io::Error), 25 | 26 | #[error("invalid config file format at '{0}'")] 27 | InvalidFormat(String), 28 | 29 | #[error("unknown directive '{0}' at '{1}'")] 30 | UnknownDirective(String, String), 31 | 32 | #[error(transparent)] 33 | ParseIntError(#[from] ParseIntError), 34 | 35 | #[error(transparent)] 36 | AddrParseError(#[from] AddrParseError), 37 | } 38 | 39 | impl ServerConfig { 40 | pub fn load_from_disk(path: &str) -> Result { 41 | let reader = BufReader::new(File::open(path)?); 42 | Self::parse(reader) 43 | } 44 | 45 | fn parse(reader: BufReader) -> Result { 46 | let mut config = ServerConfig { 47 | port: 1698, 48 | max_memory: 0, 49 | connection_buffer_size: 4096, 50 | bind: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 51 | }; 52 | for maybe_line in reader.lines() { 53 | let line = &maybe_line?; 54 | if line.trim().starts_with('#') || line.trim().is_empty() { 55 | continue; 56 | } 57 | 58 | let tokens: Vec<&str> = line.split('=').map(|token| token.trim()).collect(); 59 | 60 | if tokens.len() < 2 || tokens.len() > 2 { 61 | return Err(ServerConfigError::InvalidFormat(line.clone())); 62 | } 63 | 64 | match tokens[0] { 65 | PORT_LABEL => { 66 | let port = tokens[1].parse::()?; 67 | config.port = port; 68 | } 69 | MAX_MEMORY_LABEL => { 70 | if tokens[1].len() < 3 { 71 | return Err(ServerConfigError::InvalidFormat(line.clone())); 72 | } 73 | let unit = &tokens[1][tokens[1].len() - 2..]; 74 | let memory = tokens[1][..tokens[1].len() - 2].parse::()?; 75 | match unit { 76 | "mb" => config.max_memory = memory * 1024 * 1024, 77 | "gb" => config.max_memory = memory * 1024 * 1024 * 1024, 78 | _ => { 79 | return Err(ServerConfigError::InvalidFormat(line.clone())); 80 | } 81 | } 82 | } 83 | CONNECTION_BUFFER_SIZE_LABEL => { 84 | let connection_buffer_size = tokens[1].parse::()?; 85 | config.connection_buffer_size = connection_buffer_size; 86 | } 87 | BIND_LABEL => { 88 | let bind = IpAddr::from_str(tokens[1])?; 89 | config.bind = bind 90 | } 91 | _ => { 92 | return Err(ServerConfigError::UnknownDirective( 93 | tokens[0].to_string(), 94 | line.clone(), 95 | )) 96 | } 97 | } 98 | } 99 | 100 | Ok(config) 101 | } 102 | 103 | pub fn port(&self) -> u16 { 104 | self.port 105 | } 106 | 107 | pub fn max_memory(&self) -> u64 { 108 | self.max_memory 109 | } 110 | 111 | pub fn connection_buffer_size(&self) -> usize { 112 | self.connection_buffer_size 113 | } 114 | 115 | pub fn bind(&self) -> String { 116 | self.bind.to_string() 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /docs/protocol.v1.md: -------------------------------------------------------------------------------- 1 | # Segment Wire Protocol 2 | 3 | This document is the specification for the Segment wire protocol. 4 | 5 | ## Introduction 6 | 7 | Segment has a text based TCP protocol that is used for client-server architecture. It is inspired by [RESP](https://redis.io/docs/reference/protocol-spec/) and [Memcached](https://github.com/memcached/memcached/blob/master/doc/protocol.txt). The protocol aims to be _human readable_, _fast to parse_ and _simple to implement_. 8 | 9 | ## Request-Response Model 10 | 11 | The client connects to the Segment server on port 1698 by default. The client sends a command which is made up of arguments and flags, on receiving a request the server sends back a reply. This is a traditional client-server architecture. 12 | 13 | ## Protocol Specification 14 | 15 | ### Data Types 16 | 17 | The protocol supports 8 data types: 18 | 19 | - [String](#strings) 20 | - [Integer](#integers) 21 | - [Double](#doubles) 22 | - [Boolean](#booleans) 23 | - [Null](#null) 24 | - [Error](#errors) 25 | - [Array](#arrays) 26 | - [Map](#maps) 27 | 28 | The first byte determines the data type. 29 | 30 | - For **Strings** the first byte is `$` 31 | - For **Integers** the first byte is `%` 32 | - For **Doubles** the first byte is `.` 33 | - For **Booleans** the first byte is `^` 34 | - For **Null** the first byte is `-` 35 | - For **Errors** the first byte is `!` 36 | - For **Arrays** the first byte is `*` 37 | - For **Maps** the first byte is `#` 38 | 39 | The protocol uses CRLF (`\r\n`) as the delimiter. 40 | 41 | #### Strings 42 | 43 | Strings are encoded as follows: A `$` character followed by the length of the string followed by CRLF. After encoding the length, the actual data is appended followed by a CRLF. 44 | 45 | ``` 46 | $11\r\nhello world\r\n 47 | ``` 48 | 49 | Since the strings are prefixed with their lengths we don't need to search for any delimiter to mark the end of the string. This makes it fast to parse and it also makes the strings **binary safe**. 50 | 51 | #### Integers 52 | 53 | Integers are encoded as follows: A `%` character followed by the integer that we want to encode followed by CRLF. 54 | 55 | ``` 56 | %100\r\n 57 | ``` 58 | 59 | #### Doubles 60 | 61 | Doubles are similar to integers, they are encoded as follows: A `.` character followed by the double that we want to encode followed by CRLF. 62 | 63 | ``` 64 | .26.3\r\n 65 | ``` 66 | 67 | #### Booleans 68 | 69 | Booleans are encoded as follows: A `^` character followed by a `0` for false and a `1` for true followed by CRLF. 70 | 71 | ``` 72 | ^0\r\n // false 73 | 74 | ^1\r\n // true 75 | ``` 76 | 77 | #### Null 78 | 79 | Nulls are encoded as follows: A `-` character followed by CRLF. 80 | 81 | ``` 82 | -\r\n 83 | ``` 84 | 85 | #### Errors 86 | 87 | Errors are similar to strings and are encoded as follows: A `!` character followed by the length of the error data followed by CRLF. After encoding the length, the data is appended followed by CRLF 88 | 89 | ``` 90 | !5\r\nerror\r\n 91 | ``` 92 | 93 | #### Arrays 94 | 95 | Array is a container type, it can contain all the other data types. An array is encoded as follows: A `*` character followed by the number of items in the array followed By CRLF. After encoding the length of the array we can just encode any type into it. Arrays can contain different data types at once. 96 | 97 | ``` 98 | // empty array 99 | *0\r\n 100 | 101 | // array containing one integer 102 | *1\r\n%100\r\n 103 | 104 | // array containing an integer and a string 105 | *2\r\n%100\r\n$5\r\nhello\r\n 106 | ``` 107 | 108 | #### Maps 109 | 110 | A map is a hash map, it's similar to an array and is encoded as follows: A `#` character followed by the number of items in the map followed by CRLF. Please note that a key-value pair is considered as a single unit/item. After encoding the number of items we can encode any type as key and value. Even though a key can be of any type, segment will only send keys as strings. 111 | 112 | ``` 113 | // empty map 114 | #0\r\n 115 | 116 | // map containing one key-value pair, with key being hello and value being world 117 | #1\r\n$5\r\nhello\r\n$5\r\nworld\r\n 118 | ``` 119 | 120 | ## Sending Commands to a Segment Server 121 | 122 | Now that you are familiar with the wire protocol, you can use it to write a client to interact with a Segment server. 123 | 124 | - A client can send the command to a Segment server as an _Array of Strings_ only. Using any other data type to send the command will result in an error. 125 | - The server can respond with any of the above data type. 126 | 127 | For example, the create command will be encoded as follows: 128 | 129 | ```shell 130 | *2\r\n$6\r\nCREATE\r\n$3\r\nfoo\r\n 131 | ``` 132 | -------------------------------------------------------------------------------- /src/server.rs: -------------------------------------------------------------------------------- 1 | use crate::command; 2 | use crate::config::ServerConfig; 3 | use crate::connection::Connection; 4 | use crate::db::Db; 5 | use anyhow::Result; 6 | use crossbeam::sync::WaitGroup; 7 | use std::sync::Arc; 8 | use std::time::Duration; 9 | use sysinfo::{Pid, ProcessExt, System, SystemExt}; 10 | use tokio::net::{TcpListener, TcpStream}; 11 | use tokio::signal; 12 | use tokio::sync::broadcast; 13 | use tracing::{debug, error, info}; 14 | 15 | struct Server { 16 | ln: TcpListener, 17 | cfg: ServerConfig, 18 | wg: WaitGroup, 19 | db: Arc, 20 | done_tx: broadcast::Sender<()>, 21 | evict_tx: broadcast::Sender<()>, 22 | } 23 | 24 | struct ConnectionHandler { 25 | connection: Connection, 26 | done: broadcast::Receiver<()>, 27 | db: Arc, 28 | } 29 | 30 | pub async fn start(ln: TcpListener, cfg: ServerConfig) -> Result<()> { 31 | let srv = Server::new(ln, cfg); 32 | srv.start().await 33 | } 34 | 35 | impl Server { 36 | pub fn new(ln: TcpListener, cfg: ServerConfig) -> Self { 37 | let wg = WaitGroup::new(); 38 | let (done_tx, _) = broadcast::channel(1); 39 | let (evict_tx, _) = broadcast::channel(1); 40 | let db = Db::new(done_tx.subscribe(), wg.clone(), evict_tx.subscribe()); 41 | Server { 42 | ln, 43 | cfg, 44 | wg, 45 | done_tx, 46 | db: Arc::new(db), 47 | evict_tx, 48 | } 49 | } 50 | 51 | pub async fn start(self) -> Result<()> { 52 | info!( 53 | "server started on port {}:{}", 54 | self.cfg.bind(), 55 | self.cfg.port() 56 | ); 57 | let monitor_wg = self.wg.clone(); 58 | let mut monitor_done_rx = self.done_tx.subscribe(); 59 | let monitor_evict_tx = self.evict_tx.clone(); 60 | let server_max_memory = self.cfg.max_memory(); 61 | // FIXME: move this to a separate fn 62 | tokio::spawn(async move { 63 | let pid = std::process::id() as i32; 64 | let mut monitor = System::new(); 65 | monitor.refresh_process(Pid::from(pid)); 66 | loop { 67 | tokio::select! { 68 | _ = monitor_done_rx.recv() => { 69 | debug!("stopping system monitor, shutdown signal received"); 70 | break; 71 | } 72 | _ = tokio::time::sleep(Duration::from_millis(1000)) => { 73 | monitor.refresh_process(Pid::from(pid)); 74 | if let Some(process) = monitor.process(Pid::from(pid)) { 75 | let memory = process.memory(); 76 | if memory >= server_max_memory && server_max_memory > 0 { 77 | debug!("broadcasting evict event, server max memory (bytes) = {}, current memory usage (bytes) = {}", server_max_memory, memory); 78 | if let Err(err) = monitor_evict_tx.send(()) { 79 | error!("no listeners available for max memory eviction event, error = {:?}", err); 80 | } 81 | } 82 | }else { 83 | error!("no process found with pid {}, max memory evictors will not work", pid); 84 | break; 85 | } 86 | } 87 | } 88 | } 89 | drop(monitor_wg) 90 | }); 91 | loop { 92 | tokio::select! { 93 | maybe_connection = self.ln.accept() => { 94 | let (stream, _) = maybe_connection?; 95 | let mut handler = ConnectionHandler::new(self.done_tx.subscribe(), stream, self.cfg.connection_buffer_size(), self.db.clone()); 96 | let wg = self.wg.clone(); 97 | tokio::spawn(async move { 98 | if let Err(e) = handler.handle().await { 99 | error!("{}", e) 100 | } 101 | drop(wg); 102 | }); 103 | } 104 | _ = signal::ctrl_c() => { 105 | info!("shutdown signal received"); 106 | drop(self.ln); 107 | drop(self.done_tx); 108 | break; 109 | } 110 | } 111 | } 112 | drop(self.db); 113 | self.wg.wait(); 114 | info!("shutdown complete, bye bye :)"); 115 | Ok(()) 116 | } 117 | } 118 | 119 | impl ConnectionHandler { 120 | pub fn new( 121 | done: broadcast::Receiver<()>, 122 | stream: TcpStream, 123 | connection_buf_size: usize, 124 | db: Arc, 125 | ) -> Self { 126 | let connection = Connection::new(stream, connection_buf_size); 127 | ConnectionHandler { 128 | connection, 129 | done, 130 | db, 131 | } 132 | } 133 | 134 | pub async fn handle(&mut self) -> Result<()> { 135 | debug!("new connection started"); 136 | loop { 137 | let maybe_frame = tokio::select! { 138 | _ = self.done.recv() => { 139 | break; 140 | } 141 | res = self.connection.read_frame() => res?, 142 | }; 143 | 144 | let frame = match maybe_frame { 145 | Some(frame) => frame, 146 | None => return Ok(()), 147 | }; 148 | 149 | let maybe_cmd = match command::parse(frame) { 150 | Ok(cmd) => Some(cmd), 151 | Err(e) => { 152 | self.connection.write_error(e).await?; 153 | None 154 | } 155 | }; 156 | 157 | let cmd = match maybe_cmd { 158 | Some(cmd) => cmd, 159 | None => continue, 160 | }; 161 | 162 | let maybe_result = match self.db.execute(cmd).await { 163 | Ok(frame) => Some(frame), 164 | Err(e) => { 165 | self.connection.write_error(e).await?; 166 | None 167 | } 168 | }; 169 | 170 | match maybe_result { 171 | Some(frame) => self.connection.write_frame(&frame).await?, 172 | 173 | None => continue, 174 | } 175 | } 176 | Ok(()) 177 | } 178 | } 179 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Segment 2 | 3 | segment fox logo 4 | 5 | ![build status](https://github.com/segment-dev/segment/actions/workflows/build.yml/badge.svg) 6 | 7 | Segment is a _simple_ and _fast_ in-memory key-value database written in Rust. 8 | 9 | ### Features 10 | 11 | - Simple to use and understand. 12 | - Keys can be separated into multiple dynamic keyspaces. 13 | - Keyspace level configuration. 14 | 15 | ### Why Segment? 16 | 17 | Segment's goal to is to provide a simpler and more intuitive in-memory key-value solution. It has certain features that other solutions don't. Let's go over them one by one. 18 | 19 | #### Keyspaces 20 | 21 | Segment has a concept of keyspaces. You can think of keyspaces like tables in a relational database, except that keyspaces don't have any schema. When the segment server starts there are no keyspaces and you can create as many as you like. 22 | 23 | #### Evictors 24 | 25 | Separating keys into keyspaces comes with a benefit, we can now have keyspace level configurations. Evictors are one such configuration. 26 | There are two types of evicros in Segment, **expiring evictors** and **max memory evictors**. 27 | 28 | ##### Expiring Evictors 29 | 30 | The expiring evictor is responsible for evicting expired keys which runs for every keyspaces. 31 | 32 | ##### Max Memory Evictors 33 | 34 | The second type of evictor is max memory evictor, which is responsible for evicting keys when the server reaches the max memory specified in `segment.conf`. 35 | Currently there are 3 max memory evictors: 36 | 37 | - Nop - Stands for no-operation which doesn't evict any keys. 38 | - Random - Evicts keys in a random order. 39 | - LRU - Evicts keys in a LRU fashion. 40 | 41 | There are plans to include even more evictors out of the box in future. 42 | 43 | Max memory evitors can configured at a keyspace level, which means that you can have a keyspace that does not evict at all while some keyspaces evict. 44 | This is powerful becuase now you don't have to spin up a separate server just because you want to have a separate eviction policy. 45 | 46 | #### Multithreaded 47 | 48 | Segment is multithreaded, which means it uses locks which can be a deal breaker for some use cases. But It works for most use cases and that's what segment is aiming for. 49 | 50 | #### Ease of Use 51 | 52 | Segment aims to be easy to use and intuitive. One way we are aiming to solve this is by having only one way of doing things. There is only one command to insert data and one way to get it back, this helps reduce the stuff that a developer needs to remember. 53 | 54 | One more thing that we are doing is using simple commands, for example let's take a look at the command to insert some data in a keyspace. 55 | 56 | ```shell 57 | SET my_keyspace my_key my_value IF NOT EXISTS EXPIRE AFTER 60000 58 | ``` 59 | 60 | This commnd tells the segment server to insert the key `my_key` with value `my_value` into the keyspace `my_keyspace` if the key does not exist already and expire the key after 1 minute (60000ms). 61 | 62 | The command reads like english and the intent is clear 63 | 64 | A similar command in redis would look like this. 65 | 66 | ``` 67 | SET my_key my_value NX EX 60 68 | ``` 69 | 70 | If you are not familiar with redis you will not understand what is happeining here. and if you want to have your ttl in milliseconds there is another flag for that and I don't even remember what it is called and that's the point, to reduce the dev effort. 71 | 72 | ### Installation 73 | 74 | Segment is built using Rust, so you will need rust and it's toolchain installed on your system. To install rust you can follow the steps [here](https://rustup.rs/). 75 | 76 | After installing you can follow the steps below to build and run segment from source. 77 | 78 | 1. Clone this repository. 79 | 2. `cd /path/to/cloned/repo` 80 | 3. `cargo build --release` 81 | 4. The final binary can be found in `./target/release` 82 | 83 | ### Running the sever 84 | 85 | After building you will find the `segment` binary in the `./target/release` directory. 86 | 87 | Segment requires `segment.conf` file to start. `segment.conf` is the config file that contains several server configurations. 88 | 89 | If you have the `segment.conf` file in the same directory as the segment binary you can just run the binary however, if the config file is in some separate directory you can start the Segment server using the command below 90 | 91 | ```shell 92 | segment --config=/path/to/segment.conf 93 | ``` 94 | 95 | If the server is started successfully you will see a log similar to this in your terminal. 96 | 97 | ```shell 98 | 2022-10-29T07:23:05.308471Z INFO segment::server: server started on port 1698 99 | ``` 100 | 101 | ### Client Libraries 102 | 103 | - [Node.js](https://github.com/segment-dev/segment-node) 104 | - [Rust (Alpha)](https://github.com/segment-dev/segment-rs) 105 | 106 | ### List of Commands 107 | 108 | #### `CREATE` 109 | 110 | ##### Description 111 | 112 | Used to create a new keyspace. By defualt it doesn't take any arguments except the name of the keyspace, but you can specify the evictor you want to use for the keyspace. 113 | 114 | ##### Essential Arguments 115 | 116 | - `` - Name of the keyspace that you want to create. 117 | 118 | ##### Optional Arguments 119 | 120 | - `EVICTOR` - Indicates the evictor that you want to use for the keyspace. Possible values include `NOP`, `RANDOM` and `LRU`. 121 | 122 | ##### Optional Flags 123 | 124 | - `IF NOT EXISTS` - If a keyspace already exists and you try to create it again the server will throw an error, but if you don't want an error you can pass this flag with the create command. 125 | 126 | ##### Return Type 127 | 128 | The return type can be a boolean or an error. 129 | 130 | ##### Examples 131 | 132 | ```shell 133 | CREATE my_keyspace 134 | ``` 135 | 136 | ```shell 137 | CREATE my_keyspace EVICTOR LRU 138 | ``` 139 | 140 | ```shell 141 | CREATE my_keyspace EVICTOR LRU IF NOT EXISTS 142 | ``` 143 | 144 | #### `DROP` 145 | 146 | ##### Description 147 | 148 | Used to drop a keyspace. 149 | 150 | ##### Essential Arguments 151 | 152 | - `` - Name of the keyspace that you want to drop. 153 | 154 | ##### Optional Flags 155 | 156 | - `IF EXISTS` - If a keyspace doesn't already exists and you try to drop it the server will throw an error, but if you don't want an error you can pass this flag with the drop command. 157 | 158 | ##### Return Type 159 | 160 | The return type can be a boolean or an error. 161 | 162 | ##### Examples 163 | 164 | ```shell 165 | DROP my_keyspace 166 | ``` 167 | 168 | ```shell 169 | DROP my_keyspace IF EXISTS 170 | ``` 171 | 172 | #### `SET` 173 | 174 | ##### Description 175 | 176 | Used to insert a value in the keyspace. 177 | 178 | ##### Essential Arguments 179 | 180 | - `` - Name of the keyspace that you want to create. 181 | - `` - Key that you want to insert. 182 | - `` - Value for the key. 183 | 184 | ##### Optional Arguments 185 | 186 | - `EXPIRE AFTER` - Expiry time of the key in milliseconds after which it will expire. 187 | - `EXPIRE AT` - Unix timestamp after which the key will expire. 188 | 189 | ##### Optional Flags 190 | 191 | - `IF NOT EXISTS` - If you want to set a key only if it does not already exists. 192 | - `IF EXISTS` - If you want to set a key only if it already exists. 193 | 194 | ##### Return Type 195 | 196 | The return type can be a boolean or an error. 197 | 198 | ##### Examples 199 | 200 | ```shell 201 | SET my_keyspace my_key my_value 202 | ``` 203 | 204 | ```shell 205 | SET my_keyspace my_key my_value IF NOT EXISTS 206 | ``` 207 | 208 | ```shell 209 | SET my_keyspace my_key my_value IF EXISTS 210 | ``` 211 | 212 | ```shell 213 | SET my_keyspace my_key my_value EXPIRE AFTER 60000 214 | ``` 215 | 216 | ```shell 217 | SET my_keyspace my_key my_value EXPIRE AT 1667041052 218 | ``` 219 | 220 | #### `GET` 221 | 222 | ##### Description 223 | 224 | Used to get a key from the keyspace. 225 | 226 | ##### Essential Arguments 227 | 228 | - `` - Name of the keyspace that you want to get the key from. 229 | - `` - key that you want to get. 230 | 231 | ##### Return Type 232 | 233 | The return type can be a string, null, or error. 234 | 235 | ##### Examples 236 | 237 | ```shell 238 | GET my_keyspace my_key 239 | ``` 240 | 241 | #### `DEL` 242 | 243 | ##### Description 244 | 245 | Used to delete a key from the keyspace. 246 | 247 | ##### Essential Arguments 248 | 249 | - `` - Name of the keyspace that you want to create. 250 | - `` - Name of the keyspace that you want to create. 251 | 252 | ##### Return Type 253 | 254 | The return type can be a boolean or error. 255 | 256 | ##### Examples 257 | 258 | ```shell 259 | DEL my_keyspace my_key 260 | ``` 261 | 262 | #### `COUNT` 263 | 264 | ##### Description 265 | 266 | Returns the number of keys in a keyspace. 267 | 268 | ##### Essential Arguments 269 | 270 | - `` - Name of the keyspace. 271 | 272 | ##### Return Type 273 | 274 | The return type can be an integer or error. 275 | 276 | ##### Examples 277 | 278 | ```shell 279 | COUNT my_keyspace 280 | ``` 281 | 282 | #### `TTL` 283 | 284 | ##### Description 285 | 286 | Returns the remaining TTL of the key in milliseconds. 287 | 288 | ##### Essential Arguments 289 | 290 | - `` - Name of the keyspace. 291 | - `` - Name of the key. 292 | 293 | ##### Return Type 294 | 295 | The return type can be an integer or null (if the key doesn't have an expiry or is already expired) or an error. 296 | 297 | ##### Examples 298 | 299 | ```shell 300 | TTL my_keyspace my_key 301 | ``` 302 | 303 | #### `PING` 304 | 305 | ##### Description 306 | 307 | Used to ping the server. 308 | 309 | ##### Return Type 310 | 311 | The return type is the string `pong`. 312 | 313 | ##### Examples 314 | 315 | ```shell 316 | PING 317 | ``` 318 | 319 | #### `KEYSPACES` 320 | 321 | ##### Description 322 | 323 | Returns the list of keyspaces. 324 | 325 | ##### Return Type 326 | 327 | The return type is an array of strings. 328 | 329 | ##### Examples 330 | 331 | ```shell 332 | KEYSPACES 333 | ``` 334 | -------------------------------------------------------------------------------- /src/db.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | command::{Command, Count, Create, Del, Drop, Get, Set, Ttl}, 3 | connection::ConnectionError, 4 | frame::Frame, 5 | }; 6 | use bytes::Bytes; 7 | use crossbeam::sync::WaitGroup; 8 | use parking_lot::{Mutex, RwLock}; 9 | use std::{ 10 | collections::HashMap, 11 | str::{self, Utf8Error}, 12 | time::Duration, 13 | }; 14 | use std::{ 15 | sync::Arc, 16 | time::{Instant, SystemTime, SystemTimeError, UNIX_EPOCH}, 17 | }; 18 | use thiserror::Error; 19 | use tokio::sync::broadcast; 20 | use tokio::time; 21 | use tracing::{debug, error}; 22 | 23 | static EXPIRING_EVICTOR_SAMPLE_SIZE: u8 = 5; 24 | static MAX_MEMORY_EVICTOR_SAMPLE_SIZE: u8 = 3; 25 | 26 | #[derive(Debug)] 27 | pub struct Value { 28 | data: Bytes, 29 | last_accessed: Instant, 30 | expire_at: Option, 31 | } 32 | 33 | #[derive(Debug, Clone, Copy, PartialEq)] 34 | pub enum Evictor { 35 | Nop, 36 | Random, 37 | Lru, 38 | } 39 | 40 | #[derive(Debug)] 41 | pub struct Keyspace { 42 | store: Arc>>, 43 | expiring: Arc>>, 44 | evictor: Evictor, 45 | wg: WaitGroup, 46 | done: broadcast::Receiver<()>, 47 | drop: broadcast::Sender<()>, 48 | evict: broadcast::Receiver<()>, 49 | } 50 | 51 | #[derive(Debug)] 52 | pub struct Db { 53 | keyspaces: RwLock>, 54 | done: broadcast::Receiver<()>, 55 | wg: WaitGroup, 56 | evict: broadcast::Receiver<()>, 57 | } 58 | 59 | #[derive(Debug, Error)] 60 | pub enum ExecuteCommandError { 61 | #[error(transparent)] 62 | ConnectionError(#[from] ConnectionError), 63 | 64 | #[error("keyspace '{0}' already exists")] 65 | KeyspaceExists(String), 66 | 67 | #[error("keyspace '{0}' does not exist")] 68 | KeyspaceDoesNotExist(String), 69 | 70 | #[error(transparent)] 71 | Utf8Error(#[from] Utf8Error), 72 | 73 | #[error(transparent)] 74 | SystemTimeError(#[from] SystemTimeError), 75 | } 76 | 77 | impl Db { 78 | pub fn new( 79 | done: broadcast::Receiver<()>, 80 | wg: WaitGroup, 81 | evict: broadcast::Receiver<()>, 82 | ) -> Self { 83 | Db { 84 | keyspaces: RwLock::new(HashMap::new()), 85 | done, 86 | wg, 87 | evict, 88 | } 89 | } 90 | 91 | pub async fn execute(&self, command: Command) -> Result { 92 | match command { 93 | Command::Create(cmd) => self.exec_create(&cmd).await, 94 | Command::Drop(cmd) => self.exec_drop(&cmd), 95 | Command::Keyspaces => self.exec_keyspaces(), 96 | Command::Set(cmd) => self.exec_set(&cmd), 97 | Command::Ping => Ok(Frame::String(Bytes::from_static(b"PONG"))), 98 | Command::Get(cmd) => self.exec_get(&cmd), 99 | Command::Del(cmd) => self.exec_del(&cmd), 100 | Command::Count(cmd) => self.exec_count(&cmd), 101 | Command::Ttl(cmd) => self.exec_ttl(&cmd), 102 | } 103 | } 104 | 105 | async fn exec_create(&self, cmd: &Create) -> Result { 106 | let mut handle = self.keyspaces.write(); 107 | if handle.contains_key(&cmd.keyspace()) { 108 | if cmd.if_not_exists() { 109 | return Ok(Frame::Boolean(false)); 110 | } else { 111 | return Err(ExecuteCommandError::KeyspaceExists( 112 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 113 | )); 114 | } 115 | } 116 | 117 | let ks = Keyspace::new( 118 | self.done.resubscribe(), 119 | self.wg.clone(), 120 | cmd.evictor(), 121 | self.evict.resubscribe(), 122 | ); 123 | 124 | ks.start_expiring_evictor(); 125 | ks.start_max_memory_evictor(); 126 | 127 | handle.insert(cmd.keyspace(), ks); 128 | 129 | Ok(Frame::Boolean(true)) 130 | } 131 | 132 | fn exec_drop(&self, cmd: &Drop) -> Result { 133 | let mut handle = self.keyspaces.write(); 134 | if !handle.contains_key(&cmd.keyspace()) { 135 | if cmd.if_exists() { 136 | return Ok(Frame::Boolean(false)); 137 | } else { 138 | return Err(ExecuteCommandError::KeyspaceDoesNotExist( 139 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 140 | )); 141 | } 142 | } 143 | handle.remove(&cmd.keyspace()); 144 | Ok(Frame::Boolean(true)) 145 | } 146 | 147 | fn exec_keyspaces(&self) -> Result { 148 | let handle = self.keyspaces.read(); 149 | let mut keyspaces = Vec::with_capacity(handle.keys().count()); 150 | for key in handle.keys() { 151 | if let Some(keyspace) = handle.get(key) { 152 | let mut map = Vec::with_capacity(2); 153 | let name = Frame::String(key.clone()); 154 | let evictor = Frame::String(Bytes::copy_from_slice(keyspace.evictor().as_bytes())); 155 | map.push(Frame::String(Bytes::from_static(b"name"))); 156 | map.push(name); 157 | map.push(Frame::String(Bytes::from_static(b"evictor"))); 158 | map.push(evictor); 159 | keyspaces.push(Frame::Map(map)) 160 | } else { 161 | continue; 162 | } 163 | } 164 | Ok(Frame::Array(keyspaces)) 165 | } 166 | 167 | fn exec_set(&self, cmd: &Set) -> Result { 168 | let handle = self.keyspaces.read(); 169 | let keyspace = handle.get(&cmd.keyspace()); 170 | if let Some(ks) = keyspace { 171 | if cmd.if_exists() || cmd.if_not_exists() { 172 | if cmd.if_exists() { 173 | return ks.set_if_exists(cmd.key(), cmd.value(), cmd.expire_at()); 174 | } else { 175 | return ks.set_if_not_exists(cmd.key(), cmd.value(), cmd.expire_at()); 176 | } 177 | } else { 178 | return ks.set(cmd.key(), cmd.value(), cmd.expire_at()); 179 | } 180 | } 181 | 182 | Err(ExecuteCommandError::KeyspaceDoesNotExist( 183 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 184 | )) 185 | } 186 | 187 | fn exec_get(&self, cmd: &Get) -> Result { 188 | let handle = self.keyspaces.read(); 189 | let keyspace = handle.get(&cmd.keyspace()); 190 | if let Some(ks) = keyspace { 191 | return ks.get(cmd.key()); 192 | } 193 | 194 | Err(ExecuteCommandError::KeyspaceDoesNotExist( 195 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 196 | )) 197 | } 198 | 199 | fn exec_del(&self, cmd: &Del) -> Result { 200 | let handle = self.keyspaces.read(); 201 | let keyspace = handle.get(&cmd.keyspace()); 202 | if let Some(ks) = keyspace { 203 | return ks.del(cmd.key()); 204 | } 205 | 206 | Err(ExecuteCommandError::KeyspaceDoesNotExist( 207 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 208 | )) 209 | } 210 | 211 | fn exec_count(&self, cmd: &Count) -> Result { 212 | let handle = self.keyspaces.read(); 213 | let keyspace = handle.get(&cmd.keyspace()); 214 | if let Some(ks) = keyspace { 215 | return ks.count(); 216 | } 217 | 218 | Err(ExecuteCommandError::KeyspaceDoesNotExist( 219 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 220 | )) 221 | } 222 | 223 | fn exec_ttl(&self, cmd: &Ttl) -> Result { 224 | let handle = self.keyspaces.read(); 225 | let keyspace = handle.get(&cmd.keyspace()); 226 | if let Some(ks) = keyspace { 227 | return ks.ttl(cmd.key()); 228 | } 229 | 230 | Err(ExecuteCommandError::KeyspaceDoesNotExist( 231 | str::from_utf8(&cmd.keyspace()[..])?.to_string(), 232 | )) 233 | } 234 | } 235 | 236 | impl Keyspace { 237 | pub fn new( 238 | done: broadcast::Receiver<()>, 239 | wg: WaitGroup, 240 | evictor: Evictor, 241 | evict: broadcast::Receiver<()>, 242 | ) -> Self { 243 | let (drop_tx, _) = broadcast::channel(1); 244 | Keyspace { 245 | store: Arc::new(Mutex::new(HashMap::new())), 246 | expiring: Arc::new(Mutex::new(HashMap::new())), 247 | evictor, 248 | done, 249 | wg, 250 | drop: drop_tx, 251 | evict, 252 | } 253 | } 254 | pub fn set_if_not_exists( 255 | &self, 256 | key: Bytes, 257 | value: Bytes, 258 | expire_at: Option, 259 | ) -> Result { 260 | let handle = self.store.lock(); 261 | let val = handle.get(&key); 262 | if val.is_some() { 263 | return Ok(Frame::Boolean(false)); 264 | } 265 | drop(handle); 266 | self.set(key, value, expire_at) 267 | } 268 | 269 | pub fn set_if_exists( 270 | &self, 271 | key: Bytes, 272 | value: Bytes, 273 | expire_at: Option, 274 | ) -> Result { 275 | let handle = self.store.lock(); 276 | let val = handle.get(&key); 277 | if val.is_none() { 278 | return Ok(Frame::Boolean(false)); 279 | } 280 | drop(handle); 281 | self.set(key, value, expire_at) 282 | } 283 | 284 | pub fn set( 285 | &self, 286 | key: Bytes, 287 | value: Bytes, 288 | expire_at: Option, 289 | ) -> Result { 290 | let mut handle = self.store.lock(); 291 | let value = Value::new(value, expire_at); 292 | handle.insert(key.clone(), value); 293 | if let Some(expiry) = expire_at { 294 | let mut expring_handle = self.expiring.lock(); 295 | expring_handle.insert(key, expiry); 296 | } 297 | Ok(Frame::Boolean(true)) 298 | } 299 | 300 | pub fn get(&self, key: Bytes) -> Result { 301 | let mut handle = self.store.lock(); 302 | if let Some(val) = handle.get_mut(&key) { 303 | val.touch(); 304 | if let Some(expiry) = val.expire_at() { 305 | let current_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); 306 | if expiry < current_time { 307 | handle.remove(&key); 308 | return Ok(Frame::Null); 309 | } 310 | } 311 | return Ok(Frame::String(val.data())); 312 | } 313 | Ok(Frame::Null) 314 | } 315 | 316 | pub fn del(&self, key: Bytes) -> Result { 317 | let mut handle = self.store.lock(); 318 | let result = handle.remove(&key); 319 | Ok(Frame::Boolean(result.is_some())) 320 | } 321 | 322 | pub fn count(&self) -> Result { 323 | let handle = self.store.lock(); 324 | let count = handle.iter().count(); 325 | Ok(Frame::Integer(count as i64)) 326 | } 327 | 328 | pub fn ttl(&self, key: Bytes) -> Result { 329 | let mut handle = self.store.lock(); 330 | if let Some(val) = handle.get_mut(&key) { 331 | val.touch(); 332 | if let Some(expiry) = val.expire_at() { 333 | let current_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); 334 | if expiry <= current_time { 335 | handle.remove(&key); 336 | return Ok(Frame::Null); 337 | } else { 338 | return Ok(Frame::Integer(((expiry - current_time) * 1000) as i64)); 339 | } 340 | } 341 | return Ok(Frame::Null); 342 | } 343 | Ok(Frame::Null) 344 | } 345 | fn start_expiring_evictor(&self) { 346 | let mut done = self.done.resubscribe(); 347 | let wg = self.wg.clone(); 348 | let expiring = self.expiring.clone(); 349 | let store = self.store.clone(); 350 | let mut drop_rx = self.drop.subscribe(); 351 | tokio::spawn(async move { 352 | debug!("expiring evictor started"); 353 | loop { 354 | tokio::select! { 355 | _ = done.recv() => { 356 | drop(wg); 357 | debug!("shutting down expiring evictor, shutdown signal received"); 358 | break; 359 | } 360 | _ = drop_rx.recv() => { 361 | drop(wg); 362 | debug!("shutting down expiring evictor, keyspace is dropped"); 363 | break; 364 | } 365 | _ = time::sleep(Duration::from_millis(500)) => { 366 | let mut expring_handle = expiring.lock(); 367 | let mut store_handle = store.lock(); 368 | let mut expired_keys = Vec::with_capacity(5); 369 | 370 | for (idx, (key, expiry)) in expring_handle.iter().enumerate() { 371 | if idx >= EXPIRING_EVICTOR_SAMPLE_SIZE as usize { 372 | break; 373 | } 374 | 375 | let current_time = match SystemTime::now().duration_since(UNIX_EPOCH) { 376 | Ok(time) => time.as_secs(), 377 | Err(e) => { 378 | error!("{}", e); 379 | break; 380 | } 381 | }; 382 | if *expiry <= current_time { 383 | expired_keys.push(key.clone()); 384 | store_handle.remove(key); 385 | } 386 | } 387 | 388 | for key in expired_keys { 389 | expring_handle.remove(&key); 390 | } 391 | } 392 | } 393 | } 394 | }); 395 | } 396 | 397 | fn start_max_memory_evictor(&self) { 398 | if self.evictor == Evictor::Nop { 399 | return; 400 | } 401 | let mut done = self.done.resubscribe(); 402 | let mut drop_rx = self.drop.subscribe(); 403 | let mut evict_rx = self.evict.resubscribe(); 404 | let wg = self.wg.clone(); 405 | let store = self.store.clone(); 406 | let evictor = self.evictor; 407 | tokio::spawn(async move { 408 | debug!("max memory evictor started"); 409 | loop { 410 | tokio::select! { 411 | _ = done.recv() => { 412 | drop(wg); 413 | debug!("shutting down max memory evictor, shutdown signal received"); 414 | break; 415 | } 416 | _ = drop_rx.recv() => { 417 | drop(wg); 418 | debug!("shutting down max memory evictor, keyspace is dropped"); 419 | break; 420 | } 421 | _ = evict_rx.recv() => { 422 | match evictor { 423 | Evictor::Lru => { 424 | let mut handle = store.lock(); 425 | let mut lru = Instant::now(); 426 | let mut to_evict: Option = None; 427 | for (idx, (key, value)) in handle.iter().enumerate() { 428 | if idx >= MAX_MEMORY_EVICTOR_SAMPLE_SIZE as usize { 429 | break; 430 | } 431 | 432 | let last_accessed = value.last_accessed(); 433 | 434 | if last_accessed < lru { 435 | lru = last_accessed; 436 | to_evict = Some(key.clone()); 437 | } 438 | } 439 | 440 | if let Some(key) = to_evict { 441 | debug!("key '{:?}' evicted using lru policy", key); 442 | handle.remove(&key); 443 | } 444 | }, 445 | Evictor::Random => { 446 | let mut handle = store.lock(); 447 | let mut to_evict: Option = None; 448 | for (idx, key) in handle.keys().enumerate() { 449 | if idx >= MAX_MEMORY_EVICTOR_SAMPLE_SIZE as usize { 450 | break; 451 | } 452 | to_evict = Some(key.clone()); 453 | } 454 | 455 | if let Some(key) = to_evict { 456 | debug!("key '{:?}' evicted using random policy", key); 457 | handle.remove(&key); 458 | } 459 | }, 460 | _ => unreachable!(), 461 | } 462 | } 463 | } 464 | } 465 | }); 466 | } 467 | 468 | pub fn evictor(&self) -> Evictor { 469 | self.evictor 470 | } 471 | } 472 | 473 | impl Value { 474 | pub fn new(data: Bytes, expire_at: Option) -> Self { 475 | Value { 476 | data, 477 | last_accessed: Instant::now(), 478 | expire_at, 479 | } 480 | } 481 | 482 | pub fn touch(&mut self) { 483 | self.last_accessed = Instant::now(); 484 | } 485 | 486 | pub fn data(&self) -> Bytes { 487 | self.data.clone() 488 | } 489 | 490 | pub fn expire_at(&self) -> Option { 491 | self.expire_at 492 | } 493 | 494 | pub fn last_accessed(&self) -> Instant { 495 | self.last_accessed 496 | } 497 | } 498 | 499 | impl Evictor { 500 | pub fn as_bytes(&self) -> &[u8] { 501 | match self { 502 | Evictor::Lru => b"LRU", 503 | Evictor::Nop => b"NOP", 504 | Evictor::Random => b"RANDOM", 505 | } 506 | } 507 | } 508 | -------------------------------------------------------------------------------- /src/command/mod.rs: -------------------------------------------------------------------------------- 1 | use crate::db::Evictor; 2 | use crate::frame::Frame; 3 | use bytes::Bytes; 4 | use std::iter::Peekable; 5 | use std::ops::Add; 6 | use std::str::{self, Utf8Error}; 7 | use std::time::{Duration, SystemTime, SystemTimeError, UNIX_EPOCH}; 8 | use std::vec::IntoIter; 9 | use thiserror::Error; 10 | 11 | #[cfg(test)] 12 | mod test; 13 | 14 | #[derive(Debug)] 15 | struct Parser { 16 | tokens: Peekable>, 17 | } 18 | 19 | #[derive(Debug, PartialEq)] 20 | pub struct Create { 21 | keyspace: Bytes, 22 | evictor: Evictor, 23 | if_not_exists: bool, 24 | } 25 | 26 | #[derive(Debug, PartialEq)] 27 | pub struct Set { 28 | keyspace: Bytes, 29 | key: Bytes, 30 | value: Bytes, 31 | expire_at: Option, 32 | if_not_exists: bool, 33 | if_exists: bool, 34 | } 35 | 36 | #[derive(Debug, PartialEq)] 37 | pub struct Get { 38 | keyspace: Bytes, 39 | key: Bytes, 40 | } 41 | 42 | #[derive(Debug, PartialEq)] 43 | pub struct Del { 44 | keyspace: Bytes, 45 | key: Bytes, 46 | } 47 | 48 | #[derive(Debug, PartialEq)] 49 | pub struct Drop { 50 | keyspace: Bytes, 51 | if_exists: bool, 52 | } 53 | 54 | #[derive(Debug, PartialEq)] 55 | pub struct Count { 56 | keyspace: Bytes, 57 | } 58 | 59 | #[derive(Debug, PartialEq)] 60 | pub struct Ttl { 61 | keyspace: Bytes, 62 | key: Bytes, 63 | } 64 | 65 | #[derive(Debug, PartialEq)] 66 | pub enum Command { 67 | Create(Create), 68 | Set(Set), 69 | Get(Get), 70 | Del(Del), 71 | Drop(Drop), 72 | Count(Count), 73 | Ttl(Ttl), 74 | Ping, 75 | Keyspaces, 76 | } 77 | 78 | #[derive(Debug, Error)] 79 | pub enum ParseCommandError { 80 | #[error("invalid command format")] 81 | InvalidFormat, 82 | 83 | #[error("wrong number of arguments for '{0}' command")] 84 | WrongArgCount(String), 85 | 86 | #[error(transparent)] 87 | Utf8Error(#[from] Utf8Error), 88 | 89 | #[error("invalid argument '{0}' for '{1}' command")] 90 | InvalidArg(String, String), 91 | 92 | #[error("invalid value '{0}' for argument '{1}' for '{2}' command")] 93 | InvalidArgValue(String, String, String), 94 | 95 | #[error(transparent)] 96 | SystemTimeError(#[from] SystemTimeError), 97 | 98 | #[error("unknown command '{0}'")] 99 | UnknownCommand(String), 100 | } 101 | 102 | impl Parser { 103 | pub fn new(frame: Frame) -> Result { 104 | match frame { 105 | Frame::Array(tokens) => Ok(Parser { 106 | tokens: tokens.into_iter().peekable(), 107 | }), 108 | _ => Err(ParseCommandError::InvalidFormat), 109 | } 110 | } 111 | 112 | pub fn next(&mut self) -> Option { 113 | self.tokens.next() 114 | } 115 | 116 | pub fn has_remaining(&mut self) -> bool { 117 | self.tokens.peek().is_some() 118 | } 119 | 120 | pub fn next_as_string(&mut self) -> Result, ParseCommandError> { 121 | let frame = match self.next() { 122 | Some(frame) => frame, 123 | None => return Ok(None), 124 | }; 125 | 126 | match frame { 127 | Frame::String(data) => Ok(Some(str::from_utf8(&data[..])?.to_string())), 128 | _ => Err(ParseCommandError::InvalidFormat), 129 | } 130 | } 131 | 132 | pub fn next_as_bytes(&mut self) -> Result, ParseCommandError> { 133 | let frame = match self.next() { 134 | Some(frame) => frame, 135 | None => return Ok(None), 136 | }; 137 | 138 | match frame { 139 | Frame::String(data) => Ok(Some(data)), 140 | _ => Err(ParseCommandError::InvalidFormat), 141 | } 142 | } 143 | } 144 | 145 | impl Create { 146 | fn parse(parser: &mut Parser) -> Result { 147 | let keyspace = parser 148 | .next_as_bytes()? 149 | .ok_or_else(|| ParseCommandError::WrongArgCount("create".to_string()))?; 150 | 151 | let mut command = Create { 152 | keyspace, 153 | evictor: Evictor::Nop, 154 | if_not_exists: false, 155 | }; 156 | 157 | if !parser.has_remaining() { 158 | return Ok(command); 159 | }; 160 | 161 | while parser.has_remaining() { 162 | let token = parser 163 | .next_as_string()? 164 | .ok_or_else(|| ParseCommandError::WrongArgCount("create".to_string()))? 165 | .to_lowercase(); 166 | 167 | if matches!(token.as_str(), "evictor") { 168 | let value = parser 169 | .next_as_string()? 170 | .ok_or_else(|| ParseCommandError::WrongArgCount("create".to_string()))? 171 | .to_lowercase(); 172 | match value.as_str() { 173 | "nop" => command.evictor = Evictor::Nop, 174 | "random" => command.evictor = Evictor::Random, 175 | "lru" => command.evictor = Evictor::Lru, 176 | _ => { 177 | return Err(ParseCommandError::InvalidArgValue( 178 | value, 179 | token, 180 | "create".to_string(), 181 | )) 182 | } 183 | } 184 | } else if matches!(token.as_str(), "if") { 185 | let not_token = parser 186 | .next_as_string()? 187 | .ok_or_else(|| ParseCommandError::WrongArgCount("create".to_string()))? 188 | .to_lowercase(); 189 | let exists_token = parser 190 | .next_as_string()? 191 | .ok_or_else(|| ParseCommandError::WrongArgCount("create".to_string()))? 192 | .to_lowercase(); 193 | 194 | if !matches!(not_token.as_str(), "not") { 195 | return Err(ParseCommandError::InvalidArg( 196 | not_token, 197 | "create".to_string(), 198 | )); 199 | } 200 | 201 | if !matches!(exists_token.as_str(), "exists") { 202 | return Err(ParseCommandError::InvalidArg( 203 | exists_token, 204 | "create".to_string(), 205 | )); 206 | } 207 | 208 | if !command.if_not_exists { 209 | command.if_not_exists = true 210 | } else { 211 | return Err(ParseCommandError::InvalidFormat); 212 | } 213 | } else { 214 | return Err(ParseCommandError::InvalidArg(token, "create".to_string())); 215 | } 216 | } 217 | 218 | Ok(command) 219 | } 220 | 221 | pub fn keyspace(&self) -> Bytes { 222 | self.keyspace.clone() 223 | } 224 | pub fn evictor(&self) -> Evictor { 225 | self.evictor 226 | } 227 | pub fn if_not_exists(&self) -> bool { 228 | self.if_not_exists 229 | } 230 | } 231 | 232 | impl Set { 233 | fn parse(parser: &mut Parser) -> Result { 234 | let keyspace = parser 235 | .next_as_bytes()? 236 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))?; 237 | 238 | let key = parser 239 | .next_as_bytes()? 240 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))?; 241 | 242 | let value = parser 243 | .next_as_bytes()? 244 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))?; 245 | 246 | let mut command = Set { 247 | keyspace, 248 | key, 249 | value, 250 | expire_at: None, 251 | if_not_exists: false, 252 | if_exists: false, 253 | }; 254 | 255 | if !parser.has_remaining() { 256 | return Ok(command); 257 | } 258 | 259 | while parser.has_remaining() { 260 | let token = parser 261 | .next_as_string()? 262 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))? 263 | .to_lowercase(); 264 | 265 | if matches!(token.as_str(), "expire") { 266 | let at_or_after_token = parser 267 | .next_as_string()? 268 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))? 269 | .to_lowercase(); 270 | if matches!(at_or_after_token.as_str(), "at") { 271 | let value = parser 272 | .next_as_string()? 273 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))?; 274 | let timestamp = value.parse::().map_err(|_| { 275 | ParseCommandError::InvalidArgValue(value, token, "set".to_string()) 276 | })?; 277 | match command.expire_at { 278 | Some(_) => return Err(ParseCommandError::InvalidFormat), 279 | None => command.expire_at = Some(timestamp), 280 | } 281 | } else if matches!(at_or_after_token.as_str(), "after") { 282 | let value = parser 283 | .next_as_string()? 284 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))?; 285 | 286 | let millis = value.parse::().map_err(|_| { 287 | ParseCommandError::InvalidArgValue(value, token, "set".to_string()) 288 | })?; 289 | 290 | let timestamp = SystemTime::now() 291 | .add(Duration::from_millis(millis)) 292 | .duration_since(UNIX_EPOCH)? 293 | .as_secs(); 294 | 295 | match command.expire_at { 296 | Some(_) => return Err(ParseCommandError::InvalidFormat), 297 | None => command.expire_at = Some(timestamp), 298 | } 299 | } else { 300 | return Err(ParseCommandError::InvalidArg( 301 | at_or_after_token, 302 | "set".to_string(), 303 | )); 304 | } 305 | } else if matches!(token.as_str(), "if") { 306 | let not_or_exists_token = parser 307 | .next_as_string()? 308 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))? 309 | .to_lowercase(); 310 | 311 | if matches!(not_or_exists_token.as_str(), "not") { 312 | let exists_token = parser 313 | .next_as_string()? 314 | .ok_or_else(|| ParseCommandError::WrongArgCount("set".to_string()))? 315 | .to_lowercase(); 316 | if matches!(exists_token.as_str(), "exists") { 317 | if !command.if_not_exists && !command.if_exists { 318 | command.if_not_exists = true 319 | } else { 320 | return Err(ParseCommandError::InvalidFormat); 321 | } 322 | } else { 323 | return Err(ParseCommandError::InvalidArg( 324 | exists_token, 325 | "set".to_string(), 326 | )); 327 | } 328 | } else if matches!(not_or_exists_token.as_str(), "exists") { 329 | if !command.if_not_exists && !command.if_exists { 330 | command.if_exists = true 331 | } else { 332 | return Err(ParseCommandError::InvalidFormat); 333 | } 334 | } else { 335 | return Err(ParseCommandError::InvalidArg( 336 | not_or_exists_token, 337 | "set".to_string(), 338 | )); 339 | } 340 | } else { 341 | return Err(ParseCommandError::InvalidArg(token, "set".to_string())); 342 | } 343 | } 344 | 345 | Ok(command) 346 | } 347 | 348 | pub fn expire_at(&self) -> Option { 349 | self.expire_at 350 | } 351 | 352 | pub fn if_exists(&self) -> bool { 353 | self.if_exists 354 | } 355 | 356 | pub fn if_not_exists(&self) -> bool { 357 | self.if_not_exists 358 | } 359 | 360 | pub fn key(&self) -> Bytes { 361 | self.key.clone() 362 | } 363 | 364 | pub fn value(&self) -> Bytes { 365 | self.value.clone() 366 | } 367 | 368 | pub fn keyspace(&self) -> Bytes { 369 | self.keyspace.clone() 370 | } 371 | } 372 | 373 | impl Get { 374 | fn parse(parser: &mut Parser) -> Result { 375 | let keyspace = parser 376 | .next_as_bytes()? 377 | .ok_or_else(|| ParseCommandError::WrongArgCount("get".to_string()))?; 378 | 379 | let key = parser 380 | .next_as_bytes()? 381 | .ok_or_else(|| ParseCommandError::WrongArgCount("get".to_string()))?; 382 | 383 | let command = Get { keyspace, key }; 384 | 385 | if parser.has_remaining() { 386 | return Err(ParseCommandError::WrongArgCount("get".to_string())); 387 | } 388 | 389 | Ok(command) 390 | } 391 | 392 | pub fn keyspace(&self) -> Bytes { 393 | self.keyspace.clone() 394 | } 395 | 396 | pub fn key(&self) -> Bytes { 397 | self.key.clone() 398 | } 399 | } 400 | 401 | impl Del { 402 | fn parse(parser: &mut Parser) -> Result { 403 | let keyspace = parser 404 | .next_as_bytes()? 405 | .ok_or_else(|| ParseCommandError::WrongArgCount("del".to_string()))?; 406 | 407 | let key = parser 408 | .next_as_bytes()? 409 | .ok_or_else(|| ParseCommandError::WrongArgCount("del".to_string()))?; 410 | 411 | let command = Del { keyspace, key }; 412 | 413 | if parser.has_remaining() { 414 | return Err(ParseCommandError::WrongArgCount("del".to_string())); 415 | } 416 | 417 | Ok(command) 418 | } 419 | 420 | pub fn keyspace(&self) -> Bytes { 421 | self.keyspace.clone() 422 | } 423 | 424 | pub fn key(&self) -> Bytes { 425 | self.key.clone() 426 | } 427 | } 428 | 429 | impl Drop { 430 | fn parse(parser: &mut Parser) -> Result { 431 | let keyspace = parser 432 | .next_as_bytes()? 433 | .ok_or_else(|| ParseCommandError::WrongArgCount("drop".to_string()))?; 434 | 435 | let mut command = Drop { 436 | keyspace, 437 | if_exists: false, 438 | }; 439 | 440 | if !parser.has_remaining() { 441 | return Ok(command); 442 | } 443 | 444 | while parser.has_remaining() { 445 | let key = parser 446 | .next_as_string()? 447 | .ok_or_else(|| ParseCommandError::WrongArgCount("drop".to_string()))? 448 | .to_lowercase(); 449 | 450 | if matches!(key.as_str(), "if") { 451 | let exists_token = parser 452 | .next_as_string()? 453 | .ok_or_else(|| ParseCommandError::WrongArgCount("drop".to_string()))? 454 | .to_lowercase(); 455 | if matches!(exists_token.as_str(), "exists") { 456 | if !command.if_exists { 457 | command.if_exists = true 458 | } else { 459 | return Err(ParseCommandError::InvalidFormat); 460 | } 461 | } else { 462 | return Err(ParseCommandError::InvalidArg( 463 | exists_token, 464 | "set".to_string(), 465 | )); 466 | } 467 | } else { 468 | return Err(ParseCommandError::InvalidArg(key, "drop".to_string())); 469 | } 470 | } 471 | 472 | Ok(command) 473 | } 474 | 475 | pub fn keyspace(&self) -> Bytes { 476 | self.keyspace.clone() 477 | } 478 | 479 | pub fn if_exists(&self) -> bool { 480 | self.if_exists 481 | } 482 | } 483 | 484 | impl Count { 485 | fn parse(parser: &mut Parser) -> Result { 486 | let keyspace = parser 487 | .next_as_bytes()? 488 | .ok_or_else(|| ParseCommandError::WrongArgCount("count".to_string()))?; 489 | 490 | let command = Count { keyspace }; 491 | 492 | if parser.has_remaining() { 493 | return Err(ParseCommandError::WrongArgCount("count".to_string())); 494 | } 495 | 496 | Ok(command) 497 | } 498 | 499 | pub fn keyspace(&self) -> Bytes { 500 | self.keyspace.clone() 501 | } 502 | } 503 | 504 | impl Ttl { 505 | fn parse(parser: &mut Parser) -> Result { 506 | let keyspace = parser 507 | .next_as_bytes()? 508 | .ok_or_else(|| ParseCommandError::WrongArgCount("ttl".to_string()))?; 509 | 510 | let key = parser 511 | .next_as_bytes()? 512 | .ok_or_else(|| ParseCommandError::WrongArgCount("ttl".to_string()))?; 513 | 514 | let command = Ttl { keyspace, key }; 515 | 516 | if parser.has_remaining() { 517 | return Err(ParseCommandError::WrongArgCount("ttl".to_string())); 518 | } 519 | 520 | Ok(command) 521 | } 522 | 523 | pub fn keyspace(&self) -> Bytes { 524 | self.keyspace.clone() 525 | } 526 | 527 | pub fn key(&self) -> Bytes { 528 | self.key.clone() 529 | } 530 | } 531 | 532 | pub fn parse(frame: Frame) -> Result { 533 | let mut parser = Parser::new(frame)?; 534 | let command = match parser.next().ok_or(ParseCommandError::InvalidFormat)? { 535 | Frame::String(data) => str::from_utf8(&data[..])?.to_lowercase(), 536 | _ => return Err(ParseCommandError::InvalidFormat), 537 | }; 538 | 539 | match command.as_str() { 540 | "create" => Ok(Command::Create(Create::parse(&mut parser)?)), 541 | "set" => Ok(Command::Set(Set::parse(&mut parser)?)), 542 | "get" => Ok(Command::Get(Get::parse(&mut parser)?)), 543 | "del" => Ok(Command::Del(Del::parse(&mut parser)?)), 544 | "drop" => Ok(Command::Drop(Drop::parse(&mut parser)?)), 545 | "count" => Ok(Command::Count(Count::parse(&mut parser)?)), 546 | "ttl" => Ok(Command::Ttl(Ttl::parse(&mut parser)?)), 547 | "ping" => Ok(Command::Ping), 548 | "keyspaces" => Ok(Command::Keyspaces), 549 | _ => Err(ParseCommandError::UnknownCommand(command)), 550 | } 551 | } 552 | -------------------------------------------------------------------------------- /src/connection.rs: -------------------------------------------------------------------------------- 1 | use crate::frame::{ 2 | self, Frame, ParseFrameError, ARRAY_IDENT, BOOLEAN_IDENT, DOUBLE_IDENT, ERROR_IDENT, 3 | INTEGER_IDENT, MAP_IDENT, STRING_IDENT, 4 | }; 5 | use async_recursion::async_recursion; 6 | use bytes::{Buf, Bytes, BytesMut}; 7 | use std::io::{self, Cursor}; 8 | use thiserror::Error; 9 | use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; 10 | 11 | #[derive(Debug)] 12 | pub struct Connection 13 | where 14 | T: AsyncRead + AsyncWrite + Unpin + Send, 15 | { 16 | stream: T, 17 | buf: BytesMut, 18 | } 19 | 20 | #[derive(Debug, Error)] 21 | pub enum ConnectionError { 22 | #[error(transparent)] 23 | Io(#[from] io::Error), 24 | 25 | #[error("connection reset by peer")] 26 | Reset, 27 | 28 | #[error(transparent)] 29 | Frame(#[from] ParseFrameError), 30 | 31 | #[error("malformed frame received for write")] 32 | MalformedFrameForWrite, 33 | } 34 | 35 | impl Connection 36 | where 37 | T: AsyncRead + AsyncWrite + Unpin + Send, 38 | { 39 | pub fn new(stream: T, buf_size: usize) -> Self { 40 | Connection { 41 | stream, 42 | buf: BytesMut::with_capacity(buf_size), 43 | } 44 | } 45 | 46 | pub async fn read_frame(&mut self) -> Result, ConnectionError> { 47 | loop { 48 | if let Some(frame) = self.parse_frame()? { 49 | return Ok(Some(frame)); 50 | } 51 | 52 | if self.stream.read_buf(&mut self.buf).await? == 0 { 53 | if self.buf.is_empty() { 54 | return Ok(None); 55 | } else { 56 | return Err(ConnectionError::Reset); 57 | } 58 | } 59 | } 60 | } 61 | 62 | fn parse_frame(&mut self) -> Result, ConnectionError> { 63 | let mut cursor = Cursor::new(&self.buf[..]); 64 | match frame::parse(&mut cursor) { 65 | Ok(frame) => { 66 | self.buf.advance(cursor.position() as usize); 67 | Ok(Some(frame)) 68 | } 69 | Err(ParseFrameError::Incomplete) => Ok(None), 70 | Err(e) => Err(e.into()), 71 | } 72 | } 73 | 74 | #[async_recursion] 75 | pub async fn write_frame(&mut self, frame: &Frame) -> Result<(), ConnectionError> { 76 | match frame { 77 | Frame::String(data) => { 78 | let len = data.len(); 79 | self.stream.write_u8(STRING_IDENT).await?; 80 | self.stream 81 | .write_all(format!("{}\r\n", len).as_bytes()) 82 | .await?; 83 | self.stream.write_all(data).await?; 84 | self.stream.write_all(b"\r\n").await?; 85 | } 86 | Frame::Integer(data) => { 87 | self.stream.write_u8(INTEGER_IDENT).await?; 88 | self.stream 89 | .write_all(format!("{}\r\n", data).as_bytes()) 90 | .await?; 91 | } 92 | Frame::Boolean(data) => { 93 | self.stream.write_u8(BOOLEAN_IDENT).await?; 94 | if *data { 95 | self.stream 96 | .write_all(format!("{}\r\n", 1).as_bytes()) 97 | .await?; 98 | } else { 99 | self.stream 100 | .write_all(format!("{}\r\n", 0).as_bytes()) 101 | .await?; 102 | } 103 | } 104 | Frame::Null => { 105 | self.stream.write_all(b"-\r\n").await?; 106 | } 107 | Frame::Double(data) => { 108 | self.stream.write_u8(DOUBLE_IDENT).await?; 109 | self.stream 110 | .write_all(format!("{}\r\n", data).as_bytes()) 111 | .await?; 112 | } 113 | Frame::Error(data) => { 114 | let len = data.len(); 115 | self.stream.write_u8(ERROR_IDENT).await?; 116 | self.stream 117 | .write_all(format!("{}\r\n", len).as_bytes()) 118 | .await?; 119 | self.stream.write_all(data).await?; 120 | self.stream.write_all(b"\r\n").await?; 121 | } 122 | Frame::Array(array) => { 123 | self.stream.write_u8(ARRAY_IDENT).await?; 124 | self.stream 125 | .write_all(format!("{}\r\n", array.len()).as_bytes()) 126 | .await?; 127 | for value in array { 128 | self.write_frame(value).await?; 129 | } 130 | } 131 | Frame::Map(map) => { 132 | if map.len() % 2 != 0 { 133 | return Err(ConnectionError::MalformedFrameForWrite); 134 | } 135 | self.stream.write_u8(MAP_IDENT).await?; 136 | self.stream 137 | .write_all(format!("{}\r\n", map.len() / 2).as_bytes()) 138 | .await?; 139 | for value in map { 140 | self.write_frame(value).await?; 141 | } 142 | } 143 | } 144 | 145 | self.stream.flush().await?; 146 | Ok(()) 147 | } 148 | 149 | pub async fn write_error( 150 | &mut self, 151 | error: impl std::error::Error, 152 | ) -> Result<(), ConnectionError> { 153 | self.write_frame(&Frame::Error(Bytes::from(error.to_string()))) 154 | .await 155 | } 156 | } 157 | 158 | #[cfg(test)] 159 | mod tests { 160 | use super::*; 161 | use bytes::{BufMut, Bytes, BytesMut}; 162 | use std::fs; 163 | use std::path::{Path, PathBuf}; 164 | use tokio_test::io::Builder; 165 | 166 | fn read_file(path: PathBuf) -> Vec { 167 | fs::read(path).unwrap() 168 | } 169 | 170 | fn get_frame_from_file(data: &[u8], ident: u8) -> Bytes { 171 | let mut frame = BytesMut::new(); 172 | frame.put_u8(ident); 173 | frame.put_slice(format!("{}", data.len()).as_bytes()); 174 | frame.put_u8(b'\r'); 175 | frame.put_u8(b'\n'); 176 | frame.put(data); 177 | frame.put_u8(b'\r'); 178 | frame.put_u8(b'\n'); 179 | 180 | frame.copy_to_bytes(frame.len()) 181 | } 182 | 183 | #[tokio::test] 184 | async fn write_frame_given_string_writes_string_frame() { 185 | let mock = Builder::new().write(b"$3\r\nfoo\r\n").build(); 186 | let mut connection = Connection::new(mock, 1024); 187 | connection 188 | .write_frame(&Frame::String(Bytes::from("foo"))) 189 | .await 190 | .unwrap(); 191 | } 192 | 193 | #[tokio::test] 194 | async fn write_frame_given_string_with_crlf_writes_string_frame() { 195 | let mock = Builder::new().write(b"$5\r\nfoo\r\n\r\n").build(); 196 | let mut connection = Connection::new(mock, 1024); 197 | connection 198 | .write_frame(&Frame::String(Bytes::from("foo\r\n"))) 199 | .await 200 | .unwrap(); 201 | } 202 | 203 | #[tokio::test] 204 | async fn write_frame_given_string_with_pdf_data_writes_string_frame() { 205 | let file_data = read_file(Path::new("test_data").join("test.pdf")); 206 | let mock = Builder::new() 207 | .write(&get_frame_from_file(file_data.as_slice(), STRING_IDENT)) 208 | .build(); 209 | let mut connection = Connection::new(mock, 1024); 210 | connection 211 | .write_frame(&Frame::String(Bytes::from(file_data))) 212 | .await 213 | .unwrap(); 214 | } 215 | 216 | #[tokio::test] 217 | async fn write_frame_given_string_with_png_data_writes_string_frame() { 218 | let file_data = read_file(Path::new("test_data").join("test.png")); 219 | let mock = Builder::new() 220 | .write(&get_frame_from_file(file_data.as_slice(), STRING_IDENT)) 221 | .build(); 222 | let mut connection = Connection::new(mock, 1024); 223 | connection 224 | .write_frame(&Frame::String(Bytes::from(file_data))) 225 | .await 226 | .unwrap(); 227 | } 228 | 229 | #[tokio::test] 230 | async fn write_frame_given_string_with_jpg_data_writes_string_frame() { 231 | let file_data = read_file(Path::new("test_data").join("test.jpg")); 232 | let mock = Builder::new() 233 | .write(&get_frame_from_file(file_data.as_slice(), STRING_IDENT)) 234 | .build(); 235 | let mut connection = Connection::new(mock, 1024); 236 | connection 237 | .write_frame(&Frame::String(Bytes::from(file_data))) 238 | .await 239 | .unwrap(); 240 | } 241 | 242 | #[tokio::test] 243 | async fn write_frame_given_string_with_html_data_writes_string_frame() { 244 | let file_data = read_file(Path::new("test_data").join("test.html")); 245 | let mock = Builder::new() 246 | .write(&get_frame_from_file(file_data.as_slice(), STRING_IDENT)) 247 | .build(); 248 | let mut connection = Connection::new(mock, 1024); 249 | connection 250 | .write_frame(&Frame::String(Bytes::from(file_data))) 251 | .await 252 | .unwrap(); 253 | } 254 | 255 | #[tokio::test] 256 | async fn write_frame_given_positive_integer_writes_integer_frame() { 257 | let mock = Builder::new().write(b"%100\r\n").build(); 258 | let mut connection = Connection::new(mock, 1024); 259 | connection.write_frame(&Frame::Integer(100)).await.unwrap(); 260 | } 261 | 262 | #[tokio::test] 263 | async fn write_frame_given_negative_integer_writes_integer_frame() { 264 | let mock = Builder::new().write(b"%-100\r\n").build(); 265 | let mut connection = Connection::new(mock, 1024); 266 | connection.write_frame(&Frame::Integer(-100)).await.unwrap(); 267 | } 268 | 269 | #[tokio::test] 270 | async fn write_frame_given_zero_writes_integer_frame() { 271 | let mock = Builder::new().write(b"%0\r\n").build(); 272 | let mut connection = Connection::new(mock, 1024); 273 | connection.write_frame(&Frame::Integer(0)).await.unwrap(); 274 | } 275 | 276 | #[tokio::test] 277 | async fn write_frame_given_false_writes_boolean_frame() { 278 | let mock = Builder::new().write(b"^0\r\n").build(); 279 | let mut connection = Connection::new(mock, 1024); 280 | connection 281 | .write_frame(&Frame::Boolean(false)) 282 | .await 283 | .unwrap(); 284 | } 285 | 286 | #[tokio::test] 287 | async fn write_frame_given_true_writes_boolean_frame() { 288 | let mock = Builder::new().write(b"^1\r\n").build(); 289 | let mut connection = Connection::new(mock, 1024); 290 | connection.write_frame(&Frame::Boolean(true)).await.unwrap(); 291 | } 292 | 293 | #[tokio::test] 294 | async fn write_frame_given_null_writes_null_frame() { 295 | let mock = Builder::new().write(b"-\r\n").build(); 296 | let mut connection = Connection::new(mock, 1024); 297 | connection.write_frame(&Frame::Null).await.unwrap(); 298 | } 299 | 300 | #[tokio::test] 301 | async fn write_frame_given_positive_double_writes_double_frame() { 302 | let mock = Builder::new().write(b".100\r\n").build(); 303 | let mut connection = Connection::new(mock, 1024); 304 | connection 305 | .write_frame(&Frame::Double(100.0000)) 306 | .await 307 | .unwrap(); 308 | } 309 | 310 | #[tokio::test] 311 | async fn write_frame_given_negative_double_writes_double_frame() { 312 | let mock = Builder::new().write(b".-100\r\n").build(); 313 | let mut connection = Connection::new(mock, 1024); 314 | connection 315 | .write_frame(&Frame::Double(-100.0)) 316 | .await 317 | .unwrap(); 318 | } 319 | 320 | #[tokio::test] 321 | async fn write_frame_given_zero_writes_double_frame() { 322 | let mock = Builder::new().write(b".0\r\n").build(); 323 | let mut connection = Connection::new(mock, 1024); 324 | connection.write_frame(&Frame::Double(0.0)).await.unwrap(); 325 | } 326 | 327 | #[tokio::test] 328 | async fn write_frame_given_positive_double_with_decimal_part_writes_double_frame() { 329 | let mock = Builder::new().write(b".100.00001\r\n").build(); 330 | let mut connection = Connection::new(mock, 1024); 331 | connection 332 | .write_frame(&Frame::Double(100.00001)) 333 | .await 334 | .unwrap(); 335 | } 336 | 337 | #[tokio::test] 338 | async fn write_frame_given_negative_double_with_decimal_part_writes_double_frame() { 339 | let mock = Builder::new().write(b".-100.6789\r\n").build(); 340 | let mut connection = Connection::new(mock, 1024); 341 | connection 342 | .write_frame(&Frame::Double(-100.6789)) 343 | .await 344 | .unwrap(); 345 | } 346 | 347 | #[tokio::test] 348 | async fn write_frame_given_zero_with_decimal_part_writes_double_frame() { 349 | let mock = Builder::new().write(b".0.1\r\n").build(); 350 | let mut connection = Connection::new(mock, 1024); 351 | connection.write_frame(&Frame::Double(0.10)).await.unwrap(); 352 | } 353 | 354 | #[tokio::test] 355 | async fn write_frame_given_error_writes_error_frame() { 356 | let mock = Builder::new().write(b"!3\r\nfoo\r\n").build(); 357 | let mut connection = Connection::new(mock, 1024); 358 | connection 359 | .write_frame(&Frame::Error(Bytes::from("foo"))) 360 | .await 361 | .unwrap(); 362 | } 363 | 364 | #[tokio::test] 365 | async fn write_frame_given_error_with_crlf_writes_error_frame() { 366 | let mock = Builder::new().write(b"!5\r\nfoo\r\n\r\n").build(); 367 | let mut connection = Connection::new(mock, 1024); 368 | connection 369 | .write_frame(&Frame::Error(Bytes::from("foo\r\n"))) 370 | .await 371 | .unwrap(); 372 | } 373 | 374 | #[tokio::test] 375 | async fn write_frame_given_error_with_pdf_data_writes_error_frame() { 376 | let file_data = read_file(Path::new("test_data").join("test.pdf")); 377 | let mock = Builder::new() 378 | .write(&get_frame_from_file(file_data.as_slice(), ERROR_IDENT)) 379 | .build(); 380 | let mut connection = Connection::new(mock, 1024); 381 | connection 382 | .write_frame(&Frame::Error(Bytes::from(file_data))) 383 | .await 384 | .unwrap(); 385 | } 386 | 387 | #[tokio::test] 388 | async fn write_frame_given_error_with_png_data_writes_error_frame() { 389 | let file_data = read_file(Path::new("test_data").join("test.png")); 390 | let mock = Builder::new() 391 | .write(&get_frame_from_file(file_data.as_slice(), ERROR_IDENT)) 392 | .build(); 393 | let mut connection = Connection::new(mock, 1024); 394 | connection 395 | .write_frame(&Frame::Error(Bytes::from(file_data))) 396 | .await 397 | .unwrap(); 398 | } 399 | 400 | #[tokio::test] 401 | async fn write_frame_given_error_with_jpg_data_writes_error_frame() { 402 | let file_data = read_file(Path::new("test_data").join("test.jpg")); 403 | let mock = Builder::new() 404 | .write(&get_frame_from_file(file_data.as_slice(), ERROR_IDENT)) 405 | .build(); 406 | let mut connection = Connection::new(mock, 1024); 407 | connection 408 | .write_frame(&Frame::Error(Bytes::from(file_data))) 409 | .await 410 | .unwrap(); 411 | } 412 | 413 | #[tokio::test] 414 | async fn write_frame_given_error_with_html_data_writes_error_frame() { 415 | let file_data = read_file(Path::new("test_data").join("test.html")); 416 | let mock = Builder::new() 417 | .write(&get_frame_from_file(file_data.as_slice(), ERROR_IDENT)) 418 | .build(); 419 | let mut connection = Connection::new(mock, 1024); 420 | connection 421 | .write_frame(&Frame::Error(Bytes::from(file_data))) 422 | .await 423 | .unwrap(); 424 | } 425 | 426 | #[tokio::test] 427 | async fn write_frame_given_empty_array_writes_array_frame() { 428 | let mock = Builder::new().write(b"*0\r\n").build(); 429 | let mut connection = Connection::new(mock, 1024); 430 | connection.write_frame(&Frame::Array(vec![])).await.unwrap(); 431 | } 432 | 433 | #[tokio::test] 434 | async fn write_frame_given_array_with_one_element_writes_array_frame() { 435 | let mock = Builder::new().write(b"*1\r\n$3\r\nfoo\r\n").build(); 436 | let mut connection = Connection::new(mock, 1024); 437 | connection 438 | .write_frame(&Frame::Array(vec![Frame::String(Bytes::from("foo"))])) 439 | .await 440 | .unwrap(); 441 | } 442 | 443 | #[tokio::test] 444 | async fn write_frame_given_nested_arrays_writes_array_frame() { 445 | let mock = Builder::new().write(b"*1\r\n*1\r\n*1\r\n*0\r\n").build(); 446 | let mut connection = Connection::new(mock, 1024); 447 | connection 448 | .write_frame(&Frame::Array(vec![Frame::Array(vec![Frame::Array(vec![ 449 | Frame::Array(vec![]), 450 | ])])])) 451 | .await 452 | .unwrap(); 453 | } 454 | 455 | #[tokio::test] 456 | async fn write_frame_given_array_with_maps_writes_array_frame() { 457 | let mock = Builder::new() 458 | .write(b"*1\r\n#1\r\n$3\r\nfoo\r\n$3\r\nbar\r\n") 459 | .build(); 460 | let mut connection = Connection::new(mock, 1024); 461 | connection 462 | .write_frame(&Frame::Array(vec![Frame::Map(vec![ 463 | Frame::String(Bytes::from("foo")), 464 | Frame::String(Bytes::from("bar")), 465 | ])])) 466 | .await 467 | .unwrap(); 468 | } 469 | 470 | #[tokio::test] 471 | async fn write_frame_given_empty_map_writes_map_frame() { 472 | let mock = Builder::new().write(b"#0\r\n").build(); 473 | let mut connection = Connection::new(mock, 1024); 474 | connection.write_frame(&Frame::Map(vec![])).await.unwrap(); 475 | } 476 | 477 | #[tokio::test] 478 | async fn write_frame_given_malformed_map_returns_malformed_frame_for_write_error() { 479 | let mock = Builder::new().build(); 480 | let mut connection = Connection::new(mock, 1024); 481 | match connection 482 | .write_frame(&Frame::Map(vec![Frame::String(Bytes::from("foo"))])) 483 | .await 484 | { 485 | Err(ConnectionError::MalformedFrameForWrite) => {} 486 | _ => unreachable!(), 487 | } 488 | } 489 | 490 | #[tokio::test] 491 | async fn write_frame_given_map_with_one_element_writes_map_frame() { 492 | let mock = Builder::new() 493 | .write(b"#1\r\n$3\r\nfoo\r\n$3\r\nbar\r\n") 494 | .build(); 495 | let mut connection = Connection::new(mock, 1024); 496 | connection 497 | .write_frame(&Frame::Map(vec![ 498 | Frame::String(Bytes::from("foo")), 499 | Frame::String(Bytes::from("bar")), 500 | ])) 501 | .await 502 | .unwrap(); 503 | } 504 | } 505 | -------------------------------------------------------------------------------- /src/frame.rs: -------------------------------------------------------------------------------- 1 | use atoi::atoi; 2 | use bytes::Buf; 3 | use bytes::Bytes; 4 | use std::io::Cursor; 5 | use std::str; 6 | use thiserror::Error; 7 | 8 | pub const STRING_IDENT: u8 = b'$'; 9 | pub const INTEGER_IDENT: u8 = b'%'; 10 | pub const ARRAY_IDENT: u8 = b'*'; 11 | pub const BOOLEAN_IDENT: u8 = b'^'; 12 | pub const NULL_IDENT: u8 = b'-'; 13 | pub const MAP_IDENT: u8 = b'#'; 14 | pub const DOUBLE_IDENT: u8 = b'.'; 15 | pub const ERROR_IDENT: u8 = b'!'; 16 | 17 | #[derive(Debug, PartialEq)] 18 | pub enum Frame { 19 | String(Bytes), 20 | Integer(i64), 21 | Array(Vec), 22 | Boolean(bool), 23 | Null, 24 | Map(Vec), 25 | Double(f64), 26 | Error(Bytes), 27 | } 28 | 29 | #[derive(Debug, Error, PartialEq)] 30 | pub enum ParseFrameError { 31 | #[error("more data is required to parse the frame")] 32 | Incomplete, 33 | 34 | #[error("invalid frame format")] 35 | InvalidFormat, 36 | } 37 | 38 | pub fn parse(buf: &mut Cursor<&[u8]>) -> Result { 39 | // since our frames are CRLF delimited, we read our frames line by line. 40 | // A line here represents a CRLF delimited section of frame. This is binary 41 | // safe because when reading bytes which might contain binary data, we 42 | // don't look for a delimiter, instead of depend on the prefixed length. 43 | // This is safe to do for other data types like, booleans, null, integers, doubles etc. 44 | let line = get_line(buf)?; 45 | if line.is_empty() { 46 | return Err(ParseFrameError::InvalidFormat); 47 | } 48 | // first byte of a line determines it's type 49 | let frame_type = line[0]; 50 | // rest of the line can contain the length of the data (in case of strings and errors), 51 | // or the whole data in case of other data types 52 | let line = &line[1..]; 53 | match frame_type { 54 | STRING_IDENT => parse_string(buf, line), 55 | INTEGER_IDENT => parse_integer(line), 56 | ARRAY_IDENT => parse_array(buf, line), 57 | BOOLEAN_IDENT => parse_boolean(line), 58 | NULL_IDENT => parse_null(line), 59 | MAP_IDENT => parse_map(buf, line), 60 | DOUBLE_IDENT => parse_double(line), 61 | ERROR_IDENT => parse_error(buf, line), 62 | _ => Err(ParseFrameError::InvalidFormat), 63 | } 64 | } 65 | 66 | fn get_line<'a>(buf: &mut Cursor<&'a [u8]>) -> Result<&'a [u8], ParseFrameError> { 67 | if !buf.has_remaining() { 68 | return Err(ParseFrameError::Incomplete); 69 | } 70 | 71 | let start = buf.position() as usize; 72 | let end = buf.get_ref().len() - 1; 73 | 74 | for i in start..end { 75 | // A line is CRLF terminated and we look for that 76 | if buf.get_ref()[i] == b'\r' && buf.get_ref()[i + 1] == b'\n' { 77 | buf.set_position((i + 2) as u64); 78 | return Ok(&buf.get_ref()[start..i]); 79 | } 80 | } 81 | 82 | Err(ParseFrameError::Incomplete) 83 | } 84 | 85 | fn skip(buf: &mut Cursor<&[u8]>, n: usize) -> Result<(), ParseFrameError> { 86 | if buf.remaining() < n { 87 | return Err(ParseFrameError::Incomplete); 88 | } 89 | buf.advance(n); 90 | Ok(()) 91 | } 92 | 93 | fn parse_string(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result { 94 | // len is the length of encoded data 95 | let len = atoi::(line).ok_or(ParseFrameError::InvalidFormat)?; 96 | // add 2 to accommodate CRLF 97 | let n = len + 2; 98 | 99 | if buf.remaining() < n { 100 | return Err(ParseFrameError::Incomplete); 101 | } 102 | 103 | let data = Bytes::copy_from_slice(&buf.chunk()[..len]); 104 | 105 | skip(buf, n)?; 106 | 107 | Ok(Frame::String(data)) 108 | } 109 | 110 | fn parse_integer(line: &[u8]) -> Result { 111 | let int = atoi::(line).ok_or(ParseFrameError::InvalidFormat)?; 112 | Ok(Frame::Integer(int)) 113 | } 114 | 115 | fn parse_array(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result { 116 | let len = atoi::(line).ok_or(ParseFrameError::InvalidFormat)?; 117 | let mut vec = Vec::with_capacity(len); 118 | for _ in 0..len { 119 | vec.push(parse(buf)?); 120 | } 121 | 122 | Ok(Frame::Array(vec)) 123 | } 124 | 125 | fn parse_boolean(line: &[u8]) -> Result { 126 | if line.len() > 1 { 127 | return Err(ParseFrameError::InvalidFormat); 128 | } 129 | 130 | let val = line[0]; 131 | 132 | match val { 133 | b'0' => Ok(Frame::Boolean(false)), 134 | b'1' => Ok(Frame::Boolean(true)), 135 | _ => Err(ParseFrameError::InvalidFormat), 136 | } 137 | } 138 | 139 | fn parse_null(line: &[u8]) -> Result { 140 | if !line.is_empty() { 141 | return Err(ParseFrameError::InvalidFormat); 142 | } 143 | Ok(Frame::Null) 144 | } 145 | 146 | fn parse_map(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result { 147 | let len = atoi::(line).ok_or(ParseFrameError::InvalidFormat)?; 148 | let mut map = Vec::with_capacity(2 * len); 149 | for _ in 0..len { 150 | let key = parse(buf)?; 151 | let value = parse(buf)?; 152 | map.push(key); 153 | map.push(value); 154 | } 155 | 156 | Ok(Frame::Map(map)) 157 | } 158 | 159 | fn parse_double(line: &[u8]) -> Result { 160 | let double = str::from_utf8(line) 161 | .map_err(|_| ParseFrameError::InvalidFormat)? 162 | .parse::() 163 | .map_err(|_| ParseFrameError::InvalidFormat)?; 164 | Ok(Frame::Double(double)) 165 | } 166 | 167 | fn parse_error(buf: &mut Cursor<&[u8]>, line: &[u8]) -> Result { 168 | let len = atoi::(line).ok_or(ParseFrameError::InvalidFormat)?; 169 | let n = len + 2; 170 | 171 | if buf.remaining() < n { 172 | return Err(ParseFrameError::Incomplete); 173 | } 174 | 175 | let data = Bytes::copy_from_slice(&buf.chunk()[..len]); 176 | 177 | skip(buf, n)?; 178 | 179 | Ok(Frame::Error(data)) 180 | } 181 | 182 | #[cfg(test)] 183 | mod tests { 184 | use super::*; 185 | use bytes::{BufMut, BytesMut}; 186 | use std::fs; 187 | use std::io::Cursor; 188 | use std::path::{Path, PathBuf}; 189 | 190 | fn get_cursor_from_bytes(bytes: &[u8]) -> Cursor<&[u8]> { 191 | Cursor::new(bytes) 192 | } 193 | 194 | fn read_file(path: PathBuf) -> Vec { 195 | fs::read(path).unwrap() 196 | } 197 | 198 | fn get_frame_from_file(data: &[u8], ident: u8) -> Bytes { 199 | let mut frame = BytesMut::new(); 200 | frame.put_u8(ident); 201 | frame.put_slice(format!("{}", data.len()).as_bytes()); 202 | frame.put_u8(b'\r'); 203 | frame.put_u8(b'\n'); 204 | frame.put(data); 205 | frame.put_u8(b'\r'); 206 | frame.put_u8(b'\n'); 207 | 208 | frame.copy_to_bytes(frame.len()) 209 | } 210 | 211 | #[test] 212 | fn parse_given_empty_line_returns_invalid_format_error() { 213 | let mut buf = get_cursor_from_bytes(b"\r\n"); 214 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 215 | } 216 | 217 | #[test] 218 | fn parse_given_unknown_type_returns_invalid_format_error() { 219 | let mut buf = get_cursor_from_bytes(b"foo\r\n"); 220 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 221 | } 222 | 223 | #[test] 224 | fn parse_given_string_with_no_length_returns_invalid_format_error() { 225 | let mut buf = get_cursor_from_bytes(b"$\r\n"); 226 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 227 | } 228 | 229 | #[test] 230 | fn parse_given_string_with_invalid_length_returns_invalid_format_error() { 231 | let mut buf = get_cursor_from_bytes(b"$abc\r\n"); 232 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 233 | } 234 | 235 | #[test] 236 | fn parse_given_incomplete_string_with_zero_length_returns_incomplete_error() { 237 | let mut buf = get_cursor_from_bytes(b"$0\r\n"); 238 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 239 | } 240 | 241 | #[test] 242 | fn parse_given_incomplete_string_with_non_zero_length_returns_incomplete_error() { 243 | let mut buf = get_cursor_from_bytes(b"$1\r\n"); 244 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 245 | } 246 | 247 | #[test] 248 | fn parse_given_string_with_zero_length_returns_empty_string() { 249 | let mut buf = get_cursor_from_bytes(b"$0\r\n\r\n"); 250 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("")))) 251 | } 252 | 253 | #[test] 254 | fn parse_given_string_with_length_less_than_length_of_data_returns_data_upto_given_length() { 255 | let mut buf = get_cursor_from_bytes(b"$1\r\nfoo\r\n"); 256 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("f")))) 257 | } 258 | 259 | #[test] 260 | fn parse_given_string_with_length_greater_than_length_of_data_returns_incomplete_error() { 261 | let mut buf = get_cursor_from_bytes(b"$100\r\nfoo\r\n"); 262 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 263 | } 264 | 265 | #[test] 266 | fn parse_given_string_returns_string() { 267 | let mut buf = get_cursor_from_bytes(b"$3\r\nfoo\r\n"); 268 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("foo")))) 269 | } 270 | 271 | #[test] 272 | fn parse_given_string_with_delimiter_in_data_returns_string() { 273 | let mut buf = get_cursor_from_bytes(b"$5\r\nfoo\r\n\r\n"); 274 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from("foo\r\n")))) 275 | } 276 | 277 | #[test] 278 | fn parse_given_string_with_pdf_data_returns_pdf_data() { 279 | let file_data = read_file(Path::new("test_data").join("test.pdf")); 280 | let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT); 281 | let mut buf = get_cursor_from_bytes(&frame); 282 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data)))) 283 | } 284 | 285 | #[test] 286 | fn parse_given_string_with_png_data_returns_png_data() { 287 | let file_data = read_file(Path::new("test_data").join("test.png")); 288 | let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT); 289 | let mut buf = get_cursor_from_bytes(&frame); 290 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data)))) 291 | } 292 | 293 | #[test] 294 | fn parse_given_string_with_jpg_data_returns_jpg_data() { 295 | let file_data = read_file(Path::new("test_data").join("test.jpg")); 296 | let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT); 297 | let mut buf = get_cursor_from_bytes(&frame); 298 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data)))) 299 | } 300 | 301 | #[test] 302 | fn parse_given_string_with_html_data_returns_html_data() { 303 | let file_data = read_file(Path::new("test_data").join("test.html")); 304 | let frame = get_frame_from_file(file_data.as_slice(), STRING_IDENT); 305 | let mut buf = get_cursor_from_bytes(&frame); 306 | assert_eq!(parse(&mut buf), Ok(Frame::String(Bytes::from(file_data)))) 307 | } 308 | 309 | #[test] 310 | fn parse_given_invalid_integer_returns_invalid_format_error() { 311 | let mut buf = get_cursor_from_bytes(b"%abc\r\n"); 312 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 313 | } 314 | 315 | #[test] 316 | fn parse_given_empty_integer_returns_invalid_format_error() { 317 | let mut buf = get_cursor_from_bytes(b"%\r\n"); 318 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 319 | } 320 | 321 | #[test] 322 | fn parse_given_negative_integer_returns_given_integer() { 323 | let mut buf = get_cursor_from_bytes(b"%-1\r\n"); 324 | assert_eq!(parse(&mut buf), Ok(Frame::Integer(-1))) 325 | } 326 | 327 | #[test] 328 | fn parse_given_zero_returns_zero() { 329 | let mut buf = get_cursor_from_bytes(b"%0\r\n"); 330 | assert_eq!(parse(&mut buf), Ok(Frame::Integer(0))) 331 | } 332 | 333 | #[test] 334 | fn parse_given_positive_integer_returns_given_integer() { 335 | let mut buf = get_cursor_from_bytes(b"%1000\r\n"); 336 | assert_eq!(parse(&mut buf), Ok(Frame::Integer(1000))) 337 | } 338 | 339 | #[test] 340 | fn parse_given_out_of_range_integer_returns_format_error() { 341 | let mut buf = get_cursor_from_bytes(b"%9223372036854775808\r\n"); 342 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 343 | } 344 | 345 | #[test] 346 | fn parse_given_false_returns_false() { 347 | let mut buf = get_cursor_from_bytes(b"^0\r\n"); 348 | assert_eq!(parse(&mut buf), Ok(Frame::Boolean(false))) 349 | } 350 | 351 | #[test] 352 | fn parse_given_true_returns_true() { 353 | let mut buf = get_cursor_from_bytes(b"^1\r\n"); 354 | assert_eq!(parse(&mut buf), Ok(Frame::Boolean(true))) 355 | } 356 | 357 | #[test] 358 | fn parse_given_invalid_boolean_returns_invalid_format_error() { 359 | let mut buf = get_cursor_from_bytes(b"^foo\r\n"); 360 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 361 | } 362 | 363 | #[test] 364 | fn parse_given_null_returns_null() { 365 | let mut buf = get_cursor_from_bytes(b"-\r\n"); 366 | assert_eq!(parse(&mut buf), Ok(Frame::Null)) 367 | } 368 | 369 | #[test] 370 | fn parse_given_invalid_null_returns_invalid_format_error() { 371 | let mut buf = get_cursor_from_bytes(b"-foo\r\n"); 372 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 373 | } 374 | 375 | #[test] 376 | fn parse_given_double_with_invalid_decimal_part_returns_invalid_format_error() { 377 | let mut buf = get_cursor_from_bytes(b".20.foo\r\n"); 378 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 379 | } 380 | 381 | #[test] 382 | fn parse_given_double_with_invalid_integer_part_returns_invalid_format_error() { 383 | let mut buf = get_cursor_from_bytes(b".foo.90\r\n"); 384 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 385 | } 386 | 387 | #[test] 388 | fn parse_given_double_with_zero_decimal_part_returns_double() { 389 | let mut buf = get_cursor_from_bytes(b".10.000\r\n"); 390 | assert_eq!(parse(&mut buf), Ok(Frame::Double(10.0))) 391 | } 392 | 393 | #[test] 394 | fn parse_given_double_with_trailing_zeroes_returns_double() { 395 | let mut buf = get_cursor_from_bytes(b".10.100\r\n"); 396 | assert_eq!(parse(&mut buf), Ok(Frame::Double(10.1))) 397 | } 398 | 399 | #[test] 400 | fn parse_given_double_with_leading_zeroes_returns_double() { 401 | let mut buf = get_cursor_from_bytes(b".000010.100\r\n"); 402 | assert_eq!(parse(&mut buf), Ok(Frame::Double(10.1))) 403 | } 404 | 405 | #[test] 406 | fn parse_given_invalid_double_returns_invalid_format_error() { 407 | let mut buf = get_cursor_from_bytes(b".abc\r\n"); 408 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 409 | } 410 | 411 | #[test] 412 | fn parse_given_error_with_no_length_returns_invalid_format_error() { 413 | let mut buf = get_cursor_from_bytes(b"!\r\n"); 414 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 415 | } 416 | 417 | #[test] 418 | fn parse_given_error_with_invalid_length_returns_invalid_format_error() { 419 | let mut buf = get_cursor_from_bytes(b"!abc\r\n"); 420 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 421 | } 422 | 423 | #[test] 424 | fn parse_given_incomplete_error_with_zero_length_returns_incomplete_error() { 425 | let mut buf = get_cursor_from_bytes(b"!0\r\n"); 426 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 427 | } 428 | 429 | #[test] 430 | fn parse_given_incomplete_error_with_non_zero_length_returns_incomplete_error() { 431 | let mut buf = get_cursor_from_bytes(b"!1\r\n"); 432 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 433 | } 434 | 435 | #[test] 436 | fn parse_given_error_with_zero_length_returns_empty_error_frame() { 437 | let mut buf = get_cursor_from_bytes(b"!0\r\n\r\n"); 438 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("")))) 439 | } 440 | 441 | #[test] 442 | fn parse_given_error_with_length_less_than_length_of_data_returns_data_upto_given_length() { 443 | let mut buf = get_cursor_from_bytes(b"!1\r\nfoo\r\n"); 444 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("f")))) 445 | } 446 | 447 | #[test] 448 | fn parse_given_error_with_length_greater_than_length_of_data_returns_incomplete_error() { 449 | let mut buf = get_cursor_from_bytes(b"!100\r\nfoo\r\n"); 450 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 451 | } 452 | 453 | #[test] 454 | fn parse_given_error_returns_error_frame() { 455 | let mut buf = get_cursor_from_bytes(b"!3\r\nfoo\r\n"); 456 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("foo")))) 457 | } 458 | 459 | #[test] 460 | fn parse_given_error_with_delimiter_in_data_returns_error_frame() { 461 | let mut buf = get_cursor_from_bytes(b"!5\r\nfoo\r\n\r\n"); 462 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from("foo\r\n")))) 463 | } 464 | 465 | #[test] 466 | fn parse_given_error_with_pdf_data_returns_pdf_data() { 467 | let file_data = read_file(Path::new("test_data").join("test.pdf")); 468 | let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT); 469 | let mut buf = get_cursor_from_bytes(&frame); 470 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data)))) 471 | } 472 | 473 | #[test] 474 | fn parse_given_error_with_png_data_returns_png_data() { 475 | let file_data = read_file(Path::new("test_data").join("test.png")); 476 | let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT); 477 | let mut buf = get_cursor_from_bytes(&frame); 478 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data)))) 479 | } 480 | 481 | #[test] 482 | fn parse_given_error_with_jpg_data_returns_jpg_data() { 483 | let file_data = read_file(Path::new("test_data").join("test.jpg")); 484 | let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT); 485 | let mut buf = get_cursor_from_bytes(&frame); 486 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data)))) 487 | } 488 | 489 | #[test] 490 | fn parse_given_error_with_html_data_returns_html_data() { 491 | let file_data = read_file(Path::new("test_data").join("test.html")); 492 | let frame = get_frame_from_file(file_data.as_slice(), ERROR_IDENT); 493 | let mut buf = get_cursor_from_bytes(&frame); 494 | assert_eq!(parse(&mut buf), Ok(Frame::Error(Bytes::from(file_data)))) 495 | } 496 | 497 | #[test] 498 | fn parse_given_array_with_no_length_returns_invalid_format_error() { 499 | let mut buf = get_cursor_from_bytes(b"*\r\n"); 500 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 501 | } 502 | 503 | #[test] 504 | fn parse_given_array_with_zero_length_returns_empty_array() { 505 | let mut buf = get_cursor_from_bytes(b"*0\r\n"); 506 | assert_eq!(parse(&mut buf), Ok(Frame::Array(Vec::new()))) 507 | } 508 | 509 | #[test] 510 | fn parse_given_array_with_invalid_length_invalid_format_error() { 511 | let mut buf = get_cursor_from_bytes(b"*abc\r\n"); 512 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 513 | } 514 | 515 | #[test] 516 | fn parse_given_map_with_no_length_returns_invalid_format_error() { 517 | let mut buf = get_cursor_from_bytes(b"#\r\n"); 518 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 519 | } 520 | 521 | #[test] 522 | fn parse_given_map_with_zero_length_returns_empty_map() { 523 | let mut buf = get_cursor_from_bytes(b"#0\r\n"); 524 | assert_eq!(parse(&mut buf), Ok(Frame::Map(Vec::new()))) 525 | } 526 | 527 | #[test] 528 | fn parse_given_map_with_invalid_length_returns_invalid_format_error() { 529 | let mut buf = get_cursor_from_bytes(b"#abc\r\n"); 530 | assert_eq!(parse(&mut buf), Err(ParseFrameError::InvalidFormat)) 531 | } 532 | 533 | #[test] 534 | fn parse_given_incomplete_map_return_incomplete_error() { 535 | let mut buf = get_cursor_from_bytes(b"#2\r\n$3\r\nfoo\r\n"); 536 | assert_eq!(parse(&mut buf), Err(ParseFrameError::Incomplete)) 537 | } 538 | } 539 | -------------------------------------------------------------------------------- /src/command/test.rs: -------------------------------------------------------------------------------- 1 | use super::parse; 2 | use crate::db::Evictor; 3 | use crate::{ 4 | command::{Command, Count, Create, Del, Drop, Get, Set, Ttl}, 5 | frame::Frame, 6 | }; 7 | use bytes::Bytes; 8 | use std::ops::Add; 9 | use std::time::{Duration, SystemTime, UNIX_EPOCH}; 10 | 11 | fn get_frame_from_str(str: &'static str) -> Frame { 12 | Frame::String(Bytes::from(str)) 13 | } 14 | 15 | #[test] 16 | fn parse_given_unknown_command_returns_error() { 17 | let command = vec![get_frame_from_str("foo")]; 18 | assert!(parse(Frame::Array(command)).is_err()) 19 | } 20 | 21 | #[test] 22 | fn parse_given_command_with_non_string_tokens_returns_error() { 23 | let command = vec![get_frame_from_str("create"), Frame::Integer(100)]; 24 | assert!(parse(Frame::Array(command)).is_err()) 25 | } 26 | 27 | #[test] 28 | fn parse_given_command_with_non_array_container_returns_error() { 29 | assert!(parse(get_frame_from_str("create")).is_err()) 30 | } 31 | 32 | #[test] 33 | fn parse_given_create_command_without_keyspace_returns_error() { 34 | let command = vec![get_frame_from_str("create")]; 35 | assert!(parse(Frame::Array(command)).is_err()) 36 | } 37 | 38 | #[test] 39 | fn parse_given_create_command_with_keyspace_and_no_config_returns_create() { 40 | let command = vec![get_frame_from_str("create"), get_frame_from_str("foo")]; 41 | assert_eq!( 42 | parse(Frame::Array(command)).unwrap(), 43 | Command::Create(Create { 44 | evictor: Evictor::Nop, 45 | if_not_exists: false, 46 | keyspace: Bytes::from("foo") 47 | }) 48 | ); 49 | } 50 | 51 | #[test] 52 | fn parse_given_create_command_with_no_evictor_value_returns_error() { 53 | let command = vec![ 54 | get_frame_from_str("create"), 55 | get_frame_from_str("foo"), 56 | get_frame_from_str("evictor"), 57 | ]; 58 | assert!(parse(Frame::Array(command)).is_err()) 59 | } 60 | 61 | #[test] 62 | fn parse_given_create_command_with_invalid_evictor_value_returns_error() { 63 | let command = vec![ 64 | get_frame_from_str("create"), 65 | get_frame_from_str("foo"), 66 | get_frame_from_str("evictor"), 67 | get_frame_from_str("bar"), 68 | ]; 69 | assert!(parse(Frame::Array(command)).is_err()) 70 | } 71 | 72 | #[test] 73 | fn parse_given_create_command_with_lru_evictor_returns_create() { 74 | let command = vec![ 75 | get_frame_from_str("create"), 76 | get_frame_from_str("foo"), 77 | get_frame_from_str("evictor"), 78 | get_frame_from_str("lru"), 79 | ]; 80 | assert_eq!( 81 | parse(Frame::Array(command)).unwrap(), 82 | Command::Create(Create { 83 | evictor: Evictor::Lru, 84 | if_not_exists: false, 85 | keyspace: Bytes::from("foo") 86 | }) 87 | ); 88 | } 89 | 90 | #[test] 91 | fn parse_given_create_command_with_random_evictor_returns_create() { 92 | let command = vec![ 93 | get_frame_from_str("create"), 94 | get_frame_from_str("foo"), 95 | get_frame_from_str("evictor"), 96 | get_frame_from_str("random"), 97 | ]; 98 | assert_eq!( 99 | parse(Frame::Array(command)).unwrap(), 100 | Command::Create(Create { 101 | evictor: Evictor::Random, 102 | if_not_exists: false, 103 | keyspace: Bytes::from("foo") 104 | }) 105 | ); 106 | } 107 | 108 | #[test] 109 | fn parse_given_create_command_with_nop_evictor_returns_create() { 110 | let command = vec![ 111 | get_frame_from_str("create"), 112 | get_frame_from_str("foo"), 113 | get_frame_from_str("evictor"), 114 | get_frame_from_str("nop"), 115 | ]; 116 | assert_eq!( 117 | parse(Frame::Array(command)).unwrap(), 118 | Command::Create(Create { 119 | evictor: Evictor::Nop, 120 | if_not_exists: false, 121 | keyspace: Bytes::from("foo") 122 | }) 123 | ); 124 | } 125 | 126 | #[test] 127 | fn parse_given_create_command_with_incomplete_if_flag_1_returns_error() { 128 | let command = vec![ 129 | get_frame_from_str("create"), 130 | get_frame_from_str("foo"), 131 | get_frame_from_str("evictor"), 132 | get_frame_from_str("lru"), 133 | get_frame_from_str("if"), 134 | ]; 135 | assert!(parse(Frame::Array(command)).is_err()) 136 | } 137 | 138 | #[test] 139 | fn parse_given_create_command_with_incomplete_if_flag_2_returns_error() { 140 | let command = vec![ 141 | get_frame_from_str("create"), 142 | get_frame_from_str("foo"), 143 | get_frame_from_str("evictor"), 144 | get_frame_from_str("lru"), 145 | get_frame_from_str("if"), 146 | get_frame_from_str("not"), 147 | ]; 148 | assert!(parse(Frame::Array(command)).is_err()) 149 | } 150 | 151 | #[test] 152 | fn parse_given_create_command_with_invalid_if_flag_returns_error() { 153 | let command = vec![ 154 | get_frame_from_str("create"), 155 | get_frame_from_str("foo"), 156 | get_frame_from_str("evictor"), 157 | get_frame_from_str("lru"), 158 | get_frame_from_str("if"), 159 | get_frame_from_str("exists"), 160 | ]; 161 | assert!(parse(Frame::Array(command)).is_err()) 162 | } 163 | 164 | #[test] 165 | fn parse_given_create_command_with_if_not_exists_flag_returns_create() { 166 | let command = vec![ 167 | get_frame_from_str("create"), 168 | get_frame_from_str("foo"), 169 | get_frame_from_str("evictor"), 170 | get_frame_from_str("lru"), 171 | get_frame_from_str("if"), 172 | get_frame_from_str("not"), 173 | get_frame_from_str("exists"), 174 | ]; 175 | assert_eq!( 176 | parse(Frame::Array(command)).unwrap(), 177 | Command::Create(Create { 178 | evictor: Evictor::Lru, 179 | if_not_exists: true, 180 | keyspace: Bytes::from("foo") 181 | }) 182 | ); 183 | } 184 | 185 | #[test] 186 | fn parse_given_create_command_with_if_not_exists_flag_and_evictor_position_reversed_returns_create() 187 | { 188 | let command = vec![ 189 | get_frame_from_str("create"), 190 | get_frame_from_str("foo"), 191 | get_frame_from_str("if"), 192 | get_frame_from_str("not"), 193 | get_frame_from_str("exists"), 194 | get_frame_from_str("evictor"), 195 | get_frame_from_str("lru"), 196 | ]; 197 | assert_eq!( 198 | parse(Frame::Array(command)).unwrap(), 199 | Command::Create(Create { 200 | evictor: Evictor::Lru, 201 | if_not_exists: true, 202 | keyspace: Bytes::from("foo") 203 | }) 204 | ); 205 | } 206 | 207 | #[test] 208 | fn parse_given_create_command_with_if_not_exists_flag_and_default_evictor_returns_create() { 209 | let command = vec![ 210 | get_frame_from_str("create"), 211 | get_frame_from_str("foo"), 212 | get_frame_from_str("if"), 213 | get_frame_from_str("not"), 214 | get_frame_from_str("exists"), 215 | ]; 216 | assert_eq!( 217 | parse(Frame::Array(command)).unwrap(), 218 | Command::Create(Create { 219 | evictor: Evictor::Nop, 220 | if_not_exists: true, 221 | keyspace: Bytes::from("foo") 222 | }) 223 | ); 224 | } 225 | 226 | #[test] 227 | fn parse_given_create_command_with_invalid_if_flag_2_returns_error() { 228 | let command = vec![ 229 | get_frame_from_str("create"), 230 | get_frame_from_str("foo"), 231 | get_frame_from_str("if"), 232 | get_frame_from_str("not"), 233 | get_frame_from_str("evictor"), 234 | get_frame_from_str("exists"), 235 | get_frame_from_str("lru"), 236 | ]; 237 | assert!(parse(Frame::Array(command)).is_err()) 238 | } 239 | 240 | #[test] 241 | fn parse_given_set_command_without_keyspace_returns_error() { 242 | let command = vec![get_frame_from_str("set")]; 243 | assert!(parse(Frame::Array(command)).is_err()) 244 | } 245 | 246 | #[test] 247 | fn parse_given_set_command_without_key_returns_error() { 248 | let command = vec![get_frame_from_str("set"), get_frame_from_str("my_keyspace")]; 249 | assert!(parse(Frame::Array(command)).is_err()) 250 | } 251 | 252 | #[test] 253 | fn parse_given_set_command_without_value_returns_error() { 254 | let command = vec![ 255 | get_frame_from_str("set"), 256 | get_frame_from_str("my_keyspace"), 257 | get_frame_from_str("foo"), 258 | ]; 259 | assert!(parse(Frame::Array(command)).is_err()) 260 | } 261 | 262 | #[test] 263 | fn parse_given_set_command_without_config_returns_set() { 264 | let command = vec![ 265 | get_frame_from_str("set"), 266 | get_frame_from_str("my_keyspace"), 267 | get_frame_from_str("foo"), 268 | get_frame_from_str("bar"), 269 | ]; 270 | assert_eq!( 271 | parse(Frame::Array(command)).unwrap(), 272 | Command::Set(Set { 273 | keyspace: Bytes::from("my_keyspace"), 274 | key: Bytes::from("foo"), 275 | value: Bytes::from("bar"), 276 | if_not_exists: false, 277 | expire_at: None, 278 | if_exists: false 279 | }) 280 | ); 281 | } 282 | 283 | #[test] 284 | fn parse_given_set_command_with_incomplete_expire_at_returns_error() { 285 | let command = vec![ 286 | get_frame_from_str("set"), 287 | get_frame_from_str("my_keyspace"), 288 | get_frame_from_str("foo"), 289 | get_frame_from_str("bar"), 290 | get_frame_from_str("expire"), 291 | ]; 292 | assert!(parse(Frame::Array(command)).is_err()) 293 | } 294 | 295 | #[test] 296 | fn parse_given_set_command_without_expire_at_value_returns_error() { 297 | let command = vec![ 298 | get_frame_from_str("set"), 299 | get_frame_from_str("my_keyspace"), 300 | get_frame_from_str("foo"), 301 | get_frame_from_str("bar"), 302 | get_frame_from_str("expire"), 303 | get_frame_from_str("at"), 304 | ]; 305 | assert!(parse(Frame::Array(command)).is_err()) 306 | } 307 | 308 | #[test] 309 | fn parse_given_set_command_with_invalid_expire_at_value_returns_error() { 310 | let command = vec![ 311 | get_frame_from_str("set"), 312 | get_frame_from_str("my_keyspace"), 313 | get_frame_from_str("foo"), 314 | get_frame_from_str("bar"), 315 | get_frame_from_str("expire"), 316 | get_frame_from_str("at"), 317 | get_frame_from_str("baz"), 318 | ]; 319 | assert!(parse(Frame::Array(command)).is_err()) 320 | } 321 | 322 | #[test] 323 | fn parse_given_set_command_with_expire_at_returns_set() { 324 | let command = vec![ 325 | get_frame_from_str("set"), 326 | get_frame_from_str("my_keyspace"), 327 | get_frame_from_str("foo"), 328 | get_frame_from_str("bar"), 329 | get_frame_from_str("expire"), 330 | get_frame_from_str("at"), 331 | get_frame_from_str("1667041052"), 332 | ]; 333 | assert_eq!( 334 | parse(Frame::Array(command)).unwrap(), 335 | Command::Set(Set { 336 | keyspace: Bytes::from("my_keyspace"), 337 | key: Bytes::from("foo"), 338 | value: Bytes::from("bar"), 339 | if_not_exists: false, 340 | expire_at: Some(1667041052), 341 | if_exists: false 342 | }) 343 | ); 344 | } 345 | 346 | #[test] 347 | fn parse_given_set_command_with_incomplete_expire_after_returns_error() { 348 | let command = vec![ 349 | get_frame_from_str("set"), 350 | get_frame_from_str("my_keyspace"), 351 | get_frame_from_str("foo"), 352 | get_frame_from_str("bar"), 353 | get_frame_from_str("expire"), 354 | ]; 355 | assert!(parse(Frame::Array(command)).is_err()) 356 | } 357 | 358 | #[test] 359 | fn parse_given_set_command_without_expire_after_value_returns_error() { 360 | let command = vec![ 361 | get_frame_from_str("set"), 362 | get_frame_from_str("my_keyspace"), 363 | get_frame_from_str("foo"), 364 | get_frame_from_str("bar"), 365 | get_frame_from_str("expire"), 366 | get_frame_from_str("after"), 367 | ]; 368 | assert!(parse(Frame::Array(command)).is_err()) 369 | } 370 | 371 | #[test] 372 | fn parse_given_set_command_with_invalid_expire_after_value_returns_error() { 373 | let command = vec![ 374 | get_frame_from_str("set"), 375 | get_frame_from_str("my_keyspace"), 376 | get_frame_from_str("foo"), 377 | get_frame_from_str("bar"), 378 | get_frame_from_str("expire"), 379 | get_frame_from_str("after"), 380 | get_frame_from_str("baz"), 381 | ]; 382 | assert!(parse(Frame::Array(command)).is_err()) 383 | } 384 | 385 | // FIXME: this is a finicky test, find a better way to test the timestamp 386 | #[test] 387 | fn parse_given_set_command_with_expire_after_returns_set() { 388 | let command = vec![ 389 | get_frame_from_str("set"), 390 | get_frame_from_str("my_keyspace"), 391 | get_frame_from_str("foo"), 392 | get_frame_from_str("bar"), 393 | get_frame_from_str("expire"), 394 | get_frame_from_str("after"), 395 | get_frame_from_str("60000"), 396 | ]; 397 | 398 | let timestamp = SystemTime::now() 399 | .add(Duration::from_millis(60000)) 400 | .duration_since(UNIX_EPOCH) 401 | .unwrap() 402 | .as_secs(); 403 | 404 | assert_eq!( 405 | parse(Frame::Array(command)).unwrap(), 406 | Command::Set(Set { 407 | keyspace: Bytes::from("my_keyspace"), 408 | key: Bytes::from("foo"), 409 | value: Bytes::from("bar"), 410 | if_not_exists: false, 411 | expire_at: Some(timestamp), 412 | if_exists: false 413 | }) 414 | ); 415 | } 416 | 417 | #[test] 418 | fn parse_given_set_command_with_both_expire_at_and_expire_after_returns_error() { 419 | let command = vec![ 420 | get_frame_from_str("set"), 421 | get_frame_from_str("my_keyspace"), 422 | get_frame_from_str("foo"), 423 | get_frame_from_str("bar"), 424 | get_frame_from_str("expire"), 425 | get_frame_from_str("after"), 426 | get_frame_from_str("100"), 427 | get_frame_from_str("expire"), 428 | get_frame_from_str("at"), 429 | get_frame_from_str("1667041052"), 430 | ]; 431 | assert!(parse(Frame::Array(command)).is_err()) 432 | } 433 | 434 | #[test] 435 | fn parse_given_set_command_with_incomplete_if_flag_1_returns_error() { 436 | let command = vec![ 437 | get_frame_from_str("set"), 438 | get_frame_from_str("my_keyspace"), 439 | get_frame_from_str("foo"), 440 | get_frame_from_str("bar"), 441 | get_frame_from_str("expire"), 442 | get_frame_from_str("after"), 443 | get_frame_from_str("100"), 444 | get_frame_from_str("if"), 445 | ]; 446 | assert!(parse(Frame::Array(command)).is_err()) 447 | } 448 | 449 | #[test] 450 | fn parse_given_set_command_with_invalid_if_flag_1_returns_error() { 451 | let command = vec![ 452 | get_frame_from_str("set"), 453 | get_frame_from_str("my_keyspace"), 454 | get_frame_from_str("foo"), 455 | get_frame_from_str("bar"), 456 | get_frame_from_str("expire"), 457 | get_frame_from_str("after"), 458 | get_frame_from_str("100"), 459 | get_frame_from_str("if"), 460 | get_frame_from_str("foo"), 461 | ]; 462 | assert!(parse(Frame::Array(command)).is_err()) 463 | } 464 | 465 | #[test] 466 | fn parse_given_set_command_with_if_exists_flag_returns_set() { 467 | let command = vec![ 468 | get_frame_from_str("set"), 469 | get_frame_from_str("my_keyspace"), 470 | get_frame_from_str("foo"), 471 | get_frame_from_str("bar"), 472 | get_frame_from_str("if"), 473 | get_frame_from_str("exists"), 474 | ]; 475 | 476 | assert_eq!( 477 | parse(Frame::Array(command)).unwrap(), 478 | Command::Set(Set { 479 | keyspace: Bytes::from("my_keyspace"), 480 | key: Bytes::from("foo"), 481 | value: Bytes::from("bar"), 482 | if_not_exists: false, 483 | expire_at: None, 484 | if_exists: true 485 | }) 486 | ); 487 | } 488 | 489 | #[test] 490 | fn parse_given_set_command_with_incomplete_if_flag_2_returns_error() { 491 | let command = vec![ 492 | get_frame_from_str("set"), 493 | get_frame_from_str("my_keyspace"), 494 | get_frame_from_str("foo"), 495 | get_frame_from_str("bar"), 496 | get_frame_from_str("expire"), 497 | get_frame_from_str("after"), 498 | get_frame_from_str("100"), 499 | get_frame_from_str("if"), 500 | get_frame_from_str("not"), 501 | ]; 502 | assert!(parse(Frame::Array(command)).is_err()) 503 | } 504 | 505 | #[test] 506 | fn parse_given_set_command_with_invalid_if_flag_2_returns_error() { 507 | let command = vec![ 508 | get_frame_from_str("set"), 509 | get_frame_from_str("my_keyspace"), 510 | get_frame_from_str("foo"), 511 | get_frame_from_str("bar"), 512 | get_frame_from_str("expire"), 513 | get_frame_from_str("after"), 514 | get_frame_from_str("100"), 515 | get_frame_from_str("if"), 516 | get_frame_from_str("not"), 517 | get_frame_from_str("foo"), 518 | ]; 519 | assert!(parse(Frame::Array(command)).is_err()) 520 | } 521 | 522 | #[test] 523 | fn parse_given_set_command_with_if_not_exists_flag_returns_set() { 524 | let command = vec![ 525 | get_frame_from_str("set"), 526 | get_frame_from_str("my_keyspace"), 527 | get_frame_from_str("foo"), 528 | get_frame_from_str("bar"), 529 | get_frame_from_str("if"), 530 | get_frame_from_str("not"), 531 | get_frame_from_str("exists"), 532 | ]; 533 | 534 | assert_eq!( 535 | parse(Frame::Array(command)).unwrap(), 536 | Command::Set(Set { 537 | keyspace: Bytes::from("my_keyspace"), 538 | key: Bytes::from("foo"), 539 | value: Bytes::from("bar"), 540 | if_not_exists: true, 541 | expire_at: None, 542 | if_exists: false 543 | }) 544 | ); 545 | } 546 | 547 | #[test] 548 | fn parse_given_set_command_with_both_if_exists_and_if_not_exists_returns_error() { 549 | let command = vec![ 550 | get_frame_from_str("set"), 551 | get_frame_from_str("my_keyspace"), 552 | get_frame_from_str("foo"), 553 | get_frame_from_str("bar"), 554 | get_frame_from_str("expire"), 555 | get_frame_from_str("after"), 556 | get_frame_from_str("100"), 557 | get_frame_from_str("if"), 558 | get_frame_from_str("not"), 559 | get_frame_from_str("exists"), 560 | get_frame_from_str("if"), 561 | get_frame_from_str("exists"), 562 | ]; 563 | assert!(parse(Frame::Array(command)).is_err()) 564 | } 565 | 566 | #[test] 567 | fn parse_given_set_command_with_if_not_exists_and_tll_returns_set() { 568 | let command = vec![ 569 | get_frame_from_str("set"), 570 | get_frame_from_str("my_keyspace"), 571 | get_frame_from_str("foo"), 572 | get_frame_from_str("bar"), 573 | get_frame_from_str("expire"), 574 | get_frame_from_str("at"), 575 | get_frame_from_str("1667041052"), 576 | get_frame_from_str("if"), 577 | get_frame_from_str("not"), 578 | get_frame_from_str("exists"), 579 | ]; 580 | 581 | assert_eq!( 582 | parse(Frame::Array(command)).unwrap(), 583 | Command::Set(Set { 584 | keyspace: Bytes::from("my_keyspace"), 585 | key: Bytes::from("foo"), 586 | value: Bytes::from("bar"), 587 | if_not_exists: true, 588 | expire_at: Some(1667041052), 589 | if_exists: false 590 | }) 591 | ); 592 | } 593 | 594 | #[test] 595 | fn parse_given_set_command_with_if_exists_and_tll_returns_set() { 596 | let command = vec![ 597 | get_frame_from_str("set"), 598 | get_frame_from_str("my_keyspace"), 599 | get_frame_from_str("foo"), 600 | get_frame_from_str("bar"), 601 | get_frame_from_str("expire"), 602 | get_frame_from_str("at"), 603 | get_frame_from_str("1667041052"), 604 | get_frame_from_str("if"), 605 | get_frame_from_str("exists"), 606 | ]; 607 | 608 | assert_eq!( 609 | parse(Frame::Array(command)).unwrap(), 610 | Command::Set(Set { 611 | keyspace: Bytes::from("my_keyspace"), 612 | key: Bytes::from("foo"), 613 | value: Bytes::from("bar"), 614 | if_not_exists: false, 615 | expire_at: Some(1667041052), 616 | if_exists: true 617 | }) 618 | ); 619 | } 620 | 621 | #[test] 622 | fn parse_given_set_command_with_if_not_exists_and_tll_reverse_order_returns_set() { 623 | let command = vec![ 624 | get_frame_from_str("set"), 625 | get_frame_from_str("my_keyspace"), 626 | get_frame_from_str("foo"), 627 | get_frame_from_str("bar"), 628 | get_frame_from_str("if"), 629 | get_frame_from_str("not"), 630 | get_frame_from_str("exists"), 631 | get_frame_from_str("expire"), 632 | get_frame_from_str("at"), 633 | get_frame_from_str("1667041052"), 634 | ]; 635 | 636 | assert_eq!( 637 | parse(Frame::Array(command)).unwrap(), 638 | Command::Set(Set { 639 | keyspace: Bytes::from("my_keyspace"), 640 | key: Bytes::from("foo"), 641 | value: Bytes::from("bar"), 642 | if_not_exists: true, 643 | expire_at: Some(1667041052), 644 | if_exists: false 645 | }) 646 | ); 647 | } 648 | 649 | #[test] 650 | fn parse_given_set_command_with_if_exists_and_tll_reverse_order_returns_set() { 651 | let command = vec![ 652 | get_frame_from_str("set"), 653 | get_frame_from_str("my_keyspace"), 654 | get_frame_from_str("foo"), 655 | get_frame_from_str("bar"), 656 | get_frame_from_str("if"), 657 | get_frame_from_str("exists"), 658 | get_frame_from_str("expire"), 659 | get_frame_from_str("at"), 660 | get_frame_from_str("1667041052"), 661 | ]; 662 | 663 | assert_eq!( 664 | parse(Frame::Array(command)).unwrap(), 665 | Command::Set(Set { 666 | keyspace: Bytes::from("my_keyspace"), 667 | key: Bytes::from("foo"), 668 | value: Bytes::from("bar"), 669 | if_not_exists: false, 670 | expire_at: Some(1667041052), 671 | if_exists: true 672 | }) 673 | ); 674 | } 675 | 676 | #[test] 677 | fn parse_given_get_without_keyspace_returns_error() { 678 | let command = vec![get_frame_from_str("get")]; 679 | assert!(parse(Frame::Array(command)).is_err()) 680 | } 681 | 682 | #[test] 683 | fn parse_given_get_without_key_returns_error() { 684 | let command = vec![get_frame_from_str("get"), get_frame_from_str("foo")]; 685 | assert!(parse(Frame::Array(command)).is_err()) 686 | } 687 | 688 | #[test] 689 | fn parse_given_get_returns_get() { 690 | let command = vec![ 691 | get_frame_from_str("get"), 692 | get_frame_from_str("foo"), 693 | get_frame_from_str("bar"), 694 | ]; 695 | 696 | assert_eq!( 697 | parse(Frame::Array(command)).unwrap(), 698 | Command::Get(Get { 699 | keyspace: Bytes::from("foo"), 700 | key: Bytes::from("bar"), 701 | }) 702 | ); 703 | } 704 | 705 | #[test] 706 | fn parse_given_del_without_keyspace_returns_error() { 707 | let command = vec![get_frame_from_str("del")]; 708 | assert!(parse(Frame::Array(command)).is_err()) 709 | } 710 | 711 | #[test] 712 | fn parse_given_del_without_key_returns_error() { 713 | let command = vec![get_frame_from_str("del"), get_frame_from_str("foo")]; 714 | assert!(parse(Frame::Array(command)).is_err()) 715 | } 716 | 717 | #[test] 718 | fn parse_given_del_returns_del() { 719 | let command = vec![ 720 | get_frame_from_str("del"), 721 | get_frame_from_str("foo"), 722 | get_frame_from_str("bar"), 723 | ]; 724 | 725 | assert_eq!( 726 | parse(Frame::Array(command)).unwrap(), 727 | Command::Del(Del { 728 | keyspace: Bytes::from("foo"), 729 | key: Bytes::from("bar"), 730 | }) 731 | ); 732 | } 733 | 734 | #[test] 735 | fn parse_given_drop_without_keyspace_returns_error() { 736 | let command = vec![get_frame_from_str("drop")]; 737 | assert!(parse(Frame::Array(command)).is_err()) 738 | } 739 | 740 | #[test] 741 | fn parse_given_drop_without_if_exists_returns_drop() { 742 | let command = vec![get_frame_from_str("drop"), get_frame_from_str("foo")]; 743 | 744 | assert_eq!( 745 | parse(Frame::Array(command)).unwrap(), 746 | Command::Drop(Drop { 747 | keyspace: Bytes::from("foo"), 748 | if_exists: false 749 | }) 750 | ); 751 | } 752 | 753 | #[test] 754 | fn parse_given_drop_with_incomplete_if_returns_error() { 755 | let command = vec![ 756 | get_frame_from_str("drop"), 757 | get_frame_from_str("foo"), 758 | get_frame_from_str("if"), 759 | ]; 760 | assert!(parse(Frame::Array(command)).is_err()) 761 | } 762 | 763 | #[test] 764 | fn parse_given_drop_with_invalid_if_returns_error() { 765 | let command = vec![ 766 | get_frame_from_str("drop"), 767 | get_frame_from_str("foo"), 768 | get_frame_from_str("if"), 769 | get_frame_from_str("not"), 770 | get_frame_from_str("exists"), 771 | ]; 772 | assert!(parse(Frame::Array(command)).is_err()) 773 | } 774 | 775 | #[test] 776 | fn parse_given_drop_with_if_exists_returns_drop() { 777 | let command = vec![ 778 | get_frame_from_str("drop"), 779 | get_frame_from_str("foo"), 780 | get_frame_from_str("if"), 781 | get_frame_from_str("exists"), 782 | ]; 783 | 784 | assert_eq!( 785 | parse(Frame::Array(command)).unwrap(), 786 | Command::Drop(Drop { 787 | keyspace: Bytes::from("foo"), 788 | if_exists: true 789 | }) 790 | ); 791 | } 792 | 793 | #[test] 794 | fn parse_given_drop_with_if_exists_reverse_position_returns_error() { 795 | let command = vec![ 796 | get_frame_from_str("drop"), 797 | get_frame_from_str("if"), 798 | get_frame_from_str("exists"), 799 | get_frame_from_str("foo"), 800 | ]; 801 | 802 | assert!(parse(Frame::Array(command)).is_err()) 803 | } 804 | 805 | #[test] 806 | fn parse_given_count_without_keyspace_returns_error() { 807 | let command = vec![get_frame_from_str("count")]; 808 | assert!(parse(Frame::Array(command)).is_err()) 809 | } 810 | 811 | #[test] 812 | fn parse_given_count_returns_count() { 813 | let command = vec![get_frame_from_str("count"), get_frame_from_str("foo")]; 814 | 815 | assert_eq!( 816 | parse(Frame::Array(command)).unwrap(), 817 | Command::Count(Count { 818 | keyspace: Bytes::from("foo"), 819 | }) 820 | ); 821 | } 822 | 823 | #[test] 824 | fn parse_given_tll_without_keyspace_returns_error() { 825 | let command = vec![get_frame_from_str("ttl")]; 826 | assert!(parse(Frame::Array(command)).is_err()) 827 | } 828 | 829 | #[test] 830 | fn parse_given_ttl_without_key_returns_error() { 831 | let command = vec![get_frame_from_str("ttl"), get_frame_from_str("foo")]; 832 | assert!(parse(Frame::Array(command)).is_err()) 833 | } 834 | 835 | #[test] 836 | fn parse_given_ttl_returns_ttl() { 837 | let command = vec![ 838 | get_frame_from_str("ttl"), 839 | get_frame_from_str("foo"), 840 | get_frame_from_str("bar"), 841 | ]; 842 | 843 | assert_eq!( 844 | parse(Frame::Array(command)).unwrap(), 845 | Command::Ttl(Ttl { 846 | keyspace: Bytes::from("foo"), 847 | key: Bytes::from("bar"), 848 | }) 849 | ); 850 | } 851 | --------------------------------------------------------------------------------