Grcov report - roots.rs

1

use std::{

2

    any::Any,

3

    borrow::{Borrow, Cow},

4

    collections::HashMap,

5

    convert::Infallible,

6

    fmt::{Debug, Display},

7

fs,

8

    ops::{Deref, DerefMut, RangeBounds},

9

    path::{Path, PathBuf},

10

    sync::{

11

        atomic::{AtomicU16, Ordering},

12

        Arc,

13

},

14

};

15

16

use flume::Sender;

17

use once_cell::sync::Lazy;

18

use parking_lot::{MappedMutexGuard, Mutex, MutexGuard};

19

20

use crate::{

21

    context::Context,

22

    error::Error,

23

    io::{fs::StdFileManager, FileManager, ManagedFile, PathId},

24

    transaction::{LogEntry, ManagedTransaction, TransactionId, TransactionManager},

25

    tree::{

26

        self,

27

        root::{AnyReducer, AnyTreeRoot},

28

        state::AnyTreeState,

29

        EmbeddedIndex, KeySequence, Modification, ModificationResult, Operation, PersistenceMode,

30

        ScanEvaluation, SequenceEntry, SequenceId, SequenceIndex, State, TransactableCompaction,

31

        TreeEntry, TreeFile, TreeRoot, TreeValueIndex, VersionedTreeRoot,

32

},

33

    vault::AnyVault,

34

    ArcBytes, ChunkCache, ErrorKind,

35

};

36

37

/// A multi-tree transactional B-Tree database.

38

#[derive(Debug)]

39

pub struct Roots<File: ManagedFile> {

40

    data: Arc<Data<File>>,

41

42

43

#[derive(Debug)]

44

struct Data<File: ManagedFile> {

45

    context: Context<File::Manager>,

46

    transactions: TransactionManager<File::Manager>,

47

    thread_pool: ThreadPool<File>,

48

    path: PathBuf,

49

    tree_states: Mutex<HashMap<String, Box<dyn AnyTreeState>>>,

50

    tree_paths: Mutex<HashMap<String, PathId>>,

51

52

53

impl<File: ManagedFile> Roots<File> {

54

17

    fn open<P: Into<PathBuf> + Send>(

55

17

        path: P,

56

17

        context: Context<File::Manager>,

57

17

        thread_pool: ThreadPool<File>,

58

17

    ) -> Result<Self, Error> {

59

17

        let path = path.into();

60

17

        if !path.exists() {

61

            fs::create_dir_all(&path)?;

62

17

        } else if !path.is_dir() {

63

            return Err(Error::from(format!(

64

                "'{:?}' already exists, but is not a directory.",

65

                path

66

            )));

67

17

68

69

17

        let transactions = TransactionManager::spawn(&path, context.clone())?;

70

17

        Ok(Self {

71

17

            data: Arc::new(Data {

72

17

                context,

73

17

                path,

74

17

                transactions,

75

17

                thread_pool,

76

17

                tree_states: Mutex::default(),

77

17

                tree_paths: Mutex::default(),

78

17

}),

79

17

})

80

17

81

82

    /// Returns the path to the database directory.

83

    #[must_use]

84

19

    pub fn path(&self) -> &Path {

85

19

        &self.data.path

86

19

87

88

    /// Returns the vault used to encrypt this database.

89

39845

    pub fn context(&self) -> &Context<File::Manager> {

90

39845

        &self.data.context

91

39845

92

93

    /// Returns the transaction manager for this database.

94

    #[must_use]

95

15198

    pub fn transactions(&self) -> &TransactionManager<File::Manager> {

96

15198

        &self.data.transactions

97

15198

98

99

    /// Opens a tree named `name`.

100

///

101

    /// ## Errors

102

///

103

    /// - [`InvalidTreeName`](ErrorKind::InvalidTreeName): The name contained an

104

    ///   invalid character. For a full list of valid characters, see the

105

    ///   documentation on [`InvalidTreeName`](ErrorKind::InvalidTreeName).

106

    pub fn tree<Root: tree::Root>(

107

        &self,

108

        root: TreeRoot<Root, File>,

109

    ) -> Result<Tree<Root, File>, Error> {

110

19

        check_name(&root.name)?;

111

19

        let path = self.tree_path(&root.name);

112

19

        if !self.context().file_manager.exists(&path)? {

113

14

            self.context().file_manager.append(&path)?;

114

5

115

19

        let state = self.tree_state(root.clone());

116

19

        Ok(Tree {

117

19

            roots: self.clone(),

118

19

            path,

119

19

            state,

120

19

            vault: root.vault,

121

19

            reducer: root.reducer,

122

19

            name: root.name,

123

19

})

124

19

125

126

24614

    fn tree_path(&self, name: &str) -> PathId {

127

24614

        let mut paths = self.data.tree_paths.lock();

128

24614

        if let Some(id) = paths.get(name) {

129

24595

            id.clone()

130

        } else {

131

19

            let id = self

132

19

                .context()

133

19

                .file_manager

134

19

                .resolve_path(self.path().join(format!("{}.nebari", name)), true)

135

19

                .unwrap();

136

19

            paths.insert(name.to_owned(), id.clone());

137

19

id

138

139

24614

140

141

    /// Removes a tree. Returns true if a tree was deleted.

142

    pub fn delete_tree(&self, name: impl Into<Cow<'static, str>>) -> Result<bool, Error> {

143

        let name = name.into();

144

        let mut tree_states = self.data.tree_states.lock();

145

        self.context()

146

            .file_manager

147

            .delete(self.tree_path(name.as_ref()))?;

148

        Ok(tree_states.remove(name.as_ref()).is_some())

149

150

151

    /// Returns a list of all the names of trees contained in this database.

152

    pub fn tree_names(&self) -> Result<Vec<String>, Error> {

153

        let mut names = Vec::new();

154

        // TODO use the file manager

155

        for entry in std::fs::read_dir(self.path())? {

156

            let entry = entry?;

157

            if let Some(name) = entry.file_name().to_str() {

158

                if let Some(without_extension) = name.strip_suffix(".nebari") {

159

                    names.push(without_extension.to_string());

160

161

162

163

        Ok(names)

164

165

166

19

    fn tree_state<Root: tree::Root>(&self, root: TreeRoot<Root, File>) -> State<Root> {

167

19

        self.tree_states(&[root])

168

19

            .into_iter()

169

19

            .next()

170

19

            .unwrap()

171

19

            .as_ref()

172

19

            .as_any()

173

19

            .downcast_ref::<State<Root>>()

174

19

            .unwrap()

175

19

            .clone()

176

19

177

178

24614

    fn tree_states<R: Borrow<T>, T: AnyTreeRoot<File> + ?Sized>(

179

24614

        &self,

180

24614

        trees: &[R],

181

24614

    ) -> Vec<Box<dyn AnyTreeState>> {

182

24614

        let mut tree_states = self.data.tree_states.lock();

183

24614

        let mut output = Vec::with_capacity(trees.len());

184

49228

        for tree in trees {

185

24614

            let state = tree_states

186

24614

                .entry(tree.borrow().name().to_string())

187

24614

                .or_insert_with(|| tree.borrow().default_state())

188

24614

                .cloned();

189

24614

            output.push(state);

190

24614

191

24614

        output

192

24614

193

194

    /// Begins a transaction over `trees`. All trees will be exclusively

195

    /// accessible by the transaction. Dropping the executing transaction will

196

    /// roll the transaction back.

197

///

198

    /// ## Errors

199

///

200

    /// - [`InvalidTreeName`](ErrorKind::InvalidTreeName): A tree name contained

201

    ///   an invalid character. For a full list of valid characters, see the

202

    ///   documentation on [`InvalidTreeName`](ErrorKind::InvalidTreeName).

203

24595

    pub fn transaction<R: Borrow<T>, T: AnyTreeRoot<File> + ?Sized>(

204

24595

        &self,

205

24595

        trees: &[R],

206

24595

    ) -> Result<ExecutingTransaction<File>, Error> {

207

49190

        for tree in trees {

208

24595

            check_name(tree.borrow().name()).map(|_| tree.borrow().name().as_bytes())?;

209

210

24595

        let transaction = self

211

24595

            .data

212

24595

            .transactions

213

24595

            .new_transaction(trees.iter().map(|t| t.borrow().name().as_bytes()));

214

24595

        let states = self.tree_states(trees);

215

24595

        let trees = trees

216

24595

            .iter()

217

24595

            .zip(states)

218

24595

            .map(|(tree, state)| {

219

24595

                tree.borrow()

220

24595

                    .begin_transaction(

221

24595

                        transaction.id,

222

24595

                        &self.tree_path(tree.borrow().name()),

223

24595

                        state.as_ref(),

224

24595

                        self.context(),

225

24595

                        Some(&self.data.transactions),

226

24595

227

24595

                    .map(UnlockedTransactionTree::new)

228

24595

})

229

24595

            .collect::<Result<Vec<_>, Error>>()?;

230

24595

        Ok(ExecutingTransaction {

231

24595

            roots: self.clone(),

232

24595

            transaction: Some(transaction),

233

24595

            trees,

234

24595

})

235

24595

236

237

238

24617

fn check_name(name: &str) -> Result<(), Error> {

239

24617

    if name != "_transactions"

240

24616

        && name

241

24616

            .bytes()

242

98548

            .all(|c| matches!(c as char, 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '.' | '_'))

243

244

24615

        Ok(())

245

    } else {

246

2

        Err(Error::from(ErrorKind::InvalidTreeName))

247

248

24617

249

250

impl<File: ManagedFile> Clone for Roots<File> {

251

24646

    fn clone(&self) -> Self {

252

24646

        Self {

253

24646

            data: self.data.clone(),

254

24646

255

24646

256

257

258

/// An executing transaction. While this exists, no other transactions can

259

/// execute across the same trees as this transaction holds.

260

#[must_use]

261

pub struct ExecutingTransaction<File: ManagedFile> {

262

    roots: Roots<File>,

263

    trees: Vec<UnlockedTransactionTree<File>>,

264

    transaction: Option<ManagedTransaction<File::Manager>>,

265

266

267

/// A tree that belongs to an [`ExecutingTransaction`].

268

#[must_use]

269

pub struct UnlockedTransactionTree<File: ManagedFile>(Mutex<Box<dyn AnyTransactionTree<File>>>);

270

271

impl<File: ManagedFile> UnlockedTransactionTree<File> {

272

24595

    fn new(file: Box<dyn AnyTransactionTree<File>>) -> Self {

273

24595

        Self(Mutex::new(file))

274

24595

275

276

    /// Locks this tree so that operations can be performed against it.

277

///

278

    /// # Panics

279

///

280

    /// This function panics if `Root` does not match the type specified when

281

    /// starting the transaction.

282

24596

    pub fn lock<Root: tree::Root>(&self) -> LockedTransactionTree<'_, Root, File> {

283

24596

        LockedTransactionTree(MutexGuard::map(self.0.lock(), |tree| {

284

24596

            tree.as_mut().as_any_mut().downcast_mut().unwrap()

285

24596

}))

286

24596

287

288

289

/// A locked transaction tree. This transactional tree is exclusively available

290

/// for writing and reading to the thread that locks it.

291

#[must_use]

292

pub struct LockedTransactionTree<'transaction, Root: tree::Root, File: ManagedFile>(

293

    MappedMutexGuard<'transaction, TransactionTree<Root, File>>,

294

);

295

296

impl<'transaction, Root: tree::Root, File: ManagedFile> Deref

297

    for LockedTransactionTree<'transaction, Root, File>

298

299

    type Target = TransactionTree<Root, File>;

300

301

    fn deref(&self) -> &Self::Target {

302

        &self.0

303

304

305

306

impl<'transaction, Root: tree::Root, File: ManagedFile> DerefMut

307

    for LockedTransactionTree<'transaction, Root, File>

308

309

24596

    fn deref_mut(&mut self) -> &mut Self::Target {

310

24596

        &mut self.0

311

24596

312

313

314

impl<File: ManagedFile> ExecutingTransaction<File> {

315

    /// Returns the [`LogEntry`] for this transaction.

316

    #[must_use]

317

    #[allow(clippy::missing_panics_doc)]

318

    pub fn entry(&self) -> &LogEntry<'static> {

319

        self.transaction

320

            .as_ref()

321

            .and_then(|tx| tx.transaction.as_ref())

322

            .unwrap()

323

324

325

    /// Returns a mutable reference to the [`LogEntry`] for this transaction.

326

    #[must_use]

327

    #[allow(clippy::missing_panics_doc)]

328

1

    pub fn entry_mut(&mut self) -> &mut LogEntry<'static> {

329

1

        self.transaction

330

1

            .as_mut()

331

1

            .and_then(|tx| tx.transaction.as_mut())

332

1

            .unwrap()

333

1

334

335

    /// Commits the transaction. Once this function has returned, all data

336

    /// updates are guaranteed to be able to be accessed by all other readers as

337

    /// well as impervious to sudden failures such as a power outage.

338

    #[allow(clippy::missing_panics_doc)]

339

24592

    pub fn commit(mut self) -> Result<(), Error> {

340

24592

        let trees = std::mem::take(&mut self.trees);

341

        // Write the trees to disk

342

24592

        let trees = self.roots.data.thread_pool.commit_trees(trees)?;

343

344

        // Push the transaction to the log.

345

24592

        let transaction = self.transaction.take().unwrap();

346

24592

        let tree_locks = transaction.commit()?;

347

348

        // Publish the tree states, now that the transaction has been fully recorded

349

49184

        for tree in trees {

350

24592

            tree.state().publish();

351

24592

352

353

        // Release the locks for the trees, allowing a new transaction to begin.

354

24592

        drop(tree_locks);

355

24592

356

24592

        Ok(())

357

24592

358

359

    /// Rolls the transaction back. It is not necessary to call this function --

360

    /// transactions will automatically be rolled back when the transaction is

361

    /// dropped, if `commit()` isn't called first.

362

2

    pub fn rollback(self) {

363

2

        drop(self);

364

2

365

366

    /// Accesses a locked tree.

367

24596

    pub fn tree<Root: tree::Root>(

368

24596

        &self,

369

24596

        index: usize,

370

24596

    ) -> Option<LockedTransactionTree<'_, Root, File>> {

371

24596

        self.unlocked_tree(index).map(UnlockedTransactionTree::lock)

372

24596

373

374

    /// Accesses an unlocked tree. Note: If you clone an

375

    /// [`UnlockedTransactionTree`], you must make sure to drop all instances

376

    /// before calling commit.

377

24596

    pub fn unlocked_tree(&self, index: usize) -> Option<&UnlockedTransactionTree<File>> {

378

24596

        self.trees.get(index)

379

24596

380

381

3

    fn rollback_tree_states(&mut self) {

382

3

        for tree in self.trees.drain(..) {

383

3

            let tree = tree.0.lock();

384

3

            tree.rollback();

385

3

386

3

387

388

389

impl<File: ManagedFile> Drop for ExecutingTransaction<File> {

390

    fn drop(&mut self) {

391

24595

        if let Some(transaction) = self.transaction.take() {

392

3

            self.rollback_tree_states();

393

3

            // Now the transaction can be dropped safely, freeing up access to the trees.

394

3

            drop(transaction);

395

24592

396

24595

397

398

399

/// A tree that is modifiable during a transaction.

400

pub struct TransactionTree<Root: tree::Root, File: ManagedFile> {

401

    pub(crate) transaction_id: TransactionId,

402

    /// The underlying tree file.

403

    pub tree: TreeFile<Root, File>,

404

405

406

pub trait AnyTransactionTree<File: ManagedFile>: Any + Send + Sync {

407

    fn as_any(&self) -> &dyn Any;

408

    fn as_any_mut(&mut self) -> &mut dyn Any;

409

410

    fn state(&self) -> Box<dyn AnyTreeState>;

411

412

    fn commit(&mut self) -> Result<(), Error>;

413

    fn rollback(&self);

414

415

416

impl<Root: tree::Root, File: ManagedFile> AnyTransactionTree<File> for TransactionTree<Root, File> {

417

    fn as_any(&self) -> &dyn Any {

418

        self

419

420

24596

    fn as_any_mut(&mut self) -> &mut dyn Any {

421

24596

        self

422

24596

423

424

24592

    fn state(&self) -> Box<dyn AnyTreeState> {

425

24592

        Box::new(self.tree.state.clone())

426

24592

427

428

24592

    fn commit(&mut self) -> Result<(), Error> {

429

24592

        self.tree.commit()

430

24592

431

432

3

    fn rollback(&self) {

433

3

        let mut state = self.tree.state.lock();

434

3

        state.rollback(&self.tree.state);

435

3

436

437

438

impl<File: ManagedFile, Index> TransactionTree<VersionedTreeRoot<Index>, File>

439

where

440

    Index: Clone + EmbeddedIndex<ArcBytes<'static>> + Debug + 'static,

441

442

    /// Returns the latest sequence id.

443

    #[must_use]

444

    pub fn current_sequence_id(&self) -> SequenceId {

445

        let state = self.tree.state.lock();

446

        state.root.sequence

447

448

449

    /// Scans the tree for keys that are contained within `range`. If `forwards`

450

    /// is true, scanning starts at the lowest sort-order key and scans forward.

451

    /// Otherwise, scanning starts at the highest sort-order key and scans

452

    /// backwards. `key_evaluator` is invoked for each key as it is encountered.

453

    /// For all [`ScanEvaluation::ReadData`] results returned, `callback` will be

454

    /// invoked with the key and values. The callback may not be invoked in the

455

    /// same order as the keys are scanned.

456

    pub fn scan_sequences<CallerError, Range, KeyEvaluator, DataCallback>(

457

        &mut self,

458

        range: Range,

459

        forwards: bool,

460

        key_evaluator: &mut KeyEvaluator,

461

        data_callback: &mut DataCallback,

462

    ) -> Result<(), AbortError<CallerError>>

463

    where

464

        Range: RangeBounds<SequenceId> + Debug + 'static,

465

        KeyEvaluator: FnMut(KeySequence<Index>) -> ScanEvaluation,

466

        DataCallback:

467

            FnMut(KeySequence<Index>, ArcBytes<'static>) -> Result<(), AbortError<CallerError>>,

468

        CallerError: Display + Debug,

469

470

        self.tree

471

            .scan_sequences(range, forwards, true, key_evaluator, data_callback)

472

473

474

    /// Retrieves the keys and values associated with one or more `sequences`.

475

    /// The value retrieved is the value of the key at the given [`SequenceId`].

476

    /// If a sequence is not found, it will not appear in the result map. If

477

    /// the value was removed, None is returned for the value.

478

    pub fn get_multiple_by_sequence<Sequences>(

479

        &mut self,

480

        sequences: Sequences,

481

    ) -> Result<HashMap<SequenceId, (ArcBytes<'static>, Option<ArcBytes<'static>>)>, Error>

482

    where

483

        Sequences: Iterator<Item = SequenceId>,

484

485

        self.tree.get_multiple_by_sequence(sequences, true)

486

487

488

    /// Retrieves the keys and indexes associated with one or more `sequences`.

489

    /// The value retrieved is the value of the key at the given [`SequenceId`].

490

    /// If a sequence is not found, it will not appear in the result list.

491

    pub fn get_multiple_indexes_by_sequence<Sequences>(

492

        &mut self,

493

        sequences: Sequences,

494

    ) -> Result<Vec<SequenceIndex<Index>>, Error>

495

    where

496

        Sequences: Iterator<Item = SequenceId>,

497

498

        self.tree.get_multiple_indexes_by_sequence(sequences, true)

499

500

501

    /// Retrieves the keys, values, and indexes associated with one or more

502

    /// `sequences`. The value retrieved is the value of the key at the given

503

    /// [`SequenceId`]. If a sequence is not found, it will not appear in the

504

    /// result list.

505

    pub fn get_multiple_with_indexes_by_sequence<Sequences>(

506

        &mut self,

507

        sequences: Sequences,

508

    ) -> Result<HashMap<SequenceId, SequenceEntry<Index>>, Error>

509

    where

510

        Sequences: Iterator<Item = SequenceId>,

511

512

        self.tree

513

            .get_multiple_with_indexes_by_sequence(sequences, true)

514

515

516

517

impl<Root: tree::Root, File: ManagedFile> TransactionTree<Root, File> {

518

    /// Sets `key` to `value`. Returns the newly created index for this key.

519

16402

    pub fn set(

520

16402

        &mut self,

521

16402

        key: impl Into<ArcBytes<'static>>,

522

16402

        value: impl Into<Root::Value>,

523

16402

    ) -> Result<Root::Index, Error> {

524

16402

        self.tree.set(

525

16402

            PersistenceMode::Transactional(self.transaction_id),

526

16402

            key,

527

16402

            value,

528

16402

529

16402

530

531

    /// Executes a modification. Returns a list of all changed keys.

532

    pub fn modify<'a>(

533

        &mut self,

534

        keys: Vec<ArcBytes<'a>>,

535

        operation: Operation<'a, Root::Value, Root::Index>,

536

    ) -> Result<Vec<ModificationResult<Root::Index>>, Error> {

537

        self.tree.modify(Modification {

538

            keys,

539

            persistence_mode: PersistenceMode::Transactional(self.transaction_id),

540

            operation,

541

})

542

543

544

    /// Sets `key` to `value`. Returns a tuple containing two elements:

545

///

546

    /// - The previously stored value, if a value was already present.

547

    /// - The new/updated index for this key.

548

    pub fn replace(

549

        &mut self,

550

        key: impl Into<ArcBytes<'static>>,

551

        value: impl Into<Root::Value>,

552

    ) -> Result<(Option<Root::Value>, Root::Index), Error> {

553

        self.tree.replace(key, value, self.transaction_id)

554

555

556

    /// Returns the current value of `key`. This will return updated information

557

    /// if it has been previously updated within this transaction.

558

2

    pub fn get(&mut self, key: &[u8]) -> Result<Option<Root::Value>, Error> {

559

2

        self.tree.get(key, true)

560

2

561

562

    /// Returns the current index of `key`. This will return updated information

563

    /// if it has been previously updated within this transaction.

564

    pub fn get_index(&mut self, key: &[u8]) -> Result<Option<Root::Index>, Error> {

565

        self.tree.get_index(key, true)

566

567

568

    /// Returns the current value and index of `key`. This will return updated

569

    /// information if it has been previously updated within this transaction.

570

    pub fn get_with_index(&mut self, key: &[u8]) -> Result<Option<TreeValueIndex<Root>>, Error> {

571

        self.tree.get_with_index(key, true)

572

573

574

    /// Removes `key` and returns the existing value amd index, if present.

575

8192

    pub fn remove(&mut self, key: &[u8]) -> Result<Option<TreeValueIndex<Root>>, Error> {

576

8192

        self.tree.remove(key, self.transaction_id)

577

8192

578

579

    /// Compares the value of `key` against `old`. If the values match, key will

580

    /// be set to the new value if `new` is `Some` or removed if `new` is

581

    /// `None`.

582

    pub fn compare_and_swap<Old>(

583

        &mut self,

584

        key: &[u8],

585

        old: Option<&Old>,

586

        new: Option<Root::Value>,

587

    ) -> Result<(), CompareAndSwapError<Root::Value>>

588

    where

589

        Old: PartialEq,

590

        Root::Value: AsRef<Old> + Clone,

591

592

        self.tree

593

            .compare_and_swap(key, old, new, self.transaction_id)

594

595

596

    /// Retrieves the values of `keys`. If any keys are not found, they will be

597

    /// omitted from the results. Keys are required to be pre-sorted.

598

    pub fn get_multiple<'keys, KeysIntoIter, KeysIter>(

599

        &mut self,

600

        keys: KeysIntoIter,

601

    ) -> Result<Vec<(ArcBytes<'static>, Root::Value)>, Error>

602

    where

603

        KeysIntoIter: IntoIterator<Item = &'keys [u8], IntoIter = KeysIter>,

604

        KeysIter: Iterator<Item = &'keys [u8]> + ExactSizeIterator,

605

606

        self.tree.get_multiple(keys, true)

607

608

609

    /// Retrieves the indexes of `keys`. If any keys are not found, they will be

610

    /// omitted from the results. Keys are required to be pre-sorted.

611

    pub fn get_multiple_indexes<'keys, KeysIntoIter, KeysIter>(

612

        &mut self,

613

        keys: KeysIntoIter,

614

    ) -> Result<Vec<(ArcBytes<'static>, Root::Index)>, Error>

615

    where

616

        KeysIntoIter: IntoIterator<Item = &'keys [u8], IntoIter = KeysIter>,

617

        KeysIter: Iterator<Item = &'keys [u8]> + ExactSizeIterator,

618

619

        self.tree.get_multiple_indexes(keys, true)

620

621

622

    /// Retrieves the values and indexes of `keys`. If any keys are not found,

623

    /// they will be omitted from the results. Keys are required to be

624

    /// pre-sorted.

625

    pub fn get_multiple_with_indexes<'keys, KeysIntoIter, KeysIter>(

626

        &mut self,

627

        keys: KeysIntoIter,

628

    ) -> Result<Vec<TreeEntry<Root>>, Error>

629

    where

630

        KeysIntoIter: IntoIterator<Item = &'keys [u8], IntoIter = KeysIter>,

631

        KeysIter: Iterator<Item = &'keys [u8]> + ExactSizeIterator,

632

633

        self.tree.get_multiple_with_indexes(keys, true)

634

635

636

    /// Retrieves all of the values of keys within `range`.

637

    pub fn get_range<'keys, KeyRangeBounds>(

638

        &mut self,

639

        range: &'keys KeyRangeBounds,

640

    ) -> Result<Vec<(ArcBytes<'static>, Root::Value)>, Error>

641

    where

642

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

643

644

        self.tree.get_range(range, true)

645

646

647

    /// Retrieves all of the indexes of keys within `range`.

648

    pub fn get_range_indexes<'keys, KeyRangeBounds>(

649

        &mut self,

650

        range: &'keys KeyRangeBounds,

651

    ) -> Result<Vec<(ArcBytes<'static>, Root::Index)>, Error>

652

    where

653

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

654

655

        self.tree.get_range_indexes(range, true)

656

657

658

    /// Retrieves all of the values and indexes of keys within `range`.

659

    pub fn get_range_with_indexes<'keys, KeyRangeBounds>(

660

        &mut self,

661

        range: &'keys KeyRangeBounds,

662

    ) -> Result<Vec<TreeEntry<Root>>, Error>

663

    where

664

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

665

666

        self.tree.get_range_with_indexes(range, true)

667

668

669

    /// Scans the tree across all nodes that might contain nodes within `range`.

670

///

671

    /// If `forwards` is true, the tree is scanned in ascending order.

672

    /// Otherwise, the tree is scanned in descending order.

673

///

674

    /// `node_evaluator` is invoked for each [`Interior`](crate::tree::Interior)

675

    /// node to determine if the node should be traversed. The parameters to the

676

    /// callback are:

677

///

678

    /// - `&ArcBytes<'static>`: The maximum key stored within the all children

679

    ///   nodes.

680

    /// - `&Root::ReducedIndex`: The reduced index value stored within the node.

681

    /// - `usize`: The depth of the node. The root nodes are depth 0.

682

///

683

    /// The result of the callback is a [`ScanEvaluation`]. To read children

684

    /// nodes, return [`ScanEvaluation::ReadData`].

685

///

686

    /// `key_evaluator` is invoked for each key encountered that is contained

687

    /// within `range`. For all [`ScanEvaluation::ReadData`] results returned,

688

    /// `callback` will be invoked with the key and values. `callback` may not

689

    /// be invoked in the same order as the keys are scanned.

690

    #[cfg_attr(

691

        feature = "tracing",

692

        tracing::instrument(skip(self, node_evaluator, key_evaluator, callback))

693

)]

694

    pub fn scan<'b, 'keys, CallerError, KeyRangeBounds, NodeEvaluator, KeyEvaluator, DataCallback>(

695

        &mut self,

696

        range: &'keys KeyRangeBounds,

697

        forwards: bool,

698

        mut node_evaluator: NodeEvaluator,

699

        mut key_evaluator: KeyEvaluator,

700

        mut callback: DataCallback,

701

    ) -> Result<(), AbortError<CallerError>>

702

    where

703

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

704

        NodeEvaluator: FnMut(&ArcBytes<'static>, &Root::ReducedIndex, usize) -> ScanEvaluation,

705

        KeyEvaluator: FnMut(&ArcBytes<'static>, &Root::Index) -> ScanEvaluation,

706

        DataCallback: FnMut(

707

            ArcBytes<'static>,

708

            &Root::Index,

709

            Root::Value,

710

        ) -> Result<(), AbortError<CallerError>>,

711

        CallerError: Display + Debug,

712

713

        self.tree.scan(

714

            range,

715

            forwards,

716

            true,

717

            &mut node_evaluator,

718

            &mut key_evaluator,

719

            &mut callback,

720

721

722

723

    /// Returns the reduced index over the provided range. This is an

724

    /// aggregation function that builds atop the `scan()` operation which calls

725

    /// [`Reducer::reduce()`](crate::tree::Reducer::reduce) and

726

    /// [`Reducer::rereduce()`](crate::tree::Reducer::rereduce) on all matching

727

    /// indexes stored within the nodes of this tree, producing a single

728

    /// aggregated [`Root::ReducedIndex`](tree::Root::ReducedIndex) value.

729

///

730

    /// If no keys match, the returned result is what

731

    /// [`Reducer::rereduce()`](crate::tree::Reducer::rereduce) returns when an

732

    /// empty slice is provided.

733

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

734

    pub fn reduce<'keys, KeyRangeBounds>(

735

        &mut self,

736

        range: &'keys KeyRangeBounds,

737

    ) -> Result<Option<Root::ReducedIndex>, Error>

738

    where

739

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + Clone + ?Sized,

740

741

        self.tree.reduce(range, true)

742

743

744

    /// Returns the first key of the tree.

745

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

746

    pub fn first_key(&mut self) -> Result<Option<ArcBytes<'static>>, Error> {

747

        self.tree.first_key(true)

748

749

750

    /// Returns the first key and value of the tree.

751

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

752

    pub fn first(&mut self) -> Result<Option<(ArcBytes<'static>, Root::Value)>, Error> {

753

        self.tree.first(true)

754

755

756

    /// Returns the last key of the tree.

757

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

758

    pub fn last_key(&mut self) -> Result<Option<ArcBytes<'static>>, Error> {

759

        self.tree.last_key(true)

760

761

762

    /// Returns the last key and value of the tree.

763

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

764

    pub fn last(&mut self) -> Result<Option<(ArcBytes<'static>, Root::Value)>, Error> {

765

        self.tree.last(true)

766

767

768

769

/// An error returned from `compare_and_swap()`.

770

#[derive(Debug, thiserror::Error)]

771

pub enum CompareAndSwapError<Value: Debug> {

772

    /// The stored value did not match the conditional value.

773

    #[error("value did not match. existing value: {0:?}")]

774

    Conflict(Option<Value>),

775

    /// Another error occurred while executing the operation.

776

    #[error("error during compare_and_swap: {0}")]

777

    Error(#[from] Error),

778

779

780

/// A database configuration used to open a database.

781

#[derive(Debug)]

782

#[must_use]

783

pub struct Config<M: FileManager = StdFileManager> {

784

    path: PathBuf,

785

    vault: Option<Arc<dyn AnyVault>>,

786

    cache: Option<ChunkCache>,

787

    file_manager: Option<M>,

788

    thread_pool: Option<ThreadPool<M::File>>,

789

790

791

impl<M: FileManager> Clone for Config<M> {

792

1

    fn clone(&self) -> Self {

793

1

        Self {

794

1

            path: self.path.clone(),

795

1

            vault: self.vault.clone(),

796

1

            cache: self.cache.clone(),

797

1

            file_manager: self.file_manager.clone(),

798

1

            thread_pool: self.thread_pool.clone(),

799

1

800

1

801

802

803

impl Config<StdFileManager> {

804

    /// Creates a new config to open a database located at `path`.

805

16

    pub fn new<P: AsRef<Path>>(path: P) -> Self {

806

16

        Self {

807

16

            path: path.as_ref().to_path_buf(),

808

16

            vault: None,

809

16

            cache: None,

810

16

            thread_pool: None,

811

16

            file_manager: None,

812

16

813

16

814

815

    /// Returns a default configuration to open a database located at `path`.

816

    pub fn default_for<P: AsRef<Path>>(path: P) -> Self {

817

        Self {

818

            path: path.as_ref().to_path_buf(),

819

            vault: None,

820

            cache: Some(ChunkCache::new(2000, 65536)),

821

            thread_pool: Some(ThreadPool::default()),

822

            file_manager: None,

823

824

825

826

    /// Sets the file manager.

827

///

828

    /// ## Panics

829

///

830

    /// Panics if called after a shared thread pool has been set.

831

10

    pub fn file_manager<M: FileManager>(self, file_manager: M) -> Config<M> {

832

10

        assert!(self.thread_pool.is_none());

833

10

        Config {

834

10

            path: self.path,

835

10

            vault: self.vault,

836

10

            cache: self.cache,

837

10

            file_manager: Some(file_manager),

838

10

            thread_pool: None,

839

10

840

10

841

842

843

impl<M: FileManager> Config<M> {

844

    /// Sets the vault to use for this database.

845

2

    pub fn vault<V: AnyVault>(mut self, vault: V) -> Self {

846

2

        self.vault = Some(Arc::new(vault));

847

2

        self

848

2

849

850

    /// Sets the chunk cache to use for this database.

851

    pub fn cache(mut self, cache: ChunkCache) -> Self {

852

        self.cache = Some(cache);

853

        self

854

855

856

    /// Uses the `thread_pool` provided instead of creating its own. This will

857

    /// allow a single thread pool to manage multiple [`Roots`] instances'

858

    /// transactions.

859

    pub fn shared_thread_pool(mut self, thread_pool: &ThreadPool<M::File>) -> Self {

860

        self.thread_pool = Some(thread_pool.clone());

861

        self

862

863

864

    /// Opens the database, or creates one if the target path doesn't exist.

865

17

    pub fn open(self) -> Result<Roots<M::File>, Error> {

866

17

        Roots::open(

867

17

            self.path,

868

17

            Context {

869

17

                file_manager: self.file_manager.unwrap_or_default(),

870

17

                vault: self.vault,

871

17

                cache: self.cache,

872

17

},

873

17

            self.thread_pool.unwrap_or_default(),

874

17

875

17

876

877

878

/// A named collection of keys and values.

879

pub struct Tree<Root: tree::Root, File: ManagedFile> {

880

    roots: Roots<File>,

881

    path: PathId,

882

    state: State<Root>,

883

    reducer: Arc<dyn AnyReducer>,

884

    vault: Option<Arc<dyn AnyVault>>,

885

    name: Cow<'static, str>,

886

887

888

impl<Root: tree::Root, File: ManagedFile> Clone for Tree<Root, File> {

889

32

    fn clone(&self) -> Self {

890

32

        Self {

891

32

            roots: self.roots.clone(),

892

32

            path: self.path.clone(),

893

32

            state: self.state.clone(),

894

32

            vault: self.vault.clone(),

895

32

            reducer: self.reducer.clone(),

896

32

            name: self.name.clone(),

897

32

898

32

899

900

901

impl<Root: tree::Root, File: ManagedFile> Tree<Root, File> {

902

    /// Returns the name of the tree.

903

    #[must_use]

904

    pub fn name(&self) -> &str {

905

        &self.name

906

907

908

    /// Returns the path to the file for this tree.

909

    #[must_use]

910

    pub fn path(&self) -> &Path {

911

        self.path.path()

912

913

914

    /// Returns the number of keys stored in the tree. Does not include deleted keys.

915

    #[must_use]

916

3506

    pub fn count(&self) -> u64 {

917

3506

        let state = self.state.lock();

918

3506

        state.root.count()

919

3506

920

921

    /// Sets `key` to `value`. This is executed within its own transaction.

922

    #[allow(clippy::missing_panics_doc)]

923

16398

    pub fn set(

924

16398

        &self,

925

16398

        key: impl Into<ArcBytes<'static>>,

926

16398

        value: impl Into<Root::Value>,

927

16398

    ) -> Result<(), Error> {

928

16398

        let transaction = self.begin_transaction()?;

929

16398

        transaction.tree::<Root>(0).unwrap().set(key, value)?;

930

16398

        transaction.commit()

931

16398

932

933

24590

    fn begin_transaction(&self) -> Result<ExecutingTransaction<File>, Error> {

934

24590

        let reducer = self

935

24590

            .reducer

936

24590

            .as_ref()

937

24590

            .as_any()

938

24590

            .downcast_ref::<Root::Reducer>()

939

24590

            .unwrap()

940

24590

            .clone();

941

24590

        let mut root = Root::tree_with_reducer(self.name.clone(), reducer);

942

24590

        if let Some(vault) = &self.vault {

943

1

            root.vault = Some(vault.clone());

944

24589

945

24590

        self.roots.transaction(&[root])

946

24590

947

948

    /// Returns a [`TreeFile`] for lower-level operations within the context of

949

    /// Roots.

950

///

951

    /// Using this direct access, it is possible to circumvent some of the

952

    /// safety provided by Roots (e.g., passing an incorrect value for

953

    /// `in_transaction`). This function is provided for those who are

954

    /// implementing custom roots and wish to expose functionality through

955

    /// Roots.

956

11700

    pub fn open_for_read(&self) -> Result<TreeFile<Root, File>, Error> {

957

11700

        let context = self.vault.as_ref().map_or_else(

958

11700

            || Cow::Borrowed(self.roots.context()),

959

11700

            |vault| Cow::Owned(self.roots.context().clone().with_any_vault(vault.clone())),

960

11700

);

961

11700

962

11700

        TreeFile::<Root, File>::read(

963

11700

            &self.path,

964

11700

            self.state.clone(),

965

11700

            &context,

966

11700

            Some(self.roots.transactions()),

967

11700

968

11700

969

970

    /// Returns a [`TreeFile`] for lower-level operations within the context of

971

    /// Roots.

972

///

973

    /// Using this direct access, it is possible to circumvent some of the

974

    /// safety provided by Roots (e.g., passing an incorrect value for

975

    /// `in_transaction`). This function is provided for those who are

976

    /// implementing custom roots and wish to expose functionality through

977

    /// Roots.

978

    pub fn open_for_write(&self) -> Result<TreeFile<Root, File>, Error> {

979

        let context = self.vault.as_ref().map_or_else(

980

            || Cow::Borrowed(self.roots.context()),

981

            |vault| Cow::Owned(self.roots.context().clone().with_any_vault(vault.clone())),

982

);

983

984

        TreeFile::<Root, File>::write(

985

            &self.path,

986

            self.state.clone(),

987

            &context,

988

            Some(self.roots.transactions()),

989

990

991

992

    /// Retrieves the current value of `key`, if present. Does not reflect any

993

    /// changes in pending transactions.

994

8202

    pub fn get(&self, key: &[u8]) -> Result<Option<Root::Value>, Error> {

995

8202

        catch_compaction_and_retry(|| {

996

8202

            let mut tree = match self.open_for_read() {

997

8200

                Ok(tree) => tree,

998

2

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

999

2

                Err(err) => return Err(err),

1000

};

1001

1002

8200

            tree.get(key, false)

1003

8202

})

1004

8202

1005

1006

    /// Retrieves the current index of `key`, if present. Does not reflect any

1007

    /// changes in pending transactions.

1008

    pub fn get_index(&self, key: &[u8]) -> Result<Option<Root::Index>, Error> {

1009

        catch_compaction_and_retry(|| {

1010

            let mut tree = match self.open_for_read() {

1011

                Ok(tree) => tree,

1012

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1013

                Err(err) => return Err(err),

1014

};

1015

1016

            tree.get_index(key, false)

1017

})

1018

1019

1020

    /// Retrieves the current value and index of `key`, if present. Does not reflect any

1021

    /// changes in pending transactions.

1022

    pub fn get_with_index(&self, key: &[u8]) -> Result<Option<TreeValueIndex<Root>>, Error> {

1023

        catch_compaction_and_retry(|| {

1024

            let mut tree = match self.open_for_read() {

1025

                Ok(tree) => tree,

1026

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1027

                Err(err) => return Err(err),

1028

};

1029

1030

            tree.get_with_index(key, false)

1031

})

1032

1033

1034

    /// Sets `key` to `value`. Returns a tuple containing two elements:

1035

///

1036

    /// - The previously stored value, if a value was already present.

1037

    /// - The new/updated index for this key.

1038

    #[allow(clippy::missing_panics_doc)]

1039

    pub fn replace(

1040

        &self,

1041

        key: impl Into<ArcBytes<'static>>,

1042

        value: impl Into<Root::Value>,

1043

    ) -> Result<(Option<Root::Value>, Root::Index), Error> {

1044

        let transaction = self.begin_transaction()?;

1045

        let existing_value = transaction.tree::<Root>(0).unwrap().replace(key, value)?;

1046

        transaction.commit()?;

1047

        Ok(existing_value)

1048

1049

1050

    /// Executes a modification. Returns a list of all changed keys.

1051

    #[allow(clippy::missing_panics_doc)]

1052

    pub fn modify<'a>(

1053

        &self,

1054

        keys: Vec<ArcBytes<'a>>,

1055

        operation: Operation<'a, Root::Value, Root::Index>,

1056

    ) -> Result<Vec<ModificationResult<Root::Index>>, Error> {

1057

        let transaction = self.begin_transaction()?;

1058

        let results = transaction

1059

            .tree::<Root>(0)

1060

            .unwrap()

1061

            .modify(keys, operation)?;

1062

        transaction.commit()?;

1063

        Ok(results)

1064

1065

1066

    /// Removes `key` and returns the existing value and index, if present. This

1067

    /// is executed within its own transaction.

1068

    #[allow(clippy::missing_panics_doc)]

1069

8192

    pub fn remove(&self, key: &[u8]) -> Result<Option<TreeValueIndex<Root>>, Error> {

1070

8192

        let transaction = self.begin_transaction()?;

1071

8192

        let existing_value = transaction.tree::<Root>(0).unwrap().remove(key)?;

1072

8192

        transaction.commit()?;

1073

8192

        Ok(existing_value)

1074

8192

1075

1076

    /// Compares the value of `key` against `old`. If the values match, key will

1077

    /// be set to the new value if `new` is `Some` or removed if `new` is

1078

    /// `None`. This is executed within its own transaction.

1079

    #[allow(clippy::missing_panics_doc)]

1080

    pub fn compare_and_swap<Old>(

1081

        &self,

1082

        key: &[u8],

1083

        old: Option<&Old>,

1084

        new: Option<Root::Value>,

1085

    ) -> Result<(), CompareAndSwapError<Root::Value>>

1086

    where

1087

        Old: PartialEq,

1088

        Root::Value: AsRef<Old> + Clone,

1089

1090

        let transaction = self.begin_transaction()?;

1091

        transaction

1092

            .tree::<Root>(0)

1093

            .unwrap()

1094

            .compare_and_swap(key, old, new)?;

1095

        transaction.commit()?;

1096

        Ok(())

1097

1098

1099

    /// Retrieves the values of `keys`. If any keys are not found, they will be

1100

    /// omitted from the results. Keys are required to be pre-sorted.

1101

    #[allow(clippy::needless_pass_by_value)]

1102

    pub fn get_multiple<'keys, Keys>(

1103

        &self,

1104

        keys: Keys,

1105

    ) -> Result<Vec<(ArcBytes<'static>, Root::Value)>, Error>

1106

    where

1107

        Keys: Iterator<Item = &'keys [u8]> + ExactSizeIterator + Clone,

1108

1109

        catch_compaction_and_retry(|| {

1110

            let mut tree = match self.open_for_read() {

1111

                Ok(tree) => tree,

1112

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1113

                Err(err) => return Err(err),

1114

};

1115

1116

            tree.get_multiple(keys.clone(), false)

1117

})

1118

1119

1120

    /// Retrieves the indexes of `keys`. If any keys are not found, they will be

1121

    /// omitted from the results. Keys are required to be pre-sorted.

1122

    #[allow(clippy::needless_pass_by_value)]

1123

    pub fn get_multiple_indexes<'keys, KeysIntoIter, KeysIter>(

1124

        &self,

1125

        keys: KeysIntoIter,

1126

    ) -> Result<Vec<(ArcBytes<'static>, Root::Index)>, Error>

1127

    where

1128

        KeysIntoIter: IntoIterator<Item = &'keys [u8], IntoIter = KeysIter> + Clone,

1129

        KeysIter: Iterator<Item = &'keys [u8]> + ExactSizeIterator,

1130

1131

        catch_compaction_and_retry(|| {

1132

            let mut tree = match self.open_for_read() {

1133

                Ok(tree) => tree,

1134

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1135

                Err(err) => return Err(err),

1136

};

1137

1138

            tree.get_multiple_indexes(keys.clone(), false)

1139

})

1140

1141

1142

    /// Retrieves the values and indexes of `keys`. If any keys are not found,

1143

    /// they will be omitted from the results. Keys are required to be

1144

    /// pre-sorted.

1145

    #[allow(clippy::needless_pass_by_value)]

1146

    pub fn get_multiple_with_indexes<'keys, KeysIntoIter, KeysIter>(

1147

        &self,

1148

        keys: KeysIntoIter,

1149

    ) -> Result<Vec<TreeEntry<Root>>, Error>

1150

    where

1151

        KeysIntoIter: IntoIterator<Item = &'keys [u8], IntoIter = KeysIter> + Clone,

1152

        KeysIter: Iterator<Item = &'keys [u8]> + ExactSizeIterator,

1153

1154

        catch_compaction_and_retry(|| {

1155

            let mut tree = match self.open_for_read() {

1156

                Ok(tree) => tree,

1157

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1158

                Err(err) => return Err(err),

1159

};

1160

1161

            tree.get_multiple_with_indexes(keys.clone(), false)

1162

})

1163

1164

1165

    /// Retrieves all of the values of keys within `range`.

1166

    pub fn get_range<'keys, KeyRangeBounds>(

1167

        &self,

1168

        range: &'keys KeyRangeBounds,

1169

    ) -> Result<Vec<(ArcBytes<'static>, Root::Value)>, Error>

1170

    where

1171

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + Clone + ?Sized,

1172

1173

        catch_compaction_and_retry(|| {

1174

            let mut tree = match self.open_for_read() {

1175

                Ok(tree) => tree,

1176

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1177

                Err(err) => return Err(err),

1178

};

1179

1180

            tree.get_range(range, false)

1181

})

1182

1183

1184

    /// Retrieves all of the indexes of keys within `range`.

1185

    pub fn get_range_indexes<'keys, KeyRangeBounds>(

1186

        &self,

1187

        range: &'keys KeyRangeBounds,

1188

    ) -> Result<Vec<(ArcBytes<'static>, Root::Index)>, Error>

1189

    where

1190

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

1191

1192

        catch_compaction_and_retry(|| {

1193

            let mut tree = match self.open_for_read() {

1194

                Ok(tree) => tree,

1195

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1196

                Err(err) => return Err(err),

1197

};

1198

1199

            tree.get_range_indexes(range, false)

1200

})

1201

1202

1203

    /// Retrieves all of the values and indexes of keys within `range`.

1204

    pub fn get_range_with_indexes<'keys, KeyRangeBounds>(

1205

        &self,

1206

        range: &'keys KeyRangeBounds,

1207

    ) -> Result<Vec<TreeEntry<Root>>, Error>

1208

    where

1209

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + ?Sized,

1210

1211

        catch_compaction_and_retry(|| {

1212

            let mut tree = match self.open_for_read() {

1213

                Ok(tree) => tree,

1214

                Err(err) if err.kind.is_file_not_found() => return Ok(Vec::new()),

1215

                Err(err) => return Err(err),

1216

};

1217

1218

            tree.get_range_with_indexes(range, false)

1219

})

1220

1221

1222

    /// Scans the tree across all nodes that might contain nodes within `range`.

1223

///

1224

    /// If `forwards` is true, the tree is scanned in ascending order.

1225

    /// Otherwise, the tree is scanned in descending order.

1226

///

1227

    /// `node_evaluator` is invoked for each [`Interior`](crate::tree::Interior) node to determine if

1228

    /// the node should be traversed. The parameters to the callback are:

1229

///

1230

    /// - `&ArcBytes<'static>`: The maximum key stored within the all children

1231

    ///   nodes.

1232

    /// - `&Root::ReducedIndex`: The reduced index value stored within the node.

1233

    /// - `usize`: The depth of the node. The root nodes are depth 0.

1234

///

1235

    /// The result of the callback is a [`ScanEvaluation`]. To read children

1236

    /// nodes, return [`ScanEvaluation::ReadData`].

1237

///

1238

    /// `key_evaluator` is invoked for each key encountered that is contained

1239

    /// within `range`. For all [`ScanEvaluation::ReadData`] results returned,

1240

    /// `callback` will be invoked with the key and values. `callback` may not

1241

    /// be invoked in the same order as the keys are scanned.

1242

    #[cfg_attr(

1243

        feature = "tracing",

1244

        tracing::instrument(skip(self, node_evaluator, key_evaluator, callback))

1245

)]

1246

    pub fn scan<'keys, CallerError, KeyRangeBounds, NodeEvaluator, KeyEvaluator, DataCallback>(

1247

        &self,

1248

        range: &'keys KeyRangeBounds,

1249

        forwards: bool,

1250

        mut node_evaluator: NodeEvaluator,

1251

        mut key_evaluator: KeyEvaluator,

1252

        mut callback: DataCallback,

1253

    ) -> Result<(), AbortError<CallerError>>

1254

    where

1255

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + Clone + ?Sized,

1256

        NodeEvaluator: FnMut(&ArcBytes<'static>, &Root::ReducedIndex, usize) -> ScanEvaluation,

1257

        KeyEvaluator: FnMut(&ArcBytes<'static>, &Root::Index) -> ScanEvaluation,

1258

        DataCallback: FnMut(

1259

            ArcBytes<'static>,

1260

            &Root::Index,

1261

            Root::Value,

1262

        ) -> Result<(), AbortError<CallerError>>,

1263

        CallerError: Display + Debug,

1264

1265

        catch_compaction_and_retry_abortable(move || {

1266

            let mut tree = match self.open_for_read() {

1267

                Ok(tree) => tree,

1268

                Err(err) if err.kind.is_file_not_found() => return Ok(()),

1269

                Err(err) => return Err(AbortError::from(err)),

1270

};

1271

1272

            tree.scan(

1273

                range,

1274

                forwards,

1275

                false,

1276

                &mut node_evaluator,

1277

                &mut key_evaluator,

1278

                &mut callback,

1279

1280

})

1281

1282

1283

    /// Returns the reduced index over the provided range. This is an

1284

    /// aggregation function that builds atop the `scan()` operation which calls

1285

    /// [`Reducer::reduce()`](crate::tree::Reducer::reduce) and

1286

    /// [`Reducer::rereduce()`](crate::tree::Reducer::rereduce) on all matching

1287

    /// indexes stored within the nodes of this tree, producing a single

1288

    /// aggregated [`Root::ReducedIndex`](tree::Root::ReducedIndex) value.

1289

///

1290

    /// If no keys match, the returned result is what

1291

    /// [`Reducer::rereduce()`](crate::tree::Reducer::rereduce) returns when an

1292

    /// empty slice is provided.

1293

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

1294

    pub fn reduce<'keys, KeyRangeBounds>(

1295

        &self,

1296

        range: &'keys KeyRangeBounds,

1297

    ) -> Result<Option<Root::ReducedIndex>, Error>

1298

    where

1299

        KeyRangeBounds: RangeBounds<&'keys [u8]> + Debug + Clone + ?Sized,

1300

1301

        catch_compaction_and_retry(move || {

1302

            let mut tree = match self.open_for_read() {

1303

                Ok(tree) => tree,

1304

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1305

                Err(err) => return Err(err),

1306

};

1307

1308

            tree.reduce(range, false)

1309

})

1310

1311

1312

    /// Returns the first key of the tree.

1313

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

1314

    pub fn first_key(&self) -> Result<Option<ArcBytes<'static>>, Error> {

1315

        catch_compaction_and_retry(|| {

1316

            let mut tree = match self.open_for_read() {

1317

                Ok(tree) => tree,

1318

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1319

                Err(err) => return Err(err),

1320

};

1321

1322

            tree.first_key(false)

1323

})

1324

1325

1326

    /// Returns the first key and value of the tree.

1327

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

1328

    pub fn first(&self) -> Result<Option<(ArcBytes<'static>, Root::Value)>, Error> {

1329

        catch_compaction_and_retry(|| {

1330

            let mut tree = match self.open_for_read() {

1331

                Ok(tree) => tree,

1332

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1333

                Err(err) => return Err(err),

1334

};

1335

1336

            tree.first(false)

1337

})

1338

1339

1340

    /// Returns the last key of the tree.

1341

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

1342

    pub fn last_key(&self) -> Result<Option<ArcBytes<'static>>, Error> {

1343

        catch_compaction_and_retry(|| {

1344

            let mut tree = match self.open_for_read() {

1345

                Ok(tree) => tree,

1346

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1347

                Err(err) => return Err(err),

1348

};

1349

1350

            tree.last_key(false)

1351

})

1352

1353

1354

    /// Returns the last key and value of the tree.

1355

    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]

1356

    pub fn last(&self) -> Result<Option<(ArcBytes<'static>, Root::Value)>, Error> {

1357

        catch_compaction_and_retry(|| {

1358

            let mut tree = match self.open_for_read() {

1359

                Ok(tree) => tree,

1360

                Err(err) if err.kind.is_file_not_found() => return Ok(None),

1361

                Err(err) => return Err(err),

1362

};

1363

1364

            tree.last(false)

1365

})

1366

1367

1368

    /// Rewrites the database to remove data that is no longer current. Because

1369

    /// Nebari uses an append-only format, this is helpful in reducing disk

1370

    /// usage.

1371

///

1372

    /// See [`TreeFile::compact()`](crate::tree::TreeFile::compact) for more

1373

    /// information.

1374

3498

    pub fn compact(&self) -> Result<(), Error> {

1375

3498

        let tree = match self.open_for_read() {

1376

3498

            Ok(tree) => tree,

1377

            Err(err) if err.kind.is_file_not_found() => return Ok(()),

1378

            Err(err) => return Err(err),

1379

};

1380

3498

        tree.compact(

1381

3498

            &self.roots.context().file_manager,

1382

3498

            Some(TransactableCompaction {

1383

3498

                name: self.name.as_ref(),

1384

3498

                manager: self.roots.transactions(),

1385

3498

}),

1386

3498

)?;

1387

3498

        Ok(())

1388

3498

1389

1390

1391

impl<Root: tree::Root, File: ManagedFile> AnyTreeRoot<File> for Tree<Root, File> {

1392

    fn name(&self) -> &str {

1393

        &self.name

1394

1395

1396

    fn default_state(&self) -> Box<dyn AnyTreeState> {

1397

        Box::new(State::<Root>::new(

1398

            None,

1399

            None,

1400

            Root::default_with(

1401

                self.reducer

1402

                    .as_ref()

1403

                    .as_any()

1404

                    .downcast_ref::<Root::Reducer>()

1405

                    .unwrap()

1406

                    .clone(),

1407

),

1408

))

1409

1410

1411

    fn begin_transaction(

1412

        &self,

1413

        transaction_id: TransactionId,

1414

        file_path: &PathId,

1415

        state: &dyn AnyTreeState,

1416

        context: &Context<File::Manager>,

1417

        transactions: Option<&TransactionManager<File::Manager>>,

1418

    ) -> Result<Box<dyn AnyTransactionTree<File>>, Error> {

1419

        let context = self.vault.as_ref().map_or_else(

1420

            || Cow::Borrowed(context),

1421

            |vault| Cow::Owned(context.clone().with_any_vault(vault.clone())),

1422

);

1423

        let tree = TreeFile::write(

1424

            file_path,

1425

            state

1426

                .as_any()

1427

                .downcast_ref::<State<Root>>()

1428

                .unwrap()

1429

                .clone(),

1430

            &context,

1431

            transactions,

1432

)?;

1433

1434

        Ok(Box::new(TransactionTree {

1435

            transaction_id,

1436

            tree,

1437

}))

1438

1439

1440

1441

impl<File: ManagedFile, Index> Tree<VersionedTreeRoot<Index>, File>

1442

where

1443

    Index: EmbeddedIndex<ArcBytes<'static>> + Clone + Debug + 'static,

1444

1445

    /// Returns the latest sequence id.

1446

    #[must_use]

1447

    pub fn current_sequence_id(&self) -> SequenceId {

1448

        let state = self.state.lock();

1449

        state.root.sequence

1450

1451

1452

    /// Scans the tree for keys that are contained within `range`. If `forwards`

1453

    /// is true, scanning starts at the lowest sort-order key and scans forward.

1454

    /// Otherwise, scanning starts at the highest sort-order key and scans

1455

    /// backwards. `key_evaluator` is invoked for each key as it is encountered.

1456

    /// For all [`ScanEvaluation::ReadData`] results returned, `callback` will be

1457

    /// invoked with the key and values. The callback may not be invoked in the

1458

    /// same order as the keys are scanned.

1459

    pub fn scan_sequences<CallerError, Range, KeyEvaluator, DataCallback>(

1460

        &self,

1461

        range: Range,

1462

        forwards: bool,

1463

        mut key_evaluator: KeyEvaluator,

1464

        mut data_callback: DataCallback,

1465

    ) -> Result<(), AbortError<CallerError>>

1466

    where

1467

        Range: Clone + RangeBounds<SequenceId> + Debug + 'static,

1468

        KeyEvaluator: FnMut(KeySequence<Index>) -> ScanEvaluation,

1469

        DataCallback:

1470

            FnMut(KeySequence<Index>, ArcBytes<'static>) -> Result<(), AbortError<CallerError>>,

1471

        CallerError: Display + Debug,

1472

1473

        catch_compaction_and_retry_abortable(|| {

1474

            let mut tree = TreeFile::<VersionedTreeRoot<Index>, File>::read(

1475

                &self.path,

1476

                self.state.clone(),

1477

                self.roots.context(),

1478

                Some(self.roots.transactions()),

1479

)?;

1480

1481

            tree.scan_sequences(

1482

                range.clone(),

1483

                forwards,

1484

                false,

1485

                &mut key_evaluator,

1486

                &mut data_callback,

1487

1488

})

1489

1490

1491

    /// Retrieves the keys and values associated with one or more `sequences`.

1492

    /// The value retrieved is the value of the key at the given [`SequenceId`].

1493

    /// If a sequence is not found, it will not appear in the result map. If

1494

    /// the value was removed, None is returned for the value.

1495

    #[allow(clippy::needless_pass_by_value)]

1496

    pub fn get_multiple_by_sequence<Sequences>(

1497

        &self,

1498

        sequences: Sequences,

1499

    ) -> Result<HashMap<SequenceId, (ArcBytes<'static>, Option<ArcBytes<'static>>)>, Error>

1500

    where

1501

        Sequences: Iterator<Item = SequenceId> + Clone,

1502

1503

        catch_compaction_and_retry(|| {

1504

            let mut tree = TreeFile::<VersionedTreeRoot<Index>, File>::read(

1505

                &self.path,

1506

                self.state.clone(),

1507

                self.roots.context(),

1508

                Some(self.roots.transactions()),

1509

)?;

1510

1511

            tree.get_multiple_by_sequence(sequences.clone(), false)

1512

})

1513

1514

1515

    /// Retrieves the keys and indexes associated with one or more `sequences`.

1516

    /// The value retrieved is the value of the key at the given [`SequenceId`].

1517

    /// If a sequence is not found, it will not appear in the result list.

1518

    #[allow(clippy::needless_pass_by_value)]

1519

    pub fn get_multiple_indexes_by_sequence<Sequences>(

1520

        &self,

1521

        sequences: Sequences,

1522

    ) -> Result<Vec<SequenceIndex<Index>>, Error>

1523

    where

1524

        Sequences: Iterator<Item = SequenceId> + Clone,

1525

1526

        catch_compaction_and_retry(|| {

1527

            let mut tree = TreeFile::<VersionedTreeRoot<Index>, File>::read(

1528

                &self.path,

1529

                self.state.clone(),

1530

                self.roots.context(),

1531

                Some(self.roots.transactions()),

1532

)?;

1533

1534

            tree.get_multiple_indexes_by_sequence(sequences.clone(), false)

1535

})

1536

1537

1538

    /// Retrieves the keys, values, and indexes associated with one or more

1539

    /// `sequences`. The value retrieved is the value of the key at the given

1540

    /// [`SequenceId`]. If a sequence is not found, it will not appear in the

1541

    /// result list.

1542

    #[allow(clippy::needless_pass_by_value)]

1543

    pub fn get_multiple_with_indexes_by_sequence<Sequences>(

1544

        &self,

1545

        sequences: Sequences,

1546

    ) -> Result<HashMap<SequenceId, SequenceEntry<Index>>, Error>

1547

    where

1548

        Sequences: Iterator<Item = SequenceId> + Clone,

1549

1550

        catch_compaction_and_retry(|| {

1551

            let mut tree = TreeFile::<VersionedTreeRoot<Index>, File>::read(

1552

                &self.path,

1553

                self.state.clone(),

1554

                self.roots.context(),

1555

                Some(self.roots.transactions()),

1556

)?;

1557

1558

            tree.get_multiple_with_indexes_by_sequence(sequences.clone(), false)

1559

})

1560

1561

1562

1563

/// An error that could come from user code or Nebari.

1564

#[derive(thiserror::Error, Debug)]

1565

pub enum AbortError<CallerError: Display + Debug = Infallible> {

1566

    /// An error unrelated to Nebari occurred.

1567

    #[error("other error: {0}")]

1568

    Other(CallerError),

1569

    /// An error from Roots occurred.

1570

    #[error("database error: {0}")]

1571

    Nebari(#[from] Error),

1572

1573

1574

impl AbortError<Infallible> {

1575

    /// Unwraps the error contained within an infallible abort error.

1576

    #[must_use]

1577

    pub fn infallible(self) -> Error {

1578

        match self {

1579

            Self::Other(infallible) => match infallible {},

1580

            Self::Nebari(error) => error,

1581

1582

1583

1584

1585

/// A thread pool that commits transactions to disk in parallel.

1586

#[derive(Debug)]

1587

pub struct ThreadPool<File>

1588

where

1589

    File: ManagedFile,

1590

1591

    sender: flume::Sender<ThreadCommit<File>>,

1592

    receiver: flume::Receiver<ThreadCommit<File>>,

1593

    thread_count: Arc<AtomicU16>,

1594

    maximum_threads: usize,

1595

1596

1597

impl<File: ManagedFile> ThreadPool<File> {

1598

    /// Returns a thread pool that will spawn up to `maximum_threads` to process

1599

    /// file operations.

1600

    #[must_use]

1601

17

    pub fn new(maximum_threads: usize) -> Self {

1602

17

        let (sender, receiver) = flume::unbounded();

1603

17

        Self {

1604

17

            sender,

1605

17

            receiver,

1606

17

            thread_count: Arc::new(AtomicU16::new(0)),

1607

17

            maximum_threads,

1608

17

1609

17

1610

1611

24592

    fn commit_trees(

1612

24592

        &self,

1613

24592

        trees: Vec<UnlockedTransactionTree<File>>,

1614

24592

    ) -> Result<Vec<Box<dyn AnyTransactionTree<File>>>, Error> {

1615

24592

        // If we only have one tree, there's no reason to split IO across

1616

24592

        // threads. If we have multiple trees, we should split even with one

1617

24592

        // cpu: if one thread blocks, the other can continue executing.

1618

24592

        if trees.len() == 1 {

1619

24592

            let mut tree = trees.into_iter().next().unwrap().0.into_inner();

1620

24592

            tree.commit()?;

1621

24592

            Ok(vec![tree])

1622

        } else {

1623

            // Push the trees so that any existing threads can begin processing the queue.

1624

            let (completion_sender, completion_receiver) = flume::unbounded();

1625

            let tree_count = trees.len();

1626

            for tree in trees {

1627

                self.sender.send(ThreadCommit {

1628

                    tree: tree.0.into_inner(),

1629

                    completion_sender: completion_sender.clone(),

1630

                })?;

1631

1632

1633

            // Scale the queue if needed.

1634

            let desired_threads = tree_count.min(self.maximum_threads);

1635

            loop {

1636

                let thread_count = self.thread_count.load(Ordering::SeqCst);

1637

                if (thread_count as usize) >= desired_threads {

1638

                    break;

1639

1640

1641

                // Spawn a thread, but ensure that we don't spin up too many threads if another thread is committing at the same time.

1642

                if self

1643

                    .thread_count

1644

                    .compare_exchange(

1645

                        thread_count,

1646

                        thread_count + 1,

1647

                        Ordering::SeqCst,

1648

                        Ordering::SeqCst,

1649

1650

                    .is_ok()

1651

1652

                    let commit_receiver = self.receiver.clone();

1653

                    std::thread::Builder::new()

1654

                        .name(String::from("roots-txwriter"))

1655

                        .spawn(move || transaction_commit_thread(commit_receiver))

1656

                        .unwrap();

1657

1658

1659

1660

            // Wait for our results

1661

            let mut results = Vec::with_capacity(tree_count);

1662

            for _ in 0..tree_count {

1663

                results.push(completion_receiver.recv()??);

1664

1665

1666

            Ok(results)

1667

1668

24592

1669

1670

1671

impl<File: ManagedFile> Clone for ThreadPool<File> {

1672

    fn clone(&self) -> Self {

1673

        Self {

1674

            sender: self.sender.clone(),

1675

            receiver: self.receiver.clone(),

1676

            thread_count: self.thread_count.clone(),

1677

            maximum_threads: self.maximum_threads,

1678

1679

1680

1681

1682

impl<File: ManagedFile> Default for ThreadPool<File> {

1683

17

    fn default() -> Self {

1684

17

        static CPU_COUNT: Lazy<usize> = Lazy::new(num_cpus::get);

1685

17

        Self::new(*CPU_COUNT)

1686

17

1687

1688

1689

#[allow(clippy::needless_pass_by_value)]

1690

fn transaction_commit_thread<File: ManagedFile>(receiver: flume::Receiver<ThreadCommit<File>>) {

1691

    while let Ok(ThreadCommit {

1692

        mut tree,

1693

        completion_sender,

1694

    }) = receiver.recv()

1695

1696

        let result = tree.commit();

1697

        let result = result.map(move |_| tree);

1698

        drop(completion_sender.send(result));

1699

1700

1701

1702

struct ThreadCommit<File>

1703

where

1704

    File: ManagedFile,

1705

1706

    tree: Box<dyn AnyTransactionTree<File>>,

1707

    completion_sender: Sender<Result<Box<dyn AnyTransactionTree<File>>, Error>>,

1708

1709

1710

8202

fn catch_compaction_and_retry<R, F: Fn() -> Result<R, Error>>(func: F) -> Result<R, Error> {

1711

8202

    loop {

1712

8202

        match func() {

1713

8200

            Ok(result) => return Ok(result),

1714

2

            Err(error) => {

1715

2

                if matches!(error.kind, ErrorKind::TreeCompacted) {

1716

                    continue;

1717

2

1718

2

1719

2

                return Err(error);

1720

1721

1722

1723

8202

1724

1725

fn catch_compaction_and_retry_abortable<

1726

R,

1727

    E: Display + Debug,

1728

    F: FnMut() -> Result<R, AbortError<E>>,

1729

>(

1730

    mut func: F,

1731

) -> Result<R, AbortError<E>> {

1732

    loop {

1733

        match func() {

1734

            Ok(result) => return Ok(result),

1735

            Err(AbortError::Nebari(error)) => {

1736

                if matches!(error.kind, ErrorKind::TreeCompacted) {

1737

                    continue;

1738

1739

1740

                return Err(AbortError::Nebari(error));

1741

1742

            Err(other) => return Err(other),

1743

1744

1745

1746

1747

#[cfg(test)]

1748

mod tests {

1749

    use byteorder::{BigEndian, ByteOrder};

1750

    use tempfile::tempdir;

1751

1752

    use super::*;

1753

    use crate::{

1754

        io::{any::AnyFileManager, fs::StdFileManager, memory::MemoryFileManager},

1755

        test_util::RotatorVault,

1756

        tree::{Root, Unversioned, ValueIndex, Versioned},

1757

};

1758

1759

2

    fn basic_get_set<M: FileManager>(file_manager: M) {

1760

2

        let tempdir = tempdir().unwrap();

1761

2

        let roots = Config::new(tempdir.path())

1762

2

            .file_manager(file_manager)

1763

2

            .open()

1764

2

            .unwrap();

1765

2

1766

2

        let tree = roots.tree(Versioned::tree("test")).unwrap();

1767

2

        tree.set(b"test", b"value").unwrap();

1768

2

        let result = tree.get(b"test").unwrap().expect("key not found");

1769

2

1770

2

        assert_eq!(result, b"value");

1771

2

1772

1773

1

    #[test]

1774

1

    fn memory_basic_get_set() {

1775

1

        basic_get_set(MemoryFileManager::default());

1776

1

1777

1778

1

    #[test]

1779

1

    fn std_basic_get_set() {

1780

1

        basic_get_set(StdFileManager::default());

1781

1

1782

1783

1

    #[test]

1784

1

    fn basic_transaction_isolation_test() {

1785

1

        let tempdir = tempdir().unwrap();

1786

1

1787

1

        let roots = Config::<StdFileManager>::new(tempdir.path())

1788

1

            .open()

1789

1

            .unwrap();

1790

1

        let tree = roots.tree(Versioned::tree("test")).unwrap();

1791

1

        tree.set(b"test", b"value").unwrap();

1792

1

1793

1

        // Begin a transaction

1794

1

        let transaction = roots.transaction(&[Versioned::tree("test")]).unwrap();

1795

1

1796

1

        // Replace the key with a new value.

1797

1

        transaction

1798

1

            .tree::<Versioned>(0)

1799

1

            .unwrap()

1800

1

            .set(b"test", b"updated value")

1801

1

            .unwrap();

1802

1

1803

1

        // Check that the transaction can read the new value

1804

1

        let result = transaction

1805

1

            .tree::<Versioned>(0)

1806

1

            .unwrap()

1807

1

            .get(b"test")

1808

1

            .unwrap()

1809

1

            .expect("key not found");

1810

1

        assert_eq!(result, b"updated value");

1811

1812

        // Ensure that existing read-access doesn't see the new value

1813

1

        let result = tree.get(b"test").unwrap().expect("key not found");

1814

1

        assert_eq!(result, b"value");

1815

1816

        // Commit the transaction

1817

1

        transaction.commit().unwrap();

1818

1

1819

1

        // Ensure that the reader now sees the new value

1820

1

        let result = tree.get(b"test").unwrap().expect("key not found");

1821

1

        assert_eq!(result, b"updated value");

1822

1

1823

1824

1

    #[test]

1825

1

    fn basic_transaction_rollback_test() {

1826

1

        let tempdir = tempdir().unwrap();

1827

1

1828

1

        let roots = Config::<StdFileManager>::new(tempdir.path())

1829

1

            .open()

1830

1

            .unwrap();

1831

1

        let tree = roots.tree(Versioned::tree("test")).unwrap();

1832

1

        tree.set(b"test", b"value").unwrap();

1833

1

1834

1

        // Begin a transaction

1835

1

        let transaction = roots.transaction(&[Versioned::tree("test")]).unwrap();

1836

1

1837

1

        // Replace the key with a new value.

1838

1

        transaction

1839

1

            .tree::<Versioned>(0)

1840

1

            .unwrap()

1841

1

            .set(b"test", b"updated value")

1842

1

            .unwrap();

1843

1

1844

1

        // Roll the transaction back

1845

1

        transaction.rollback();

1846

1

1847

1

        // Ensure that the reader still sees the old value

1848

1

        let result = tree.get(b"test").unwrap().expect("key not found");

1849

1

        assert_eq!(result, b"value");

1850

1851

        // Begin a new transaction

1852

1

        let transaction = roots.transaction(&[Versioned::tree("test")]).unwrap();

1853

1

        // Check that the transaction has the original value

1854

1

        let result = transaction

1855

1

            .tree::<Versioned>(0)

1856

1

            .unwrap()

1857

1

            .get(b"test")

1858

1

            .unwrap()

1859

1

            .expect("key not found");

1860

1

        assert_eq!(result, b"value");

1861

1

1862

1863

1

    #[test]

1864

1

    fn std_compact_test_versioned() {

1865

1

        compact_test::<Versioned, _>(StdFileManager::default());

1866

1

1867

1868

1

    #[test]

1869

1

    fn std_compact_test_unversioned() {

1870

1

        compact_test::<Unversioned, _>(StdFileManager::default());

1871

1

1872

1873

1

    #[test]

1874

1

    fn memory_compact_test_versioned() {

1875

1

        compact_test::<Versioned, _>(MemoryFileManager::default());

1876

1

1877

1878

1

    #[test]

1879

1

    fn memory_compact_test_unversioned() {

1880

1

        compact_test::<Unversioned, _>(MemoryFileManager::default());

1881

1

1882

1883

1

    #[test]

1884

1

    fn any_compact_test_versioned() {

1885

1

        compact_test::<Versioned, _>(AnyFileManager::std());

1886

1

        compact_test::<Versioned, _>(AnyFileManager::memory());

1887

1

1888

1889

1

    #[test]

1890

1

    fn any_compact_test_unversioned() {

1891

1

        compact_test::<Unversioned, _>(AnyFileManager::std());

1892

1

        compact_test::<Unversioned, _>(AnyFileManager::memory());

1893

1

1894

1895

8

    fn compact_test<R: Root<Value = ArcBytes<'static>>, M: FileManager>(file_manager: M)

1896

8

    where

1897

8

        R::Reducer: Default,

1898

8

1899

8

        const OPERATION_COUNT: usize = 256;

1900

8

        const WORKER_COUNT: usize = 4;

1901

8

        let tempdir = tempdir().unwrap();

1902

8

1903

8

        let roots = Config::new(tempdir.path())

1904

8

            .file_manager(file_manager)

1905

8

            .open()

1906

8

            .unwrap();

1907

8

        let tree = roots.tree(R::tree("test")).unwrap();

1908

8

        tree.set("foo", b"bar").unwrap();

1909

8

1910

8

        // Spawn a pool of threads that will perform a series of operations

1911

8

        let mut threads = Vec::new();

1912

40

        for worker in 0..WORKER_COUNT {

1913

32

            let tree = tree.clone();

1914

32

            threads.push(std::thread::spawn(move || {

1915

8224

                for relative_id in 0..OPERATION_COUNT {

1916

8192

                    let absolute_id = (worker * OPERATION_COUNT + relative_id) as u64;

1917

8192

                    tree.set(absolute_id.to_be_bytes(), absolute_id.to_be_bytes())

1918

8192

                        .unwrap();

1919

8192

                    let ValueIndex { value, .. } = tree

1920

8192

                        .remove(&absolute_id.to_be_bytes())

1921

8192

                        .unwrap()

1922

8192

                        .ok_or_else(|| panic!("value not found: {:?}", absolute_id))

1923

8192

                        .unwrap();

1924

8192

                    assert_eq!(BigEndian::read_u64(&value), absolute_id);

1925

8192

                    tree.set(absolute_id.to_be_bytes(), absolute_id.to_be_bytes())

1926

8192

                        .unwrap();

1927

8192

                    let newer_value = tree

1928

8192

                        .get(&absolute_id.to_be_bytes())

1929

8192

                        .unwrap()

1930

8192

                        .expect("couldn't find found");

1931

8192

                    assert_eq!(value, newer_value);

1932

1933

32

            }));

1934

32

1935

1936

8

        threads.push(std::thread::spawn(move || {

1937

            // While those workers are running, this thread is going to continually

1938

            // execute compaction.

1939

3506

            while tree.count() < (OPERATION_COUNT * WORKER_COUNT) as u64 {

1940

3498

                tree.compact().unwrap();

1941

3498

1942

8

        }));

1943

1944

48

        for thread in threads {

1945

40

            thread.join().unwrap();

1946

40

1947

8

1948

1949

1

    #[test]

1950

1

    fn name_tests() {

1951

1

        assert!(check_name("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_-.").is_ok());

1952

1

        assert!(check_name("=").is_err());

1953

1

        assert!(check_name("_transactions").is_err());

1954

1

1955

1956

1

    #[test]

1957

1

    fn context_encryption_tests() {

1958

1

        let tempdir = tempdir().unwrap();

1959

1

1960

1

        // Encrypt a tree using default encryption via the context

1961

1

1962

1

            let roots = Config::<StdFileManager>::new(tempdir.path())

1963

1

                .vault(RotatorVault::new(13))

1964

1

                .open()

1965

1

                .unwrap();

1966

1

            let tree = roots.tree(Versioned::tree("test")).unwrap();

1967

1

            tree.set(b"test", b"value").unwrap();

1968

1

            let other_tree = roots

1969

1

                .tree(Versioned::tree("test-otherkey").with_vault(RotatorVault::new(42)))

1970

1

                .unwrap();

1971

1

            other_tree.set(b"test", b"other").unwrap();

1972

1

1973

1

        // Try to access the tree with the vault again.

1974

1

1975

1

            let roots = Config::<StdFileManager>::new(tempdir.path())

1976

1

                .vault(RotatorVault::new(13))

1977

1

                .open()

1978

1

                .unwrap();

1979

1

            let tree = roots.tree(Versioned::tree("test")).unwrap();

1980

1

            let value = tree.get(b"test").unwrap();

1981

1

            assert_eq!(value.as_deref(), Some(&b"value"[..]));

1982

1983

            // Verify we can't read the other tree without the right vault

1984

1

            let bad_tree = roots.tree(Versioned::tree("test-otherkey")).unwrap();

1985

1

            assert!(bad_tree.get(b"test").is_err());

1986

1987

            // And test retrieving the other key with the correct vault

1988

1

            let tree = roots

1989

1

                .tree(Versioned::tree("test-otherkey").with_vault(RotatorVault::new(42)))

1990

1

                .unwrap();

1991

1

            let value = tree.get(b"test").unwrap();

1992

1

            assert_eq!(value.as_deref(), Some(&b"other"[..]));

1993

1994

1995

1

            let roots = Config::<StdFileManager>::new(tempdir.path())

1996

1

                .open()

1997

1

                .unwrap();

1998

1

            // Try to access roots without the vault.

1999

1

            let bad_tree = roots.tree(Versioned::tree("test")).unwrap();

2000

1

            assert!(bad_tree.get(b"test").is_err());

2001

2002

            // Try to access roots with the vault specified. In this situation, the transaction log will be unreadable, causing itself to not consider any transactions valid.

2003

1

            let bad_tree = roots

2004

1

                .tree(Versioned::tree("test").with_vault(RotatorVault::new(13)))

2005

1

                .unwrap();

2006

1

            assert_eq!(bad_tree.get(b"test").unwrap(), None);

2007

2008

1

2009

2010

1

    #[test]

2011

1

    fn too_large_transaction() {

2012

1

        let tempdir = tempdir().unwrap();

2013

1

2014

1

        let config = Config::<StdFileManager>::new(tempdir.path());

2015

1

2016

1

            let roots = config.clone().open().unwrap();

2017

1

2018

1

            let mut transaction = roots.transaction(&[Versioned::tree("test")]).unwrap();

2019

1

2020

1

            // Write some data to the tree.

2021

1

            transaction

2022

1

                .tree::<Versioned>(0)

2023

1

                .unwrap()

2024

1

                .set(b"test", vec![0; 16 * 1024 * 1024])

2025

1

                .unwrap();

2026

2027

            // Issue a transaction that's too large.

2028

1

            assert!(matches!(

2029

1

                transaction

2030

1

                    .entry_mut()

2031

1

                    .set_data(vec![0; 16 * 1024 * 1024 - 7])

2032

1

                    .unwrap_err()

2033

1

                    .kind,

2034

                ErrorKind::ValueTooLarge

2035

));

2036

            // Roll the transaction back

2037

1

            transaction.rollback();

2038

1

2039

1

        // Ensure that we can still write to the tree.

2040

1

2041

1

            let roots = config.open().unwrap();

2042

1

2043

1

            let transaction = roots.transaction(&[Versioned::tree("test")]).unwrap();

2044

1

2045

1

            // Write some data to the tree

2046

1

            transaction

2047

1

                .tree::<Versioned>(0)

2048

1

                .unwrap()

2049

1

                .set(b"test", b"updated value")

2050

1

                .unwrap();

2051

1

2052

1

            transaction.commit().unwrap();

2053

1

2054

1

2055