File size: 4,451 Bytes
84d2a97
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
use std::collections::HashSet;
use std::mem;
use std::sync::Arc;

use parking_lot::{Mutex, RwLock};
use rocksdb::DB;

use super::rocksdb_wrapper::DatabaseColumnIterator;
use crate::common::operation_error::OperationResult;
use crate::common::rocksdb_wrapper::{DatabaseColumnWrapper, LockedDatabaseColumnWrapper};
use crate::common::Flusher;

/// Wrapper around `DatabaseColumnWrapper` that ensures, that keys that were removed from the
/// database are only persisted on flush explicitly.
///
/// This might be required to guarantee consistency of the database component.
/// E.g. copy-on-write implementation should guarantee that data in the `write` component is
/// persisted before it is removed from the `copy` component.
#[derive(Debug)]
pub struct DatabaseColumnScheduledDeleteWrapper {
    db: DatabaseColumnWrapper,
    deleted_pending_persistence: Arc<Mutex<HashSet<Vec<u8>>>>,
}

impl Clone for DatabaseColumnScheduledDeleteWrapper {
    fn clone(&self) -> Self {
        Self {
            db: self.db.clone(),
            deleted_pending_persistence: self.deleted_pending_persistence.clone(),
        }
    }
}

impl DatabaseColumnScheduledDeleteWrapper {
    pub fn new(db: DatabaseColumnWrapper) -> Self {
        Self {
            db,
            deleted_pending_persistence: Arc::new(Mutex::new(HashSet::new())),
        }
    }

    pub fn put<K, V>(&self, key: K, value: V) -> OperationResult<()>
    where
        K: AsRef<[u8]>,
        V: AsRef<[u8]>,
    {
        self.deleted_pending_persistence.lock().remove(key.as_ref());
        self.db.put(key, value)
    }

    pub fn remove<K>(&self, key: K) -> OperationResult<()>
    where
        K: AsRef<[u8]>,
    {
        self.deleted_pending_persistence
            .lock()
            .insert(key.as_ref().to_vec());
        Ok(())
    }

    fn is_pending_removal<K>(&self, key: K) -> bool
    where
        K: AsRef<[u8]>,
    {
        self.deleted_pending_persistence
            .lock()
            .contains(key.as_ref())
    }

    pub fn flusher(&self) -> Flusher {
        let ids_to_delete = mem::take(&mut *self.deleted_pending_persistence.lock());
        let wrapper = self.db.clone();
        Box::new(move || {
            for id in ids_to_delete {
                wrapper.remove(id)?;
            }
            wrapper.flusher()()
        })
    }

    pub fn lock_db(&self) -> LockedDatabaseColumnScheduledDeleteWrapper<'_> {
        LockedDatabaseColumnScheduledDeleteWrapper {
            base: self.db.lock_db(),
            deleted_pending_persistence: &self.deleted_pending_persistence,
        }
    }

    pub fn get_pinned<T, F>(&self, key: &[u8], f: F) -> OperationResult<Option<T>>
    where
        F: FnOnce(&[u8]) -> T,
    {
        if self.is_pending_removal(key) {
            return Ok(None);
        }
        self.db.get_pinned(key, f)
    }

    pub fn recreate_column_family(&self) -> OperationResult<()> {
        self.db.recreate_column_family()
    }

    pub fn get_database(&self) -> Arc<RwLock<DB>> {
        self.db.get_database()
    }

    pub fn get_column_name(&self) -> &str {
        self.db.get_column_name()
    }

    pub fn has_column_family(&self) -> OperationResult<bool> {
        self.db.has_column_family()
    }

    pub fn remove_column_family(&self) -> OperationResult<()> {
        self.db.remove_column_family()
    }
}

pub struct LockedDatabaseColumnScheduledDeleteWrapper<'a> {
    base: LockedDatabaseColumnWrapper<'a>,
    deleted_pending_persistence: &'a Mutex<HashSet<Vec<u8>>>,
}

impl LockedDatabaseColumnScheduledDeleteWrapper<'_> {
    pub fn iter(&self) -> OperationResult<DatabaseColumnScheduledDeleteIterator<'_>> {
        Ok(DatabaseColumnScheduledDeleteIterator {
            base: self.base.iter()?,
            deleted_pending_persistence: self.deleted_pending_persistence,
        })
    }
}

pub struct DatabaseColumnScheduledDeleteIterator<'a> {
    base: DatabaseColumnIterator<'a>,
    deleted_pending_persistence: &'a Mutex<HashSet<Vec<u8>>>,
}

impl<'a> Iterator for DatabaseColumnScheduledDeleteIterator<'a> {
    type Item = (Box<[u8]>, Box<[u8]>);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let (key, value) = self.base.next()?;
            if !self
                .deleted_pending_persistence
                .lock()
                .contains(key.as_ref())
            {
                return Some((key, value));
            }
        }
    }
}