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

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

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

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

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

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

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

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

            
235
                            // Call the schema map function
236
134809
                            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
150554
                        let keys: HashSet<OwnedBytes> = map_result
242
150554
                            .iter()
243
150554
                            .map(|map| OwnedBytes::from(map.key.as_slice()))
244
150554
                            .collect();
245
150554
                        let new_keys = ArcBytes::from(bincode::serialize(&keys)?);
246

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

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

            
275
114956
    fn update_document_map(
276
114956
        document_ids: Vec<ArcBytes<'static>>,
277
114956
        document_map: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
278
114956
        document_maps: &BTreeMap<ArcBytes<'static>, ArcBytes<'static>>,
279
114956
        mut document_keys: BTreeMap<ArcBytes<'static>, HashSet<OwnedBytes>>,
280
114956
        all_keys: &mut BTreeSet<ArcBytes<'static>>,
281
114956
    ) -> Result<BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>, Error> {
282
114956
        // We need to store a record of all the mappings this document produced.
283
114956
        let mut maps_to_clear = Vec::new();
284
114956
        document_map.modify(
285
114956
            document_ids,
286
114956
            nebari::tree::Operation::CompareSwap(CompareSwap::new(&mut |key, value| {
287
150554
                if let Some(existing_map) = value {
288
17253
                    maps_to_clear.push((key.to_owned(), existing_map));
289
133301
                }
290
150554
                let new_map = document_maps.get(key).unwrap();
291
150554
                KeyOperation::Set(new_map.clone())
292
150554
            })),
293
114956
        )?;
294
114956
        let mut view_entries_to_clean = BTreeMap::new();
295
132209
        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
114956
        Ok(view_entries_to_clean)
307
114956
    }
308

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

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

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

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

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

            
401
114428
        Ok(())
402
114956
    }
403
}
404

            
405
114956
#[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
167705
    fn key(&self) -> Task {
416
167705
        Task::ViewMap(self.map.clone())
417
167705
    }
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
143333
    fn compare_swap_view_entry(
431
143333
        &mut self,
432
143333
        key: &ArcBytes<'_>,
433
143333
        view_entries: Option<ArcBytes<'static>>,
434
143333
    ) -> KeyOperation<ArcBytes<'static>> {
435
143333
        let mut view_entry = view_entries
436
143333
            .and_then(|view_entries| bincode::deserialize::<ViewEntry>(&view_entries).ok())
437
143333
            .unwrap_or_else(|| ViewEntry {
438
119508
                key: Bytes::from(key.to_vec()),
439
119508
                view_version: self.view.version(),
440
119508
                mappings: vec![],
441
119508
                reduced_value: Bytes::default(),
442
143333
            });
443
143333
        let key = key.to_owned();
444
143333
        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
126844
        }
472

            
473
127340
        if let Some(new_mappings) = self.new_mappings.remove(&key[..]) {
474
261548
            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
135108
                if self.view.unique()
477
85375
                    && !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
134550
                }
487
134550
                let entry_mapping = EntryMapping { source, value };
488
134550

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

            
501
                // If an existing mapping wasn't
502
                // found, add it
503
134550
                if !found {
504
133750
                    view_entry.mappings.push(entry_mapping);
505
133750
                }
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
126440
            if self.has_reduce {
520
126188
                let mappings = view_entry
521
126188
                    .mappings
522
126188
                    .iter()
523
144296
                    .map(|m| (key.as_slice(), m.value.as_slice()))
524
126188
                    .collect::<Vec<_>>();
525
126188

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

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