1
use std::path::Path;
2
use std::time::Duration;
3

            
4
use bonsaidb::client::url::Url;
5
use bonsaidb::client::{AsyncClient, AsyncRemoteDatabase};
6
use bonsaidb::core::connection::AsyncStorageConnection;
7
use bonsaidb::core::keyvalue::AsyncKeyValue;
8
use bonsaidb::local::config::{Builder, KeyValuePersistence, PersistenceThreshold};
9
use bonsaidb::server::{DefaultPermissions, Server, ServerConfiguration};
10
use criterion::measurement::WallTime;
11
use criterion::{BenchmarkGroup, BenchmarkId};
12
use tokio::runtime::Runtime;
13
use ubyte::ToByteUnit;
14

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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