eventstream fixes (#206)

- fix `setup.py` so that eventstream modules get installed
- fix crash when `eventstream.rpc.ClientConnection` was GC'd during `on_connection_shutdown` callback
- change default args to always be `=None`, so that it's easy to pass along defaults from wrapper functions
- centralize logic that sanitizes message args going in each direction between C <-> Python

Author: graebm

awscrt/eventstream/rpc.py

@@ -185,6 +185,35 @@ def on_protocol_message(
         pass
 
 
+def _to_binding_msg_args(headers, payload, message_type, flags):
+    """
+    Transform args that a python send-msg function would take,
+    into args that a native send-msg function would take.
+    """
+    # python functions for sending messages
+    if headers is None:
+        headers = []
+    else:
+        headers = [i._as_binding_tuple() for i in headers]
+    if payload is None:
+        payload = b''
+    if flags is None:
+        flags = MessageFlag.NONE
+    return (headers, payload, message_type, flags)
+
+
+def _from_binding_msg_args(headers, payload, message_type, flags):
+    """
+    Transform msg-received args that came from native,
+    into msg-received args presented to python users.
+    """
+    headers = [Header._from_binding_tuple(i) for i in headers]
+    if payload is None:
+        payload = b''
+    message_type = MessageType(message_type)
+    return (headers, payload, message_type, flags)
+
+
 class ClientConnection(NativeResource):
     """A client connection for the event-stream RPC protocol.
 
@@ -323,8 +352,7 @@ def _on_protocol_message(bound_weak_handler, headers, payload, message_type, fla
         handler = bound_weak_handler()
         if handler:
             # transform from simple types to actual classes
-            headers = [Header._from_binding_tuple(i) for i in headers]
-            message_type = MessageType(message_type)
+            headers, payload, message_type, flags = _from_binding_msg_args(headers, payload, message_type, flags)
             handler.on_protocol_message(
                 headers=headers,
                 payload=payload,
@@ -336,8 +364,7 @@ def _on_continuation_message(bound_weak_handler, headers, payload, message_type,
         handler = bound_weak_handler()
         if handler:
             # transform from simple types to actual classes
-            headers = [Header._from_binding_tuple(i) for i in headers]
-            message_type = MessageType(message_type)
+            headers, payload, message_type, flags = _from_binding_msg_args(headers, payload, message_type, flags)
             handler.on_continuation_message(
                 headers=headers,
                 payload=payload,
@@ -358,7 +385,7 @@ def _on_flush(bound_future, bound_callback, error_code):
             else:
                 bound_future.set_result(None)
 
-    def close(self, reason=None):
+    def close(self):
         """Close the connection.
 
         Shutdown is asynchronous. This call has no effect if the connection is already
@@ -384,10 +411,10 @@ def is_open(self):
     def send_protocol_message(
             self,
             *,
-            headers: Optional[Sequence[Header]] = [],
-            payload: Optional[ByteString] = b'',
+            headers: Optional[Sequence[Header]] = None,
+            payload: Optional[ByteString] = None,
             message_type: MessageType,
-            flags: int = MessageFlag.NONE,
+            flags: Optional[int] = None,
             on_flush: Callable = None) -> Future:
         """Send a protocol message.
 
@@ -428,7 +455,7 @@ def send_protocol_message(
         future = Future()
 
         # native code deals with simplified types
-        headers = [i._as_binding_tuple() for i in headers]
+        headers, payload, message_type, flags = _to_binding_msg_args(headers, payload, message_type, flags)
 
         _awscrt.event_stream_rpc_client_connection_send_protocol_message(
             self._binding,
@@ -488,10 +515,10 @@ def activate(
             self,
             *,
             operation: str,
-            headers: Sequence[Header] = [],
-            payload: ByteString = b'',
+            headers: Sequence[Header] = None,
+            payload: ByteString = None,
             message_type: MessageType,
-            flags: int = MessageFlag.NONE,
+            flags: int = None,
             on_flush: Callable = None):
         """
         Activate the stream by sending its first message.
@@ -537,7 +564,7 @@ def activate(
         flush_future = Future()
 
         # native code deals with simplified types
-        headers = [i._as_binding_tuple() for i in headers]
+        headers, payload, message_type, flags = _to_binding_msg_args(headers, payload, message_type, flags)
 
         _awscrt.event_stream_rpc_client_continuation_activate(
             self._binding,
@@ -553,10 +580,10 @@ def activate(
     def send_message(
             self,
             *,
-            headers: Sequence[Header] = [],
-            payload: ByteString = b'',
+            headers: Sequence[Header] = None,
+            payload: ByteString = None,
             message_type: MessageType,
-            flags: int = MessageFlag.NONE,
+            flags: int = None,
             on_flush: Callable = None) -> Future:
         """
         Send a continuation message.
@@ -600,7 +627,7 @@ def send_message(
         """
         future = Future()
         # native code deals with simplified types
-        headers = [i._as_binding_tuple() for i in headers]
+        headers, payload, message_type, flags = _to_binding_msg_args(headers, payload, message_type, flags)
 
         _awscrt.event_stream_rpc_client_continuation_send_message(
             self._binding,

setup.py

@@ -304,7 +304,8 @@ def awscrt_ext():
     author_email="aws-sdk-common-runtime@amazon.com",
     description="A common runtime for AWS Python projects",
     url="https://github.com/awslabs/aws-crt-python",
-    packages=['awscrt'],
+    # Note: find_packages() without extra args will end up installing test/
+    packages=setuptools.find_packages(include=['awscrt*']),
     classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: Apache Software License",

source/event_stream_rpc_client_connection.c

@@ -187,7 +187,6 @@ static void s_on_connection_shutdown(
 
     AWS_FATAL_ASSERT(connection->native && "Illegal for event-stream connection shutdown to fire before setup");
     AWS_FATAL_ASSERT(!connection->shutdown_complete && "illegal for event-stream connection shutdown to fire twice");
-    connection->shutdown_complete = true;
 
     PyGILState_STATE state;
     if (aws_py_gilstate_ensure(&state)) {
@@ -202,6 +201,7 @@ static void s_on_connection_shutdown(
         PyErr_WriteUnraisable(PyErr_Occurred());
     }
 
+    connection->shutdown_complete = true;
     s_connection_destroy_if_ready(connection);
     PyGILState_Release(state);
 }
@@ -257,20 +257,12 @@ static void s_on_protocol_message(
         return; /* Python has shut down. Nothing matters anymore, but don't crash */
     }
 
-    /* We always want to deliver bytes to python user, even if the length is 0.
-     * But PyObject_CallFunction() with "y#" will convert a NULL ptr to None instead of 0-length bytes.
-     * Therefore, if message_args->payload_buffer is NULL, pass some other valid ptr instead. */
-    const char *payload_ptr = (void *)message_args->payload->buffer;
-    if (payload_ptr == NULL) {
-        payload_ptr = "";
-    }
-
     PyObject *result = PyObject_CallFunction(
         connection->on_protocol_message,
         "(Oy#iI)",
         /* NOTE: if headers_create() returns NULL, then PyObject_CallFunction() fails too, which is convenient */
         aws_py_event_stream_python_headers_create(message_args->headers, message_args->headers_count),
-        payload_ptr,
+        message_args->payload->buffer,
         message_args->payload->len,
         message_args->message_type,
         message_args->message_flags);

source/event_stream_rpc_client_continuation.c

@@ -162,20 +162,12 @@ static void s_on_continuation_message(
         return; /* Python has shut down. Nothing matters anymore, but don't crash */
     }
 
-    /* We always want to deliver bytes to python user, even if the length is 0.
-     * But PyObject_CallFunction() with "y#" will convert a NULL ptr to None instead of 0-length bytes.
-     * Therefore, if message_args->payload_buffer is NULL, pass some other valid ptr instead. */
-    const char *payload_ptr = (void *)message_args->payload->buffer;
-    if (payload_ptr == NULL) {
-        payload_ptr = "";
-    }
-
     PyObject *result = PyObject_CallFunction(
         continuation->on_message,
         "(Oy#iI)",
         /* NOTE: if headers_create() returns NULL, then PyObject_CallFunction() fails too, which is convenient */
         aws_py_event_stream_python_headers_create(message_args->headers, message_args->headers_count),
-        payload_ptr,
+        message_args->payload->buffer,
         message_args->payload->len,
         message_args->message_type,
         message_args->message_flags);