Emit drain event, add comments, etc

src/WebSocket.h

@@ -88,7 +88,7 @@ public:
         return true;
     }
 
-    /* Emit close event, stat passive timeout */
+    /* Emit close event, start passive timeout */
     void close(int code, std::string_view message = {}) {
         static const int MAX_CLOSE_PAYLOAD = 123;
         int length = std::min<size_t>(MAX_CLOSE_PAYLOAD, message.length());
@@ -96,11 +96,17 @@ public:
         // todo: here we start a timeout and handle it accordingly in the timeout handler
 
         WebSocketData *webSocketData = (WebSocketData *) static_dispatch(us_ssl_socket_ext, us_socket_ext)((SOCKET_TYPE *) this);
+
+        /* We postpone any FIN sending to either drainage or uncorking */
         webSocketData->isShuttingDown = true;
 
         /* Format and send the close frame */
         char closePayload[MAX_CLOSE_PAYLOAD + 2];
         int closePayloadLength = (int) WebSocketProtocol<isServer, WebSocketContext<SSL, isServer>>::formatClosePayload(closePayload, code, message.data(), length);
+
+        // but what if we are NOT corked, THEN we can FIN here if we succeeded
+
+        // if we are corked and send returns true we cannot know for sure if we can fin
         send(std::string_view(closePayload, closePayloadLength), OpCode::CLOSE);
 
         // why should we fin here?
@@ -110,7 +116,9 @@ public:
         WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)(
             (SOCKET_CONTEXT_TYPE *) static_dispatch(us_ssl_socket_get_context, us_socket_get_context)((SOCKET_TYPE *) this)
         );
-        webSocketContextData->closeHandler(this, code, message);
+        if (webSocketContextData->closeHandler) {
+            webSocketContextData->closeHandler(this, code, message);
+        }
     }
 };
 

src/WebSocketContext.h

@@ -21,7 +21,7 @@
 #include "WebSocketContextData.h"
 #include "WebSocketProtocol.h"
 #include "WebSocketData.h"
-#include "AsyncSocket.h"
+#include "WebSocket.h"
 
 namespace uWS {
 
@@ -99,9 +99,11 @@ private:
                 }
 
                 /* Emit message event & break if we are closed or shut down when returning */
-                webSocketContextData->messageHandler((WebSocket<SSL, isServer> *) s, std::string_view(data, length), (uWS::OpCode) opCode);
+                if (webSocketContextData->messageHandler) {
-                if (us_socket_is_closed((us_socket *) s) || webSocketData->isShuttingDown) {
+                    webSocketContextData->messageHandler((WebSocket<SSL, isServer> *) s, std::string_view(data, length), (uWS::OpCode) opCode);
-                    return true;
+                    if (us_socket_is_closed((us_socket *)s) || webSocketData->isShuttingDown) {
+                        return true;
+                    }
                 }
             } else {
                 /* Allocate fragment buffer up front first time */
@@ -150,9 +152,11 @@ private:
                     }
 
                     /* Emit message and check for shutdown or close */
-                    webSocketContextData->messageHandler((WebSocket<SSL, isServer> *) s, std::string_view(data, length), (uWS::OpCode) opCode);
+                    if (webSocketContextData->messageHandler) {
-                    if (us_socket_is_closed((us_socket *) s) || webSocketData->isShuttingDown) {
+                        webSocketContextData->messageHandler((WebSocket<SSL, isServer> *) s, std::string_view(data, length), (uWS::OpCode) opCode);
-                        return true;
+                        if (us_socket_is_closed((us_socket *)s) || webSocketData->isShuttingDown) {
+                            return true;
+                        }
                     }
 
                     /* If we shutdown or closed, this will be taken care of elsewhere */
@@ -247,18 +251,20 @@ private:
             webSocket->cork();
 
             /* We need the websocket data */
-            WebSocketData *wsState = (WebSocketData *) (static_dispatch(us_ssl_socket_ext, us_socket_ext)(s));
+            WebSocketData *webSocketData = (WebSocketData *) (static_dispatch(us_ssl_socket_ext, us_socket_ext)(s));
 
             /* This parser has virtually no overhead */
-            uWS::WebSocketProtocol<isServer, WebSocketContext<SSL, isServer>>::consume(data, length, wsState, s);
+            uWS::WebSocketProtocol<isServer, WebSocketContext<SSL, isServer>>::consume(data, length, (WebSocketState<isServer> *) webSocketData, s);
+
+            // todo: we need to check for close and shutdown here? as we just emitted a bunch of message/close events!
 
             // todo: check for failures here just like for HTTP
             webSocket->uncork();
 
-            // I guess we need to check drain here
+            // I guess we need to check drain here - emit drain if we had to poll for writable
 
-            // are we shutdown?
+            // are we shutdown? can onnly call this if we did succeed uncork!
-            if (wsState->isShuttingDown) {
+            if (webSocketData->isShuttingDown) {
                 webSocket->shutdown();
             }
 
@@ -268,18 +274,34 @@ private:
         /* Handle HTTP write out (note: SSL_read may trigger this spuriously, the app need to handle spurious calls) */
         static_dispatch(us_ssl_socket_context_on_writable, us_socket_context_on_writable)(getSocketContext(), [](auto *s) {
 
-            std::cout << "websocket writable" << std::endl;
+            // check if we already shut down the us socket and return?
-
-            // we need to drain here!
 
             AsyncSocket<SSL> *webSocket = (AsyncSocket<SSL> *) s;
+            WebSocketData *webSocketData = (WebSocketData *)(static_dispatch(us_ssl_socket_ext, us_socket_ext)(s));
 
-            // check for failures and shutdown just like in data event
+            /* Drain as much as possible */
-            webSocket->write(nullptr, 0); // drainage - also check for shutdown!
+            webSocket->write(nullptr, 0);
-
-            // call drain here
 
+            /* Are we in (WebSocket) shutdown mode? As in, have we called WebSocket::close? */
+            if (webSocketData->isShuttingDown) {
+                /* Check if we just now drained completely */
+                if (webSocket->getBufferedAmount() == 0) {
+                    /* Now perform the actual TCP/TLS shutdown which was postponed due to backpressure */
+                    webSocket->shutdown();
 
+                    /* Set us to not shutting down so to avoid any spurious extra calls */
+                    webSocketData->isShuttingDown = false;
+                }
+            } else {
+                /* Call drain event even though nothing might actually changed */
+                WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)(
+                    static_dispatch(us_ssl_socket_get_context, us_socket_get_context)((SOCKET_TYPE *)s)
+                    );
+                if (webSocketContextData->drainHandler) {
+                    webSocketContextData->drainHandler((WebSocket<SSL, isServer> *) s);
+                }
+                /* No need to check for closed here as we leave the handler immediately*/
+            }
 
             return s;
         });
@@ -315,8 +337,7 @@ private:
     }
 
 public:
-
+    /* WebSocket contexts are always child contexts to a HTTP context so no SSL options are needed as they are inherited */
-    // we do not need SSL options as we come from adoptions
     static WebSocketContext *create(Loop *loop, SOCKET_CONTEXT_TYPE *parentSocketContext) {
         WebSocketContext *webSocketContext = (WebSocketContext *)static_dispatch(us_create_child_ssl_socket_context, us_create_child_socket_context)(parentSocketContext, sizeof(WebSocketContextData<SSL>));
         if (!webSocketContext) {
@@ -327,7 +348,6 @@ public:
         new ((WebSocketContextData<SSL> *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)((SOCKET_CONTEXT_TYPE *)webSocketContext)) WebSocketContextData<SSL>;
         return webSocketContext->init();
     }
-
 };
 
 }

src/WebSocketContextData.h

@@ -29,9 +29,9 @@ template <bool, bool> struct WebSocket;
 template <bool SSL>
 struct WebSocketContextData {
     /* The callbacks for this context */
-    std::function<void(WebSocket<SSL, true> *, std::string_view, uWS::OpCode)> messageHandler;
+    std::function<void(WebSocket<SSL, true> *, std::string_view, uWS::OpCode)> messageHandler = nullptr;
-    std::function<void(uWS::WebSocket<SSL, true> *)> drainHandler;
+    std::function<void(uWS::WebSocket<SSL, true> *)> drainHandler = nullptr;
-    std::function<void(uWS::WebSocket<SSL, true> *, int, std::string_view)> closeHandler;
+    std::function<void(uWS::WebSocket<SSL, true> *, int, std::string_view)> closeHandler = nullptr;
 };
 
 }