add ref-counting to aws_websocket (#405)

Author: graebm

include/aws/http/websocket.h

@@ -351,11 +351,21 @@ AWS_HTTP_API
 int aws_websocket_client_connect(const struct aws_websocket_client_connection_options *options);
 
 /**
+ * Increment the websocket's ref-count, preventing it from being destroyed.
+ * @return Always returns the same pointer that is passed in.
+ */
+AWS_HTTP_API
+struct aws_websocket *aws_websocket_acquire(struct aws_websocket *websocket);
+
+/**
+ * Decrement the websocket's ref-count.
+ * When the ref-count reaches zero, the connection will shut down, if it hasn't already.
  * Users must release the websocket when they are done with it.
  * The websocket's memory cannot be reclaimed until this is done.
- * If the websocket connection was not already shutting down, it will be shut down.
  * Callbacks may continue firing after this is called, with "shutdown" being the final callback.
  * This function may be called from any thread.
+ *
+ * It is safe to pass NULL, nothing will happen.
  */
 AWS_HTTP_API
 void aws_websocket_release(struct aws_websocket *websocket);

source/websocket.c

@@ -9,6 +9,7 @@
 #include <aws/common/device_random.h>
 #include <aws/common/encoding.h>
 #include <aws/common/mutex.h>
+#include <aws/common/ref_count.h>
 #include <aws/http/private/websocket_decoder.h>
 #include <aws/http/private/websocket_encoder.h>
 #include <aws/http/request_response.h>
@@ -41,6 +42,7 @@ struct outgoing_frame {
 
 struct aws_websocket {
     struct aws_allocator *alloc;
+    struct aws_ref_count ref_count;
     struct aws_channel_handler channel_handler;
     struct aws_channel_slot *channel_slot;
     size_t initial_window_size;
@@ -137,9 +139,6 @@ struct aws_websocket {
 
         /* Mirrors variable from thread_data */
         bool is_midchannel_handler;
-
-        /* Whether aws_websocket_release() has been called */
-        bool is_released;
     } synced_data;
 };
 
@@ -168,6 +167,7 @@ static int s_handler_shutdown(
 static size_t s_handler_initial_window_size(struct aws_channel_handler *handler);
 static size_t s_handler_message_overhead(struct aws_channel_handler *handler);
 static void s_handler_destroy(struct aws_channel_handler *handler);
+static void s_websocket_on_refcount_zero(void *user_data);
 
 static int s_encoder_stream_outgoing_payload(struct aws_byte_buf *out_buf, void *user_data);
 
@@ -271,6 +271,7 @@ struct aws_websocket *aws_websocket_handler_new(const struct aws_websocket_handl
     }
 
     websocket->alloc = options->allocator;
+    aws_ref_count_init(&websocket->ref_count, websocket, s_websocket_on_refcount_zero);
     websocket->channel_handler.vtable = &s_channel_handler_vtable;
     websocket->channel_handler.alloc = options->allocator;
     websocket->channel_handler.impl = websocket;
@@ -357,30 +358,28 @@ static void s_handler_destroy(struct aws_channel_handler *handler) {
     aws_mem_release(websocket->alloc, websocket);
 }
 
-void aws_websocket_release(struct aws_websocket *websocket) {
-    AWS_ASSERT(websocket);
-    AWS_ASSERT(websocket->channel_slot);
-
-    bool was_already_released;
-
-    /* BEGIN CRITICAL SECTION */
-    s_lock_synced_data(websocket);
-    if (websocket->synced_data.is_released) {
-        was_already_released = true;
-    } else {
-        was_already_released = false;
-        websocket->synced_data.is_released = true;
-    }
-    s_unlock_synced_data(websocket);
-    /* END CRITICAL SECTION */
+struct aws_websocket *aws_websocket_acquire(struct aws_websocket *websocket) {
+    AWS_PRECONDITION(websocket);
+    AWS_LOGF_TRACE(AWS_LS_HTTP_WEBSOCKET, "id=%p: Acquiring websocket ref-count.", (void *)websocket);
+    aws_ref_count_acquire(&websocket->ref_count);
+    return websocket;
+}
 
-    if (was_already_released) {
-        AWS_LOGF_TRACE(
-            AWS_LS_HTTP_WEBSOCKET, "id=%p: Ignoring multiple calls to websocket release.", (void *)websocket);
+void aws_websocket_release(struct aws_websocket *websocket) {
+    if (!websocket) {
         return;
     }
 
-    AWS_LOGF_TRACE(AWS_LS_HTTP_WEBSOCKET, "id=%p: Websocket released, shut down if necessary.", (void *)websocket);
+    AWS_LOGF_TRACE(AWS_LS_HTTP_WEBSOCKET, "id=%p: Releasing websocket ref-count.", (void *)websocket);
+    aws_ref_count_release(&websocket->ref_count);
+}
+
+static void s_websocket_on_refcount_zero(void *user_data) {
+    struct aws_websocket *websocket = user_data;
+    AWS_ASSERT(websocket->channel_slot);
+
+    AWS_LOGF_TRACE(
+        AWS_LS_HTTP_WEBSOCKET, "id=%p: Websocket ref-count is zero, shut down if necessary.", (void *)websocket);
 
     /* Channel might already be shut down, but make sure */
     s_schedule_channel_shutdown(websocket, AWS_ERROR_SUCCESS);
@@ -422,26 +421,6 @@ int aws_websocket_convert_to_midchannel_handler(struct aws_websocket *websocket)
         return aws_raise_error(AWS_ERROR_INVALID_STATE);
     }
 
-    bool was_released = false;
-
-    /* BEGIN CRITICAL SECTION */
-    s_lock_synced_data(websocket);
-    if (websocket->synced_data.is_released) {
-        was_released = true;
-    } else {
-        websocket->synced_data.is_midchannel_handler = true;
-    }
-    s_unlock_synced_data(websocket);
-    /* END CRITICAL SECTION */
-
-    if (was_released) {
-        AWS_LOGF_ERROR(
-            AWS_LS_HTTP_WEBSOCKET,
-            "id=%p: Cannot convert websocket to midchannel handler because it was already released.",
-            (void *)websocket);
-        return aws_raise_error(AWS_ERROR_HTTP_CONNECTION_CLOSED);
-    }
-
     websocket->thread_data.is_midchannel_handler = true;
 
     return AWS_OP_SUCCESS;

tests/CMakeLists.txt

@@ -180,6 +180,7 @@ add_test_case(websocket_encoder_fragmented_message)
 add_test_case(websocket_encoder_fragmentation_failure_checks)
 add_test_case(websocket_encoder_payload_callback_can_fail_encoder)
 add_test_case(websocket_handler_sanity_check)
+add_test_case(websocket_handler_refcounting)
 add_test_case(websocket_handler_send_frame)
 add_test_case(websocket_handler_send_frame_off_thread)
 add_test_case(websocket_handler_send_multiple_frames)

tests/test_websocket_handler.c

@@ -650,6 +650,25 @@ TEST_CASE(websocket_handler_sanity_check) {
     return AWS_OP_SUCCESS;
 }
 
+TEST_CASE(websocket_handler_refcounting) {
+    (void)ctx;
+    struct tester tester;
+    ASSERT_SUCCESS(s_tester_init(&tester, allocator));
+
+    /* acquire() and then release() a refcount. The websocket should not shut down yet */
+    ASSERT_PTR_EQUALS(tester.websocket, aws_websocket_acquire(tester.websocket));
+    aws_websocket_release(tester.websocket);
+
+    testing_channel_drain_queued_tasks(&tester.testing_channel);
+    ASSERT_FALSE(tester.testing_channel.channel_shutdown_completed);
+
+    /* should be safe to call release() on NULL */
+    aws_websocket_release(NULL);
+
+    ASSERT_SUCCESS(s_tester_clean_up(&tester));
+    return AWS_OP_SUCCESS;
+}
+
 static int s_websocket_handler_send_frame_common(struct aws_allocator *allocator, bool on_thread) {
     struct tester tester;
     ASSERT_SUCCESS(s_tester_init(&tester, allocator));