test(benchmarking)

.github/workflows/ci_tests.yaml

@@ -62,6 +62,20 @@ jobs:
         env:
           TOXENV: ${{ matrix.category }}
         run: tox -- -vv
+  benchmarks:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v1
+        with:
+          python-version: 3.x
+      - run: |
+          python -m pip install --upgrade pip
+          pip install --upgrade -r ci-requirements.txt
+      - name: run test
+        env:
+          TOXENV: benchmark-nokms
+        run: tox -- -vv
   upstream-py3:
     runs-on: ubuntu-latest
     strategy:

dev_requirements/test-requirements.txt

@@ -2,3 +2,4 @@ mock==4.0.3
 pytest==7.0.1
 pytest-cov==3.0.0
 pytest-mock==3.6.1
+pytest-benchmark>=3.2.3

setup.cfg

@@ -20,6 +20,7 @@ markers =
     integ: mark a test as an integration test (requires network access)
     accept: mark a test as an acceptance test (requires network access)
     examples: mark a test as an examples test (requires network access)
+    benchmark: mark a test as a performance benchmark test
 
 # Flake8 Configuration
 [flake8]

test/benchmark/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Performance benchmarking tests."""

test/benchmark/benchmark_test_utils.py

@@ -0,0 +1,89 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Helper utilities for benchmark tests."""
+import copy
+
+import pytest
+
+import aws_encryption_sdk
+from aws_encryption_sdk.caches.local import LocalCryptoMaterialsCache
+from aws_encryption_sdk.identifiers import AlgorithmSuite
+from aws_encryption_sdk.keyrings.base import Keyring
+from aws_encryption_sdk.materials_managers.caching import CachingCryptoMaterialsManager
+from aws_encryption_sdk.materials_managers.default import DefaultCryptoMaterialsManager
+
+ENCRYPTION_CONTEXT = {
+    "encryption": "context",
+    "is not": "secret",
+    "but adds": "useful metadata",
+    "that can help you": "be confident that",
+    "the data you are handling": "is what you think it is",
+}
+
+
+def all_operations():
+    return pytest.mark.parametrize(
+        "operation",
+        (
+            pytest.param(aws_encryption_sdk.encrypt, id="encrypt only"),
+            pytest.param(aws_encryption_sdk.decrypt, id="decrypt only"),
+            pytest.param(encrypt_decrypt_cycle, id="encrypt decrypt cycle"),
+        ),
+    )
+
+
+def encrypt_decrypt_cycle(**kwargs):
+    encrypt_kwargs = copy.copy(kwargs)
+    decrypt_kwargs = copy.copy(kwargs)
+    for param in ("encryption_context", "frame_length", "source", "algorithm"):
+        try:
+            del decrypt_kwargs[param]
+        except KeyError:
+            pass
+
+    encrypted = aws_encryption_sdk.encrypt(**encrypt_kwargs)
+    decrypt_kwargs["source"] = encrypted.result
+    aws_encryption_sdk.decrypt(**decrypt_kwargs)
+
+
+def build_cmm(provider_builder, cache_messages):
+    provider = provider_builder()
+    if isinstance(provider, Keyring):
+        provider_param = "keyring"
+    else:
+        provider_param = "master_key_provider"
+
+    if cache_messages == 0:
+        cmm = DefaultCryptoMaterialsManager(**{provider_param: provider})
+    else:
+        cmm = CachingCryptoMaterialsManager(
+            max_age=6000.0,
+            max_messages_encrypted=cache_messages,
+            cache=LocalCryptoMaterialsCache(capacity=10),
+            **{provider_param: provider}
+        )
+
+    return cmm
+
+
+def run_benchmark(
+    benchmark,
+    provider_builder,
+    operation,
+    cache_messages=0,
+    plaintext="foo",
+    frame_length=1024,
+    algorithm=AlgorithmSuite.AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384,
+):
+    cmm = build_cmm(provider_builder, cache_messages)
+
+    kwargs = dict(
+        source=plaintext,
+        materials_manager=cmm,
+        encryption_context=copy.copy(ENCRYPTION_CONTEXT),
+        frame_length=frame_length,
+        algorithm=algorithm,
+    )
+    if operation is aws_encryption_sdk.decrypt:
+        kwargs = dict(source=aws_encryption_sdk.encrypt(**kwargs).result, materials_manager=cmm,)
+    benchmark.pedantic(target=operation, kwargs=kwargs, iterations=10, rounds=10)

test/benchmark/test_client_performance.py

@@ -0,0 +1,121 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Functional performance test suite for ``aws_encryption_sdk``."""
+import os
+
+import pytest
+
+from aws_encryption_sdk.identifiers import AlgorithmSuite
+
+from ..integration.integration_test_utils import build_aws_kms_keyring, setup_kms_master_key_provider
+from ..unit.unit_test_utils import (
+    ephemeral_raw_aes_keyring,
+    ephemeral_raw_aes_master_key,
+    ephemeral_raw_rsa_keyring,
+    ephemeral_raw_rsa_master_key,
+)
+from .benchmark_test_utils import all_operations, run_benchmark
+
+pytestmark = [pytest.mark.benchmark]
+
+PLAINTEXTS = {
+    "SMALL": os.urandom(32),  # 32B
+    "LARGE": os.urandom(1024 * 1024),  # 1MiB
+    "VERY_LARGE": os.urandom(10 * 1024 * 1024),  # 10MiB
+}
+
+
+@pytest.mark.parametrize("algorithm_suite", AlgorithmSuite)
+@all_operations()
+def test_compare_algorithm_suite_performance(benchmark, algorithm_suite, operation):
+    """Compare the affect of algorithm suite on performance.
+    Use the Raw AES keyring as a baseline keyring.
+    """
+    run_benchmark(
+        benchmark=benchmark, provider_builder=ephemeral_raw_aes_keyring, operation=operation, algorithm=algorithm_suite
+    )
+
+
+@pytest.mark.parametrize(
+    "cache_messages",
+    (
+        pytest.param(0, id="no cache"),
+        pytest.param(1000000, id="cache and only miss once"),
+        pytest.param(10, id="cache and miss every 10"),
+    ),
+)
+@all_operations()
+def test_compare_caching_performance(benchmark, operation, cache_messages):
+    """Compare the affect of caching on performance.
+    Use the Raw AES keyring as a baseline keyring.
+    """
+    run_benchmark(
+        benchmark=benchmark,
+        provider_builder=ephemeral_raw_aes_keyring,
+        operation=operation,
+        cache_messages=cache_messages,
+    )
+
+
+@pytest.mark.parametrize(
+    "plaintext, frame_length",
+    (
+        pytest.param("SMALL", 0, id="small message, unframed"),
+        pytest.param("SMALL", 128, id="small message, single frame"),
+        pytest.param("LARGE", 1024 * 1024 * 1024, id="large message, single frame"),
+        pytest.param("LARGE", 102400, id="large message, few large frames"),
+        pytest.param("LARGE", 1024, id="large message, many small frames"),
+    ),
+)
+@all_operations()
+def test_compare_framing_performance(benchmark, operation, plaintext, frame_length):
+    """Compare the affect of framing and on performance.
+    Use the Raw AES keyring as a baseline keyring.
+    """
+    run_benchmark(
+        benchmark=benchmark,
+        provider_builder=ephemeral_raw_aes_keyring,
+        operation=operation,
+        plaintext=PLAINTEXTS[plaintext],
+        frame_length=frame_length,
+    )
+
+
+def _frame_sizes():
+    for frame_kb in (2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 10240):
+        yield pytest.param(frame_kb * 1024, id="{} kiB frame".format(frame_kb))
+
+
+@pytest.mark.parametrize(
+    "plaintext", (pytest.param("LARGE", id="1MiB plaintext"), pytest.param("VERY_LARGE", id="10MiB plaintext"),),
+)
+@pytest.mark.parametrize("frame_length", _frame_sizes())
+@all_operations()
+def test_compare_frame_size_performance(benchmark, operation, plaintext, frame_length):
+    """Compare the affect of framing and on performance.
+    Use the Raw AES keyring as a baseline keyring.
+    """
+    run_benchmark(
+        benchmark=benchmark,
+        provider_builder=ephemeral_raw_aes_keyring,
+        operation=operation,
+        plaintext=PLAINTEXTS[plaintext],
+        frame_length=frame_length,
+    )
+
+
+@pytest.mark.parametrize(
+    "provider_builder",
+    (
+        pytest.param(ephemeral_raw_aes_keyring, id="Raw AES keyring"),
+        pytest.param(ephemeral_raw_aes_master_key, id="Raw AES master key"),
+        pytest.param(ephemeral_raw_rsa_keyring, id="Raw RSA keyring"),
+        pytest.param(ephemeral_raw_rsa_master_key, id="Raw RSA master key"),
+        pytest.param(build_aws_kms_keyring, id="AWS KMS keyring", marks=pytest.mark.integ),
+        pytest.param(setup_kms_master_key_provider, id="AWS KMS master key provider", marks=pytest.mark.integ),
+    ),
+)
+@all_operations()
+def test_compare_keyring_performance(benchmark, provider_builder, operation):
+    """Compare the performance of different keyrings and master key providers."""
+    run_benchmark(benchmark=benchmark, provider_builder=provider_builder, operation=operation)

test/functional/test_f_aws_encryption_sdk_client.py

@@ -10,7 +10,7 @@
 # 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.
-"""Functional test suite for aws_encryption_sdk.kms_thick_client"""
+"""Functional test suite for aws_encryption_sdk"""
 from __future__ import division
 
 import io

tox.ini

@@ -1,4 +1,5 @@
 [tox]
+minversion = 3.4.0
 envlist =
     py{36,37,38,39,310}-{local,integ,accept,examples}, nocmk,
     bandit, doc8, readme, docs,
@@ -62,7 +63,12 @@ passenv =
     PIP_CONFIG_FILE
 sitepackages = False
 deps = -rdev_requirements/test-requirements.txt
+# 'download' forces tox to always upgrade pip to the latest
+download = true
 commands =
+    benchmark-full: {[testenv:base-command]commands} test/ -m benchmark
+    benchmark-kms: {[testenv:base-command]commands} test/ -m "benchmark and integ"
+    benchmark-nokms: {[testenv:base-command]commands} test/ -m "benchmark and not integ"
     local: {[testenv:base-command]commands} test/ -m local
     integ: {[testenv:base-command]commands} test/ -m integ
     accept: {[testenv:base-command]commands} test/ -m accept