title: mini-redis项目-2-存储层 toc: true cover: 'https://img.paulzzh.com/touhou/random?2' date: 2022-12-05 19:58:43 categories: Rust tags: [Rust, 数据库, Redis]
本篇接上一篇《mini-redis项目-1-简介》，讲解mini-redis存储层的实现；
源代码：
系列文章：
<br/>
<!--more-->在 mini-redis 中为了简单起见，数据是直接通过 HashMap 中存储了；
对于实际生产环境中使用的数据库来说，通常会使用：
<br/>
数据操作接口被定义在 traits.rs 文件的 KvStore Trait 中；
src/storage/traits.rs
pub trait KvStore {
fn get(&self, key: &str) -> Option<Bytes>;
/// Set the value associated with a key along with an optional expiration
/// Duration.
///
/// If a value is already associated with the key, it is removed.
fn set(&self, key: String, value: Bytes, expire: Option<Duration>);
/// Returns a `Receiver` for the requested channel.
///
/// The returned `Receiver` is used to receive values broadcast by `PUBLISH`
/// commands.
fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes>;
/// Publish a message to the channel. Returns the number of subscribers
/// listening on the channel.
fn publish(&self, key: &str, value: Bytes) -> usize;
}
主要是 get、set、pub、sub 四个接口，非常的简单；
其中 set 可以设置键的过期时间；
<br/>
Store 结构体中存放实际的数据，定义如下：
src/storage/store.rs
#[derive(Debug)]
pub(crate) struct Store {
pub(crate) entries: HashMap<String, Entry>,
pub(crate) pub_sub: HashMap<String, broadcast::Sender<Bytes>>,
pub(crate) expirations: BTreeMap<(Instant, u64), String>,
pub(crate) next_id: u64,
pub(crate) shutdown: bool,
}
#[derive(Debug)]
pub(crate) struct Entry {
pub(crate) id: u64,
pub(crate) data: Bytes,
pub(crate) expires_at: Option<Instant>,
}
impl Store {
pub(crate) fn new() -> Store {
Store {
entries: HashMap::new(),
pub_sub: HashMap::new(),
expirations: BTreeMap::new(),
next_id: 0,
shutdown: false,
}
}
pub(crate) fn next_expiration(&self) -> Option<Instant> {
self.expirations.keys().next().map(|expire| expire.0)
}
}
Store 的逻辑也非常简单：
这里的 Store 层用了
pub(crate)即整个 crate 可见，是一个简化操作；个人感觉应当将 Store 进行封装，这样如果实际上底层的存储结构变化，对于上次的操作是无感知的！
<br/>
Db 封装了上面的 Store，并实现了 KvStore Trait，对外提供了一组数据操作接口，下面来看具体实现；
Db实现：
src/storage/db.rs
#[derive(Debug, Clone)]
pub(crate) struct Db {
shared: Arc<SharedDb>,
}
impl Db {
pub(crate) fn new() -> Db {
let shared = Arc::new(SharedDb::new());
tokio::spawn(Db::purge_expired_tasks(shared.clone()));
Db { shared }
}
async fn purge_expired_tasks(shared: Arc<SharedDb>) {
while !shared.is_shutdown() {
if let Some(when) = shared.purge_expired_keys() {
tokio::select! {
_ = time::sleep_until(when) => {}
_ = shared.background_task.notified() => {}
}
} else {
shared.background_task.notified().await;
}
}
info!("Purge background task shut down")
}
fn shutdown_purge_task(&self) {
let mut store = self.shared.store.lock().unwrap();
store.shutdown = true;
drop(store);
self.shared.background_task.notify_one();
}
}
由于可能会存在多个异步任务同时获取数据库的场景（例如：键过期、Channel订阅都会用到数据库）；
因此需要将 Store 进行封装至 Arc 中，同时封装了 SharedDb（下文会讲）；
Db 中定义了 new 方法，用于初始化数据库，并且启动 purge_expired_tasks 任务；
当服务器未关闭时：
purge_expired_tasks 会调用 shared.purge_expired_keys 清除最近的过期的键，并获取下一个最新过期的键的时间并等待；shared.background_task 再次通知了当前异步任务，此时大概率是加入了更早的过期键，需要重新调用 shared.purge_expired_keys 获取最新的最近的过期时间；shared.background_task.notified().await；这样，就完成了过期键的处理！
同时，Db 定义了 shutdown_purge_task 方法，提供给 DbDropGuard，用于在服务器停止后清理资源以及其他任务：
drop(store) 就显式释放了数据库资源，但是这个方法不能由在 Drop Db 的时候调用，因为还需要停止其他异步任务；self.shared.background_task.notify_one() 通知异步过期键任务，此时由于 shutdown 变为了 true，异步任务会直接退出！<br/>
前面提到了可能会存在多个异步任务同时获取数据库，因此将 Db 封装至了SharedDb 中：
src/storage/db.rs
#[derive(Debug)]
struct SharedDb {
store: Mutex<Store>,
background_task: Notify,
}
impl SharedDb {
fn new() -> Self {
SharedDb {
store: Mutex::new(Store::new()),
background_task: Notify::new(),
}
}
fn purge_expired_keys(&self) -> Option<Instant> {
let mut store = self.store.lock().unwrap();
if store.shutdown {
return None;
}
let store = &mut *store;
let now = Instant::now();
while let Some((&(when, id), key)) = store.expirations.iter().next() {
if when > now {
return Some(when);
}
store.entries.remove(key);
store.expirations.remove(&(when, id));
}
None
}
fn is_shutdown(&self) -> bool {
self.store.lock().unwrap().shutdown
}
}
SharedDb 中的内容非常简单：
store: Mutex<Store>：互斥锁保护底层数据不会并发访问；background_task: Notify：和前面提到的异步任务交互；同时定义了几个方法，new、is_shutdown 很简单不说了；
和上文对应，purge_expired_keys 方法由 Db 中的异步任务调用，在服务器没有关闭时，通过 while let Some((&(when, id), key)) = store.expirations.iter().next() 按照过期时间获取过期键：
<br/>
下面是 Db 实现的 KvStore Trait：
src/storage/db.rs
impl KvStore for Db {
fn get(&self, key: &str) -> Option<Bytes> {
let store = self.shared.store.lock().unwrap();
store.entries.get(key).map(|entry| entry.data.clone())
}
fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
let mut store = self.shared.store.lock().unwrap();
let id = store.next_id;
store.next_id += 1;
let mut notify = false;
let expires_at = expire.map(|duration| {
let when = Instant::now() + duration;
notify = store
.next_expiration()
.map(|expiration| expiration > when)
.unwrap_or(true);
store.expirations.insert((when, id), key.clone());
when
});
let prev = store.entries.insert(
key,
Entry {
id,
data: value,
expires_at,
},
);
if let Some(prev) = prev {
if let Some(when) = prev.expires_at {
// clear expiration
store.expirations.remove(&(when, prev.id));
}
}
drop(store);
if notify {
self.shared.background_task.notify_one();
}
}
fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes> {
use std::collections::hash_map::Entry;
let mut store = self.shared.store.lock().unwrap();
match store.pub_sub.entry(key) {
Entry::Occupied(e) => e.get().subscribe(),
Entry::Vacant(e) => {
let (tx, rx) = broadcast::channel(1024);
e.insert(tx);
rx
}
}
}
fn publish(&self, key: &str, value: Bytes) -> usize {
debug!("publish: (key={}, len(value)={})", key, value.len());
let state = self.shared.store.lock().unwrap();
state
.pub_sub
.get(key)
.map(|tx| tx.send(value).unwrap_or(0))
.unwrap_or(0)
}
}
下面分别来看；
<br/>
Get 方法比较简单：
fn get(&self, key: &str) -> Option<Bytes> {
// Acquire the lock, get the entry and clone the value.
//
// Because data is stored using `Bytes`, a clone here is a shallow
// clone. Data is not copied.
let store = self.shared.store.lock().unwrap();
store.entries.get(key).map(|entry| entry.data.clone())
}
首先获取存储的锁，然后通过 store 中的 HashMap entries 直接获取键值即可；
<font color="#f00">**这里注意：我们使用的是 `bytes::Bytes`，因此这里 clone 的实际上是一个胖指针，而非将整个数据copy了一份！**</font>
<br/>
Set方法如下：
fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
let mut store = self.shared.store.lock().unwrap();
let id = store.next_id;
store.next_id += 1;
let mut notify = false;
let expires_at = expire.map(|duration| {
let when = Instant::now() + duration;
notify = store
.next_expiration()
.map(|expiration| expiration > when)
.unwrap_or(true);
store.expirations.insert((when, id), key.clone());
when
});
let prev = store.entries.insert(
key,
Entry {
id,
data: value,
expires_at,
},
);
if let Some(prev) = prev {
if let Some(when) = prev.expires_at {
store.expirations.remove(&(when, prev.id));
}
}
drop(store);
if notify {
self.shared.background_task.notify_one();
}
}
我们为当前 set 的 key：
purge_expired_keys 方法，重新刷新过期任务！随后，将键加入 store.entries 中，如果是过期键还会加入 store.expirations 中；
然后，需要判断如果 set 的 key 之前在数据库中存在了，则需要将前一个过期任务删除；
最后，显式 drop 掉我们获取的 store 的锁，这样最后通知刷新过期键异步任务后，这个异步任务不用再等待这个函数释放锁了～
<br/>
subscribe方法如下：
fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes> {
use std::collections::hash_map::Entry;
let mut store = self.shared.store.lock().unwrap();
match store.pub_sub.entry(key) {
Entry::Occupied(e) => e.get().subscribe(),
Entry::Vacant(e) => {
let (tx, rx) = broadcast::channel(1024);
e.insert(tx);
rx
}
}
}
subscribe 方法的逻辑也非常简单：
subscribe 了，则直接返回一个新的 Receiver 即可；subscribe 则创建一个新的channel，并将 Receiver 加入 store.pub_sub 存储后返回Receiver即可；<br/>
有了 Receiver，则 publish 方法的实现更加简单：
fn publish(&self, key: &str, value: Bytes) -> usize {
let state = self.shared.store.lock().unwrap();
state
.pub_sub
.get(key)
.map(|tx| tx.send(value).unwrap_or(0))
.unwrap_or(0)
}
直接获取 state.pub_sub 中 key 对应的 channel 并调用 send 方法即可；
需要注意的是：如果不存在这个 key，或者发生失败，则返回 0 大小的消息，表示未发送成功；
<br/>
最后就是 DbDropGuard 的实现了，DbDropGuard 实际上是 Db 的一个封装；
当通过 DbDropGuard 获取 Db 时，实际上是一个 Db 的 clone，而 Db 内部是一个 Arc 计数智能指针；
因此，当所有使用 Db 都被回收，仍然需要 DbDropGuard 被 drop 后才会对数据库进行回收（DbDropGuard 中持有第一个 Db）：
src/storage/db.rs
/// A wrapper around a `Db` instance. This exists to allow orderly cleanup
/// of the `Db` by signalling the background purge task to shut down when
/// this struct is dropped.
#[derive(Debug)]
pub(crate) struct DbDropGuard {
/// The `Db` instance that will be shut down when this `DbHolder` struct
/// is dropped.
db: Db,
}
impl DbDropGuard {
/// Create a new `DbHolder`, wrapping a `Db` instance. When this is dropped
/// the `Db`'s purge task will be shut down.
pub(crate) fn new() -> DbDropGuard {
DbDropGuard { db: Db::new() }
}
/// Get the shared database. Internally, this is an `Arc`,
/// so a clone only increments the ref count.
pub(crate) fn db(&self) -> Db {
self.db.clone()
}
}
impl Drop for DbDropGuard {
fn drop(&mut self) {
// Signal the 'Db' instance to shut down the task that purges expired keys
self.db.shutdown_purge_task();
}
}
同时，为 DbDropGuard 实现了 Drop Trait，当 DbDropGuard 被回收后，会调用前文中提到的 db 中的 shutdown_purge_task 方法进行资源回收；
<br/>
本小节完成了对 mini-redis 存储层的实现，对外封装了 DbDropGuard；
在使用时只需创建 DbDropGuard 对象，然后通过其内部封装的智能指针 Db 操作底层的 Store 即可！
实际上，真实的存储结构远远要比这里的 HashMap 要复杂，但是本文也能起到一定抛砖引玉的作用！
<br/>
源代码：
系列文章：
<br/>