From 891380bdf509c479df7af886f4358cfc1cbd8fb6 Mon Sep 17 00:00:00 2001
From: Euripedes Rocha <rocha.euripedes@espressif.com>
Date: Tue, 12 Mar 2024 08:49:36 +0100
Subject: [PATCH] feat: Make state and size atomic

This makes the mqtt client state atomic to reduce the scope of locking in
some parts of the code.
---
 lib/include/mqtt_client_priv.h |  2 +-
 lib/mqtt_outbox.c              |  2 +-
 mqtt_client.c                  | 14 +++++++-------
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/lib/include/mqtt_client_priv.h b/lib/include/mqtt_client_priv.h
index 593f3d7..0b152b0 100644
--- a/lib/include/mqtt_client_priv.h
+++ b/lib/include/mqtt_client_priv.h
@@ -114,7 +114,7 @@ struct esp_mqtt_client {
     esp_transport_handle_t transport;
     mqtt_config_storage_t *config;
     mqtt_state_t  mqtt_state;
-    mqtt_client_state_t state;
+    _Atomic mqtt_client_state_t state;
     uint64_t refresh_connection_tick;
     int64_t keepalive_tick;
     uint64_t reconnect_tick;
diff --git a/lib/mqtt_outbox.c b/lib/mqtt_outbox.c
index 2623904..ab44e82 100644
--- a/lib/mqtt_outbox.c
+++ b/lib/mqtt_outbox.c
@@ -24,7 +24,7 @@ typedef struct outbox_item {
 STAILQ_HEAD(outbox_list_t, outbox_item);
 
 struct outbox_t {
-    uint64_t size;
+    _Atomic uint64_t size;
     struct outbox_list_t *list;
 };
 
diff --git a/mqtt_client.c b/mqtt_client.c
index 22f6f15..bca83cd 100644
--- a/mqtt_client.c
+++ b/mqtt_client.c
@@ -1585,7 +1585,8 @@ static void esp_mqtt_task(void *pv)
     while (client->run) {
         MQTT_API_LOCK(client);
         run_event_loop(client);
-        switch (client->state) {
+        mqtt_client_state_t state = client->state;
+        switch (state) {
         case MQTT_STATE_DISCONNECTED:
             break;
         case MQTT_STATE_INIT:
@@ -1874,12 +1875,13 @@ int esp_mqtt_client_subscribe_multiple(esp_mqtt_client_handle_t client,
     if (client->config->outbox_limit > 0 && outbox_get_size(client->outbox) > client->config->outbox_limit) {
         return -2;
     }
-    MQTT_API_LOCK(client);
+
     if (client->state != MQTT_STATE_CONNECTED) {
         ESP_LOGE(TAG, "Client has not connected");
-        MQTT_API_UNLOCK(client);
         return -1;
     }
+
+    MQTT_API_LOCK(client);
     if (client->mqtt_state.connection.information.protocol_ver == MQTT_PROTOCOL_V_5) {
 #ifdef MQTT_PROTOCOL_5
         int max_qos = topic_list[0].qos;
@@ -1942,12 +1944,11 @@ int esp_mqtt_client_unsubscribe(esp_mqtt_client_handle_t client, const char *top
         ESP_LOGE(TAG, "Client was not initialized");
         return -1;
     }
-    MQTT_API_LOCK(client);
     if (client->state != MQTT_STATE_CONNECTED) {
-        MQTT_API_UNLOCK(client);
         ESP_LOGE(TAG, "Client has not connected");
         return -1;
     }
+    MQTT_API_LOCK(client);
     if (client->mqtt_state.connection.information.protocol_ver == MQTT_PROTOCOL_V_5) {
 #ifdef MQTT_PROTOCOL_5
         mqtt5_msg_unsubscribe(&client->mqtt_state.connection,
@@ -2049,15 +2050,14 @@ int esp_mqtt_client_publish(esp_mqtt_client_handle_t client, const char *topic,
         ESP_LOGE(TAG, "Client was not initialized");
         return -1;
     }
-    MQTT_API_LOCK(client);
 #if MQTT_SKIP_PUBLISH_IF_DISCONNECTED
     if (client->state != MQTT_STATE_CONNECTED) {
         ESP_LOGI(TAG, "Publishing skipped: client is not connected");
-        MQTT_API_UNLOCK(client);
         return -1;
     }
 #endif
 
+MQTT_API_LOCK(client);
 #ifdef MQTT_PROTOCOL_5
     if (client->mqtt_state.connection.information.protocol_ver == MQTT_PROTOCOL_V_5) {
         if (esp_mqtt5_client_publish_check(client, qos, retain) != ESP_OK) {