1
use std::{path::Path, time::Duration};
2

            
3
use bonsaidb::{
4
    client::{url::Url, Client, RemoteDatabase},
5
    core::{connection::StorageConnection, keyvalue::KeyValue},
6
    local::config::{Builder, KeyValuePersistence, PersistenceThreshold},
7
    server::{DefaultPermissions, Server, ServerConfiguration},
8
};
9
use criterion::{measurement::WallTime, BenchmarkGroup, BenchmarkId};
10
use tokio::runtime::Runtime;
11
use ubyte::ToByteUnit;
12

            
13
4
pub fn read_blobs(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
14
4
    read_blobs_local(c, data);
15
4
    read_blobs_networked(c, data);
16
4
}
17

            
18
4
fn read_blobs_local(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
19
4
    let runtime = Runtime::new().unwrap();
20
4
    let server = runtime.block_on(initialize_server(true));
21
4
    let database = runtime.block_on(async { server.database::<()>("key-value").await.unwrap() });
22
4
    runtime.block_on(set_blob(&database, data));
23
4

            
24
4
    c.bench_function(
25
4
        BenchmarkId::new("bonsaidb-local", data.len().bytes()),
26
4
        |b| {
27
4
            b.to_async(&runtime).iter(|| get_blob(&database));
28
4
        },
29
4
    );
30
4
}
31

            
32
4
fn read_blobs_networked(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
33
4
    let runtime = Runtime::new().unwrap();
34
4
    let (quic_database, ws_database) = initialize_networked_server(&runtime, true);
35
4
    runtime.block_on(set_blob(&quic_database, data));
36
4

            
37
4
    c.bench_function(BenchmarkId::new("bonsaidb-quic", data.len().bytes()), |b| {
38
4
        b.iter(|| {
39
4
            runtime.block_on(get_blob(&quic_database));
40
4
        });
41
4
    });
42
4

            
43
4
    c.bench_function(BenchmarkId::new("bonsaidb-ws", data.len().bytes()), |b| {
44
4
        b.iter(|| {
45
4
            runtime.block_on(get_blob(&ws_database));
46
4
        });
47
4
    });
48
4
}
49

            
50
4
pub fn write_blobs(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
51
4
    write_blobs_local(c, data);
52
4
    write_blobs_networked(c, data);
53
4
}
54

            
55
4
fn write_blobs_local(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
56
4
    let runtime = Runtime::new().unwrap();
57
4
    let server = runtime.block_on(initialize_server(false));
58
4
    let database = runtime.block_on(async { server.database::<()>("key-value").await.unwrap() });
59
4

            
60
4
    c.bench_function(
61
4
        BenchmarkId::new("bonsaidb-local", data.len().bytes()),
62
4
        |b| {
63
4
            b.to_async(&runtime).iter(|| set_blob(&database, data));
64
4
        },
65
4
    );
66
4
}
67

            
68
4
fn write_blobs_networked(c: &mut BenchmarkGroup<WallTime>, data: &[u8]) {
69
4
    let runtime = Runtime::new().unwrap();
70
4
    let (quic_database, ws_database) = initialize_networked_server(&runtime, false);
71
4

            
72
4
    c.bench_function(BenchmarkId::new("bonsaidb-quic", data.len().bytes()), |b| {
73
4
        b.iter(|| {
74
4
            runtime.block_on(set_blob(&quic_database, data));
75
4
        });
76
4
    });
77
4

            
78
4
    c.bench_function(BenchmarkId::new("bonsaidb-ws", data.len().bytes()), |b| {
79
4
        b.iter(|| {
80
4
            runtime.block_on(set_blob(&ws_database, data));
81
4
        });
82
4
    });
83
4
}
84

            
85
12
async fn get_blob<C: KeyValue>(connection: &C) {
86
12
    // The set_key API provides serialization. Uisng this API, we can skip
87
12
    // serialization.
88
12
    connection.get_key("blob").await.unwrap().unwrap();
89
12
}
90

            
91
20
async fn set_blob<C: KeyValue>(connection: &C, blob: &[u8]) {
92
20
    // The set_key API provides serialization. Uisng this API, we can skip
93
20
    // serialization.
94
20
    connection.set_binary_key("blob", blob).await.unwrap();
95
20
}
96

            
97
1
pub fn increment(c: &mut BenchmarkGroup<WallTime>) {
98
1
    increment_local(c);
99
1
    increment_networked(c);
100
1
}
101

            
102
1
fn increment_local(c: &mut BenchmarkGroup<WallTime>) {
103
1
    let runtime = Runtime::new().unwrap();
104
1
    let server = runtime.block_on(initialize_server(false));
105
1
    let database = runtime.block_on(async { server.database::<()>("key-value").await.unwrap() });
106
1

            
107
1
    c.bench_function("bonsaidb-local", |b| {
108
1
        b.to_async(&runtime).iter(|| increment_key(&database));
109
1
    });
110
1
}
111

            
112
1
fn increment_networked(c: &mut BenchmarkGroup<WallTime>) {
113
1
    let runtime = Runtime::new().unwrap();
114
1
    let (quic_database, ws_database) = initialize_networked_server(&runtime, false);
115
1

            
116
1
    c.bench_function("bonsaidb-quic", |b| {
117
1
        b.iter(|| {
118
1
            runtime.block_on(increment_key(&quic_database));
119
1
        });
120
1
    });
121
1

            
122
1
    c.bench_function("bonsaidb-ws", |b| {
123
1
        b.iter(|| {
124
1
            runtime.block_on(increment_key(&ws_database));
125
1
        });
126
1
    });
127
1
}
128

            
129
3
async fn increment_key<C: KeyValue>(connection: &C) {
130
3
    connection.increment_key_by("u64", 1_u64).await.unwrap();
131
3
}
132

            
133
18
async fn initialize_server(persist_changes: bool) -> Server {
134
18
    let path = Path::new("key-value-benchmarks.bonsaidb");
135
18
    if path.exists() {
136
17
        std::fs::remove_dir_all(path).unwrap();
137
17
    }
138
18
    let server = Server::open(
139
18
        ServerConfiguration::new(path)
140
18
            .default_permissions(DefaultPermissions::AllowAll)
141
18
            .key_value_persistence(if persist_changes {
142
8
                KeyValuePersistence::immediate()
143
            } else {
144
10
                KeyValuePersistence::lazy([PersistenceThreshold::after_changes(usize::MAX)])
145
            })
146
18
            .with_schema::<()>()
147
18
            .unwrap(),
148
336
    )
149
336
    .await
150
18
    .unwrap();
151
139
    server.install_self_signed_certificate(false).await.unwrap();
152
18
    server
153
34
        .create_database::<()>("key-value", false)
154
34
        .await
155
18
        .unwrap();
156
18
    server
157
18
}
158

            
159
9
fn initialize_networked_server(
160
9
    runtime: &Runtime,
161
9
    persist_changes: bool,
162
9
) -> (RemoteDatabase, RemoteDatabase) {
163
9
    let server = runtime.block_on(initialize_server(persist_changes));
164
9
    let certificate = runtime
165
9
        .block_on(async { server.certificate_chain().await.unwrap() })
166
9
        .into_end_entity_certificate();
167
9
    let quic_server = server.clone();
168
9
    runtime.spawn(async move {
169
27
        quic_server.listen_on(7022).await.unwrap();
170
9
    });
171
9
    runtime.spawn(async move {
172
9
        server
173
9
            .listen_for_websockets_on("0.0.0.0:7023", false)
174
9
            .await
175
            .unwrap();
176
9
    });
177
9
    let quic_database = runtime.block_on(async {
178
9
        // Allow the server time to start listening
179
9
        tokio::time::sleep(Duration::from_millis(1000)).await;
180
9
        let client = Client::build(Url::parse("bonsaidb://localhost:7022").unwrap())
181
9
            .with_certificate(certificate)
182
9
            .finish()
183
            .await
184
9
            .unwrap();
185
9
        client.database::<()>("key-value").await.unwrap()
186
9
    });
187
9
    let ws_database = runtime.block_on(async {
188
        // Allow the server time to start listening
189
9
        let client = Client::build(Url::parse("ws://localhost:7023").unwrap())
190
9
            .finish()
191
            .await
192
9
            .unwrap();
193
9
        client.database::<()>("key-value").await.unwrap()
194
9
    });
195
9
    (quic_database, ws_database)
196
9
}