linkrs_sync/batch/
transaction_buffer.rs

1use dashmap::DashMap;
2
3use super::trait_def::BatchBuffer;
4use crate::batch::error::BatchResult;
5use crate::types::{IndexOpKey, IndexOperation};
6use linkrs_core::types::TransactionId;
7
8#[derive(Debug, Default)]
9pub struct TransactionBufferEntry {
10    pub operations: Vec<SequencedIndexOperation>,
11}
12
13#[derive(Debug, Clone)]
14pub struct SequencedIndexOperation {
15    pub sequence: u64,
16    pub operation: IndexOperation,
17}
18
19pub struct TransactionBatchBuffer {
20    pending: DashMap<TransactionId, DashMap<IndexOpKey, TransactionBufferEntry>>,
21}
22
23impl std::fmt::Debug for TransactionBatchBuffer {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("TransactionBatchBuffer")
26            .field("pending_count", &self.pending.len())
27            .finish_non_exhaustive()
28    }
29}
30
31impl Default for TransactionBatchBuffer {
32    fn default() -> Self {
33        Self::new()
34    }
35}
36
37impl TransactionBatchBuffer {
38    pub fn new() -> Self {
39        Self {
40            pending: DashMap::new(),
41        }
42    }
43
44    pub fn take_operations(
45        &self,
46        txn_id: TransactionId,
47    ) -> BatchResult<Vec<(IndexOpKey, Vec<IndexOperation>)>> {
48        if let Some((_, txn_buffer)) = self.pending.remove(&txn_id) {
49            let mut result = Vec::new();
50            for entry in txn_buffer.iter() {
51                let key = entry.key().clone();
52                let ops = entry
53                    .value()
54                    .operations
55                    .iter()
56                    .cloned()
57                    .map(|op| op.operation)
58                    .collect::<Vec<_>>();
59                if !ops.is_empty() {
60                    result.push((key, ops));
61                }
62            }
63            Ok(result)
64        } else {
65            Ok(Vec::new())
66        }
67    }
68
69    pub fn prepare(&self, txn_id: TransactionId, operation: IndexOperation) -> BatchResult<()> {
70        self.prepare_with_sequence(txn_id, 0, operation)
71    }
72
73    pub fn prepare_with_sequence(
74        &self,
75        txn_id: TransactionId,
76        sequence: u64,
77        operation: IndexOperation,
78    ) -> BatchResult<()> {
79        let txn_buffer = self.pending.entry(txn_id).or_default();
80
81        let key = operation.key.clone();
82
83        let mut entry = txn_buffer.entry(key).or_default();
84        entry.operations.push(SequencedIndexOperation {
85            sequence,
86            operation,
87        });
88        Ok(())
89    }
90
91    pub fn peek_operations(
92        &self,
93        txn_id: TransactionId,
94    ) -> BatchResult<Vec<(IndexOpKey, Vec<IndexOperation>)>> {
95        if let Some(txn_buffer) = self.pending.get(&txn_id) {
96            let mut result = Vec::new();
97            for entry in txn_buffer.iter() {
98                let key = entry.key().clone();
99                let ops = entry
100                    .value()
101                    .operations
102                    .iter()
103                    .cloned()
104                    .map(|op| op.operation)
105                    .collect::<Vec<_>>();
106                if !ops.is_empty() {
107                    result.push((key, ops));
108                }
109            }
110            Ok(result)
111        } else {
112            Ok(Vec::new())
113        }
114    }
115
116    pub fn rollback(&self, txn_id: TransactionId) -> BatchResult<()> {
117        self.pending.remove(&txn_id);
118        Ok(())
119    }
120
121    pub fn truncate_operations(&self, txn_id: TransactionId, sequence: u64) -> BatchResult<()> {
122        if let Some(txn_buffer) = self.pending.get_mut(&txn_id) {
123            let keys: Vec<IndexOpKey> =
124                txn_buffer.iter().map(|entry| entry.key().clone()).collect();
125
126            for key in keys {
127                if let Some(mut entry) = txn_buffer.get_mut(&key) {
128                    entry.operations.retain(|op| op.sequence <= sequence);
129                }
130            }
131        }
132        Ok(())
133    }
134
135    pub fn pending_sequence(&self, txn_id: TransactionId) -> u64 {
136        self.pending
137            .get(&txn_id)
138            .and_then(|txn_buffer| {
139                txn_buffer
140                    .iter()
141                    .flat_map(|entry| {
142                        let ops = entry.value().operations.clone();
143                        ops.into_iter().map(|op| op.sequence)
144                    })
145                    .max()
146            })
147            .unwrap_or(0)
148    }
149
150    pub fn pending_count(&self, txn_id: TransactionId) -> usize {
151        self.pending
152            .get(&txn_id)
153            .map(|txn_buffer| txn_buffer.iter().map(|e| e.value().operations.len()).sum())
154            .unwrap_or(0)
155    }
156}
157
158impl TransactionBatchBuffer {
159    pub fn add_operation(&self, txn_id: TransactionId, key: &IndexOpKey, value: IndexOperation) {
160        let txn_buffer = self.pending.entry(txn_id).or_default();
161        let mut entry = txn_buffer.entry(key.clone()).or_default();
162        entry.operations.push(SequencedIndexOperation {
163            sequence: 0,
164            operation: value,
165        });
166    }
167}
168
169impl BatchBuffer<IndexOpKey, IndexOperation> for TransactionBatchBuffer {
170    fn add(&self, key: &IndexOpKey, value: IndexOperation) {
171        self.add_operation(TransactionId::new(0), key, value);
172    }
173
174    fn drain(&self, key: &IndexOpKey) -> Vec<IndexOperation> {
175        if let Some(txn_buffer) = self.pending.get(&TransactionId::new(0)) {
176            if let Some(entry) = txn_buffer.get(key) {
177                return entry
178                    .operations
179                    .iter()
180                    .map(|op| op.operation.clone())
181                    .collect();
182            }
183        }
184        Vec::new()
185    }
186
187    fn peek(&self, key: &IndexOpKey) -> Vec<IndexOperation> {
188        if let Some(txn_buffer) = self.pending.get(&TransactionId::new(0)) {
189            if let Some(entry) = txn_buffer.get(key) {
190                return entry
191                    .operations
192                    .iter()
193                    .map(|op| op.operation.clone())
194                    .collect();
195            }
196        }
197        Vec::new()
198    }
199
200    fn count(&self, key: &IndexOpKey) -> usize {
201        if let Some(txn_buffer) = self.pending.get(&TransactionId::new(0)) {
202            if let Some(entry) = txn_buffer.get(key) {
203                return entry.operations.len();
204            }
205        }
206        0
207    }
208
209    fn is_empty(&self, key: &IndexOpKey) -> bool {
210        self.count(key) == 0
211    }
212
213    fn keys(&self) -> Vec<IndexOpKey> {
214        if let Some(txn_buffer) = self.pending.get(&TransactionId::new(0)) {
215            return txn_buffer.iter().map(|e| e.key().clone()).collect();
216        }
217        Vec::new()
218    }
219
220    fn clear(&self) {
221        self.pending.clear();
222    }
223
224    fn total_count(&self) -> usize {
225        self.pending
226            .iter()
227            .map(|txn_buffer| {
228                txn_buffer
229                    .value()
230                    .iter()
231                    .map(|e| e.value().operations.len())
232                    .sum::<usize>()
233            })
234            .sum()
235    }
236}
linkrs_sync/batch/transaction_buffer.rs

linkrs_sync/batch/
transaction_buffer.rs
