Protect MQTT publishing headroom and credit replies before disconnect

src/helpers/MQTTConnectionAdmission.h

@@ -11,6 +11,11 @@ namespace MQTTConnectionAdmission {
 
 static const size_t kDmaReserveBytes = 16384;
 static const size_t kDmaLargestBytes = 4096;
+// Keep the packet-publishing guard and connection admission on the same floor.
+// A new optional broker must not consume the contiguous block that existing
+// brokers need to continue forwarding packets.
+static const size_t kPsramPublishLargestBytes = 16000;
+static const size_t kInternalPublishLargestBytes = 8000;
 // A stopped SDK client needs a new internal task stack/TCB. First starts also
 // allocate receive/transmit/reassembly buffers and SDK transport/config objects
 // after admission. Budget that ordinary allocation separately from AES headroom.
@@ -33,26 +38,46 @@ enum class ClientMemoryState : uint8_t {
   Started,
 };
 
+static inline size_t publishLargestBytes(bool psram) {
+  return psram ? kPsramPublishLargestBytes : kInternalPublishLargestBytes;
+}
+
+static inline bool hasPublishHeadroom(bool psram, size_t largest) {
+  return largest >= publishLargestBytes(psram);
+}
+
 static inline bool canStart(bool psram, bool tls, ClientMemoryState state,
-                            size_t dma_free, size_t dma_largest) {
+                            size_t dma_free, size_t dma_largest,
+                            bool protect_publish = false) {
   size_t allocation_allowance;
-  size_t largest_required;
+  size_t sdk_largest;
   switch (state) {
     case ClientMemoryState::Uninitialized:
       allocation_allowance = kColdStartAllowanceBytes;
-      largest_required = kColdStartLargestBytes;
+      sdk_largest = kColdStartLargestBytes;
       break;
     case ClientMemoryState::Stopped:
       allocation_allowance = kRestartAllowanceBytes;
-      largest_required = kRestartLargestBytes;
+      sdk_largest = kRestartLargestBytes;
       break;
     case ClientMemoryState::Started:
       allocation_allowance = 0;
-      largest_required = kDmaLargestBytes;
+      sdk_largest = kDmaLargestBytes;
       break;
     default:
       return false;
   }
+  const size_t publication_floor = publishLargestBytes(psram);
+  const size_t tls_record = tls && !psram ? kTlsRecordAllocationBytes : 0;
+  // Bootstrap/recover one broker without requiring spare capacity for another
+  // live connection: retained SDK buffers can fragment the heap after an outage.
+  // With a live peer, conservatively allow startup/TLS allocations to consume
+  // the largest block while leaving existing forwarding's publication floor.
+  // Other tasks and TLS allocations can still change/fragment this snapshot;
+  // these are admission bounds, not an allocation guarantee.
+  size_t largest_required = publication_floor > sdk_largest ? publication_floor : sdk_largest;
+  if (tls_record > largest_required) largest_required = tls_record;
+  if (protect_publish) largest_required = publication_floor + allocation_allowance + tls_record;
   const size_t free_required = kDmaReserveBytes + allocation_allowance;
   if (dma_free < free_required || dma_largest < largest_required) {
     return false;

src/helpers/bridges/MQTTBridge.cpp

@@ -1671,10 +1671,10 @@ bool MQTTBridge::serviceNetworkPower(uint32_t now) {
         && acknowledgments != _network_publish_acks[i];
     _network_publish_acks[i] = acknowledgments;
     // A flapping CONNACK must not buy another full outage window forever.
-    // A PUBACK proves delivery immediately; a quiet connection must stay up
-    // for 30 seconds, leaving time to prove recovery within the probe window.
-    if (_slots[i].enabled && _slots[i].connected
-        && (broker_replied || (now - static_cast<uint32_t>(_slots[i].connected_at_ms)
+    // A PUBACK proves delivery even if DISCONNECTED arrives before this pass.
+    // Consume it once; a quiet connection must still stay up for 30 seconds.
+    if (_slots[i].enabled && (broker_replied || (_slots[i].connected
+        && now - static_cast<uint32_t>(_slots[i].connected_at_ms)
                 >= mesh::wifi::kNetworkHealthyMs))) {
       useful_network = true;
     }
@@ -2381,9 +2381,16 @@ bool MQTTBridge::canStartSlotConnection(int index) const {
       ? ClientMemoryState::Uninitialized
       : (slot.client->isStarted() ? ClientMemoryState::Started : ClientMemoryState::Stopped);
   const uint32_t caps = MALLOC_CAP_INTERNAL | MALLOC_CAP_DMA;
+  bool protect_publish = false;
+  for (int i = 0; i < RUNTIME_MQTT_SLOTS; ++i) {
+    if (i != index && _slots[i].enabled && _slots[i].connected) {
+      protect_publish = true;
+      break;
+    }
+  }
   return MQTTConnectionAdmission::canStart(
       ESP.getPsramSize() > 0, tls, memory_state, heap_caps_get_free_size(caps),
-      heap_caps_get_largest_free_block(caps));
+      heap_caps_get_largest_free_block(caps), protect_publish);
 #else
   return true;
 #endif
@@ -4503,11 +4510,6 @@ bool MQTTBridge::publishPacket(mesh::Packet* packet, bool is_tx,
   // at PUBLISH_JSON_BUFFER_SIZE (2 KB); the threshold includes frame overhead
   // and transient TCP buffers.
   #ifdef ESP32
-  #if defined(BOARD_HAS_PSRAM)
-  static const size_t PUBLISH_SKIP_MAX_ALLOC_THRESHOLD = 16000;
-  #else
-  static const size_t PUBLISH_SKIP_MAX_ALLOC_THRESHOLD = 8000;
-  #endif
   unsigned long now = millis();
   // Re-sample max-alloc at most once per interval and cache the verdict.
   // getMaxAllocHeap() walks the heap free-list, so it must not run per packet.
@@ -4519,10 +4521,13 @@ bool MQTTBridge::publishPacket(mesh::Packet* packet, bool is_tx,
   if (now - _last_memory_check > 5000) {
     _last_memory_check = now;
     size_t max_alloc = ESP.getMaxAllocHeap();
-    _memory_pressure = (max_alloc < PUBLISH_SKIP_MAX_ALLOC_THRESHOLD);
+    // Match connection admission's detected-PSRAM decision, including a board
+    // whose optional PSRAM failed initialization.
+    const bool psram = ESP.getPsramSize() > 0;
+    _memory_pressure = !MQTTConnectionAdmission::hasPublishHeadroom(psram, max_alloc);
     if (_memory_pressure) {
       MQTT_DEBUG_PRINTLN("MQTT: memory pressure, skipping publishes (Max alloc: %d, threshold: %d, skipped: %d)",
-                         (int)max_alloc, (int)PUBLISH_SKIP_MAX_ALLOC_THRESHOLD, _skipped_publishes);
+                         (int)max_alloc, (int)MQTTConnectionAdmission::publishLargestBytes(psram), _skipped_publishes);
     }
   }
   if (_memory_pressure) {