Parse JSON to a document and let a callback unsubscribe

esphome/components/mqtt/mqtt_client.cpp

@@ -468,6 +468,7 @@ void MQTTClientComponent::add_subscription_(std::string &&topic, mqtt_subscripti
       .qos = qos,
       .subscribed = false,
       .boxed = boxed,
+      .removed = false,
       .callback = callback,
       .resubscribe_timeout = 0,
   });
@@ -475,11 +476,6 @@ void MQTTClientComponent::add_subscription_(std::string &&topic, mqtt_subscripti
 }
 
 void MQTTClientComponent::unsubscribe(const std::string &topic) {
-  if (this->dispatching_) {
-    // Would erase, and for a boxed callback free, the subscription that may be running right now.
-    ESP_LOGE(TAG, "Cannot unsubscribe from '%s' inside a subscription callback", topic.c_str());
-    return;
-  }
   bool ret = this->mqtt_backend_.unsubscribe(topic.c_str());
   yield();
   if (ret) {
@@ -490,12 +486,22 @@ void MQTTClientComponent::unsubscribe(const std::string &topic) {
     this->status_momentary_warning(1000);
   }
 
-  auto it = subscriptions_.begin();
-  while (it != subscriptions_.end()) {
-    if (it->topic == topic) {
+  for (auto &subscription : this->subscriptions_) {
+    if (subscription.topic == topic)
+      subscription.removed = true;
+  }
+  // A callback may be running; the outermost dispatch erases them once it returns.
+  if (!this->dispatching_)
+    this->erase_removed_subscriptions_();
+}
+
+void MQTTClientComponent::erase_removed_subscriptions_() {
+  auto it = this->subscriptions_.begin();
+  while (it != this->subscriptions_.end()) {
+    if (it->removed) {
       if (it->boxed)
         it->callback.free_boxed();
-      it = subscriptions_.erase(it);
+      it = this->subscriptions_.erase(it);
     } else {
       ++it;
     }
@@ -652,16 +658,18 @@ void MQTTClientComponent::on_message(const std::string &topic, const std::string
 #endif
     // Indexed with the count taken up front: a callback may subscribe, which can reallocate the
     // vector, and the new entry waits for the next message.
-    // Restored rather than cleared, so a nested dispatch cannot reopen unsubscribe() for the outer one.
+    // Restored rather than cleared, so only the outermost dispatch erases what was unsubscribed meanwhile.
     const bool was_dispatching = this->dispatching_;
     this->dispatching_ = true;
     const size_t count = this->subscriptions_.size();
     for (size_t i = 0; i < count; i++) {
       const auto &subscription = this->subscriptions_[i];
-      if (topic_match(topic.c_str(), subscription.topic.c_str()))
+      if (!subscription.removed && topic_match(topic.c_str(), subscription.topic.c_str()))
         subscription.callback.call(topic, payload);
     }
     this->dispatching_ = was_dispatching;
+    if (!was_dispatching)
+      this->erase_removed_subscriptions_();
 #ifdef USE_ESP8266
   });
 #endif

esphome/components/mqtt/mqtt_client.h

@@ -46,7 +46,8 @@ struct MQTTSubscription {
   std::string topic;
   uint8_t qos;
   bool subscribed;
-  bool boxed;  // made by create_boxed(), freed on unsubscribe
+  bool boxed;    // made by create_boxed(), freed on unsubscribe
+  bool removed;  // unsubscribed while a callback ran, erased once dispatch returns
   mqtt_subscription_callback_t callback;
   uint32_t resubscribe_timeout;
 };
@@ -208,7 +209,7 @@ class MQTTClientComponent final : public Component {
   /** Unsubscribe from an MQTT topic.
    *
    * If multiple existing subscriptions to the same topic exist, all of them will be removed.
-   * Not allowed from inside a subscription callback; such a call is logged and ignored.
+   * From inside a subscription callback they stop receiving messages at once and are erased once it returns.
    *
    * @param topic The topic to unsubscribe from.
    * Must match the topic in the original subscribe or subscribe_json call exactly.
@@ -308,13 +309,13 @@ class MQTTClientComponent final : public Component {
 
   bool subscribe_(const char *topic, uint8_t qos);
   void resubscribe_subscription_(MQTTSubscription *sub);
-  // parse_json runs its callback before returning, so the captures can be references.
   template<typename F> static void call_json(const F &callback, const std::string &topic, const std::string &payload) {
-    json::parse_json(payload, [&topic, &callback](JsonObject root) -> bool {
-      callback(topic, root);
-      return true;
-    });
+    JsonDocument doc = json::parse_json(payload);
+    if (doc.overflowed() || doc.isNull())
+      return;
+    callback(topic, doc.as<JsonObject>());
   }
+  void erase_removed_subscriptions_();
   void add_subscription_(std::string &&topic, mqtt_subscription_callback_t callback, bool boxed, uint8_t qos);
   void resubscribe_subscriptions_();
 

tests/components/core/test_callback.cpp

@@ -60,16 +60,6 @@ TEST(CallbackCreateBoxed, MutableStateIsKeptAcrossCalls) {
   cb.free_boxed();
 }
 
-TEST(CallbackCreateBoxed, LargeCaptureIsCalled) {
-  int a = 0;
-  int b = 0;
-  auto cb = IntCallback::create_boxed([&a, &b](int x) { a = b = x; });
-  cb.call(7);
-  EXPECT_EQ(a, 7);
-  EXPECT_EQ(b, 7);
-  cb.free_boxed();
-}
-
 TEST(CallbackCreateBoxed, FreeDestroysCaptureExactlyOnce) {
   int destroyed = 0;
   int calls = 0;

tests/components/mqtt/common-triggers.yaml

@@ -1,5 +1,5 @@
 # on_message with a payload filter, a subscription callback larger than a pointer with its unsubscribe,
-# and a mutable callback, which keeps its state instead of being copied per call.
+# a callback that unsubscribes itself, and a mutable callback, which keeps its state instead of being copied per call.
 mqtt:
   id: mqtt_client
   on_message:
@@ -17,6 +17,9 @@ esphome:
           ESP_LOGD("test", "%s %s", tag.c_str(), payload.c_str());
         });
         id(mqtt_client).unsubscribe("some/topic/boxed");
+        id(mqtt_client).subscribe("some/topic/once", [](const std::string &topic, const std::string &payload) {
+          id(mqtt_client).unsubscribe(topic);
+        });
         id(mqtt_client).subscribe("some/topic/counted", [n = 0](const std::string &topic, const std::string &payload) mutable {
           ESP_LOGD("test", "message %d: %s", ++n, payload.c_str());
         });