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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
|
# Copyright 2023-present MongoDB, Inc.
#
# 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.
"""Benchmark pymongocrypt performance."""
from __future__ import annotations
import os
import sys
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
try:
import simplejson as json
except ImportError:
import json # type: ignore[no-redef]
import bson
from bson import json_util
sys.path[0:0] = [""]
from test.test_mongocrypt import MockCallback
from pymongocrypt.binding import lib, libmongocrypt_version
from pymongocrypt.mongocrypt import MongoCrypt, MongoCryptOptions
from pymongocrypt.synchronous.explicit_encrypter import ExplicitEncrypter
from pymongocrypt.version import __version__
NUM_ITERATIONS = 10
MAX_TIME = 1
NUM_FIELDS = 1500
LOCAL_MASTER_KEY = (
b"\x9d\x94K\r\x93\xd0\xc5D\xa5r\xfd2\x1b\x940\x90#5s|\xf0\xf6\xc2\xf4\xda#V\xe7\x8f\x04"
b"\xcc\xfa\xdeu\xb4Q\x87\xf3\x8b\x97\xd7KD;\xac9\xa2\xc6M\x91\x00>\xd1\xfaJ0\xc1\xd2"
b"\xc6^\xfb\xacA\xf2H\x13<\x9bP\xfc\xa7$z.\x02c\xa3\xc6\x16%QPx>\x0f\xd8n\x84\xa6\xec"
b"\x8d-$G\xe5\xaf"
)
OUTPUT_FILE = os.environ.get("OUTPUT_FILE")
result_data: list = []
def read(filename, **kwargs):
with open(os.path.join(os.path.dirname(__file__), filename), **kwargs) as fp:
return fp.read()
def json_data(filename):
return json_util.loads(read(filename))
def bson_data(filename):
return bson.encode(json_data(filename))
def tearDownModule():
output = json.dumps(result_data, indent=4)
if OUTPUT_FILE:
with open(OUTPUT_FILE, "w") as opf:
opf.write(output)
else:
print(output)
class TestBulkDecryption(unittest.TestCase):
def setUp(self):
opts = MongoCryptOptions({"local": {"key": LOCAL_MASTER_KEY}})
callback = MockCallback(key_docs=[bson_data("keyDocument.json")])
self.mongocrypt = MongoCrypt(opts, callback)
self.encrypter = ExplicitEncrypter(callback, opts)
self.addCleanup(self.mongocrypt.close)
self.addCleanup(self.encrypter.close)
def do_task(self, encrypted, duration=MAX_TIME):
start = time.monotonic()
ops = 0
while time.monotonic() - start < duration:
with self.mongocrypt.decryption_context(encrypted) as ctx:
if ctx.state == lib.MONGOCRYPT_CTX_NEED_MONGO_KEYS:
# Key is requested on the first operation, then expected to be cached for one minute.
ctx.add_mongo_operation_result(bson_data("keyDocument.json"))
ctx.complete_mongo_operation()
self.assertEqual(ctx.state, lib.MONGOCRYPT_CTX_READY)
decrypted = ctx.finish()
ops += 1
# Assert that decryption actually occurred.
self.assertGreater(ops, 0)
doc = bson.decode(decrypted)
for val in doc.values():
self.assertIsInstance(val, str)
return ops
def percentile(self, percentile):
sorted_results = sorted(self.results)
percentile_index = int(len(sorted_results) * percentile / 100) - 1
return sorted_results[percentile_index]
def runTest(self):
doc = {}
key_id = json_data("keyDocument.json")["_id"]
for i in range(NUM_FIELDS):
val = f"value {i:04}"
val_encrypted = bson.decode(
self.encrypter.encrypt(
bson.encode({"v": val}),
"AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic",
key_id=key_id,
)
)["v"]
doc[f"key{i:04}"] = val_encrypted
encrypted = bson.encode(doc)
# Warm up benchmark and discard the result.
self.do_task(encrypted, duration=2)
for n_threads in [1, 2, 8, 64]:
with ThreadPoolExecutor(max_workers=n_threads) as executor:
self.results = []
for _ in range(NUM_ITERATIONS):
start = time.monotonic()
thread_results = list(
executor.map(self.do_task, [encrypted] * n_threads)
)
interval = time.monotonic() - start
self.results.append(sum(thread_results) / interval)
median = self.percentile(50)
print(
f"Finished {self.__class__.__name__}, threads={n_threads}, median ops_per_second={median:.2f}"
)
# Remove "Test" so that TestBulkDecryption is reported as "BulkDecryption".
name = self.__class__.__name__[4:]
result_data.append(
{
"info": {
"test_name": name,
"args": {
"threads": n_threads,
},
},
"metrics": [
{"name": "ops_per_second", "type": "MEDIAN", "value": median},
],
}
)
if __name__ == "__main__":
print(
f"Running benchmark with pymongocrypt: {__version__} libmongocrypt: {libmongocrypt_version()}"
)
unittest.main()
|
