linkrs_sync/batch/
buffer.rs

1use dashmap::DashMap;
2use std::time::Instant;
3
4use super::trait_def::BatchBuffer;
5use crate::types::IndexOperation;
6
7type IndexKey = (u64, String, String);
8
9#[derive(Debug, Default, Clone)]
10pub struct BufferEntry {
11    pub inserts: Vec<IndexOperation>,
12    pub deletes: Vec<String>,
13}
14
15#[derive(Debug)]
16pub struct OpBatchBuffer {
17    buffers: DashMap<IndexKey, BufferEntry>,
18    last_commit: DashMap<IndexKey, Instant>,
19}
20
21impl Default for OpBatchBuffer {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl OpBatchBuffer {
28    pub fn new() -> Self {
29        Self {
30            buffers: DashMap::new(),
31            last_commit: DashMap::new(),
32        }
33    }
34
35    pub fn add_insert(&self, key: &IndexKey, operation: IndexOperation) {
36        let mut entry = self.buffers.entry(key.clone()).or_default();
37        entry.inserts.push(operation);
38        self.last_commit
39            .entry(key.clone())
40            .or_insert_with(Instant::now);
41    }
42
43    pub fn add_delete(&self, key: &IndexKey, id: String) {
44        let mut entry = self.buffers.entry(key.clone()).or_default();
45        entry.deletes.push(id);
46        self.last_commit
47            .entry(key.clone())
48            .or_insert_with(Instant::now);
49    }
50
51    pub fn drain_inserts(&self, key: &IndexKey) -> Vec<IndexOperation> {
52        if let Some(mut entry) = self.buffers.get_mut(key) {
53            std::mem::take(&mut entry.inserts)
54        } else {
55            Vec::new()
56        }
57    }
58
59    pub fn drain_deletes(&self, key: &IndexKey) -> Vec<String> {
60        if let Some(mut entry) = self.buffers.get_mut(key) {
61            std::mem::take(&mut entry.deletes)
62        } else {
63            Vec::new()
64        }
65    }
66
67    pub fn drain_all(&self, key: &IndexKey) -> BufferEntry {
68        self.buffers.remove(key).map(|(_, v)| v).unwrap_or_default()
69    }
70
71    pub fn peek_entry(&self, key: &IndexKey) -> BufferEntry {
72        self.buffers
73            .get(key)
74            .map(|entry| BufferEntry {
75                inserts: entry.inserts.clone(),
76                deletes: entry.deletes.clone(),
77            })
78            .unwrap_or_default()
79    }
80
81    pub fn re_enqueue(&self, key: &IndexKey, entry: BufferEntry) {
82        if entry.inserts.is_empty() && entry.deletes.is_empty() {
83            return;
84        }
85        let mut buffer = self.buffers.entry(key.clone()).or_default();
86        buffer.inserts.extend(entry.inserts);
87        buffer.deletes.extend(entry.deletes);
88    }
89
90    pub fn count(&self, key: &IndexKey) -> usize {
91        self.buffers
92            .get(key)
93            .map(|e| e.inserts.len() + e.deletes.len())
94            .unwrap_or(0)
95    }
96
97    pub fn insert_count(&self, key: &IndexKey) -> usize {
98        self.buffers.get(key).map(|e| e.inserts.len()).unwrap_or(0)
99    }
100
101    pub fn delete_count(&self, key: &IndexKey) -> usize {
102        self.buffers.get(key).map(|e| e.deletes.len()).unwrap_or(0)
103    }
104
105    pub fn is_timeout(&self, key: &IndexKey, timeout: std::time::Duration) -> bool {
106        self.last_commit
107            .get(key)
108            .map(|last| last.elapsed() >= timeout)
109            .unwrap_or(false)
110    }
111
112    pub fn update_commit_time(&self, key: &IndexKey) {
113        self.last_commit.insert(key.clone(), Instant::now());
114    }
115
116    pub fn keys(&self) -> Vec<IndexKey> {
117        self.buffers.iter().map(|e| e.key().clone()).collect()
118    }
119
120    pub fn total_count(&self) -> usize {
121        self.buffers
122            .iter()
123            .map(|e| e.inserts.len() + e.deletes.len())
124            .sum()
125    }
126
127    pub fn clear(&self) {
128        self.buffers.clear();
129        self.last_commit.clear();
130    }
131
132    pub fn remove(&self, key: &IndexKey) -> Option<BufferEntry> {
133        self.buffers.remove(key).map(|(_, v)| v)
134    }
135}
136
137impl BatchBuffer<IndexKey, IndexOperation> for OpBatchBuffer {
138    fn add(&self, key: &IndexKey, value: IndexOperation) {
139        self.add_insert(key, value);
140    }
141
142    fn drain(&self, key: &IndexKey) -> Vec<IndexOperation> {
143        self.drain_inserts(key)
144    }
145
146    fn peek(&self, key: &IndexKey) -> Vec<IndexOperation> {
147        self.peek_entry(key).inserts
148    }
149
150    fn count(&self, key: &IndexKey) -> usize {
151        self.insert_count(key)
152    }
153
154    fn is_empty(&self, key: &IndexKey) -> bool {
155        self.insert_count(key) == 0
156    }
157
158    fn keys(&self) -> Vec<IndexKey> {
159        self.keys()
160    }
161
162    fn clear(&self) {
163        self.clear();
164    }
165
166    fn total_count(&self) -> usize {
167        self.total_count()
168    }
169}
linkrs_sync/batch/buffer.rs

linkrs_sync/batch/
buffer.rs
