215 lines
10 KiB
Python
215 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Run real MQTT allocation/lifetime and multipart delivery with heap failures."""
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
import test_mqtt_publish_ack as ack
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
MEMORY = r'''
|
|
#include <map>
|
|
#include <string>
|
|
#include "helpers/MQTTConnectionAdmission.h"
|
|
enum { MALLOC_CAP_INTERNAL=1, MALLOC_CAP_DMA=2, MALLOC_CAP_SPIRAM=4, MALLOC_CAP_8BIT=8 };
|
|
struct Allocation { size_t bytes; unsigned caps; };
|
|
static std::map<void*,Allocation> allocations;
|
|
static std::vector<Allocation> heap_requests;
|
|
static bool psram_ok=true, internal_ok=true, ordinary_ok=true;
|
|
static size_t dma_free=65536, dma_largest=16384;
|
|
static unsigned ordinary_allocs=0, ordinary_frees=0, caps_frees=0;
|
|
static int required_destroy_count=0;
|
|
static void* ordinary_malloc(size_t bytes) {
|
|
++ordinary_allocs;
|
|
if(!ordinary_ok)return nullptr;
|
|
void* p=std::malloc(bytes);assert(p);allocations[p]={bytes,0};return p;
|
|
}
|
|
static void ordinary_free(void* p) {
|
|
assert(destroy_calls>=required_destroy_count);
|
|
assert(allocations.count(p)&&allocations[p].caps==0);
|
|
allocations.erase(p);++ordinary_frees;std::free(p);
|
|
}
|
|
void* heap_caps_malloc(size_t bytes,unsigned caps) {
|
|
heap_requests.push_back({bytes,caps});
|
|
if((caps&MALLOC_CAP_SPIRAM)?!psram_ok:!internal_ok)return nullptr;
|
|
if(caps&MALLOC_CAP_INTERNAL){assert(bytes<=dma_free&&bytes<=dma_largest);dma_free-=bytes;}
|
|
void* p=std::malloc(bytes);assert(p);allocations[p]={bytes,caps};return p;
|
|
}
|
|
void heap_caps_free(void* p) {
|
|
assert(destroy_calls>=required_destroy_count);
|
|
assert(allocations.count(p)&&allocations[p].caps);
|
|
if(allocations[p].caps&MALLOC_CAP_INTERNAL)dma_free+=allocations[p].bytes;
|
|
allocations.erase(p);++caps_frees;std::free(p);
|
|
}
|
|
size_t heap_caps_get_free_size(unsigned caps) {
|
|
assert(caps==(MALLOC_CAP_INTERNAL|MALLOC_CAP_DMA));return dma_free;
|
|
}
|
|
size_t heap_caps_get_largest_free_block(unsigned caps) {
|
|
assert(caps==(MALLOC_CAP_INTERNAL|MALLOC_CAP_DMA));return dma_largest;
|
|
}
|
|
#define malloc ordinary_malloc
|
|
#define free ordinary_free
|
|
'''
|
|
|
|
MAIN = r'''
|
|
static void configure(PsychicMqttClient& client,int size=896) {
|
|
client._mqtt_cfg.uri="mqtt://test";
|
|
client._mqtt_cfg.broker.address.uri="mqtt://test";
|
|
client._mqtt_cfg.buffer_size=size;
|
|
client._mqtt_cfg.buffer.size=size;
|
|
}
|
|
static size_t allocated(){return ordinary_allocs+heap_requests.size();}
|
|
static void multipart(PsychicMqttClient& client,const std::string& text) {
|
|
char topic[]="test/data";
|
|
const size_t first=text.size()/3,second=2*text.size()/3;
|
|
for(size_t part=0;part<3;++part) {
|
|
size_t offset=part==0?0:part==1?first:second;
|
|
size_t end=part==0?first:part==1?second:text.size();
|
|
esp_mqtt_event_t event{MQTT_EVENT_DATA,0,false};
|
|
event.total_data_len=text.size();event.current_data_offset=offset;
|
|
event.data_len=end-offset;event.data=const_cast<char*>(text.data()+offset);
|
|
event.topic=topic;event.topic_len=sizeof(topic)-1;
|
|
esp_mqtt_event_handle_t pointer=&event;
|
|
client._onMessage(pointer);
|
|
}
|
|
}
|
|
int main() {
|
|
{
|
|
PsychicMqttClient client;configure(client);
|
|
// A successful PSRAM allocation needs no internal-heap headroom.
|
|
dma_free=0;dma_largest=0;
|
|
assert(client.connect()==ESP_OK&&client._buffer_capacity==896);
|
|
assert(allocations[client._buffer].bytes==897);
|
|
#if defined(BOARD_HAS_PSRAM)
|
|
assert(allocations[client._buffer].caps==(MALLOC_CAP_SPIRAM|MALLOC_CAP_8BIT));
|
|
assert(ordinary_allocs==0&&heap_requests.size()==1);
|
|
#else
|
|
assert(allocations[client._buffer].caps==0&&heap_requests.empty());
|
|
#endif
|
|
char* retained=client._buffer;
|
|
size_t count=allocated();
|
|
unsigned received=0;
|
|
std::string expected(896,'Q');
|
|
client.onMessage([&](char* topic,char* payload,int,int,bool){
|
|
assert(!strcmp(topic,"test/data")&&expected==payload);++received;
|
|
});
|
|
multipart(client,expected);assert(received==1&&allocated()==count);
|
|
multipart(client,std::string(897,'X'));assert(received==1&&allocated()==count);
|
|
client.disconnect();
|
|
assert(client._buffer==retained&&allocations.count(retained)&&allocated()==count);
|
|
assert(client.connect()==ESP_OK&&client._buffer==retained&&allocated()==count);
|
|
expected="after restart";multipart(client,expected);assert(received==2);
|
|
// Existing clients retain their original fixed reassembly allocation.
|
|
configure(client,3072);
|
|
assert(client.connect()==ESP_OK&&client._buffer==retained&&client._buffer_capacity==896);
|
|
assert(allocated()==count);
|
|
required_destroy_count=destroy_calls+1; // SDK ownership ends before buffer release.
|
|
}
|
|
assert(allocations.empty()&&destroy_calls==1);
|
|
#if defined(BOARD_HAS_PSRAM)
|
|
assert(caps_frees==1&&ordinary_frees==0);
|
|
psram_ok=false;
|
|
// Exact fallback boundary preserves the full DMA reserve after allocation.
|
|
dma_free=MQTTConnectionAdmission::kDmaReserveBytes+897;dma_largest=897;
|
|
{
|
|
PsychicMqttClient client;configure(client);
|
|
assert(client.connect()==ESP_OK&&client._buffer_capacity==896);
|
|
assert(allocations[client._buffer].caps==(MALLOC_CAP_INTERNAL|MALLOC_CAP_DMA|MALLOC_CAP_8BIT));
|
|
assert(dma_free==MQTTConnectionAdmission::kDmaReserveBytes);
|
|
}
|
|
assert(allocations.empty()&&caps_frees==2);
|
|
for(unsigned scenario=0;scenario<3;++scenario) {
|
|
dma_free=MQTTConnectionAdmission::kDmaReserveBytes+(scenario==0?896:897);
|
|
dma_largest=scenario==1?896:897;
|
|
internal_ok=scenario!=2;
|
|
PsychicMqttClient client;configure(client);
|
|
const int before=init_calls;
|
|
const size_t requests=heap_requests.size();
|
|
assert(client.connect()==ESP_ERR_NO_MEM);
|
|
assert(client._buffer==nullptr&&client._buffer_capacity==0&&!client.isInitialized());
|
|
assert(init_calls==before&&allocations.empty());
|
|
assert(heap_requests.size()==requests+(scenario==2?2:1));
|
|
psram_ok=true;
|
|
assert(client.connect()==ESP_OK&&client._buffer_capacity==896);
|
|
psram_ok=false;
|
|
}
|
|
internal_ok=true;psram_ok=true;
|
|
#else
|
|
assert(ordinary_frees==1&&caps_frees==0);
|
|
{
|
|
ordinary_ok=false;
|
|
PsychicMqttClient client;configure(client);
|
|
const int before=init_calls;
|
|
assert(client.connect()==ESP_ERR_NO_MEM&&!client._buffer&&!client._buffer_capacity);
|
|
assert(init_calls==before&&heap_requests.empty());
|
|
ordinary_ok=true;
|
|
assert(client.connect()==ESP_OK&&client._buffer_capacity==896);
|
|
}
|
|
#endif
|
|
assert(allocations.empty());
|
|
// Default and explicitly larger desired capacities still include one NUL.
|
|
for(int desired:{0,3072}) {
|
|
PsychicMqttClient client;configure(client,desired);
|
|
const size_t capacity=desired?desired:1024;
|
|
assert(client.connect()==ESP_OK&&client._buffer_capacity==capacity);
|
|
assert(allocations[client._buffer].bytes==capacity+1);
|
|
}
|
|
assert(allocations.empty());
|
|
}
|
|
'''
|
|
|
|
|
|
class MqttReassemblyBufferTest(unittest.TestCase):
|
|
def test_real_allocator_lifetime_recovery_and_multipart_capacity(self):
|
|
compiler = shutil.which("g++") or shutil.which("clang++")
|
|
self.assertIsNotNone(compiler)
|
|
event = """struct esp_mqtt_event_t { int event_id, msg_id; bool session_present;
|
|
int total_data_len=0,data_len=0,current_data_offset=0,topic_len=0,retain=0,qos=0;
|
|
char* data=nullptr;char* topic=nullptr;bool dup=false; };"""
|
|
mqtt = ack.MQTT_HEADER.replace(
|
|
"struct esp_mqtt_event_t { int event_id, msg_id; bool session_present; };", event)
|
|
preamble = ack.PREAMBLE.replace(
|
|
"void PsychicMqttClient::_onMessage(esp_mqtt_event_handle_t&) { ++receive_events; }", "")
|
|
source = preamble + MEMORY + "\n".join(ack.method(ack.ADAPTER, signature) for signature in (
|
|
"PsychicMqttClient::PsychicMqttClient()", "PsychicMqttClient::~PsychicMqttClient()",
|
|
"PsychicMqttClient &PsychicMqttClient::onMessage(", "int PsychicMqttClient::subscribe(",
|
|
"esp_err_t PsychicMqttClient::applyConfig()", "esp_err_t PsychicMqttClient::connect()",
|
|
"void PsychicMqttClient::disconnect()", "void PsychicMqttClient::_onMqttEventStatic(",
|
|
"void PsychicMqttClient::_onMqttEvent(", "void PsychicMqttClient::_onConnect(",
|
|
"void PsychicMqttClient::_onDisconnect(", "void PsychicMqttClient::_onSubscribe(",
|
|
"void PsychicMqttClient::_onUnsubscribe(", "void PsychicMqttClient::_onPublish(",
|
|
"size_t PsychicMqttClient::refreshOutboxSize()",
|
|
"void PsychicMqttClient::_onMessage(", "bool PsychicMqttClient::_isTopicMatch(",
|
|
)) + MAIN
|
|
with tempfile.TemporaryDirectory(prefix="meshcore-mqtt-buffer-") as directory:
|
|
directory = Path(directory)
|
|
(directory / "Arduino.h").write_text(ack.ARDUINO_HEADER, encoding="ascii")
|
|
(directory / "mqtt_client.h").write_text(mqtt, encoding="ascii")
|
|
(directory / "esp_crt_bundle.h").write_text("#pragma once\n", encoding="ascii")
|
|
path = directory / "buffer.cpp"
|
|
path.write_text(source, encoding="ascii")
|
|
for idf in (4, 5):
|
|
for psram in (False, True):
|
|
with self.subTest(idf=idf, psram=psram):
|
|
binary = directory / "buffer"
|
|
sanitize = (["-fsanitize=address,undefined", "-fno-omit-frame-pointer", "-fno-pie", "-no-pie"]
|
|
if sys.platform.startswith("linux") else [])
|
|
command = [compiler, "-std=c++17", "-Wall", "-Wextra", "-Werror",
|
|
"-Wno-unused-function", "-Wno-unused-variable", *sanitize,
|
|
"-DARDUINO_ARCH_ESP32=1", f"-DESP_IDF_VERSION_MAJOR={idf}",
|
|
*(["-DBOARD_HAS_PSRAM=1"] if psram else []),
|
|
"-I", str(directory), "-I", str(ROOT / "lib/PsychicMqttClient/src"),
|
|
"-I", str(ROOT / "src"), str(path), "-o", str(binary)]
|
|
result = subprocess.run(command, capture_output=True, text=True, timeout=60)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
result = subprocess.run([str(binary)], capture_output=True, text=True, timeout=10)
|
|
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|