loolwsd: call poll, to check the status of the websocket

2 files changed, 32 insertions, 24 deletions

diff --git a/loolwsd/LOOLWSD.cpp b/loolwsd/LOOLWSD.cpp
index 2caec3087..0ff4915da 100644
--- a/loolwsd/LOOLWSD.cpp
+++ b/loolwsd/LOOLWSD.cpp
@@ -354,6 +354,7 @@ public:
         BasicTileQueue queue;
         Thread queueHandlerThread;
         QueueHandler handler(queue);
+        Poco::Timespan waitTime(LOOLWSD::POLL_TIMEOUT);
 
         try
         {
@@ -383,46 +384,52 @@ public:
                 // process (to be forwarded to the client).
                 int flags;
                 int n;
+                bool pollTimeout = true;
                 ws->setReceiveTimeout(0);
+
                 do
                 {
                     char buffer[200000];
-                    n = ws->receiveFrame(buffer, sizeof(buffer), flags);
 
-                    if (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE)
+                    if ((pollTimeout = ws->poll(waitTime, Socket::SELECT_READ)))
                     {
-                        std::string firstLine = getFirstLine(buffer, n);
-                        StringTokenizer tokens(firstLine, " ", StringTokenizer::TOK_IGNORE_EMPTY | StringTokenizer::TOK_TRIM);
+                        n = ws->receiveFrame(buffer, sizeof(buffer), flags);
 
-                        if (kind == LOOLSession::Kind::ToClient && firstLine.size() == static_cast<std::string::size_type>(n))
-                        {
-                            queue.put(firstLine);
-                        }
-                        else
+                        if (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE)
                         {
-                            // Check if it is a "nextmessage:" and in that case read the large
-                            // follow-up message separately, and handle that only.
-                            int size;
-                            if (tokens.count() == 2 && tokens[0] == "nextmessage:" && getTokenInteger(tokens[1], "size", size) && size > 0)
-                            {
-                                char largeBuffer[size];
+                            std::string firstLine = getFirstLine(buffer, n);
+                            StringTokenizer tokens(firstLine, " ", StringTokenizer::TOK_IGNORE_EMPTY | StringTokenizer::TOK_TRIM);
 
-                                n = ws->receiveFrame(largeBuffer, size, flags);
-                                if (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE)
-                                {
-                                    if (!session->handleInput(largeBuffer, n))
-                                        n = 0;
-                                }
+                            if (kind == LOOLSession::Kind::ToClient && firstLine.size() == static_cast<std::string::size_type>(n))
+                            {
+                                queue.put(firstLine);
                             }
                             else
                             {
-                                if (!session->handleInput(buffer, n))
-                                    n = 0;
+                                // Check if it is a "nextmessage:" and in that case read the large
+                                // follow-up message separately, and handle that only.
+                                int size;
+                                if (tokens.count() == 2 && tokens[0] == "nextmessage:" && getTokenInteger(tokens[1], "size", size) && size > 0)
+                                {
+                                    char largeBuffer[size];
+
+                                    n = ws->receiveFrame(largeBuffer, size, flags);
+                                    if (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE)
+                                    {
+                                        if (!session->handleInput(largeBuffer, n))
+                                            n = 0;
+                                    }
+                                }
+                                else
+                                {
+                                    if (!session->handleInput(buffer, n))
+                                        n = 0;
+                                }
                             }
                         }
                     }
                 }
-                while (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE);
+                while (!pollTimeout || (n > 0 && (flags & WebSocket::FRAME_OP_BITMASK) != WebSocket::FRAME_OP_CLOSE));
 
                 queue.clear();
                 queue.put("eof");
diff --git a/loolwsd/LOOLWSD.hpp b/loolwsd/LOOLWSD.hpp
index 964115fa2..f0a3aec46 100644
--- a/loolwsd/LOOLWSD.hpp
+++ b/loolwsd/LOOLWSD.hpp
@@ -47,6 +47,7 @@ public:
     static const int MASTER_PORT_NUMBER = 9981;
     static const int INTERVAL_PROBES = 10;
     static const int MAINTENANCE_INTERVAL = 1;
+    static const int POLL_TIMEOUT = 1000000;
     static const std::string CHILD_URI;
     static const std::string PIDLOG;
     static const std::string LOKIT_PIDLOG;