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

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

            
16
use crate::database::{deserialize_document, document_tree_name, Database};
17
use crate::tasks::{Job, Keyed, Task};
18
use crate::views::{
19
    view_document_map_tree_name, view_entries_tree_name, view_invalidated_docs_tree_name,
20
    EntryMapping, ViewEntry,
21
};
22
use crate::Error;
23

            
24
#[derive(Debug)]
25
pub struct Mapper {
26
    pub database: Database,
27
    pub map: Map,
28
}
29

            
30
673029
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
31
pub struct Map {
32
    pub database: Arc<Cow<'static, str>>,
33
    pub collection: CollectionName,
34
    pub view_name: ViewName,
35
}
36

            
37
impl Job for Mapper {
38
    type Error = Error;
39
    type Output = u64;
40

            
41
    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
42
    #[allow(clippy::too_many_lines)]
43
221052
    fn execute(&mut self) -> Result<Self::Output, Error> {
44
221052
        let documents =
45
221052
            self.database
46
221052
                .roots()
47
221052
                .tree(self.database.collection_tree::<Versioned, _>(
48
221052
                    &self.map.collection,
49
221052
                    document_tree_name(&self.map.collection),
50
221052
                )?)?;
51

            
52
221052
        let view_entries =
53
221052
            self.database
54
221052
                .roots()
55
221052
                .tree(self.database.collection_tree::<Unversioned, _>(
56
221052
                    &self.map.collection,
57
221052
                    view_entries_tree_name(&self.map.view_name),
58
221052
                )?)?;
59

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

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

            
76
221052
        let transaction_id = self
77
221052
            .database
78
221052
            .last_transaction_id()?
79
221052
            .expect("no way to have documents without a transaction");
80
221052

            
81
221052
        let storage = self.database.clone();
82
221052
        let map_request = self.map.clone();
83
221052

            
84
221052
        map_view(
85
221052
            &invalidated_entries,
86
221052
            &document_map,
87
221052
            &documents,
88
221052
            &view_entries,
89
221052
            &storage,
90
221052
            &map_request,
91
221052
        )?;
92

            
93
221052
        self.database.storage.instance.tasks().mark_view_updated(
94
221052
            self.map.database.clone(),
95
221052
            self.map.collection.clone(),
96
221052
            self.map.view_name.clone(),
97
221052
            transaction_id,
98
221052
        );
99
221052

            
100
221052
        Ok(transaction_id)
101
221052
    }
102
}
103

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

            
135
22612
            let document_ids = invalidated_ids
136
22612
                .drain(invalidated_ids.len().saturating_sub(CHUNK_SIZE)..)
137
22612
                .collect::<Vec<_>>();
138
22612
            let document_map = transaction.unlocked_tree(1).unwrap();
139
22612
            let documents = transaction.unlocked_tree(2).unwrap();
140
22612
            let view_entries = transaction.unlocked_tree(3).unwrap();
141
22612
            DocumentRequest {
142
22612
                document_ids: document_ids.clone(),
143
22612
                map_request,
144
22612
                database,
145
22612
                document_map,
146
22612
                documents,
147
22612
                view_entries,
148
22612
                view,
149
22612
            }
150
22612
            .map()?;
151

            
152
22612
            let mut invalidated_entries = transaction.tree::<Unversioned>(0).unwrap();
153
22612
            invalidated_entries.modify(document_ids, nebari::tree::Operation::Remove)?;
154
        }
155
22612
        transaction.commit()?;
156
    }
157

            
158
221052
    Ok(())
159
221052
}
160

            
161
pub struct DocumentRequest<'a> {
162
    pub document_ids: Vec<ArcBytes<'static>>,
163
    pub map_request: &'a Map,
164
    pub database: &'a Database,
165

            
166
    pub document_map: &'a UnlockedTransactionTree<AnyFile>,
167
    pub documents: &'a UnlockedTransactionTree<AnyFile>,
168
    pub view_entries: &'a UnlockedTransactionTree<AnyFile>,
169
    pub view: &'a dyn Serialized,
170
}
171

            
172
type DocumentIdPayload = (ArcBytes<'static>, Option<ArcBytes<'static>>);
173
type BatchPayload = (Vec<ArcBytes<'static>>, flume::Receiver<DocumentIdPayload>);
174

            
175
impl<'a> DocumentRequest<'a> {
176
461075
    fn generate_batches(
177
461075
        batch_sender: flume::Sender<BatchPayload>,
178
461075
        document_ids: &[ArcBytes<'static>],
179
461075
        documents: &UnlockedTransactionTree<AnyFile>,
180
461075
    ) -> Result<(), Error> {
181
461075
        // Generate batches
182
461075
        let mut documents = documents.lock::<Versioned>();
183
461075
        for chunk in document_ids.chunks(1024) {
184
461075
            let (document_id_sender, document_id_receiver) = flume::bounded(chunk.len());
185
461075
            batch_sender
186
461075
                .send((chunk.to_vec(), document_id_receiver))
187
461075
                .unwrap();
188
461075
            let mut documents = documents.get_multiple(chunk.iter().map(ArcBytes::as_slice))?;
189
461075
            documents.sort_by(|a, b| a.0.cmp(&b.0));
190

            
191
526903
            for document_id in chunk.iter().rev() {
192
526903
                let document = documents
193
526903
                    .last()
194
526903
                    .map_or(false, |(key, _)| (key == document_id))
195
526903
                    .then(|| documents.pop().unwrap().1);
196
526903

            
197
526903
                document_id_sender
198
526903
                    .send((document_id.clone(), document))
199
526903
                    .unwrap();
200
526903
            }
201

            
202
461075
            drop(document_id_sender);
203
        }
204
461075
        drop(batch_sender);
205
461075
        Ok(())
206
461075
    }
207

            
208
461075
    fn map_batches(
209
461075
        batch_receiver: &flume::Receiver<BatchPayload>,
210
461075
        mapped_sender: flume::Sender<Batch>,
211
461075
        view: &dyn Serialized,
212
461075
        parallelization: usize,
213
461075
    ) -> Result<(), Error> {
214
        // Process batches
215
922150
        while let Ok((document_ids, document_id_receiver)) = batch_receiver.recv() {
216
461075
            let mut batch = Batch {
217
461075
                document_ids,
218
461075
                ..Batch::default()
219
461075
            };
220
922150
            for result in Parallel::new()
221
922150
                .each(1..=parallelization, |_| -> Result<_, Error> {
222
922150
                    let mut results = Vec::new();
223
1449053
                    while let Ok((document_id, document)) = document_id_receiver.recv() {
224
526903
                        let map_result = if let Some(document) = document {
225
485222
                            let document = deserialize_document(&document)?;
226

            
227
                            // Call the schema map function
228
485222
                            view.map(&document).map_err(bonsaidb_core::Error::from)?
229
                        } else {
230
                            // Get multiple didn't return this document ID.
231
41681
                            Vec::new()
232
                        };
233
526903
                        let keys: HashSet<OwnedBytes> = map_result
234
526903
                            .iter()
235
526903
                            .map(|map| OwnedBytes::from(map.key.as_slice()))
236
526903
                            .collect();
237
526903
                        let new_keys = ArcBytes::from(bincode::serialize(&keys)?);
238

            
239
526903
                        results.push((document_id, new_keys, keys, map_result));
240
                    }
241

            
242
922150
                    Ok(results)
243
922150
                })
244
461075
                .run()
245
            {
246
922150
                for (document_id, new_keys, keys, map_result) in result? {
247
1012492
                    for key in &keys {
248
485589
                        batch.all_keys.insert(key.0.clone());
249
485589
                    }
250
526903
                    batch.document_maps.insert(document_id.clone(), new_keys);
251
526903
                    batch.document_keys.insert(document_id.clone(), keys);
252
1012492
                    for mapping in map_result {
253
485589
                        let key_mappings = batch
254
485589
                            .new_mappings
255
485589
                            .entry(ArcBytes::from(mapping.key.to_vec()))
256
485589
                            .or_insert_with(Vec::default);
257
485589
                        key_mappings.push(mapping);
258
485589
                    }
259
                }
260
            }
261
461075
            mapped_sender.send(batch).unwrap();
262
        }
263
461075
        drop(mapped_sender);
264
461075
        Ok(())
265
461075
    }
266

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

            
301
461075
    fn update_view_entries(
302
461075
        view: &dyn Serialized,
303
461075
        map_request: &Map,
304
461075
        view_entries: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
305
461075
        all_keys: BTreeSet<ArcBytes<'static>>,
306
461075
        view_entries_to_clean: BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>,
307
461075
        new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
308
461075
    ) -> Result<(), Error> {
309
461075
        let mut updater = ViewEntryUpdater {
310
461075
            view,
311
461075
            map_request,
312
461075
            view_entries_to_clean,
313
461075
            new_mappings,
314
461075
            result: Ok(()),
315
461075
            has_reduce: true,
316
461075
        };
317
461075
        view_entries
318
461075
            .modify(
319
461075
                all_keys.into_iter().collect(),
320
491660
                Operation::CompareSwap(CompareSwap::new(&mut |key, view_entries| {
321
491660
                    updater.compare_swap_view_entry(key, view_entries)
322
491660
                })),
323
461075
            )
324
461075
            .map_err(Error::from)
325
461075
            .and(updater.result)
326
461075
    }
327

            
328
461075
    fn save_mappings(
329
461075
        mapped_receiver: &flume::Receiver<Batch>,
330
461075
        view: &dyn Serialized,
331
461075
        map_request: &Map,
332
461075
        document_map: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
333
461075
        view_entries: &mut LockedTransactionTree<'_, Unversioned, AnyFile>,
334
461075
    ) -> Result<(), Error> {
335
        while let Ok(Batch {
336
461075
            document_ids,
337
461075
            document_maps,
338
461075
            document_keys,
339
461075
            new_mappings,
340
461075
            mut all_keys,
341
921429
        }) = mapped_receiver.recv()
342
        {
343
461075
            let view_entries_to_clean = Self::update_document_map(
344
461075
                document_ids,
345
461075
                document_map,
346
461075
                &document_maps,
347
461075
                document_keys,
348
461075
                &mut all_keys,
349
461075
            )?;
350

            
351
461075
            Self::update_view_entries(
352
461075
                view,
353
461075
                map_request,
354
461075
                view_entries,
355
461075
                all_keys,
356
461075
                view_entries_to_clean,
357
461075
                new_mappings,
358
461075
            )?;
359
        }
360
460354
        Ok(())
361
461075
    }
362

            
363
461075
    pub fn map(&mut self) -> Result<(), Error> {
364
461075
        let (batch_sender, batch_receiver) = flume::bounded(1);
365
461075
        let (mapped_sender, mapped_receiver) = flume::bounded(1);
366

            
367
1383225
        for result in Parallel::new()
368
461075
            .add(|| Self::generate_batches(batch_sender, &self.document_ids, self.documents))
369
461075
            .add(|| {
370
461075
                Self::map_batches(
371
461075
                    &batch_receiver,
372
461075
                    mapped_sender,
373
461075
                    self.view,
374
461075
                    self.database.storage().parallelization(),
375
461075
                )
376
461075
            })
377
461075
            .add(|| {
378
461075
                let mut document_map = self.document_map.lock();
379
461075
                let mut view_entries = self.view_entries.lock();
380
461075
                Self::save_mappings(
381
461075
                    &mapped_receiver,
382
461075
                    self.view,
383
461075
                    self.map_request,
384
461075
                    &mut document_map,
385
461075
                    &mut view_entries,
386
461075
                )
387
461075
            })
388
461075
            .run()
389
        {
390
1383225
            result?;
391
        }
392

            
393
460354
        Ok(())
394
461075
    }
395
}
396

            
397
461075
#[derive(Default)]
398
struct Batch {
399
    document_ids: Vec<ArcBytes<'static>>,
400
    document_maps: BTreeMap<ArcBytes<'static>, ArcBytes<'static>>,
401
    document_keys: BTreeMap<ArcBytes<'static>, HashSet<OwnedBytes>>,
402
    new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
403
    all_keys: BTreeSet<ArcBytes<'static>>,
404
}
405

            
406
impl Keyed<Task> for Mapper {
407
230925
    fn key(&self) -> Task {
408
230925
        Task::ViewMap(self.map.clone())
409
230925
    }
410
}
411

            
412
struct ViewEntryUpdater<'a> {
413
    view: &'a dyn Serialized,
414
    map_request: &'a Map,
415
    view_entries_to_clean: BTreeMap<ArcBytes<'static>, HashSet<ArcBytes<'static>>>,
416
    new_mappings: BTreeMap<ArcBytes<'static>, Vec<map::Serialized>>,
417
    result: Result<(), Error>,
418
    has_reduce: bool,
419
}
420

            
421
impl<'a> ViewEntryUpdater<'a> {
422
491660
    fn compare_swap_view_entry(
423
491660
        &mut self,
424
491660
        key: &ArcBytes<'_>,
425
491660
        view_entries: Option<ArcBytes<'static>>,
426
491660
    ) -> KeyOperation<ArcBytes<'static>> {
427
491660
        let mut view_entry = view_entries
428
491660
            .and_then(|view_entries| bincode::deserialize::<ViewEntry>(&view_entries).ok())
429
491660
            .unwrap_or_else(|| ViewEntry {
430
170200
                key: Bytes::from(key.to_vec()),
431
170200
                view_version: self.view.version(),
432
170200
                mappings: vec![],
433
170200
                reduced_value: Bytes::default(),
434
491660
            });
435
491660
        let key = key.to_owned();
436
491660
        if let Some(document_ids) = self.view_entries_to_clean.remove(&key) {
437
40932
            view_entry
438
40932
                .mappings
439
44327
                .retain(|m| !document_ids.contains(m.source.id.as_ref()));
440
40932

            
441
40932
            if view_entry.mappings.is_empty() && !self.new_mappings.contains_key(&key[..]) {
442
39872
                return KeyOperation::Remove;
443
1060
            } else if self.has_reduce {
444
1060
                let mappings = view_entry
445
1060
                    .mappings
446
1060
                    .iter()
447
1060
                    .map(|m| (&key[..], m.value.as_slice()))
448
1060
                    .collect::<Vec<_>>();
449
1060

            
450
1060
                match self.view.reduce(&mappings, false) {
451
1060
                    Ok(reduced) => {
452
1060
                        view_entry.reduced_value = Bytes::from(reduced);
453
1060
                    }
454
                    Err(view::Error::Core(bonsaidb_core::Error::ReduceUnimplemented)) => {
455
                        self.has_reduce = false;
456
                    }
457
                    Err(other) => {
458
                        self.result = Err(Error::from(other));
459
                        return KeyOperation::Skip;
460
                    }
461
                }
462
            }
463
450728
        }
464

            
465
451788
        if let Some(new_mappings) = self.new_mappings.remove(&key[..]) {
466
935748
            for map::Serialized { source, value, .. } in new_mappings {
467
                // Before altering any data, verify that the key is unique if this is a unique view.
468
485589
                if self.view.update_policy() == ViewUpdatePolicy::Unique
469
92603
                    && !view_entry.mappings.is_empty()
470
1699
                    && view_entry.mappings[0].source.id != source.id
471
                {
472
721
                    self.result = Err(Error::Core(bonsaidb_core::Error::UniqueKeyViolation {
473
721
                        view: self.map_request.view_name.clone(),
474
721
                        conflicting_document: Box::new(source),
475
721
                        existing_document: Box::new(view_entry.mappings[0].source.clone()),
476
721
                    }));
477
721
                    return KeyOperation::Skip;
478
484868
                }
479
484868
                let entry_mapping = EntryMapping { source, value };
480
484868

            
481
484868
                // attempt to update an existing
482
484868
                // entry for this document, if
483
484868
                // present
484
484868
                let mut found = false;
485
76939021
                for mapping in &mut view_entry.mappings {
486
76494066
                    if mapping.source.id == entry_mapping.source.id {
487
39913
                        found = true;
488
39913
                        mapping.source.revision = entry_mapping.source.revision;
489
39913
                        mapping.value = entry_mapping.value.clone();
490
39913
                        break;
491
76454153
                    }
492
                }
493

            
494
                // If an existing mapping wasn't
495
                // found, add it
496
484868
                if !found {
497
444955
                    view_entry.mappings.push(entry_mapping);
498
444955
                }
499
            }
500

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

            
519
449851
                match self.view.reduce(&mappings, false) {
520
354835
                    Ok(reduced) => {
521
354835
                        view_entry.reduced_value = Bytes::from(reduced);
522
354835
                    }
523
95016
                    Err(view::Error::Core(bonsaidb_core::Error::ReduceUnimplemented)) => {
524
95016
                        self.has_reduce = false;
525
95016
                    }
526
                    Err(other) => {
527
                        self.result = Err(Error::from(other));
528
                        return KeyOperation::Skip;
529
                    }
530
                }
531
308
            }
532
908
        }
533

            
534
451067
        let value = bincode::serialize(&view_entry).unwrap();
535
451067
        KeyOperation::Set(ArcBytes::from(value))
536
491660
    }
537
}