Simple keep-alive mechanism

Wiebe Cazemier
1 parent fce0beff
Showing 5 changed files with 54 additions and 1 deletions
client.cpp
client.h
mainapp.cpp
threaddata.cpp
threaddata.h
@@ -17,6 +17,7 @@ Client::Client(int fd, ThreadData_p threadData) :
  
 Client::~Client()
 {
+    std::cout << "Removing client: " << repr() << std::endl;
     close(fd);
 }
  
@@ -74,6 +75,8 @@ bool Client::readFdIntoBuffer()
         return false;
     }
  
+    lastActivity = time(NULL);
+
     return true;
 }
  
@@ -218,6 +221,15 @@ std::string Client::repr()
     return a.str();
 }
  
+bool Client::keepAliveExpired()
+{
+    if (!authenticated)
+        return lastActivity + 20 < time(NULL);
+
+    bool result = (lastActivity + (keepalive*10/5)) < time(NULL);
+    return result;
+}
+
 void Client::setReadyForWriting(bool val)
 {
     if (disconnecting)
@@ -6,6 +6,7 @@
 #include <vector>
 #include <mutex>
 #include <iostream>
+#include <time.h>
  
 #include "forward_declarations.h"
  
@@ -35,6 +36,7 @@ class Client
     bool readyForReading = true;
     bool disconnectWhenBytesWritten = false;
     bool disconnecting = false;
+    time_t lastActivity = time(NULL);
  
     std::string clientid;
     std::string username;
@@ -80,6 +82,7 @@ public:
     void setReadyForDisconnect() { disconnectWhenBytesWritten = true; }
  
     std::string repr();
+    bool keepAliveExpired();
  
 };
  
@@ -15,6 +15,7 @@ void do_thread_work(ThreadData *threadData)
     memset(&events, 0, sizeof (struct epoll_event)*MAX_EVENTS);
  
     std::vector<MqttPacket> packetQueueIn;
+    time_t lastKeepAliveCheck = time(NULL);
  
     while (threadData->running)
     {
@@ -46,7 +47,6 @@ void do_thread_work(ThreadData *threadData)
  
                             if (!readSuccess)
                             {
-                                std::cout << "Disconnect: " << client->repr() << std::endl;
                                 threadData->removeClient(client);
                                 continue;
                             }
@@ -92,6 +92,19 @@ void do_thread_work(ThreadData *threadData)
             }
         }
         packetQueueIn.clear();
+
+        try
+        {
+            if (lastKeepAliveCheck + 30 < time(NULL))
+            {
+                if (threadData->doKeepAliveCheck())
+                    lastKeepAliveCheck = time(NULL);
+            }
+        }
+        catch (std::exception &ex)
+        {
+            std::cerr << "Error handling keep-alives: " << ex.what() << std::endl;
+        }
     }
 }
  
@@ -72,5 +72,28 @@ std::shared_ptr&lt;SubscriptionStore&gt; &amp;ThreadData::getSubscriptionStore()
     return subscriptionStore;
 }
  
+// TODO: profile how fast hash iteration is. Perhaps having a second list/vector is beneficial?
+bool ThreadData::doKeepAliveCheck()
+{
+    std::unique_lock<std::mutex> lock(clients_by_fd_mutex, std::try_to_lock);
+    if (!lock.owns_lock())
+        return false;
+
+    auto it = clients_by_fd.begin();
+    while (it != clients_by_fd.end())
+    {
+        Client_p &client = it->second;
+        if (client->keepAliveExpired())
+        {
+            subscriptionStore->removeClient(client);
+            it = clients_by_fd.erase(it);
+        }
+        else
+            it++;
+    }
+
+    return true;
+}
+
  
  
@@ -40,6 +40,8 @@ public:
     void removeClient(Client_p client);
     void removeClient(int fd);
     std::shared_ptr<SubscriptionStore> &getSubscriptionStore();
+
+    bool doKeepAliveCheck();
 };
  
 #endif // THREADDATA_H