1 files changed, 48 insertions, 8 deletions

diff --git a/ot_http.c b/ot_http.c
index 66f0102..b0e1de3 100644
--- a/ot_http.c
+++ b/ot_http.c
@@ -22,6 +22,7 @@
 #include "ot_iovec.h"
 #include "ot_ip.h"
 #include "ot_mutex.h"
+#include "ot_workqueue.h"
 #include "ot_scan.h"
 #include "ot_event.h"
 #include "ot_fdstate.h"
@@ -39,9 +40,33 @@ ssize_t      g_stats_path_len;
 
 enum { SUCCESS_HTTP_HEADER_LENGTH = 80, SUCCESS_HTTP_SIZE_OFF = 17 };
 
-/* Give up a connection: the one place that knows everything hanging off it.
-   The slot is released before anything is freed, so that a late event or the
-   expiry sweep finds a closed descriptor rather than a freed cookie. */
+struct http_data *http_adopt(const int64_t sock) {
+  struct http_data *cookie = fdstate_cookie(sock), *heap;
+  int               interest;
+
+  if (!cookie || !(cookie->flag & STRUCT_HTTP_FLAG_ONSTACK))
+    return cookie;
+
+  if (!(heap = malloc(sizeof(*heap))))
+    return NULL;
+
+  *heap = *cookie;
+  heap->flag &= ~STRUCT_HTTP_FLAG_ONSTACK;
+
+  /* Publishing the cookie is what lets another thread reach this connection,
+     so nothing may touch it afterwards, not even to read what it asked for */
+  interest = heap->interest;
+  fdstate_set_cookie((int)sock, heap);
+
+  if (event_add((int)sock, interest, fdstate_generation((int)sock)) == -1) {
+    fdstate_set_cookie((int)sock, cookie);
+    free(heap);
+    return NULL;
+  }
+  return heap;
+}
+
+/* Give up a connection: the one place that knows everything hanging off it. */
 void http_close(const int64_t sock) {
   struct http_data *cookie = fdstate_cookie(sock);
 
@@ -53,8 +78,9 @@ void http_close(const int64_t sock) {
     free(cookie->batch);
     buf_reset(&cookie->request);
     if (cookie->flag & (STRUCT_HTTP_FLAG_WAITINGFORTASK | STRUCT_HTTP_FLAG_CHUNKED_IN_TRANSFER))
-      mutex_workqueue_canceltask(sock);
-    free(cookie);
+      workqueue_canceltask(sock);
+    if (!(cookie->flag & STRUCT_HTTP_FLAG_ONSTACK))
+      free(cookie);
   }
 
   event_close((int)sock);
@@ -140,6 +166,10 @@ void http_interest_set(const int64_t sock, struct http_data *cookie, int what) {
   if (cookie->interest == what)
     return;
   cookie->interest = what;
+  /* Still on a stack, so not in the event loop yet: whoever adopts this
+     connection registers whatever it asked for last */
+  if (cookie->flag & STRUCT_HTTP_FLAG_ONSTACK)
+    return;
   event_modify((int)sock, what, fdstate_generation((int)sock));
 }
 
@@ -327,13 +357,17 @@ static ssize_t http_handle_stats(const int64_t sock, struct ot_workstruct *ws, c
     }
 #endif
 #endif
-    /* Pass this task to the worker thread */
     cookie->flag |= STRUCT_HTTP_FLAG_WAITINGFORTASK | STRUCT_HTTP_FLAG_CHUNKED;
 
     /* Clients waiting for us should not easily timeout */
     fdstate_set_expiry(sock, 0);
-    fullscrape_deliver(sock, format);
     http_interest_set(sock, cookie, 0);
+
+    /* Everything this connection needs has been written, so it can be given
+       a home of its own and only then be made reachable by a worker thread */
+    if (!http_adopt(sock))
+      HTTPERROR_500;
+    fullscrape_deliver(sock, format);
     return ws->reply_size = -2;
   }
 #endif
@@ -345,6 +379,9 @@ static ssize_t http_handle_stats(const int64_t sock, struct ot_workstruct *ws, c
 
     /* Complex stats also include expensive memory debugging tools */
     fdstate_set_expiry(sock, 0);
+
+    if (!http_adopt(sock))
+      HTTPERROR_500;
     stats_deliver(sock, mode);
     return ws->reply_size = -2;
   }
@@ -427,8 +464,11 @@ static ssize_t http_handle_fullscrape(const int64_t sock, struct ot_workstruct *
   cookie->flag |= STRUCT_HTTP_FLAG_WAITINGFORTASK | STRUCT_HTTP_FLAG_CHUNKED;
   /* Clients waiting for us should not easily timeout */
   fdstate_set_expiry(sock, 0);
-  fullscrape_deliver(sock, format);
   http_interest_set(sock, cookie, 0);
+
+  if (!http_adopt(sock))
+    HTTPERROR_500;
+  fullscrape_deliver(sock, format);
   return ws->reply_size = -2;
 }
 #endif