1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
// Copyright 2018-2019 Mozilla
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
// this file except in compliance with the License. You may obtain a copy of the
// License at http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
//! A command-line utility to create an LMDB environment containing random data.
//! It requires one flag, `-s path/to/environment`, which specifies the location
//! where the tool should create the environment. Optionally, you may specify
//! the number of key/value pairs to create via the `-n <number>` flag
//! (for which the default value is 50).
use std::{env::args, fs, fs::File, io::Read, path::Path};
use rkv::{
backend::{BackendEnvironmentBuilder, Lmdb},
Rkv, StoreOptions, Value,
};
fn main() {
let mut args = args();
let mut database = None;
let mut path = None;
let mut num_pairs = 50;
// The first arg is the name of the program, which we can ignore.
args.next();
while let Some(arg) = args.next() {
if &arg[0..1] == "-" {
match &arg[1..] {
"s" => {
database = match args.next() {
None => panic!("-s must be followed by database arg"),
Some(str) => Some(str),
};
}
"n" => {
num_pairs = match args.next() {
None => panic!("-s must be followed by number of pairs"),
Some(str) => str.parse().expect("number"),
};
}
str => panic!("arg -{} not recognized", str),
}
} else {
if path.is_some() {
panic!("must provide only one path to the LMDB environment");
}
path = Some(arg);
}
}
if path.is_none() {
panic!("must provide a path to the LMDB environment");
}
let path = path.unwrap();
fs::create_dir_all(&path).expect("dir created");
let mut builder = Rkv::environment_builder::<Lmdb>();
builder.set_max_dbs(2);
// Allocate enough map to accommodate the largest random collection.
// We currently do this by allocating twice the maximum possible size
// of the pairs (assuming maximum key and value sizes).
builder.set_map_size((511 + 65535) * num_pairs * 2);
let rkv = Rkv::from_builder(Path::new(&path), builder).expect("Rkv");
let store = rkv
.open_single(database.as_deref(), StoreOptions::create())
.expect("opened");
let mut writer = rkv.write().expect("writer");
// Generate random values for the number of keys and key/value lengths.
// On Linux, "Just use /dev/urandom!" <https://www.2uo.de/myths-about-urandom/>.
// On macOS it doesn't matter (/dev/random and /dev/urandom are identical).
let mut random = File::open("/dev/urandom").unwrap();
let mut nums = [0u8; 4];
random.read_exact(&mut nums).unwrap();
// Generate 0–255 pairs.
for _ in 0..num_pairs {
// Generate key and value lengths. The key must be 1–511 bytes long.
// The value length can be 0 and is essentially unbounded; we generate
// value lengths of 0–0xffff (65535).
// NB: the modulus method for generating a random number within a range
// introduces distribution skew, but we don't need it to be perfect.
let key_len = ((u16::from(nums[0]) + (u16::from(nums[1]) << 8)) % 511 + 1) as usize;
let value_len = (u16::from(nums[2]) + (u16::from(nums[3]) << 8)) as usize;
let mut key: Vec<u8> = vec![0; key_len];
random.read_exact(&mut key[0..key_len]).unwrap();
let mut value: Vec<u8> = vec![0; value_len];
random.read_exact(&mut value[0..value_len]).unwrap();
store
.put(&mut writer, key, &Value::Blob(&value))
.expect("wrote");
}
writer.commit().expect("committed");
}
|
