GODRIVER-2907 Add bulk decryption benchmarks

testdata/mongocrypt/command.json

@@ -3,4 +3,4 @@
  "filter": {
  "ssn": "457-55-5462"
  }
-}
+}

testdata/mongocrypt/local-key-document.json

@@ -0,0 +1,24 @@
+{
+ "_id": {
+ "$binary": {
+ "base64": "YWFhYWFhYWFhYWFhYWFhYQ==",
+ "subType": "04"
+ }
+ },
+ "keyMaterial": {
+ "$binary": {
+ "base64": "ACR7Hm33dDOAAD7l2ubZhSpSUWK8BkALUY+qW3UgBAEcTV8sBwZnaAWnzDsmrX55dgmYHWfynDlJogC/e33u6pbhyXvFTs5ow9OLCuCWBJ39T/Ivm3kMaZJybkejY0V+uc4UEdHvVVz/SbitVnzs2WXdMGmo1/HmDRrxGYZjewFslquv8wtUHF5pyB+QDlQBd/al9M444/8bJZFbMSmtIg==",
+ "subType": "00"
+ }
+ },
+ "creationDate": {
+ "$date": "2023-08-21T14:28:20.875Z"
+ },
+ "updateDate": {
+ "$date": "2023-08-21T14:28:20.875Z"
+ },
+ "status": 0,
+ "masterKey": {
+ "provider": "local"
+ }
+}

x/mongo/driver/mongocrypt/binary_test.go

@@ -0,0 +1,218 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// 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
+
+//go:build cse
+// +build cse
+
+package mongocrypt
+
+import (
+ "fmt"
+ "os"
+ "path"
+ "runtime"
+ "sort"
+ "testing"
+ "time"
+
+ "go.mongodb.org/mongo-driver/v2/bson"
+ "go.mongodb.org/mongo-driver/v2/internal/assert"
+ "go.mongodb.org/mongo-driver/v2/internal/require"
+ "go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
+ "go.mongodb.org/mongo-driver/v2/x/mongo/driver/mongocrypt/options"
+)
+
+// load JSON for benchmark
+func resourceToDocumentForBench(b *testing.B, filename string) bsoncore.Document {
+ b.Helper()
+
+ content, err := os.ReadFile(path.Join(resourcesDir, filename))
+ require.NoError(b, err)
+
+ var doc bsoncore.Document
+
+ err = bson.UnmarshalExtJSON(content, false, &doc)
+ require.NoError(b, err)
+
+ return doc
+}
+
+// Add encryption key to the encryption context.
+func addMongoKeysForBench(b *testing.B, encryptCtx *Context) {
+ b.Helper()
+
+ if encryptCtx.State() != NeedMongoKeys {
+ return
+ }
+
+ _, err := encryptCtx.NextOperation()
+ require.NoError(b, err)
+
+ // feed result and finish op
+ err = encryptCtx.AddOperationResult(resourceToDocumentForBench(b, "local-key-document.json"))
+ require.NoError(b, err)
+
+ err = encryptCtx.CompleteOperation()
+ require.NoError(b, err)
+}
+
+// encrypt a document for benchmarking
+func createEncryptedDocForBench(b *testing.B, crypt *MongoCrypt, iter int) bsoncore.Document {
+ b.Helper()
+
+ const algorithm = "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic"
+
+ // create explicit encryption context and check initial state
+ keyID := bson.Binary{
+ Subtype: 0x04, // 0x04 is UUID subtype
+ Data: []byte("aaaaaaaaaaaaaaaa"),
+ }
+
+ encryptOpts := options.ExplicitEncryption().SetAlgorithm(algorithm).SetKeyID(keyID)
+ doc := bsoncore.NewDocumentBuilder().AppendString("v", fmt.Sprintf("value %04v", iter)).Build()
+
+ encryptCtx, err := crypt.CreateExplicitEncryptionContext(doc, encryptOpts)
+ require.NoError(b, err)
+
+ defer encryptCtx.Close()
+
+ addMongoKeysForBench(b, encryptCtx)
+
+ // perform final encryption
+ encryptedDoc, err := encryptCtx.Finish()
+ require.NoError(b, err)
+
+ return encryptedDoc
+}
+
+// decrypt an encrypted document for benchmarking.
+func decryptDocForBench(b *testing.B, crypt *MongoCrypt, encryptedDoc bsoncore.Document) bsoncore.Document {
+ b.Helper()
+
+ // create explicit decryption context and check initial state
+ decryptCtx, err := crypt.CreateDecryptionContext(encryptedDoc)
+ require.NoError(b, err)
+
+ defer decryptCtx.Close()
+
+ // perform final decryption
+ decryptedDoc, err := decryptCtx.Finish()
+ require.NoError(b, err)
+
+ return decryptedDoc
+}
+
+// create a document of the form:
+// { "key0001": <encrypted "value 0001">, "key0002": <encrypted "value 0002">, ... }
+func createEncryptedDocForBulkDecryptionBench(b *testing.B, crypt *MongoCrypt, count int) bsoncore.Document {
+ bldr := bsoncore.NewDocumentBuilder()
+
+ for i := 0; i < count; i++ {
+ encDoc := createEncryptedDocForBench(b, crypt, i)
+ bldr.AppendValue(fmt.Sprintf("key%04v", i), encDoc.Lookup("v"))
+ }
+
+ return bldr.Build()
+}
+
+// Create a MongoCrypt object for benchmarking bulk decryption.
+func newCryptForBench(b *testing.B) *MongoCrypt {
+ key := []byte{
+ 0x9d, 0x94, 0x4b, 0x0d, 0x93, 0xd0, 0xc5, 0x44,
+ 0xa5, 0x72, 0xfd, 0x32, 0x1b, 0x94, 0x30, 0x90,
+ 0x23, 0x35, 0x73, 0x7c, 0xf0, 0xf6, 0xc2, 0xf4,
+ 0xda, 0x23, 0x56, 0xe7, 0x8f, 0x04, 0xcc, 0xfa,
+ 0xde, 0x75, 0xb4, 0x51, 0x87, 0xf3, 0x8b, 0x97,
+ 0xd7, 0x4b, 0x44, 0x3b, 0xac, 0x39, 0xa2, 0xc6,
+ 0x4d, 0x91, 0x00, 0x3e, 0xd1, 0xfa, 0x4a, 0x30,
+ 0xc1, 0xd2, 0xc6, 0x5e, 0xfb, 0xac, 0x41, 0xf2,
+ 0x48, 0x13, 0x3c, 0x9b, 0x50, 0xfc, 0xa7, 0x24,
+ 0x7a, 0x2e, 0x02, 0x63, 0xa3, 0xc6, 0x16, 0x25,
+ 0x51, 0x50, 0x78, 0x3e, 0x0f, 0xd8, 0x6e, 0x84,
+ 0xa6, 0xec, 0x8d, 0x2d, 0x24, 0x47, 0xe5, 0xaf}
+
+ localProvider := bsoncore.NewDocumentBuilder().
+ AppendBinary("key", 0, key).
+ Build()
+
+ kmsProviders := bsoncore.NewDocumentBuilder().
+ AppendDocument("local", localProvider).
+ Build()
+
+ cryptOpts := options.MongoCrypt().SetKmsProviders(kmsProviders)
+
+ crypt, err := NewMongoCrypt(cryptOpts)
+
+ require.NoError(b, err)
+ assert.NotNil(b, crypt)
+
+ return crypt
+}
+
+func calcMedian(values []float64) float64 {
+ sort.Float64s(values)
+
+ n := len(values)
+ if n%2 == 0 {
+ return (values[n/2-1] + values[n/2]) / 2.0
+ }
+
+ return values[(n-1)/2]
+}
+
+func BenchmarkBulkDecryption(b *testing.B) {
+ const numKeys = 1500
+ const repeatCount = 10
+
+ // Create crypt that uses a data key with the "local" KMS provider.
+ crypt := newCryptForBench(b)
+ defer crypt.Close()
+
+ // Encrypt 1500 string values of the form value 0001, value 0002, value 0003,
+ // ... with the algorithm AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic.
+ encryptedDoc := createEncryptedDocForBulkDecryptionBench(b, crypt, numKeys)
+
+ // decrypt the document repeatedly for one second to warm up the benchmark.
+ for start := time.Now(); time.Since(start) < time.Second; {
+ decryptDocForBench(b, crypt, encryptedDoc)
+ }
+
+ benchmarks := []struct {
+ threads int
+ }{
+ {threads: 1},
+ {threads: 2},
+ {threads: 8},
+ {threads: 64},
+ }
+
+ // Run the benchmark. Repeat benchmark for thread counts: (1, 2, 8, 64).
+ // Repeat 10 times.
+ for _, bench := range benchmarks {
+ var opsPerSec []float64
+
+ b.Run(fmt.Sprintf("threadCount=%v", bench.threads), func(b *testing.B) {
+ runtime.GOMAXPROCS(bench.threads)
+
+ for i := 0; i < repeatCount; i++ {
+ b.ResetTimer()
+ b.ReportAllocs()
+
+ startTime := time.Now()
+
+ b.RunParallel(func(pb *testing.PB) {
+ for pb.Next() {
+ decryptDocForBench(b, crypt, encryptedDoc)
+ }
+ })
+
+ opsPerSec = append(opsPerSec, float64(b.N)/time.Now().Sub(startTime).Seconds())
+ }
+ })
+
+ b.Logf("thread count: %v, median ops/sec: %.2f\n", bench.threads, calcMedian(opsPerSec))
+ }
+}