bootstrap-shutdown-cb: python edition (#114)

awscrt/io.py

@@ -16,6 +16,7 @@
 from awscrt import NativeResource, isinstance_str
 from enum import IntEnum
 import io
+import threading
 
 
 class LogLevel(IntEnum):
@@ -75,15 +76,21 @@ def __init__(self, event_loop_group, max_hosts=16):
 
 
 class ClientBootstrap(NativeResource):
-    __slots__ = ()
+    __slots__ = ('shutdown_event')
 
     def __init__(self, event_loop_group, host_resolver):
         assert isinstance(event_loop_group, EventLoopGroup)
         assert isinstance(host_resolver, HostResolverBase)
 
         super(ClientBootstrap, self).__init__()
 
-        self._binding = _awscrt.client_bootstrap_new(event_loop_group, host_resolver)
+        shutdown_event = threading.Event()
+
+        def on_shutdown():
+            shutdown_event.set()
+
+        self.shutdown_event = shutdown_event
+        self._binding = _awscrt.client_bootstrap_new(event_loop_group, host_resolver, on_shutdown)
 
 
 def _read_binary_file(filepath):

source/io.c

@@ -249,16 +249,43 @@ struct client_bootstrap_binding {
     /* Dependencies that must outlive this */
     PyObject *event_loop_group;
     PyObject *host_resolver;
+    PyObject *shutdown_complete;
 };
 
-static void s_client_bootstrap_destructor(PyObject *bootstrap_capsule) {
+/* Fires after the native client bootstrap finishes shutting down. */
+static void s_client_bootstrap_on_shutdown_complete(void *user_data) {
+    struct client_bootstrap_binding *bootstrap = user_data;
+    PyObject *shutdown_complete = bootstrap->shutdown_complete;
+
+    /*************** GIL ACQUIRE ***************/
+    PyGILState_STATE state = PyGILState_Ensure();
+
+    Py_XDECREF(bootstrap->host_resolver);
+    Py_XDECREF(bootstrap->event_loop_group);
+
+    aws_mem_release(aws_py_get_allocator(), bootstrap);
+
+    if (shutdown_complete) {
+        PyObject *result = PyObject_CallFunction(shutdown_complete, "()");
+        if (result) {
+            Py_DECREF(result);
+        } else {
+            PyErr_WriteUnraisable(PyErr_Occurred());
+        }
+        Py_DECREF(shutdown_complete);
+    }
+
+    PyGILState_Release(state);
+    /*************** GIL RELEASE ***************/
+}
+
+/* Fires when python capsule is GC'd.
+ * Note that bootstrap shutdown is async, we can't release dependencies until it completes */
+static void s_client_bootstrap_capsule_destructor(PyObject *bootstrap_capsule) {
     struct client_bootstrap_binding *bootstrap =
         PyCapsule_GetPointer(bootstrap_capsule, s_capsule_name_client_bootstrap);
-    assert(bootstrap);
-    Py_DECREF(bootstrap->host_resolver);
-    Py_DECREF(bootstrap->event_loop_group);
+
     aws_client_bootstrap_release(bootstrap->native);
-    aws_mem_release(aws_py_get_allocator(), bootstrap);
 }
 
 PyObject *aws_py_client_bootstrap_new(PyObject *self, PyObject *args) {
@@ -268,8 +295,9 @@ PyObject *aws_py_client_bootstrap_new(PyObject *self, PyObject *args) {
 
     PyObject *elg_py;
     PyObject *host_resolver_py;
+    PyObject *shutdown_complete_py;
 
-    if (!PyArg_ParseTuple(args, "OO", &elg_py, &host_resolver_py)) {
+    if (!PyArg_ParseTuple(args, "OOO", &elg_py, &host_resolver_py, &shutdown_complete_py)) {
         return NULL;
     }
 
@@ -291,15 +319,22 @@ PyObject *aws_py_client_bootstrap_new(PyObject *self, PyObject *args) {
 
     /* From hereon, we need to clean up if errors occur */
 
-    bootstrap->native = aws_client_bootstrap_new(allocator, elg, host_resolver, NULL);
-    if (!bootstrap->native) {
-        PyErr_SetAwsLastError();
-        goto bootstrap_new_failed;
+    PyObject *capsule =
+        PyCapsule_New(bootstrap, s_capsule_name_client_bootstrap, s_client_bootstrap_capsule_destructor);
+    if (!capsule) {
+        goto error;
     }
 
-    PyObject *capsule = PyCapsule_New(bootstrap, s_capsule_name_client_bootstrap, s_client_bootstrap_destructor);
-    if (!capsule) {
-        goto capsule_new_failed;
+    struct aws_client_bootstrap_options bootstrap_options = {
+        .event_loop_group = elg,
+        .host_resolver = host_resolver,
+        .on_shutdown_complete = s_client_bootstrap_on_shutdown_complete,
+        .user_data = bootstrap,
+    };
+    bootstrap->native = aws_client_bootstrap_new(allocator, &bootstrap_options);
+    if (!bootstrap->native) {
+        PyErr_SetAwsLastError();
+        goto error;
     }
 
     /* From hereon, nothing will fail */
@@ -310,12 +345,17 @@ PyObject *aws_py_client_bootstrap_new(PyObject *self, PyObject *args) {
     bootstrap->host_resolver = host_resolver_py;
     Py_INCREF(host_resolver_py);
 
+    bootstrap->shutdown_complete = shutdown_complete_py;
+    Py_INCREF(bootstrap->shutdown_complete);
+
     return capsule;
 
-capsule_new_failed:
-    aws_client_bootstrap_release(bootstrap->native);
-bootstrap_new_failed:
-    aws_mem_release(allocator, bootstrap);
+error:
+    if (capsule) {
+        Py_DECREF(capsule);
+    } else {
+        aws_mem_release(allocator, bootstrap);
+    }
     return NULL;
 }
 

test/__init__.py

@@ -20,6 +20,8 @@
 import types
 import unittest
 
+TIMEOUT = 10.0
+
 
 class NativeResourceTest(unittest.TestCase):
     """
@@ -33,8 +35,9 @@ def tearDown(self):
         gc.collect()
 
         # Native resources might need a few more ticks to finish cleaning themselves up.
-        if NativeResource._living:
-            time.sleep(1)
+        wait_until = time.time() + TIMEOUT
+        while NativeResource._living and time.time() < wait_until:
+            time.sleep(0.1)
 
         # Print out debugging info on leaking resources
         if NativeResource._living:

test/test_auth.py

@@ -16,7 +16,7 @@
 import awscrt.io
 import datetime
 import os
-from test import NativeResourceTest
+from test import NativeResourceTest, TIMEOUT
 
 EXAMPLE_ACCESS_KEY_ID = 'example_access_key_id'
 EXAMPLE_SECRET_ACCESS_KEY = 'example_secret_access_key'
@@ -65,7 +65,7 @@ def test_static_provider(self):
             EXAMPLE_SESSION_TOKEN)
 
         future = provider.get_credentials()
-        credentials = future.result()
+        credentials = future.result(TIMEOUT)
 
         self.assertEqual(EXAMPLE_ACCESS_KEY_ID, credentials.access_key_id)
         self.assertEqual(EXAMPLE_SECRET_ACCESS_KEY, credentials.secret_access_key)
@@ -80,7 +80,7 @@ def test_static_provider(self):
     #         self.example_secret_access_key)
 
     #     future = provider.get_credentials()
-    #     credentials = future.result()
+    #     credentials = future.result(TIMEOUT)
 
     #     self.assertEqual(self.example_access_key_id, credentials.access_key_id)
     #     self.assertEqual(self.example_secret_access_key, credentials.secret_access_key)
@@ -96,7 +96,7 @@ def test_default_provider(self):
         provider = awscrt.auth.AwsCredentialsProvider.new_default_chain(bootstrap)
 
         future = provider.get_credentials()
-        credentials = future.result()
+        credentials = future.result(TIMEOUT)
 
         self.assertEqual('credentials_test_access_key_id', credentials.access_key_id)
         self.assertEqual('credentials_test_secret_access_key', credentials.secret_access_key)
@@ -234,7 +234,7 @@ def test_signing_sigv4_headers(self):
 
         signing_future = awscrt.auth.aws_sign_request(http_request, signing_config)
 
-        signing_result = signing_future.result(10)
+        signing_result = signing_future.result(TIMEOUT)
 
         self.assertIs(http_request, signing_result)  # should be same object
 

test/test_io.py

@@ -13,7 +13,7 @@
 
 from __future__ import absolute_import
 from awscrt.io import ClientBootstrap, ClientTlsContext, DefaultHostResolver, EventLoopGroup, TlsConnectionOptions, TlsContextOptions
-from test import NativeResourceTest
+from test import NativeResourceTest, TIMEOUT
 import unittest
 
 
@@ -32,11 +32,16 @@ def test_init(self):
 
 
 class ClientBootstrapTest(NativeResourceTest):
-    def test_init(self):
+    def test_create_destroy(self):
         event_loop_group = EventLoopGroup()
         host_resolver = DefaultHostResolver(event_loop_group)
         bootstrap = ClientBootstrap(event_loop_group, host_resolver)
 
+        # ensure shutdown_event fires
+        bootstrap_shutdown_event = bootstrap.shutdown_event
+        del bootstrap
+        self.assertTrue(bootstrap_shutdown_event.wait(TIMEOUT))
+
 
 class ClientTlsContextTest(NativeResourceTest):
     def test_init_defaults(self):

test/test_mqtt.py

@@ -23,6 +23,7 @@
 import unittest
 import boto3
 import botocore.exceptions
+import shutil
 import tempfile
 import time
 import uuid
@@ -192,20 +193,25 @@ def test_mtls_from_path(self):
         bootstrap = ClientBootstrap(elg, resolver)
 
         # test "from path" builder by writing secrets to tempfiles
-        with tempfile.NamedTemporaryFile() as cert_file:
-            with tempfile.NamedTemporaryFile() as key_file:
+        tmp_dirpath = tempfile.mkdtemp()
+        try:
+            cert_filepath = os.path.join(tmp_dirpath, 'cert')
+            with open(cert_filepath, 'wb') as cert_file:
                 cert_file.write(config.cert)
-                cert_file.flush()
 
+            key_filepath = os.path.join(tmp_dirpath, 'key')
+            with open(key_filepath, 'wb') as key_file:
                 key_file.write(config.key)
-                key_file.flush()
-
-                connection = awsiot_mqtt_connection_builder.mtls_from_path(
-                    cert_filepath=cert_file.name,
-                    pri_key_filepath=key_file.name,
-                    endpoint=config.endpoint,
-                    client_id=create_client_id(),
-                    client_bootstrap=bootstrap)
+
+            connection = awsiot_mqtt_connection_builder.mtls_from_path(
+                cert_filepath=cert_filepath,
+                pri_key_filepath=key_filepath,
+                endpoint=config.endpoint,
+                client_id=create_client_id(),
+                client_bootstrap=bootstrap)
+
+        finally:
+            shutil.rmtree(tmp_dirpath)
 
         self._test_connection(connection)