1
use std::{
2
    borrow::Cow,
3
    collections::{hash_map::RandomState, BTreeMap, BTreeSet, HashSet},
4
    sync::Arc,
5
};
6

            
7
use bonsaidb_core::{
8
    arc_bytes::{serde::Bytes, ArcBytes, OwnedBytes},
9
    connection::Connection,
10
    schema::{
11
        view::{self, map, Serialized},
12
        CollectionName, ViewName,
13
    },
14
};
15
use easy_parallel::Parallel;
16
use nebari::{
17
    io::any::AnyFile,
18
    tree::{AnyTreeRoot, CompareSwap, KeyOperation, Operation, Unversioned, Versioned},
19
    LockedTransactionTree, Tree, UnlockedTransactionTree,
20
};
21

            
22
use crate::{
23
    database::{deserialize_document, document_tree_name, Database},
24
    tasks::{Job, Keyed, Task},
25
    views::{
26
        view_document_map_tree_name, view_entries_tree_name, view_invalidated_docs_tree_name,
27
        EntryMapping, ViewEntry,
28
    },
29
    Error,
30
};
31

            
32
#[derive(Debug)]
33
pub struct Mapper {
34
    pub database: Database,
35
    pub map: Map,
36
}
37

            
38
492945
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
39
pub struct Map {
40
    pub database: Arc<Cow<'static, str>>,
41
    pub collection: CollectionName,
42
    pub view_name: ViewName,
43
}
44

            
45
impl Job for Mapper {
46
    type Output = u64;
47
    type Error = Error;
48

            
49
323788
    #[cfg_attr(feature = "tracing", tracing::instrument)]
50
161894
    #[allow(clippy::too_many_lines)]
51
    fn execute(&mut self) -> Result<Self::Output, Error> {
52
        let documents =
53
            self.database
54
                .roots()
55
                .tree(self.database.collection_tree::<Versioned, _>(
56
                    &self.map.collection,
57
                    document_tree_name(&self.map.collection),
58
                )?)?;
59

            
60
        let view_entries =
61
            self.database
62
                .roots()
63
                .tree(self.database.collection_tree::<Unversioned, _>(
64
                    &self.map.collection,
65
                    view_entries_tree_name(&self.map.view_name),
66
                )?)?;
67

            
68
        let document_map =
69
            self.database
70
                .roots()
71
                .tree(self.database.collection_tree::<Unversioned, _>(
72
                    &self.map.collection,
73
                    view_document_map_tree_name(&self.map.view_name),
74
                )?)?;
75

            
76
        let invalidated_entries =
77
            self.database
78
                .roots()
79
                .tree(self.database.collection_tree::<Unversioned, _>(
80
                    &self.map.collection,
81
                    view_invalidated_docs_tree_name(&self.map.view_name),
82
                )?)?;
83

            
84
        let transaction_id = self
85
            .database
86
            .last_transaction_id()?
87
            .expect("no way to have documents without a transaction");
88

            
89
        let storage = self.database.clone();
90
        let map_request = self.map.clone();
91

            
92
        map_view(
93
            &invalidated_entries,
94
            &document_map,
95
            &documents,
96
            &view_entries,
97
            &storage,
98
            &map_request,
99
        )?;
100

            
101
        self.database.storage.instance.tasks().mark_view_updated(
102
            self.map.database.clone(),
103
            self.map.collection.clone(),
104
            self.map.view_name.clone(),
105
            transaction_id,
106
        );
107

            
108
        Ok(transaction_id)
109
    }
110
}
111

            
112
161894
fn map_view(
113
161894
    invalidated_entries: &Tree<Unversioned, AnyFile>,
114
161894
    document_map: &Tree<Unversioned, AnyFile>,
115
161894
    documents: &Tree<Versioned, AnyFile>,
116
161894
    view_entries: &Tree<Unversioned, AnyFile>,
117
161894
    database: &Database,
118
161894
    map_request: &Map,
119
161894
) -> Result<(), Error> {
120
    const CHUNK_SIZE: usize = 100_000;
121
    // Only do any work if there are invalidated documents to process
122
161894
    let mut invalidated_ids = invalidated_entries
123
161894
        .get_range(&(..))?
124
161894
        .into_iter()
125
161894
        .map(|(key, _)| key)
126
161894
        .collect::<Vec<_>>();
127
175165
    while !invalidated_ids.is_empty() {
128
13271
        let transaction = database
129
13271
            .roots()
130
13271
            .transaction::<_, dyn AnyTreeRoot<AnyFile>>(&[
131
13271
                Box::new(invalidated_entries.clone()) as Box<dyn AnyTreeRoot<AnyFile>>,
132
13271
                Box::new(document_map.clone()),
133
13271
                Box::new(documents.clone()),
134
13271
                Box::new(view_entries.clone()),
135
13271
            ])?;
136
        {
137
13271
            let view = database
138
13271
                .data
139
13271
                .schema
140
13271
                .view_by_name(&map_request.view_name)
141
13271
                .unwrap();
142
13271

            
143
13271
            let document_ids = invalidated_ids
144
13271
                .drain(invalidated_ids.len().saturating_sub(CHUNK_SIZE)..)
145
13271
                .collect::<Vec<_>>();
146
13271
            let document_map = transaction.unlocked_tree(1).unwrap();
147
13271
            let documents = transaction.unlocked_tree(2).unwrap();
148
13271
            let view_entries = transaction.unlocked_tree(3).unwrap();
149
13271
            DocumentRequest {
150
13271
                document_ids: document_ids.clone(),
151
13271
                map_request,
152
13271
                database,
153
13271
                document_map,
154
13271
                documents,
155
13271
                view_entries,
156
13271
                view,
157
13271
            }
158
13271
            .map()?;
159

            
160
13271
            let mut invalidated_entries = transaction.tree::<Unversioned>(0).unwrap();
161
13271
            invalidated_entries.modify(document_ids, nebari::tree::Operation::Remove)?;
162
        }
163
13271
        transaction.commit()?;
164
    }
165

            
166
161894
    Ok(())
167
161894
}
168

            
169
pub struct DocumentRequest<'a> {
170
    pub document_ids: Vec<ArcBytes<'static>>,
171
    pub map_request: &'a Map,
172
    pub database: &'a Database,
173

            
174
    pub document_map: &'a UnlockedTransactionTree<AnyFile>,
175
    pub documents: &'a UnlockedTransactionTree<AnyFile>,
176
    pub view_entries: &'a UnlockedTransactionTree<AnyFile>,
177
    pub view: &'a dyn Serialized,
178
}
179

            
180
type DocumentIdPayload = (ArcBytes<'static>, Option<ArcBytes<'static>>);
181
type BatchPayload = (Vec<ArcBytes<'static>>, flume::Receiver<DocumentIdPayload>);
182

            
183
impl<'a> DocumentRequest<'a> {
184
111356
    fn generate_batches(
185
111356
        batch_sender: flume::Sender<BatchPayload>,
186
111356
        document_ids: &[ArcBytes<'static>],
187
111356
        documents: &UnlockedTransactionTree<AnyFile>,
188
111356
    ) -> Result<(), Error> {
189
111356
        // Generate batches
190
111356
        let mut documents = documents.lock::<Versioned>();
191
111356
        for chunk in document_ids.chunks(1024) {
192
111356
            let (document_id_sender, document_id_receiver) = flume::bounded(chunk.len());
193
111356
            batch_sender
194
111356
                .send((chunk.to_vec(), document_id_receiver))
195
111356
                .unwrap();
196
111356
            let mut documents = documents.get_multiple(chunk.iter().map(ArcBytes::as_slice))?;
197
111356
            documents.sort_by(|a, b| a.0.cmp(&b.0));
198

            
199
181994
            for document_id in chunk.iter().rev() {
200
181994
                let document = documents
201
181994
                    .last()
202
181994
                    .map_or(false, |(key, _)| (key == document_id))
203
181994
                    .then(|| documents.pop().unwrap().1);
204
181994

            
205
181994
                document_id_sender
206
181994
                    .send((document_id.clone(), document))
207
181994
                    .unwrap();
208
181994
            }
209

            
210
111356
            drop(document_id_sender);
211
        }
212
111356
        drop(batch_sender);
213
111356
        Ok(())
214
111356
    }
215

            
216
111356
    fn map_batches(
217
111356
        batch_receiver: &flume::Receiver<BatchPayload>,
218
111356
        mapped_sender: flume::Sender<Batch>,
219
111356
        view: &dyn Serialized,
220
111356
        parallelization: usize,
221
111356
    ) -> Result<(), Error> {
222
        // Process batches
223
222712
        while let Ok((document_ids, document_id_receiver)) = batch_receiver.recv() {
224
111356
            let mut batch = Batch {
225
111356
                document_ids,
226
111356
                ..Batch::default()
227
111356
            };
228
222712
            for result in Parallel::new()
229
222712
                .each(1..=parallelization, |_| -> Result<_, Error> {
230
222712
                    let mut results = Vec::new();
231
404706
                    while let Ok((document_id, document)) = document_id_receiver.recv() {
232
182024
                        let map_result = if let Some(document) = document {
233
166280
                            let document = deserialize_document(&document)?;
234

            
235
                            // Call the schema map function
236
166280
                            view.map(&document).map_err(bonsaidb_core::Error::from)?
237
                        } else {
238
                            // Get multiple didn't return this document ID.
239
15744
                            Vec::new()
240
                        };
241
181994
                        let keys: HashSet<OwnedBytes> = map_result
242
181994
                            .iter()
243
181994
                            .map(|map| OwnedBytes::from(map.key.as_slice()))
244
181994
                            .collect();
245
181994
                        let new_keys = ArcBytes::from(bincode::serialize(&keys)?);
246

            
247
181994
                        results.push((document_id, new_keys, keys, map_result));
248
                    }
249

            
250
222682
                    Ok(results)
251
222682
                })
252
111356
                .run()
253
            {
254
222712
                for (document_id, new_keys, keys, map_result) in result? {
255
348542
                    for key in &keys {
256
166548
                        batch.all_keys.insert(key.0.clone());
257
166548
                    }
258
181994
                    batch.document_maps.insert(document_id.clone(), new_keys);
259
181994
                    batch.document_keys.insert(document_id.clone(), keys);
260
348542
                    for mapping in map_result {
261
166548
                        let key_mappings = batch
262
166548
                            .new_mappings
263
166548
                            .entry(ArcBytes::from(mapping.key.to_vec()))
264
166548
                            .or_insert_with(Vec::default);
265
166548
                        key_mappings.push(mapping);
266
166548
                    }
267
                }
268
            }
269
111356
            mapped_sender.send(batch).unwrap();
270
        }
271
111356
        drop(mapped_sender);
272
111356
        Ok(())
273
111356
    }
274

            
275
111356
    fn update_document_map(
276
111356
        document_ids: Vec<ArcBytes<'static>>,
277
111356
        document_map: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
278
111356
        document_maps: &BTreeMap<ArcBytes<'static>, ArcBytes<'static>>,
279
111356
        mut document_keys: BTreeMap<ArcBytes<'static>, HashSet<OwnedBytes>>,
280
111356
        all_keys: &mut BTreeSet<ArcBytes<'static>>,
281
111356
    ) -> Result<BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>, Error> {
282
111356
        // We need to store a record of all the mappings this document produced.
283
111356
        let mut maps_to_clear = Vec::new();
284
111356
        document_map.modify(
285
111356
            document_ids,
286
111356
            nebari::tree::Operation::CompareSwap(CompareSwap::new(&mut |key, value| {
287
181994
                if let Some(existing_map) = value {
288
17253
                    maps_to_clear.push((key.to_owned(), existing_map));
289
164741
                }
290
181994
                let new_map = document_maps.get(key).unwrap();
291
181994
                KeyOperation::Set(new_map.clone())
292
181994
            })),
293
111356
        )?;
294
111356
        let mut view_entries_to_clean = BTreeMap::new();
295
128609
        for (document_id, existing_map) in maps_to_clear {
296
17253
            let existing_keys = bincode::deserialize::<HashSet<OwnedBytes>>(&existing_map)?;
297
17253
            let new_keys = document_keys.remove(&document_id).unwrap();
298
17253
            for key in existing_keys.difference(&new_keys) {
299
16613
                all_keys.insert(key.clone().0);
300
16613
                let key_documents = view_entries_to_clean
301
16613
                    .entry(key.clone().0)
302
16613
                    .or_insert_with(HashSet::<_, RandomState>::default);
303
16613
                key_documents.insert(document_id.clone());
304
16613
            }
305
        }
306
111356
        Ok(view_entries_to_clean)
307
111356
    }
308

            
309
111356
    fn update_view_entries(
310
111356
        view: &dyn Serialized,
311
111356
        map_request: &Map,
312
111356
        view_entries: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
313
111356
        all_keys: BTreeSet<ArcBytes<'static>>,
314
111356
        view_entries_to_clean: BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>,
315
111356
        new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
316
111356
    ) -> Result<(), Error> {
317
111356
        let mut updater = ViewEntryUpdater {
318
111356
            view,
319
111356
            map_request,
320
111356
            view_entries_to_clean,
321
111356
            new_mappings,
322
111356
            result: Ok(()),
323
111356
            has_reduce: true,
324
111356
        };
325
111356
        view_entries
326
111356
            .modify(
327
111356
                all_keys.into_iter().collect(),
328
147683
                Operation::CompareSwap(CompareSwap::new(&mut |key, view_entries| {
329
147683
                    updater.compare_swap_view_entry(key, view_entries)
330
147683
                })),
331
111356
            )
332
111356
            .map_err(Error::from)
333
111356
            .and(updater.result)
334
111356
    }
335

            
336
111356
    fn save_mappings(
337
111356
        mapped_receiver: &flume::Receiver<Batch>,
338
111356
        view: &dyn Serialized,
339
111356
        map_request: &Map,
340
111356
        document_map: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
341
111356
        view_entries: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
342
111356
    ) -> Result<(), Error> {
343
        while let Ok(Batch {
344
111356
            document_ids,
345
111356
            document_maps,
346
111356
            document_keys,
347
111356
            new_mappings,
348
111356
            mut all_keys,
349
222184
        }) = mapped_receiver.recv()
350
        {
351
111356
            let view_entries_to_clean = Self::update_document_map(
352
111356
                document_ids,
353
111356
                document_map,
354
111356
                &document_maps,
355
111356
                document_keys,
356
111356
                &mut all_keys,
357
111356
            )?;
358

            
359
111356
            Self::update_view_entries(
360
111356
                view,
361
111356
                map_request,
362
111356
                view_entries,
363
111356
                all_keys,
364
111356
                view_entries_to_clean,
365
111356
                new_mappings,
366
111356
            )?;
367
        }
368
110828
        Ok(())
369
111356
    }
370

            
371
111356
    pub fn map(&mut self) -> Result<(), Error> {
372
111356
        let (batch_sender, batch_receiver) = flume::bounded(1);
373
111356
        let (mapped_sender, mapped_receiver) = flume::bounded(1);
374

            
375
334068
        for result in Parallel::new()
376
111356
            .add(|| Self::generate_batches(batch_sender, &self.document_ids, self.documents))
377
111356
            .add(|| {
378
111356
                Self::map_batches(
379
111356
                    &batch_receiver,
380
111356
                    mapped_sender,
381
111356
                    self.view,
382
111356
                    self.database.storage().parallelization(),
383
111356
                )
384
111356
            })
385
111356
            .add(|| {
386
111356
                let mut document_map = self.document_map.lock();
387
111356
                let mut view_entries = self.view_entries.lock();
388
111356
                Self::save_mappings(
389
111356
                    &mapped_receiver,
390
111356
                    self.view,
391
111356
                    self.map_request,
392
111356
                    &mut document_map,
393
111356
                    &mut view_entries,
394
111356
                )
395
111356
            })
396
111356
            .run()
397
        {
398
334068
            result?;
399
        }
400

            
401
110828
        Ok(())
402
111356
    }
403
}
404

            
405
111356
#[derive(Default)]
406
struct Batch {
407
    document_ids: Vec<ArcBytes<'static>>,
408
    document_maps: BTreeMap<ArcBytes<'static>, ArcBytes<'static>>,
409
    document_keys: BTreeMap<ArcBytes<'static>, HashSet<OwnedBytes>>,
410
    new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
411
    all_keys: BTreeSet<ArcBytes<'static>>,
412
}
413

            
414
impl Keyed<Task> for Mapper {
415
169157
    fn key(&self) -> Task {
416
169157
        Task::ViewMap(self.map.clone())
417
169157
    }
418
}
419

            
420
struct ViewEntryUpdater<'a> {
421
    view: &'a dyn Serialized,
422
    map_request: &'a Map,
423
    view_entries_to_clean: BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>,
424
    new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
425
    result: Result<(), Error>,
426
    has_reduce: bool,
427
}
428

            
429
impl<'a> ViewEntryUpdater<'a> {
430
147683
    fn compare_swap_view_entry(
431
147683
        &mut self,
432
147683
        key: &ArcBytes<'_>,
433
147683
        view_entries: Option<ArcBytes<'static>>,
434
147683
    ) -> KeyOperation<ArcBytes<'static>> {
435
147683
        let mut view_entry = view_entries
436
147683
            .and_then(|view_entries| bincode::deserialize::<ViewEntry>(&view_entries).ok())
437
147683
            .unwrap_or_else(|| ViewEntry {
438
122148
                key: Bytes::from(key.to_vec()),
439
122148
                view_version: self.view.version(),
440
122148
                mappings: vec![],
441
122148
                reduced_value: Bytes::default(),
442
147683
            });
443
147683
        let key = key.to_owned();
444
147683
        if let Some(document_ids) = self.view_entries_to_clean.remove(&key) {
445
16489
            view_entry
446
16489
                .mappings
447
16985
                .retain(|m| !document_ids.contains(m.source.id.as_ref()));
448
16489

            
449
16489
            if view_entry.mappings.is_empty() && !self.new_mappings.contains_key(&key[..]) {
450
15993
                return KeyOperation::Remove;
451
496
            } else if self.has_reduce {
452
496
                let mappings = view_entry
453
496
                    .mappings
454
496
                    .iter()
455
496
                    .map(|m| (&key[..], m.value.as_slice()))
456
496
                    .collect::<Vec<_>>();
457
496

            
458
496
                match self.view.reduce(&mappings, false) {
459
496
                    Ok(reduced) => {
460
496
                        view_entry.reduced_value = Bytes::from(reduced);
461
496
                    }
462
                    Err(view::Error::Core(bonsaidb_core::Error::ReduceUnimplemented)) => {
463
                        self.has_reduce = false;
464
                    }
465
                    Err(other) => {
466
                        self.result = Err(Error::from(other));
467
                        return KeyOperation::Skip;
468
                    }
469
                }
470
            }
471
131194
        }
472

            
473
131690
        if let Some(new_mappings) = self.new_mappings.remove(&key[..]) {
474
297338
            for map::Serialized { source, value, .. } in new_mappings {
475
                // Before altering any data, verify that the key is unique if this is a unique view.
476
166548
                if self.view.unique()
477
82855
                    && !view_entry.mappings.is_empty()
478
1204
                    && view_entry.mappings[0].source.id != source.id
479
                {
480
528
                    self.result = Err(Error::Core(bonsaidb_core::Error::UniqueKeyViolation {
481
528
                        view: self.map_request.view_name.clone(),
482
528
                        conflicting_document: Box::new(source),
483
528
                        existing_document: Box::new(view_entry.mappings[0].source.clone()),
484
528
                    }));
485
528
                    return KeyOperation::Skip;
486
166020
                }
487
166020
                let entry_mapping = EntryMapping { source, value };
488
166020

            
489
166020
                // attempt to update an existing
490
166020
                // entry for this document, if
491
166020
                // present
492
166020
                let mut found = false;
493
265758
                for mapping in &mut view_entry.mappings {
494
100568
                    if mapping.source.id == entry_mapping.source.id {
495
830
                        found = true;
496
830
                        mapping.value = entry_mapping.value.clone();
497
830
                        break;
498
99738
                    }
499
                }
500

            
501
                // If an existing mapping wasn't
502
                // found, add it
503
166020
                if !found {
504
165190
                    view_entry.mappings.push(entry_mapping);
505
165190
                }
506
            }
507

            
508
            // There was a choice to be made here of whether to call
509
            // reduce()  with all of the existing values, or call it with
510
            // rereduce=true passing only the new value and the old stored
511
            // value. In this implementation, it's technically less
512
            // efficient, but we can guarantee that every value has only
513
            // been reduced once, and potentially re-reduced a single-time.
514
            // If we constantly try to update the value to optimize the size
515
            // of `mappings`, the fear is that the value computed may lose
516
            // precision in some contexts over time. Thus, the decision was
517
            // made to always call reduce() with all the mappings within a
518
            // single ViewEntry.
519
130790
            if self.has_reduce {
520
130538
                let mappings = view_entry
521
130538
                    .mappings
522
130538
                    .iter()
523
184916
                    .map(|m| (key.as_slice(), m.value.as_slice()))
524
130538
                    .collect::<Vec<_>>();
525
130538

            
526
130538
                match self.view.reduce(&mappings, false) {
527
45912
                    Ok(reduced) => {
528
45912
                        view_entry.reduced_value = Bytes::from(reduced);
529
45912
                    }
530
84626
                    Err(view::Error::Core(bonsaidb_core::Error::ReduceUnimplemented)) => {
531
84626
                        self.has_reduce = false;
532
84626
                    }
533
                    Err(other) => {
534
                        self.result = Err(Error::from(other));
535
                        return KeyOperation::Skip;
536
                    }
537
                }
538
252
            }
539
372
        }
540

            
541
131162
        let value = bincode::serialize(&view_entry).unwrap();
542
131162
        KeyOperation::Set(ArcBytes::from(value))
543
147683
    }
544
}