前言#

1
[dev-dependencies]
2
assert_cmd = "0.11.0"
3
predicates = "1.0.0"
4
tempfile = "3.0.7"
5
walkdir = "2.2.7"

1
/// 若仅让tests能够运行，则需要修改的地方很少
2

3
// 添加`PathBuf`导入
4
use std::{collections::HashMap, path::PathBuf};
5

6
// 添加Result<T>类型
7
pub type Result<T> = std::result::Result<T, ()>;
8

9
// 修改`KvStore`中函数的返回类型
10
impl KvStore {
11
    //...
12
    pub fn open(path: impl Into<PathBuf>) -> Result<KvStore> {
13
        Ok(Self::new())
14
    }
15
    //...
16
    pub fn get(&self, key: String) -> Result<Option<String>> {
17
        match self.db.get(&key) {
18
            Some(value) => Ok(Some(value.to_string())),
19
            None => Err(())
20
        }
21
    }
22
    //...
23
    pub fn set(&mut self, key: String, value: String) -> Result<()> {
24
        self.db.insert(key, value);
25
        Ok(())
26
    }
27
    //...
28
    pub fn remove(&mut self, key: String) -> Result<()>{
29
        self.db.remove(&key);
30
        Ok(())
31
    }

1
...
2
running 17 tests
3
test cli_invalid_subcommand ... ok
4
test cli_get_non_existent_key ... FAILED
5
test cli_rm_non_existent_key ... FAILED
6
test cli_get_stored ... FAILED
7
test cli_no_args ... ok
8
test get_non_existent_value ... FAILED
9
test get_stored_value ... FAILED
10
test compaction ... FAILED
11
test cli_invalid_get ... ok
12
test overwrite_value ... FAILED
13
test remove_non_existent_key ... FAILED
14
test cli_version ... ok
15
test remove_key ... FAILED
16
test cli_invalid_rm ... ok
17
test cli_set ... FAILED
18
test cli_rm_stored ... FAILED
19
test cli_invalid_set ... ok
20
...

1
//...
2
mod error;
3
pub use error::{Result, KvsError};
4
//...
5
impl KvStore {
6
    //...
7
    pub fn get(&self, key: String) -> Result<Option<String>> {
8
        match self.db.get(&key) {
9
            Some(value) => Ok(Some(value.to_string())),
10
            None => Err(KvsError::KeyNotFound(key))
11
        }
12
    }
13
    //...
14
}

1
use failure::Fail;
2
use std::io;
3

4
/// custom error type of kvs
5
#[derive(Fail, Debug)]
6
pub enum KvsError {
7
    /// Key not found in kvs index
8
    #[fail(display="Key `{}` not found", _0)]
9
    KeyNotFound(String),
10
}
11

12
/// simplify `Result<T, KvsError>` to `Result<T>`
13
pub type Result<T> = std::result::Result<T, KvsError>;

1
.
2
├── Cargo.lock
3
├── Cargo.toml
4
├── README.md
5
├── src
6
│  ├── bin
7
│  │  └── kvs.rs
8
│  ├── error.rs
9
│  ├── kv.rs
10
│  └── lib.rs
11
└── tests
12
   └── tests.rs

1
#![deny(missing_docs)]
2
//! A simple key/value store.
3

4
pub use error::{KvsError, Result};
5
pub use kv::KvStore;
6

7
mod error;
8
mod kv;

1
use std::collections::{BTreeMap, HashMap};
2
use std::fs::{self, File, OpenOptions};
3
use std::io::{self, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
4
use std::ops::Range;
5
use std::path::{Path, PathBuf};
6

7
use serde::{Deserialize, Serialize};
8
use serde_json::Deserializer;
9

10
use crate::{KvsError, Result};
11
use std::ffi::OsStr;
12

13
const COMPACTION_THRESHOLD: u64 = 1024 * 1024;
14

15
/// The `KvStore` stores string key/value pairs.
16
///
17
/// Key/value pairs are persisted to disk in log files. Log files are named after
18
/// monotonically increasing generation numbers with a `log` extension name.
19
/// A `BTreeMap` in memory stores the keys and the value locations for fast query.
20
///
21
/// ```rust
22
/// # use kvs::{KvStore, Result};
23
/// # fn try_main() -> Result<()> {
24
/// use std::env::current_dir;
25
/// let mut store = KvStore::open(current_dir()?)?;
26
/// store.set("key".to_owned(), "value".to_owned())?;
27
/// let val = store.get("key".to_owned())?;
28
/// assert_eq!(val, Some("value".to_owned()));
29
/// # Ok(())
30
/// # }
31
/// ```
32
pub struct KvStore {
33
    // directory for the log and other data.
34
    path: PathBuf,
35
    // map generation number to the file reader.
36
    readers: HashMap<u64, BufReaderWithPos<File>>,
37
    // writer of the current log.
38
    writer: BufWriterWithPos<File>,
39
    current_gen: u64,
40
    index: BTreeMap<String, CommandPos>,
41
    // the number of bytes representing "stale" commands that could be
42
    // deleted during a compaction.
43
    uncompacted: u64,
44
}
45

46
impl KvStore {
47
    /// Opens a `KvStore` with the given path.
48
    ///
49
    /// This will create a new directory if the given one does not exist.
50
    ///
51
    /// # Errors
52
    ///
53
    /// It propagates I/O or deserialization errors during the log replay.
54
    pub fn open(path: impl Into<PathBuf>) -> Result<KvStore> {
55
        let path = path.into();
56
        fs::create_dir_all(&path)?;
57

58
        let mut readers = HashMap::new();
59
        let mut index = BTreeMap::new();
60

61
        let gen_list = sorted_gen_list(&path)?;
62
        let mut uncompacted = 0;
63

64
        for &gen in &gen_list {
65
            let mut reader = BufReaderWithPos::new(File::open(log_path(&path, gen))?)?;
66
            uncompacted += load(gen, &mut reader, &mut index)?;
67
            readers.insert(gen, reader);
68
        }
69

70
        let current_gen = gen_list.last().unwrap_or(&0) + 1;
71
        let writer = new_log_file(&path, current_gen, &mut readers)?;
72

73
        Ok(KvStore {
74
            path,
75
            readers,
76
            writer,
77
            current_gen,
78
            index,
79
            uncompacted,
80
        })
81
    }
82

83
    /// Sets the value of a string key to a string.
84
    ///
85
    /// If the key already exists, the previous value will be overwritten.
86
    ///
87
    /// # Errors
88
    ///
89
    /// It propagates I/O or serialization errors during writing the log.
90
    pub fn set(&mut self, key: String, value: String) -> Result<()> {
91
        let cmd = Command::set(key, value);
92
        let pos = self.writer.pos;
93
        serde_json::to_writer(&mut self.writer, &cmd)?;
94
        self.writer.flush()?;
95
        if let Command::Set { key, .. } = cmd {
96
            if let Some(old_cmd) = self
97
                .index
98
                .insert(key, (self.current_gen, pos..self.writer.pos).into())
99
            {
100
                self.uncompacted += old_cmd.len;
101
            }
102
        }
103

104
        if self.uncompacted > COMPACTION_THRESHOLD {
105
            self.compact()?;
106
        }
107
        Ok(())
108
    }
109

110
    /// Gets the string value of a given string key.
111
    ///
112
    /// Returns `None` if the given key does not exist.
113
    ///
114
    /// # Errors
115
    ///
116
    /// It returns `KvsError::UnexpectedCommandType` if the given command type unexpected.
117
    pub fn get(&mut self, key: String) -> Result<Option<String>> {
118
        if let Some(cmd_pos) = self.index.get(&key) {
119
            let reader = self
120
                .readers
121
                .get_mut(&cmd_pos.gen)
122
                .expect("Cannot find log reader");
123
            reader.seek(SeekFrom::Start(cmd_pos.pos))?;
124
            let cmd_reader = reader.take(cmd_pos.len);
125
            if let Command::Set { value, .. } = serde_json::from_reader(cmd_reader)? {
126
                Ok(Some(value))
127
            } else {
128
                Err(KvsError::UnexpectedCommandType)
129
            }
130
        } else {
131
            Ok(None)
132
        }
133
    }
134

135
    /// Removes a given key.
136
    ///
137
    /// # Errors
138
    ///
139
    /// It returns `KvsError::KeyNotFound` if the given key is not found.
140
    ///
141
    /// It propagates I/O or serialization errors during writing the log.
142
    pub fn remove(&mut self, key: String) -> Result<()> {
143
        if self.index.contains_key(&key) {
144
            let cmd = Command::remove(key);
145
            serde_json::to_writer(&mut self.writer, &cmd)?;
146
            self.writer.flush()?;
147
            if let Command::Remove { key } = cmd {
148
                let old_cmd = self.index.remove(&key).expect("key not found");
149
                self.uncompacted += old_cmd.len;
150
            }
151
            Ok(())
152
        } else {
153
            Err(KvsError::KeyNotFound)
154
        }
155
    }
156

157
    /// Clears stale entries in the log.
158
    pub fn compact(&mut self) -> Result<()> {
159
        // increase current gen by 2. current_gen + 1 is for the compaction file.
160
        let compaction_gen = self.current_gen + 1;
161
        self.current_gen += 2;
162
        self.writer = self.new_log_file(self.current_gen)?;
163

164
        let mut compaction_writer = self.new_log_file(compaction_gen)?;
165

166
        let mut new_pos = 0; // pos in the new log file.
167
        for cmd_pos in &mut self.index.values_mut() {
168
            let reader = self
169
                .readers
170
                .get_mut(&cmd_pos.gen)
171
                .expect("Cannot find log reader");
172
            if reader.pos != cmd_pos.pos {
173
                reader.seek(SeekFrom::Start(cmd_pos.pos))?;
174
            }
175

176
            let mut entry_reader = reader.take(cmd_pos.len);
177
            let len = io::copy(&mut entry_reader, &mut compaction_writer)?;
178
            *cmd_pos = (compaction_gen, new_pos..new_pos + len).into();
179
            new_pos += len;
180
        }
181
        compaction_writer.flush()?;
182

183
        // remove stale log files.
184
        let stale_gens: Vec<_> = self
185
            .readers
186
            .keys()
187
            .filter(|&&gen| gen < compaction_gen)
188
            .cloned()
189
            .collect();
190
        for stale_gen in stale_gens {
191
            self.readers.remove(&stale_gen);
192
            fs::remove_file(log_path(&self.path, stale_gen))?;
193
        }
194
        self.uncompacted = 0;
195

196
        Ok(())
197
    }
198

199
    /// Create a new log file with given generation number and add the reader to the readers map.
200
    ///
201
    /// Returns the writer to the log.
202
    fn new_log_file(&mut self, gen: u64) -> Result<BufWriterWithPos<File>> {
203
        new_log_file(&self.path, gen, &mut self.readers)
204
    }
205
}
206

207
/// Create a new log file with given generation number and add the reader to the readers map.
208
///
209
/// Returns the writer to the log.
210
fn new_log_file(
211
    path: &Path,
212
    gen: u64,
213
    readers: &mut HashMap<u64, BufReaderWithPos<File>>,
214
) -> Result<BufWriterWithPos<File>> {
215
    let path = log_path(&path, gen);
216
    let writer = BufWriterWithPos::new(
217
        OpenOptions::new()
218
            .create(true)
219
            .write(true)
220
            .append(true)
221
            .open(&path)?,
222
    )?;
223
    readers.insert(gen, BufReaderWithPos::new(File::open(&path)?)?);
224
    Ok(writer)
225
}
226

227
/// Returns sorted generation numbers in the given directory.
228
fn sorted_gen_list(path: &Path) -> Result<Vec<u64>> {
229
    let mut gen_list: Vec<u64> = fs::read_dir(&path)?
230
        .flat_map(|res| -> Result<_> { Ok(res?.path()) })
231
        .filter(|path| path.is_file() && path.extension() == Some("log".as_ref()))
232
        .flat_map(|path| {
233
            path.file_name()
234
                .and_then(OsStr::to_str)
235
                .map(|s| s.trim_end_matches(".log"))
236
                .map(str::parse::<u64>)
237
        })
238
        .flatten()
239
        .collect();
240
    gen_list.sort_unstable();
241
    Ok(gen_list)
242
}
243

244
/// Load the whole log file and store value locations in the index map.
245
///
246
/// Returns how many bytes can be saved after a compaction.
247
fn load(
248
    gen: u64,
249
    reader: &mut BufReaderWithPos<File>,
250
    index: &mut BTreeMap<String, CommandPos>,
251
) -> Result<u64> {
252
    // To make sure we read from the beginning of the file.
253
    let mut pos = reader.seek(SeekFrom::Start(0))?;
254
    let mut stream = Deserializer::from_reader(reader).into_iter::<Command>();
255
    let mut uncompacted = 0; // number of bytes that can be saved after a compaction.
256
    while let Some(cmd) = stream.next() {
257
        let new_pos = stream.byte_offset() as u64;
258
        match cmd? {
259
            Command::Set { key, .. } => {
260
                if let Some(old_cmd) = index.insert(key, (gen, pos..new_pos).into()) {
261
                    uncompacted += old_cmd.len;
262
                }
263
            }
264
            Command::Remove { key } => {
265
                if let Some(old_cmd) = index.remove(&key) {
266
                    uncompacted += old_cmd.len;
267
                }
268
                // the "remove" command itself can be deleted in the next compaction.
269
                // so we add its length to `uncompacted`.
270
                uncompacted += new_pos - pos;
271
            }
272
        }
273
        pos = new_pos;
274
    }
275
    Ok(uncompacted)
276
}
277

278
fn log_path(dir: &Path, gen: u64) -> PathBuf {
279
    dir.join(format!("{}.log", gen))
280
}
281

282
/// Struct representing a command.
283
#[derive(Serialize, Deserialize, Debug)]
284
enum Command {
285
    Set { key: String, value: String },
286
    Remove { key: String },
287
}
288

289
impl Command {
290
    fn set(key: String, value: String) -> Command {
291
        Command::Set { key, value }
292
    }
293

294
    fn remove(key: String) -> Command {
295
        Command::Remove { key }
296
    }
297
}
298

299
/// Represents the position and length of a json-serialized command in the log.
300
struct CommandPos {
301
    gen: u64,
302
    pos: u64,
303
    len: u64,
304
}
305

306
impl From<(u64, Range<u64>)> for CommandPos {
307
    fn from((gen, range): (u64, Range<u64>)) -> Self {
308
        CommandPos {
309
            gen,
310
            pos: range.start,
311
            len: range.end - range.start,
312
        }
313
    }
314
}
315

316
struct BufReaderWithPos<R: Read + Seek> {
317
    reader: BufReader<R>,
318
    pos: u64,
319
}
320

321
impl<R: Read + Seek> BufReaderWithPos<R> {
322
    fn new(mut inner: R) -> Result<Self> {
323
        let pos = inner.seek(SeekFrom::Current(0))?;
324
        Ok(BufReaderWithPos {
325
            reader: BufReader::new(inner),
326
            pos,
327
        })
328
    }
329
}
330

331
impl<R: Read + Seek> Read for BufReaderWithPos<R> {
332
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
333
        let len = self.reader.read(buf)?;
334
        self.pos += len as u64;
335
        Ok(len)
336
    }
337
}
338

339
impl<R: Read + Seek> Seek for BufReaderWithPos<R> {
340
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
341
        self.pos = self.reader.seek(pos)?;
342
        Ok(self.pos)
343
    }
344
}
345

346
struct BufWriterWithPos<W: Write + Seek> {
347
    writer: BufWriter<W>,
348
    pos: u64,
349
}
350

351
impl<W: Write + Seek> BufWriterWithPos<W> {
352
    fn new(mut inner: W) -> Result<Self> {
353
        let pos = inner.seek(SeekFrom::Current(0))?;
354
        Ok(BufWriterWithPos {
355
            writer: BufWriter::new(inner),
356
            pos,
357
        })
358
    }
359
}
360

361
impl<W: Write + Seek> Write for BufWriterWithPos<W> {
362
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
363
        let len = self.writer.write(buf)?;
364
        self.pos += len as u64;
365
        Ok(len)
366
    }
367

368
    fn flush(&mut self) -> io::Result<()> {
369
        self.writer.flush()
370
    }
371
}
372

373
impl<W: Write + Seek> Seek for BufWriterWithPos<W> {
374
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
375
        self.pos = self.writer.seek(pos)?;
376
        Ok(self.pos)
377
    }
378
}

1
use failure::Fail;
2
use std::io;
3

4
/// Error type for kvs.
5
#[derive(Fail, Debug)]
6
pub enum KvsError {
7
    /// IO error.
8
    #[fail(display = "{}", _0)]
9
    Io(#[cause] io::Error),
10
    /// Serialization or deserialization error.
11
    #[fail(display = "{}", _0)]
12
    Serde(#[cause] serde_json::Error),
13
    /// Removing non-existent key error.
14
    #[fail(display = "Key not found")]
15
    KeyNotFound,
16
    /// Unexpected command type error.
17
    /// It indicated a corrupted log or a program bug.
18
    #[fail(display = "Unexpected command type")]
19
    UnexpectedCommandType,
20
}
21

22
impl From<io::Error> for KvsError {
23
    fn from(err: io::Error) -> KvsError {
24
        KvsError::Io(err)
25
    }
26
}
27

28
impl From<serde_json::Error> for KvsError {
29
    fn from(err: serde_json::Error) -> KvsError {
30
        KvsError::Serde(err)
31
    }
32
}
33

34
/// Result type for kvs.
35
pub type Result<T> = std::result::Result<T, KvsError>;