From 4df02bf45f827bd4d3d885eeb089c08e43317346 Mon Sep 17 00:00:00 2001
From: Mathieu Carbou <mathieu.carbou@gmail.com>
Date: Mon, 11 Dec 2023 18:30:45 +0100
Subject: [PATCH] Improve memory usage related to websocket

---
 src/ESPDash.cpp | 21 ++++++++++++++++++++-
 src/ESPDash.h   |  8 ++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/ESPDash.cpp b/src/ESPDash.cpp
index c98c3e7e..e6e91414 100644
--- a/src/ESPDash.cpp
+++ b/src/ESPDash.cpp
@@ -166,8 +166,21 @@ void ESPDash::remove(Statistic *statistic) {
 
 // generates the layout JSON string to the frontend
 size_t ESPDash::generateLayoutJSON(AsyncWebSocketClient *client, bool changes_only) {
+  // https://github.com/me-no-dev/ESPAsyncWebServer#limiting-the-number-of-web-socket-clients
+  // Browsers sometimes do not correctly close the websocket connection, even when the close() function is called in javascript.
+  // This will eventually exhaust the web server's resources and will cause the server to crash.
+  // Use DEFAULT_MAX_WS_CLIENTS or DASH_MAX_WS_CLIENTS (specific to ESP-DASH) to set the maximum number of clients
+  _ws->cleanupClients();
+
+  const size_t clients = _ws->count();
+  
+  if (clients == 0) {
+    // do not consume cpu and memory if no client is connected
+    return 0;
+  }
+
   String buf = "";
-  buf.reserve(DASH_LAYOUT_JSON_SIZE);
+  buf.reserve(changes_only ? DASH_PARTIAL_UPDATE_JSON_SIZE : DASH_LAYOUT_JSON_SIZE);
 
   if (changes_only) {
     buf += "{\"command\":\"update:components\",";
@@ -317,6 +330,9 @@ size_t ESPDash::generateLayoutJSON(AsyncWebSocketClient *client, bool changes_on
   // Store the length of the JSON string
   size_t total = buf.length();
   // Send resp
+  #ifdef DASH_DEBUG
+  Serial.printf("client=%d, count=%d, changes_only=%d, total=%d\n%s\n", (client == nullptr ? -1 : client->id()), clients, changes_only, total, buf.c_str());
+  #endif
   if (client != nullptr) {
     _ws->text(client->id(), buf.c_str(), total);
   } else {
@@ -463,6 +479,9 @@ void ESPDash::refreshStatistics() {
   generateLayoutJSON(nullptr, true);
 }
 
+bool ESPDash::hasClient() {
+  return _ws->count() > 0;
+}
 
 /*
   Destructor
diff --git a/src/ESPDash.h b/src/ESPDash.h
index 02f4b918..ab94c509 100644
--- a/src/ESPDash.h
+++ b/src/ESPDash.h
@@ -44,8 +44,8 @@ Github URL: https://github.com/ayushsharma82/ESP-DASH
   #define DASH_LAYOUT_JSON_SIZE 1024 * 5
 #endif
 
-#ifndef DASH_STATISTICS_JSON_SIZE
-  #define DASH_STATISTICS_JSON_SIZE 1024 * 1
+#ifndef DASH_PARTIAL_UPDATE_JSON_SIZE
+  #define DASH_PARTIAL_UPDATE_JSON_SIZE DASH_LAYOUT_JSON_SIZE
 #endif
 
 #ifndef DASH_CARD_JSON_SIZE
@@ -60,6 +60,10 @@ Github URL: https://github.com/ayushsharma82/ESP-DASH
   #define DASH_USE_LEGACY_CHART_STORAGE 0
 #endif
 
+#ifndef DASH_MAX_WS_CLIENTS
+  #define DASH_MAX_WS_CLIENTS DEFAULT_MAX_WS_CLIENTS
+#endif
+
 // Forward Declaration
 class Card;
 class Chart;