Fix mpmcq use-after-free bug (#4541)

according to address sanitizer, prior to this commit, the mpmcq node
free could cause a use-after-free error from another thread that
might still be reading from `tail->next`. This was technically safe
in this context since the logic will end up looping due to the CAS
failing in the while condition due to the aba counter being out of
sync.

This commit defers freeing of the nodes and instead recycles them
for future pushes onto the mpmcq from the same thread to ensure
that there are no use-after-free concerns any longer.

src/libponyrt/sched/mpmcq.c

@@ -1,6 +1,7 @@
 #define PONY_WANT_ATOMIC_DEFS
 
 #include "mpmcq.h"
+#include "ponyassert.h"
 #include "../mem/pool.h"
 #include "../sched/cpu.h"
 
@@ -16,9 +17,19 @@ struct mpmcq_node_t
   PONY_ATOMIC(void*) data;
 };
 
+static __pony_thread_local mpmcq_node_t* nodes_to_recycle;
+
 static mpmcq_node_t* node_alloc(void* data)
 {
-  mpmcq_node_t* node = POOL_ALLOC(mpmcq_node_t);
+  mpmcq_node_t* node = nodes_to_recycle;
+
+  if (NULL != node)
+  {
+    nodes_to_recycle = atomic_load_explicit(&node->next, memory_order_relaxed);
+  } else {
+    node = POOL_ALLOC(mpmcq_node_t);
+  }
+
   atomic_store_explicit(&node->next, NULL, memory_order_relaxed);
   atomic_store_explicit(&node->data, data, memory_order_relaxed);
 
@@ -55,6 +66,7 @@ static size_t mpmcq_size_debug(mpmcq_t* q)
 
 void ponyint_mpmcq_init(mpmcq_t* q)
 {
+  nodes_to_recycle = NULL;
   mpmcq_node_t* node = node_alloc(NULL);
 
   atomic_store_explicit(&q->head, node, memory_order_relaxed);
@@ -66,8 +78,20 @@ void ponyint_mpmcq_init(mpmcq_t* q)
 #endif
 }
 
+void ponyint_mpmcq_cleanup()
+{
+  while(nodes_to_recycle != NULL)
+  {
+    mpmcq_node_t* node = nodes_to_recycle;
+    nodes_to_recycle = atomic_load_explicit(&node->next, memory_order_relaxed);
+    node_free(node);
+  }
+}
+
 void ponyint_mpmcq_destroy(mpmcq_t* q)
 {
+  pony_assert(atomic_load_explicit(&q->head, memory_order_relaxed) == q->tail.object);
+
   atomic_store_explicit(&q->head, NULL, memory_order_relaxed);
   node_free(q->tail.object);
   q->tail.object = NULL;
@@ -177,7 +201,14 @@ void* ponyint_mpmcq_pop(mpmcq_t* q)
   ANNOTATE_HAPPENS_AFTER(&tail->data);
 #endif
 
-  node_free(tail);
+  // if tail has already been consumed and freed by this thread, the read of
+  // `tail->next` above in another thread would be a use-after-free error but is
+  // technically safe in this context since we will end up looping due to the CAS
+  // failing in the while condition due to the aba counter being out of sync
+  // we cannot free the tail node if we want to ensure there are no use-after-free
+  // concerns so we instead add it to the list of nodes to recycle
+  atomic_store_explicit(&tail->next, nodes_to_recycle, memory_order_relaxed);
+  nodes_to_recycle = tail;
 
   return data;
 }

src/libponyrt/sched/mpmcq.h

@@ -24,6 +24,8 @@ void ponyint_mpmcq_init(mpmcq_t* q);
 
 void ponyint_mpmcq_destroy(mpmcq_t* q);
 
+void ponyint_mpmcq_cleanup();
+
 void ponyint_mpmcq_push(mpmcq_t* q, void* data);
 
 void ponyint_mpmcq_push_single(mpmcq_t* q, void* data);

src/libponyrt/sched/scheduler.c

@@ -1125,6 +1125,8 @@ static DECLARE_THREAD_FN(run_thread)
 #endif
 
   run(sched);
+
+  ponyint_mpmcq_cleanup();
   ponyint_pool_thread_cleanup();
 
   return 0;