Add locking around _messageQueue and _controlQueue in AsyncWebSocketClient

Based on a combination of commits 344d31c1ffbffd845823e1f8c369fcbf2439e47d
and 9b98550f64000b937db9567b69e1275feccf1f0f of of dumbfixes branch of
0xFEEDC0DE64 fork of ESPAsyncWebServer.

This patch additionally introduces a PolymorphMessageContainer class,
containing either a AsyncWebSocketBasicMessage or an
AsyncWebSocketMultiMessage, and intented to be contained inside a
std::deque (instead of a std::queue as on the original commits, as
done in a previous commit). This patch therefore eliminates one more use
of a LinkedList<T>.

Some code formatting done on the side too.

src/AsyncWebSocket.cpp

@@ -423,7 +423,6 @@ AsyncWebSocketMultiMessage::AsyncWebSocketMultiMessage(AsyncWebSocketMessageBuff
   ,_acked(0)
   ,_WSbuffer(nullptr)
 {
-
   _opcode = opcode & 0x07;
   _mask = mask;
 
@@ -522,8 +521,7 @@ AsyncWebSocketMultiMessage::~AsyncWebSocketMultiMessage() {
  const size_t AWSC_PING_PAYLOAD_LEN = 22;
 
 AsyncWebSocketClient::AsyncWebSocketClient(AsyncWebServerRequest *request, AsyncWebSocket *server)
-  : _messageQueue(LinkedList<AsyncWebSocketMessage *>([](AsyncWebSocketMessage *m){ delete  m; }))
+  : _tempObject(NULL)
-  , _tempObject(NULL)
 {
     _client = request->client();
     _server = server;
@@ -544,23 +542,30 @@ AsyncWebSocketClient::AsyncWebSocketClient(AsyncWebServerRequest *request, Async
     memset(&_pinfo,0,sizeof(_pinfo));
 }
 
-AsyncWebSocketClient::~AsyncWebSocketClient(){
+AsyncWebSocketClient::~AsyncWebSocketClient()
-  // Serial.printf("%u FREE Q\n", id());
+{
-  _messageQueue.free();
+    {
+        AsyncWebLockGuard l(_lock);
+
+        _messageQueue.clear();
         _controlQueue.clear();
+    }
     _server->_cleanBuffers();
     _server->_handleEvent(this, WS_EVT_DISCONNECT, NULL, NULL, 0);
 }
 
 void AsyncWebSocketClient::_clearQueue(){
-  while(!_messageQueue.isEmpty() && _messageQueue.front()->finished()){
+    while (!_messageQueue.empty() && _messageQueue.front().get().finished()){
-    _messageQueue.remove(_messageQueue.front());
+      _messageQueue.pop_front();
     }
 }
 
 void AsyncWebSocketClient::_onAck(size_t len, uint32_t time){
-  // Serial.printf("%u onAck\n", id());
     _lastMessageTime = millis();
+
+    {
+        AsyncWebLockGuard l(_lock);
+
         if (!_controlQueue.empty()){
             auto &head = _controlQueue.front();
             if(head.finished()){
@@ -575,81 +580,126 @@ void AsyncWebSocketClient::_onAck(size_t len, uint32_t time){
             }
         }
 
-  if(len && !_messageQueue.isEmpty()){
+        if(len && !_messageQueue.empty()){
-    _messageQueue.front()->ack(len, time);
+          _messageQueue.front().get().ack(len, time);
         }
 
         _clearQueue();
+    }
 
     _server->_cleanBuffers();
-  // Serial.println("RUN 1");
     _runQueue();
 }
 
 void AsyncWebSocketClient::_onPoll(){
-  if(_client->canSend() && (!_controlQueue.empty() || !_messageQueue.isEmpty())){
+    if(_client->canSend() && [this](){ AsyncWebLockGuard l(_lock); return !_controlQueue.empty() || !_messageQueue.empty(); }())
+    {
         _runQueue();
-  } else if(_keepAlivePeriod > 0 && _controlQueue.empty() && _messageQueue.isEmpty() && (millis() - _lastMessageTime) >= _keepAlivePeriod){
+    } else if(_keepAlivePeriod > 0 && (millis() - _lastMessageTime) >= _keepAlivePeriod &&
+        [this](){ AsyncWebLockGuard l(_lock); return _controlQueue.empty() && _messageQueue.empty(); }())
+    {
         ping((uint8_t *)AWSC_PING_PAYLOAD, AWSC_PING_PAYLOAD_LEN);
     }
 }
 
-void AsyncWebSocketClient::_runQueue(){
+void AsyncWebSocketClient::_runQueue()
+{
+    AsyncWebLockGuard l(_lock);
+
     _clearQueue();
 
-  //size_t m0 = _messageQueue.isEmpty()? 0 : _messageQueue.length();
+    if(!_controlQueue.empty() && (_messageQueue.empty() || _messageQueue.front().get().betweenFrames()) && webSocketSendFrameWindow(_client) > (size_t)(_controlQueue.front().len() - 1))
-  //size_t m1 = _messageQueue.isEmpty()? 0 : _messageQueue.front()->betweenFrames();
+    {
-  // Serial.printf("%u R C = %u %u\n", _clientId, m0, m1);
+        l.unlock();
-  if(!_controlQueue.empty() && (_messageQueue.isEmpty() || _messageQueue.front()->betweenFrames()) && webSocketSendFrameWindow(_client) > (size_t)(_controlQueue.front().len() - 1)){
-    //  Serial.printf("%u R S C\n", _clientId);
         _controlQueue.front().send(_client);
-  } else if(!_messageQueue.isEmpty() && _messageQueue.front()->betweenFrames() && webSocketSendFrameWindow(_client)){
-    //  Serial.printf("%u R S M = ", _clientId);
-    _messageQueue.front()->send(_client);
     }
-
+    else if (!_messageQueue.empty() && _messageQueue.front().get().betweenFrames() && webSocketSendFrameWindow(_client))
-  _clearQueue();
+    {
-}
+        l.unlock();
-
+        _messageQueue.front().get().send(_client);
-bool AsyncWebSocketClient::queueIsFull() const {
-  return (_messageQueue.length() >= WS_MAX_QUEUED_MESSAGES) || (_status != WS_CONNECTED);
-}
-
-size_t AsyncWebSocketClient::queueLen() const {
-  return _messageQueue.length() + _controlQueue.size();
-}
-
-void AsyncWebSocketClient::_queueMessage(AsyncWebSocketMessage *dataMessage){
-  if(dataMessage == NULL){
-    // Serial.printf("%u Q1\n", _clientId);
-    return;
-  }
-  if(_status != WS_CONNECTED){
-    // Serial.printf("%u Q2\n", _clientId);
-    delete dataMessage;
-    return;
-  }
-  if(_messageQueue.length() >= WS_MAX_QUEUED_MESSAGES){
-      ets_printf(String(F("ERROR: Too many messages queued\n")).c_str());
-      // Serial.printf("%u Q3\n", _clientId);
-      delete dataMessage;
-  } else {
-      _messageQueue.add(dataMessage);
-      // Serial.printf("%u Q A %u\n", _clientId, _messageQueue.length());
-  }
-  if(_client->canSend()) {
-    // Serial.printf("%u Q S\n", _clientId);
-    // Serial.println("RUN 3");
-    _runQueue();
     }
 }
 
-void AsyncWebSocketClient::_queueControl(uint8_t opcode, uint8_t *data, size_t len, bool mask){
+bool AsyncWebSocketClient::queueIsFull() const
+{
+    size_t size;
+    {
+        AsyncWebLockGuard l(_lock);
+        size = _messageQueue.size();
+    }
+    return (size >= WS_MAX_QUEUED_MESSAGES) || (_status != WS_CONNECTED);
+}
+
+size_t AsyncWebSocketClient::queueLen() const
+{
+    AsyncWebLockGuard l(_lock);
+
+    return _messageQueue.size() + _controlQueue.size();
+}
+
+bool AsyncWebSocketClient::canSend() const
+{
+    size_t size;
+    {
+        AsyncWebLockGuard l(_lock);
+        size = _messageQueue.size();
+    }
+    return size < WS_MAX_QUEUED_MESSAGES;
+}
+
+void AsyncWebSocketClient::_queueControl(uint8_t opcode, uint8_t *data, size_t len, bool mask)
+{
+    {
+        AsyncWebLockGuard l(_lock);
         _controlQueue.emplace_back(opcode, data, len, mask);
-  if (_client->canSend()) {
+    }
-    // Serial.println("RUN 4");
+
+    if (_client->canSend())
         _runQueue();
 }
+
+void AsyncWebSocketClient::_queueMessage(const char *data, size_t len, uint8_t opcode, bool mask)
+{
+    if(_status != WS_CONNECTED)
+        return;
+
+    {
+        AsyncWebLockGuard l(_lock);
+        if (_messageQueue.size() >= WS_MAX_QUEUED_MESSAGES)
+        {
+            l.unlock();
+            ets_printf("ERROR: Too many messages queued\n");
+        }
+        else
+        {
+            _messageQueue.emplace_back(data, len, opcode, mask);
+        }
+    }
+
+    if(_client->canSend())
+        _runQueue();
+}
+
+void AsyncWebSocketClient::_queueMessage(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode, bool mask)
+{
+    if(_status != WS_CONNECTED)
+        return;
+
+    {
+        AsyncWebLockGuard l(_lock);
+        if (_messageQueue.size() >= WS_MAX_QUEUED_MESSAGES)
+        {
+            l.unlock();
+            ets_printf("ERROR: Too many messages queued\n");
+        }
+        else
+        {
+            _messageQueue.emplace_back(buffer, opcode, mask);
+        }
+    }
+
+    if(_client->canSend())
+        _runQueue();
 }
 
 void AsyncWebSocketClient::close(uint16_t code, const char * message){
@@ -855,7 +905,7 @@ size_t AsyncWebSocketClient::printf_P(PGM_P formatP, ...) {
 #endif
 
 void AsyncWebSocketClient::text(const char * message, size_t len){
-  _queueMessage(new AsyncWebSocketBasicMessage(message, len));
+  _queueMessage(message, len);
 }
 void AsyncWebSocketClient::text(const char * message){
   text(message, strlen(message));
@@ -874,11 +924,11 @@ void AsyncWebSocketClient::text(const __FlashStringHelper *data){
 }
 void AsyncWebSocketClient::text(AsyncWebSocketMessageBuffer * buffer)
 {
-  _queueMessage(new AsyncWebSocketMultiMessage(buffer));
+  _queueMessage(buffer);
 }
 
 void AsyncWebSocketClient::binary(const char * message, size_t len){
-  _queueMessage(new AsyncWebSocketBasicMessage(message, len, WS_BINARY));
+  _queueMessage(message, len, WS_BINARY);
 }
 void AsyncWebSocketClient::binary(const char * message){
   binary(message, strlen(message));
@@ -904,7 +954,7 @@ void AsyncWebSocketClient::binary(const __FlashStringHelper *data, size_t len){
 }
 void AsyncWebSocketClient::binary(AsyncWebSocketMessageBuffer * buffer)
 {
-  _queueMessage(new AsyncWebSocketMultiMessage(buffer, WS_BINARY));
+  _queueMessage(buffer, WS_BINARY);
 }
 
 IPAddress AsyncWebSocketClient::remoteIP() const {
@@ -1024,15 +1074,19 @@ void AsyncWebSocket::text(uint32_t id, const char * message, size_t len){
     c->text(message, len);
 }
 
-void AsyncWebSocket::textAll(AsyncWebSocketMessageBuffer * buffer){
+void AsyncWebSocket::textAll(AsyncWebSocketMessageBuffer * buffer)
-  if (!buffer) return;
+{
+    if (!buffer)
+        return;
+
     buffer->lock();
-  for(auto& c: _clients){
+
-    if(c.status() == WS_CONNECTED){
+    for(auto &c : _clients)
+        if (c.status() == WS_CONNECTED)
             c.text(buffer);
-    }
+
-  }
     buffer->unlock();
+
     _cleanBuffers();
 }
 
@@ -1056,26 +1110,38 @@ void AsyncWebSocket::binaryAll(const char * message, size_t len){
 
 void AsyncWebSocket::binaryAll(AsyncWebSocketMessageBuffer * buffer)
 {
-  if (!buffer) return;
+    if (!buffer)
+        return;
+
     buffer->lock();
-    for(auto& c: _clients){
+
+    for (auto &c : _clients)
         if (c.status() == WS_CONNECTED)
             c.binary(buffer);
-  }
+
     buffer->unlock();
+
     _cleanBuffers();
 }
 
-void AsyncWebSocket::message(uint32_t id, AsyncWebSocketMessage *message){
+void AsyncWebSocket::message(uint32_t id, const char *data, size_t len, uint8_t opcode, bool mask)
+{
     AsyncWebSocketClient *c = client(id);
     if (c)
-    c->message(message);
+        c->message(data, len, opcode, mask);
 }
 
-void AsyncWebSocket::messageAll(AsyncWebSocketMultiMessage *message){
+void AsyncWebSocket::message(uint32_t id, AsyncWebSocketMessageBuffer *buffer, uint8_t opcode, bool mask)
+{
+    AsyncWebSocketClient *c = client(id);
+    if (c)
+        c->message(buffer, opcode, mask);
+}
+
+void AsyncWebSocket::messageAll(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode, bool mask){
     for (auto &c : _clients){
         if (c.status() == WS_CONNECTED)
-      c.message(message);
+            c.message(buffer, opcode, mask);
     }
     _cleanBuffers();
 }
@@ -1298,11 +1364,13 @@ void AsyncWebSocket::handleRequest(AsyncWebServerRequest *request){
 AsyncWebSocketMessageBuffer *AsyncWebSocket::makeBuffer(size_t size)
 {
     AsyncWebSocketMessageBuffer * buffer{};
+
     {
         AsyncWebLockGuard l(_lock);
         _buffers.emplace_back(size);
         buffer = &_buffers.back();
     }
+
     return buffer;
 }
 

src/AsyncWebSocket.h

@@ -157,6 +157,63 @@ public:
     virtual size_t send(AsyncClient *client) override ;
 };
 
+class PolymorphMessageContainer
+{
+    union {
+        AsyncWebSocketBasicMessage basicMessage;
+        AsyncWebSocketMultiMessage multiMessage;
+    };
+
+    enum class Type : uint8_t { Basic, Multi };
+    const Type type;
+
+public:
+    PolymorphMessageContainer() = delete;
+    PolymorphMessageContainer(const PolymorphMessageContainer &) = delete;
+    PolymorphMessageContainer &operator=(const PolymorphMessageContainer &) = delete;
+
+    PolymorphMessageContainer(const char *data, size_t len, uint8_t opcode=WS_TEXT, bool mask=false) :
+        type{Type::Basic}
+    {
+        new (&basicMessage) AsyncWebSocketBasicMessage{data, len, opcode, mask};
+    }
+
+    PolymorphMessageContainer(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode=WS_TEXT, bool mask=false) :
+        type{Type::Multi}
+    {
+        new (&multiMessage) AsyncWebSocketMultiMessage{buffer, opcode, mask};
+    }
+
+    ~PolymorphMessageContainer()
+    {
+        switch (type)
+        {
+        case Type::Basic: basicMessage.~AsyncWebSocketBasicMessage(); break;
+        case Type::Multi: multiMessage.~AsyncWebSocketMultiMessage(); break;
+        }
+    }
+
+    AsyncWebSocketMessage &get()
+    {
+        switch (type)
+        {
+        case Type::Basic: return basicMessage;
+        case Type::Multi: return multiMessage;
+        }
+        __builtin_unreachable();
+    }
+
+    const AsyncWebSocketMessage &get() const
+    {
+        switch (type)
+        {
+        case Type::Basic: return basicMessage;
+        case Type::Multi: return multiMessage;
+        }
+        __builtin_unreachable();
+    }
+};
+
 class AsyncWebSocketClient {
   private:
     AsyncClient *_client;
@@ -164,8 +221,10 @@ class AsyncWebSocketClient {
     uint32_t _clientId;
     AwsClientStatus _status;
 
+    AsyncWebLock _lock;
+
     std::deque<AsyncWebSocketControl> _controlQueue;
-    LinkedList<AsyncWebSocketMessage *> _messageQueue;
+    std::deque<PolymorphMessageContainer> _messageQueue;
 
     uint8_t _pstate;
     AwsFrameInfo _pinfo;
@@ -173,8 +232,9 @@ class AsyncWebSocketClient {
     uint32_t _lastMessageTime;
     uint32_t _keepAlivePeriod;
 
-    void _queueMessage(AsyncWebSocketMessage *dataMessage);
     void _queueControl(uint8_t opcode, uint8_t *data=NULL, size_t len=0, bool mask=false);
+    void _queueMessage(const char *data, size_t len, uint8_t opcode=WS_TEXT, bool mask=false);
+    void _queueMessage(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode=WS_TEXT, bool mask=false);
     void _runQueue();
     void _clearQueue();
 
@@ -209,7 +269,8 @@ class AsyncWebSocketClient {
     }
 
     //data packets
-    void message(AsyncWebSocketMessage *message){ _queueMessage(message); }
+    void message(const char *data, size_t len, uint8_t opcode=WS_TEXT, bool mask=false) { _queueMessage(data, len, opcode, mask); }
+    void message(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode=WS_TEXT, bool mask=false) { _queueMessage(buffer, opcode, mask); }
     bool queueIsFull() const;
     size_t queueLen() const;
 
@@ -233,7 +294,7 @@ class AsyncWebSocketClient {
     void binary(const __FlashStringHelper *data, size_t len);
     void binary(AsyncWebSocketMessageBuffer *buffer);
 
-    bool canSend() { return _messageQueue.length() < WS_MAX_QUEUED_MESSAGES; }
+    bool canSend() const;
 
     //system callbacks (do not call)
     void _onAck(size_t len, uint32_t time);
@@ -308,8 +369,9 @@ class AsyncWebSocket: public AsyncWebHandler {
     void binaryAll(const __FlashStringHelper *message, size_t len);
     void binaryAll(AsyncWebSocketMessageBuffer *buffer);
 
-    void message(uint32_t id, AsyncWebSocketMessage *message);
+    void message(uint32_t id, const char *data, size_t len, uint8_t opcode=WS_TEXT, bool mask=false);
-    void messageAll(AsyncWebSocketMultiMessage *message);
+    void message(uint32_t id, AsyncWebSocketMessageBuffer *buffer, uint8_t opcode=WS_TEXT, bool mask=false);
+    void messageAll(AsyncWebSocketMessageBuffer *buffer, uint8_t opcode=WS_TEXT, bool mask=false);
 
     size_t printf(uint32_t id, const char *format, ...)  __attribute__ ((format (printf, 3, 4)));
     size_t printfAll(const char *format, ...)  __attribute__ ((format (printf, 2, 3)));

src/AsyncWebSynchronization.h

@@ -121,6 +121,13 @@ public:
       _lock->unlock();
     }
   }
+
+  void unlock() {
+    if (_lock) {
+      _lock->unlock();
+      _lock = NULL;
+    }
+  }
 };
 
 #endif // ASYNCWEBSYNCHRONIZATION_H_