From e8f26aff6566d97c496183409bf0f19849a3b34a Mon Sep 17 00:00:00 2001
From: Yurii <laxilef@gmail.com>
Date: Sat, 9 Dec 2023 08:14:17 +0300
Subject: [PATCH] fix exception due to mqtt client

---
 lib/HomeAssistantHelper/HomeAssistantHelper.h | 20 ++++++++++++++-
 src/MqttTask.h                                | 25 +++++++------------
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/lib/HomeAssistantHelper/HomeAssistantHelper.h b/lib/HomeAssistantHelper/HomeAssistantHelper.h
index 41f39a0..81cae9f 100644
--- a/lib/HomeAssistantHelper/HomeAssistantHelper.h
+++ b/lib/HomeAssistantHelper/HomeAssistantHelper.h
@@ -8,6 +8,16 @@ public:
     this->client = &client;
   }
 
+  void setYieldCallback(void(*yieldCallback)(void*)) {
+    this->yieldCallback = yieldCallback;
+    this->yieldArg = nullptr;
+  }
+
+  void setYieldCallback(void(*yieldCallback)(void*), void* arg) {
+    this->yieldCallback = yieldCallback;
+    this->yieldArg = arg;
+  }
+
   void setBufferedClient() {
     this->bClient = nullptr;
   }
@@ -68,7 +78,13 @@ public:
     } else {
       serializeJson(doc, *client);
     }
-    return client->endPublish();
+
+    int pubResult = client->endPublish();
+    if (this->yieldCallback != nullptr) {
+      this->yieldCallback(yieldArg);
+    }
+
+    return pubResult;
   }
 
   bool publish(const char* topic) {
@@ -89,6 +105,8 @@ public:
   }
 
 protected:
+  void(*yieldCallback)(void*) = nullptr;
+  void* yieldArg = nullptr;
   PubSubClient* client;
   BufferingPrint* bClient = nullptr;
   String prefix = "homeassistant";
diff --git a/src/MqttTask.h b/src/MqttTask.h
index 3d03a86..1272b0d 100644
--- a/src/MqttTask.h
+++ b/src/MqttTask.h
@@ -43,6 +43,14 @@ protected:
     client.setCallback(std::bind(&MqttTask::__callback, this, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3));
     client.setBufferSize(1024);
 
+    haHelper.setYieldCallback([](void* self) {
+      MqttTask* task = static_cast<MqttTask*>(self);
+      task->delay(10);
+
+      if (client.connected()) {
+        client.loop();
+      }
+    }, this);
     haHelper.setBufferedClient(this->bClient);
     haHelper.setDevicePrefix(settings.mqtt.prefix);
     haHelper.setDeviceVersion(PROJECT_VERSION);
@@ -64,7 +72,6 @@ protected:
         client.subscribe(getTopicPath("settings/set").c_str());
         client.subscribe(getTopicPath("state/set").c_str());
         publishHaEntities();
-        publishNonStaticHaEntities(true);
 
         firstFailConnect = 0;
         lastReconnectAttempt = 0;
@@ -85,6 +92,7 @@ protected:
 
         lastReconnectAttempt = millis();
       }
+      delay(100);
     }
 
 
@@ -399,13 +407,11 @@ protected:
     haHelper.publishNumberOutdoorSensorOffset(false);
     haHelper.publishNumberIndoorSensorOffset(false);
     haHelper.publishSwitchDebug(false);
-    yield();
 
     // emergency
     haHelper.publishSwitchEmergency();
     haHelper.publishNumberEmergencyTarget();
     haHelper.publishSwitchEmergencyUseEquitherm();
-    yield();
 
     // heating
     haHelper.publishSwitchHeating(false);
@@ -413,12 +419,10 @@ protected:
     haHelper.publishNumberHeatingHysteresis();
     haHelper.publishSensorHeatingSetpoint(false);
     haHelper.publishSensorCurrentHeatingMinTemp(false);
-    yield();
     haHelper.publishSensorCurrentHeatingMaxTemp(false);
     haHelper.publishNumberHeatingMinTemp(false);
     haHelper.publishNumberHeatingMaxTemp(false);
     haHelper.publishNumberHeatingMaxModulation(false);
-    yield();
 
     // pid
     haHelper.publishSwitchPID();
@@ -427,19 +431,16 @@ protected:
     haHelper.publishNumberPIDFactorD();
     haHelper.publishNumberPIDMinTemp(false);
     haHelper.publishNumberPIDMaxTemp(false);
-    yield();
 
     // equitherm
     haHelper.publishSwitchEquitherm();
     haHelper.publishNumberEquithermFactorN();
     haHelper.publishNumberEquithermFactorK();
     haHelper.publishNumberEquithermFactorT();
-    yield();
 
     // tuning
     haHelper.publishSwitchTuning();
     haHelper.publishSelectTuningRegulator();
-    yield();
 
     // states
     haHelper.publishBinSensorStatus();
@@ -448,7 +449,6 @@ protected:
     haHelper.publishBinSensorFlame();
     haHelper.publishBinSensorFault();
     haHelper.publishBinSensorDiagnostic();
-    yield();
 
     // sensors
     haHelper.publishSensorModulation(false);
@@ -456,12 +456,10 @@ protected:
     haHelper.publishSensorFaultCode();
     haHelper.publishSensorRssi(false);
     haHelper.publishSensorUptime(false);
-    yield();
 
     // temperatures
     haHelper.publishNumberIndoorTemp();
     haHelper.publishSensorHeatingTemp();
-    yield();
 
     // buttons
     haHelper.publishButtonRestart(false);
@@ -507,7 +505,6 @@ protected:
       }
 
       published = true;
-      yield();
     }
 
     if (force || _heatingMinTemp != heatingMinTemp || _heatingMaxTemp != heatingMaxTemp) {
@@ -523,14 +520,12 @@ protected:
       haHelper.publishClimateHeating(heatingMinTemp, heatingMaxTemp, isStupidMode ? HaHelper::TEMP_SOURCE_HEATING : HaHelper::TEMP_SOURCE_INDOOR);
 
       published = true;
-      yield();
 
     } else if (_isStupidMode != isStupidMode) {
       _isStupidMode = isStupidMode;
       haHelper.publishClimateHeating(heatingMinTemp, heatingMaxTemp, isStupidMode ? HaHelper::TEMP_SOURCE_HEATING : HaHelper::TEMP_SOURCE_INDOOR);
 
       published = true;
-      yield();
     }
 
     if (_dhwPresent && (force || _dhwMinTemp != settings.dhw.minTemp || _dhwMaxTemp != settings.dhw.maxTemp)) {
@@ -541,7 +536,6 @@ protected:
       haHelper.publishClimateDhw(settings.dhw.minTemp, settings.dhw.maxTemp);
 
       published = true;
-      yield();
     }
 
     if (force || _editableOutdoorTemp != editableOutdoorTemp) {
@@ -556,7 +550,6 @@ protected:
       }
 
       published = true;
-      yield();
     }
 
     if (force || _editableIndoorTemp != editableIndoorTemp) {