Added FixedPool abstraction for use in optimizing various types (#7303)

Author: alex

src/cryptography/hazmat/bindings/_rust/__init__.pyi

@@ -2,6 +2,8 @@
 # 2.0, and the BSD License. See the LICENSE file in the root of this repository
 # for complete details.
 
+import typing
+
 def check_pkcs7_padding(data: bytes) -> bool: ...
 def check_ansix923_padding(data: bytes) -> bool: ...
 
@@ -11,3 +13,17 @@ class ObjectIdentifier:
     def dotted_string(self) -> str: ...
     @property
     def _name(self) -> str: ...
+
+T = typing.TypeVar("T")
+
+class FixedPool(typing.Generic[T]):
+    def __init__(
+        self,
+        create: typing.Callable[[], T],
+        destroy: typing.Callable[[T], None],
+    ) -> None: ...
+    def acquire(self) -> PoolAcquisition[T]: ...
+
+class PoolAcquisition(typing.Generic[T]):
+    def __enter__(self) -> T: ...
+    def __exit__(self, exc_type, exc_value, exc_tb) -> None: ...

src/rust/src/lib.rs

@@ -7,6 +7,7 @@
 mod asn1;
 mod intern;
 pub(crate) mod oid;
+mod pool;
 mod x509;
 
 use std::convert::TryInto;
@@ -77,6 +78,7 @@ fn _rust(py: pyo3::Python<'_>, m: &pyo3::types::PyModule) -> pyo3::PyResult<()>
     m.add_function(pyo3::wrap_pyfunction!(check_pkcs7_padding, m)?)?;
     m.add_function(pyo3::wrap_pyfunction!(check_ansix923_padding, m)?)?;
     m.add_class::<oid::ObjectIdentifier>()?;
+    m.add_class::<pool::FixedPool>()?;
 
     m.add_submodule(asn1::create_submodule(py)?)?;
 

src/rust/src/pool.rs

@@ -0,0 +1,96 @@
+// This file is dual licensed under the terms of the Apache License, Version
+// 2.0, and the BSD License. See the LICENSE file in the root of this repository
+// for complete details.
+
+use std::cell::Cell;
+
+// An object pool that can contain a single object and will dynamically
+// allocate new objects to fulfill requests if the pool'd object is already in
+// use.
+#[pyo3::prelude::pyclass]
+pub(crate) struct FixedPool {
+    create_fn: pyo3::PyObject,
+    destroy_fn: pyo3::PyObject,
+
+    value: Cell<Option<pyo3::PyObject>>,
+}
+
+#[pyo3::prelude::pyclass]
+struct PoolAcquisition {
+    pool: pyo3::Py<FixedPool>,
+
+    value: pyo3::PyObject,
+    fresh: bool,
+}
+
+#[pyo3::pymethods]
+impl FixedPool {
+    #[new]
+    fn new(
+        py: pyo3::Python<'_>,
+        create: pyo3::PyObject,
+        destroy: pyo3::PyObject,
+    ) -> pyo3::PyResult<Self> {
+        let value = create.call0(py)?;
+
+        Ok(FixedPool {
+            create_fn: create,
+            destroy_fn: destroy,
+
+            value: Cell::new(Some(value)),
+        })
+    }
+
+    fn acquire(slf: pyo3::Py<Self>, py: pyo3::Python<'_>) -> pyo3::PyResult<PoolAcquisition> {
+        let v = slf.as_ref(py).borrow().value.replace(None);
+        if let Some(value) = v {
+            Ok(PoolAcquisition {
+                pool: slf,
+                value,
+                fresh: false,
+            })
+        } else {
+            let value = slf.as_ref(py).borrow().create_fn.call0(py)?;
+            Ok(PoolAcquisition {
+                pool: slf,
+                value,
+                fresh: true,
+            })
+        }
+    }
+}
+
+impl Drop for FixedPool {
+    fn drop(&mut self) {
+        if let Some(value) = self.value.replace(None) {
+            let gil = pyo3::Python::acquire_gil();
+            let py = gil.python();
+            self.destroy_fn
+                .call1(py, (value,))
+                .expect("FixedPool destroy function failed in destructor");
+        }
+    }
+}
+
+#[pyo3::pymethods]
+impl PoolAcquisition {
+    fn __enter__(&self, py: pyo3::Python<'_>) -> pyo3::PyObject {
+        self.value.clone_ref(py)
+    }
+
+    fn __exit__(
+        &self,
+        py: pyo3::Python<'_>,
+        _exc_type: &pyo3::PyAny,
+        _exc_value: &pyo3::PyAny,
+        _exc_tb: &pyo3::PyAny,
+    ) -> pyo3::PyResult<()> {
+        let pool = self.pool.as_ref(py).borrow();
+        if self.fresh {
+            pool.destroy_fn.call1(py, (self.value.clone_ref(py),))?;
+        } else {
+            pool.value.replace(Some(self.value.clone_ref(py)));
+        }
+        Ok(())
+    }
+}

tests/test_rust_utils.py

@@ -0,0 +1,71 @@
+# This file is dual licensed under the terms of the Apache License, Version
+# 2.0, and the BSD License. See the LICENSE file in the root of this repository
+# for complete details.
+
+import gc
+import threading
+
+from cryptography.hazmat.bindings._rust import FixedPool
+
+
+class TestFixedPool:
+    def test_basic(self):
+        c = 0
+        events = []
+
+        def create():
+            nonlocal c
+            c += 1
+            events.append(("create", c))
+            return c
+
+        def destroy(c):
+            events.append(("destroy", c))
+
+        pool = FixedPool(create, destroy)
+        assert events == [("create", 1)]
+        with pool.acquire() as c:
+            assert c == 1
+            assert events == [("create", 1)]
+
+            with pool.acquire() as c:
+                assert c == 2
+                assert events == [("create", 1), ("create", 2)]
+
+            assert events == [("create", 1), ("create", 2), ("destroy", 2)]
+
+        assert events == [("create", 1), ("create", 2), ("destroy", 2)]
+
+        del pool
+        gc.collect()
+        gc.collect()
+        gc.collect()
+
+        assert events == [
+            ("create", 1),
+            ("create", 2),
+            ("destroy", 2),
+            ("destroy", 1),
+        ]
+
+    def test_thread_stress(self):
+        def create():
+            return None
+
+        def destroy(c):
+            pass
+
+        pool = FixedPool(create, destroy)
+
+        def thread_fn():
+            with pool.acquire():
+                pass
+
+        threads = []
+        for i in range(1024):
+            t = threading.Thread(target=thread_fn)
+            t.start()
+            threads.append(t)
+
+        for t in threads:
+            t.join()