a3s_vec/collection/
mutation.rs

1//! Serialized document mutations and atomic generation publication.
2
3use super::checkpoint::maybe_checkpoint;
4use super::query_engine::{matches_filter, parse_filter_expression};
5use super::validation::{merge_patch, prepare_mutation_batch};
6use super::{ensure_same_generation, ensure_writable, next_revision, Collection};
7use crate::doc::Doc;
8use crate::error::{Error, ErrorCode, Result};
9use crate::index::IndexRegistry;
10use crate::storage::WalOperation;
11use std::collections::{BTreeSet, HashSet};
12use std::sync::Arc;
13
14/// Per-document outcome in a batch write.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct DocWriteResult {
17    pub success: bool,
18    pub code: ErrorCode,
19    pub message: String,
20}
21
22impl DocWriteResult {
23    pub fn is_success(&self) -> bool {
24        self.success
25    }
26}
27
28/// Result of a batch write.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct WriteResult {
31    pub success_count: u64,
32    pub error_count: u64,
33    pub results: Vec<DocWriteResult>,
34}
35
36impl Collection {
37    pub fn insert(&self, docs: &[&Doc]) -> Result<WriteResult> {
38        self.mutate_documents(docs, Mutation::Insert)
39    }
40
41    pub fn update(&self, docs: &[&Doc]) -> Result<WriteResult> {
42        self.mutate_documents(docs, Mutation::Update)
43    }
44
45    pub fn upsert(&self, docs: &[&Doc]) -> Result<WriteResult> {
46        self.mutate_documents(docs, Mutation::Upsert)
47    }
48
49    pub fn delete(&self, pks: &[&str]) -> Result<WriteResult> {
50        self.ensure_open()?;
51        let _writer = self
52            .inner
53            .writer
54            .lock()
55            .map_err(|_| Error::internal("writer lock poisoned"))?;
56        let current = self
57            .inner
58            .state
59            .read()
60            .map_err(|_| Error::internal("collection state lock poisoned"))?
61            .clone();
62        ensure_writable(&current.options)?;
63        if let Err(error) = current
64            .options
65            .resource_limits
66            .enforce_write_batch(pks.len())
67        {
68            current.stats.record_resource_limit_rejection();
69            return Err(error);
70        }
71        let mut results = Vec::with_capacity(pks.len());
72        let mut ids = Vec::new();
73        let mut seen = HashSet::new();
74        for pk in pks {
75            if pk.is_empty() || pk.contains('\0') {
76                results.push(write_error(
77                    ErrorCode::InvalidArgument,
78                    "primary key must be non-empty and contain no NUL byte",
79                ));
80            } else if !seen.insert(*pk) {
81                results.push(write_error(
82                    ErrorCode::AlreadyExists,
83                    "duplicate primary key in delete batch",
84                ));
85            } else if current.docs.contains_key(*pk) {
86                ids.push((*pk).to_string());
87                results.push(write_success());
88            } else {
89                results.push(write_error(
90                    ErrorCode::NotFound,
91                    format!("document '{pk}' not found"),
92                ));
93            }
94        }
95        if !ids.is_empty() {
96            self.publish_deletion(&current, &ids)?;
97        }
98        Ok(write_result(results))
99    }
100
101    pub fn delete_by_filter(&self, filter: &str) -> Result<()> {
102        self.ensure_open()?;
103        let _writer = self
104            .inner
105            .writer
106            .lock()
107            .map_err(|_| Error::internal("writer lock poisoned"))?;
108        let current = self
109            .inner
110            .state
111            .read()
112            .map_err(|_| Error::internal("collection state lock poisoned"))?
113            .clone();
114        ensure_writable(&current.options)?;
115        let parsed_filter = parse_filter_expression(filter)?;
116        let indexed = current
117            .indexes
118            .scalar_candidates(current.revision, &parsed_filter);
119        let ids: Vec<String> = if let Some(indexed) = indexed {
120            indexed
121                .ids()
122                .filter_map(|id| current.docs.get(id))
123                .filter(|doc| matches_filter(doc, Some(&parsed_filter)))
124                .filter_map(|doc| doc.get_pk().map(str::to_string))
125                .collect()
126        } else {
127            current
128                .docs
129                .values()
130                .filter(|doc| matches_filter(doc, Some(&parsed_filter)))
131                .filter_map(|doc| doc.get_pk().map(str::to_string))
132                .collect()
133        };
134        if ids.is_empty() {
135            return Ok(());
136        }
137        if let Err(error) = current
138            .options
139            .resource_limits
140            .enforce_write_batch(ids.len())
141        {
142            current.stats.record_resource_limit_rejection();
143            return Err(error);
144        }
145        self.publish_deletion(&current, &ids)
146    }
147
148    fn publish_deletion(&self, current: &super::CollectionState, ids: &[String]) -> Result<()> {
149        let config = current.config.clone();
150        let revision = next_revision(current.revision)?;
151        let mut next_docs = current.docs.as_ref().clone();
152        for id in ids {
153            next_docs.remove(id);
154        }
155        let changed_ids = ids.iter().cloned().collect::<BTreeSet<_>>();
156        let incremental_indexes = current.indexes.apply_document_changes(
157            &current.schema,
158            current.docs.as_ref(),
159            &next_docs,
160            revision,
161            &changed_ids,
162        )?;
163        let (next_indexes, next_resource_usage) =
164            match account_published(current, &next_docs, &changed_ids, &incremental_indexes) {
165                Ok(usage) => (incremental_indexes, usage),
166                Err(error) if error.code == ErrorCode::ResourceExhausted => {
167                    // A tombstone overlay can be larger than the removed payload.
168                    // Compact before rejecting so deletion remains a practical way
169                    // to return the collection below its configured budget.
170                    let compacted = IndexRegistry::build(&current.schema, &next_docs, revision)?;
171                    match account_published(current, &next_docs, &changed_ids, &compacted) {
172                        Ok(usage) => (compacted, usage),
173                        Err(error) => {
174                            current.stats.record_resource_limit_rejection();
175                            return Err(error);
176                        }
177                    }
178                }
179                Err(error) => return Err(error),
180            };
181        self.commit_visible_revision(
182            current,
183            revision,
184            WalOperation::Delete { ids: ids.to_vec() },
185            &config,
186            next_docs,
187            next_indexes,
188            next_resource_usage,
189        )
190    }
191
192    fn mutate_documents(&self, docs: &[&Doc], mutation: Mutation) -> Result<WriteResult> {
193        self.ensure_open()?;
194        let _writer = self
195            .inner
196            .writer
197            .lock()
198            .map_err(|_| Error::internal("writer lock poisoned"))?;
199        let current = self
200            .inner
201            .state
202            .read()
203            .map_err(|_| Error::internal("collection state lock poisoned"))?
204            .clone();
205        ensure_writable(&current.options)?;
206
207        if let Err(error) = current
208            .options
209            .resource_limits
210            .enforce_write_batch(docs.len())
211        {
212            current.stats.record_resource_limit_rejection();
213            return Err(error);
214        }
215
216        let (accepted, outcomes) = prepare_mutation_batch(&current, docs, mutation);
217        if accepted.is_empty() {
218            return Ok(write_result(outcomes));
219        }
220        let operation = match mutation {
221            Mutation::Insert => WalOperation::Insert {
222                docs: accepted.clone(),
223            },
224            Mutation::Update => WalOperation::Update {
225                docs: accepted.clone(),
226            },
227            Mutation::Upsert => WalOperation::Upsert {
228                docs: accepted.clone(),
229            },
230        };
231        let config = current.config.clone();
232        let revision = next_revision(current.revision)?;
233        let mut next_docs = current.docs.as_ref().clone();
234        for doc in &accepted {
235            let Some(pk) = doc.get_pk().map(str::to_string) else {
236                continue;
237            };
238            match mutation {
239                Mutation::Insert | Mutation::Upsert => {
240                    next_docs.insert(pk, Arc::new(doc.clone()));
241                }
242                Mutation::Update => {
243                    if let Some(existing) = next_docs.get_mut(&pk) {
244                        merge_patch(Arc::make_mut(existing), doc)?;
245                    }
246                }
247            }
248        }
249        let changed_ids = accepted
250            .iter()
251            .filter_map(|doc| doc.get_pk().map(str::to_string))
252            .collect::<BTreeSet<_>>();
253        let next_indexes = current.indexes.apply_document_changes(
254            &current.schema,
255            current.docs.as_ref(),
256            &next_docs,
257            revision,
258            &changed_ids,
259        )?;
260        let next_resource_usage =
261            match account_published(&current, &next_docs, &changed_ids, &next_indexes) {
262                Ok(usage) => usage,
263                Err(error) => {
264                    current.stats.record_resource_limit_rejection();
265                    return Err(error);
266                }
267            };
268        self.commit_visible_revision(
269            &current,
270            revision,
271            operation,
272            &config,
273            next_docs,
274            next_indexes,
275            next_resource_usage,
276        )?;
277        Ok(write_result(outcomes))
278    }
279
280    /// Syncs the WAL before taking the published-state lock, then publishes.
281    #[allow(clippy::too_many_arguments)]
282    fn commit_visible_revision(
283        &self,
284        current: &super::CollectionState,
285        revision: u64,
286        operation: WalOperation,
287        config: &crate::config::ConfigBuilder,
288        next_docs: crate::doc::DocumentMap,
289        next_indexes: IndexRegistry,
290        next_resource_usage: super::resource::ResourceUsage,
291    ) -> Result<()> {
292        {
293            let mut storage = self
294                .inner
295                .storage
296                .lock()
297                .map_err(|_| Error::internal("storage lock poisoned"))?;
298            storage.append(revision, operation, config)?;
299        }
300        let mut state = self
301            .inner
302            .state
303            .write()
304            .map_err(|_| Error::internal("collection state lock poisoned"))?;
305        ensure_same_generation(&state, current)?;
306        state.docs = Arc::new(next_docs);
307        state.indexes = Arc::new(next_indexes);
308        state.revision = revision;
309        state.resource_usage = next_resource_usage;
310        let mut storage = self
311            .inner
312            .storage
313            .lock()
314            .map_err(|_| Error::internal("storage lock poisoned"))?;
315        maybe_checkpoint(&mut storage, &state, config)?;
316        Ok(())
317    }
318}
319
320#[derive(Debug, Clone, Copy)]
321pub(super) enum Mutation {
322    Insert,
323    Update,
324    Upsert,
325}
326
327pub(super) fn write_success() -> DocWriteResult {
328    DocWriteResult {
329        success: true,
330        code: ErrorCode::Unknown,
331        message: String::new(),
332    }
333}
334
335pub(super) fn write_error(code: ErrorCode, message: impl Into<String>) -> DocWriteResult {
336    DocWriteResult {
337        success: false,
338        code,
339        message: message.into(),
340    }
341}
342
343fn account_published(
344    current: &super::CollectionState,
345    next_docs: &crate::doc::DocumentMap,
346    changed_ids: &BTreeSet<String>,
347    indexes: &IndexRegistry,
348) -> Result<super::resource::ResourceUsage> {
349    let usage = current.resource_usage.after_documents(
350        current.docs.as_ref(),
351        next_docs,
352        changed_ids,
353        indexes,
354    )?;
355    let document_count = u64::try_from(next_docs.len())
356        .map_err(|_| Error::resource_exhausted("collection exceeds u64 documents"))?;
357    current.options.resource_limits.admit(document_count, usage)
358}
359
360fn write_result(results: Vec<DocWriteResult>) -> WriteResult {
361    let success_count = results.iter().filter(|result| result.success).count() as u64;
362    WriteResult {
363        success_count,
364        error_count: results.len() as u64 - success_count,
365        results,
366    }
367}
a3s_vec/collection/mutation.rs

a3s_vec/collection/
mutation.rs
