esp_websocket_client: Add Kconfig options for PSRAM allocation

shardt68 · shardt68 · commit 9d9f1c97d413 · 2025-12-29T18:04:05.000+01:00
- Added CONFIG_ESP_WS_CLIENT_TASK_STACK_IN_EXT_RAM to allow allocating
  WebSocket task stack in PSRAM.
- Added CONFIG_ESP_WS_CLIENT_ENABLE_DYNAMIC_BUFFER to allow dynamic
  allocation of RX/TX buffers.
- Fixed memory leak in esp_websocket_client_destroy by deferring destruction
  to a separate task when called from the WebSocket task itself.
- Fixed stack size calculation for xTaskCreateStaticPinnedToCore (expecting words).
- Fixed race conditions and use-after-free in task destruction by using
  vTaskSuspend + vTaskDelete pattern instead of self-deletion.
- Fixed double-free race condition in destroy_on_exit where manual destroy
  could race with deferred destruction.
diff --git a/components/esp_websocket_client/Kconfig b/components/esp_websocket_client/Kconfig
@@ -20,4 +20,20 @@ menu "ESP WebSocket client"
         default 2000
         help
             Timeout for acquiring the TX lock when using separate TX lock.
+
+    config ESP_WS_CLIENT_TASK_STACK_IN_EXT_RAM
+        bool "Allocate websocket task stack in PSRAM"
+        depends on SPIRAM
+        default n
+        help
+            If enabled, the websocket task stack will be allocated in PSRAM using xTaskCreateStatic.
+            This saves internal RAM but requires PSRAM to be enabled and configured.
+
+    config ESP_WS_CLIENT_ALLOC_IN_EXT_RAM
+        bool "Allocate websocket client structure in PSRAM"
+        depends on SPIRAM
+        default n
+        help
+            If enabled, the websocket client structure and its configuration will be allocated in PSRAM.
+            This saves internal RAM but requires PSRAM to be enabled and configured.
 endmenu
diff --git a/components/esp_websocket_client/esp_websocket_client.c b/components/esp_websocket_client/esp_websocket_client.c
@@ -75,6 +75,7 @@ const static int STOPPED_BIT = BIT0;
 const static int CLOSE_FRAME_SENT_BIT = BIT1;   // Indicates that a close frame was sent by the client
 // and we are waiting for the server to continue with clean close
 const static int REQUESTED_STOP_BIT = BIT2;     // Indicates that a client stop has been requested
+const static int DESTRUCTION_IN_PROGRESS_BIT = BIT3; // Indicates that the client is being destroyed
 
 ESP_EVENT_DEFINE_BASE(WEBSOCKET_EVENTS);
 
@@ -156,6 +157,8 @@ struct esp_websocket_client {
     int                         payload_offset;
     esp_transport_keep_alive_t  keep_alive_cfg;
     struct ifreq                *if_name;
+    StackType_t                 *task_stack_buffer;
+    StaticTask_t                *task_buffer;
 };
 
 static uint64_t _tick_get_ms(void)
@@ -497,6 +500,15 @@ static void destroy_and_free_resources(esp_websocket_client_handle_t client)
     free(client->errormsg_buffer);
     if (client->status_bits) {
         vEventGroupDelete(client->status_bits);
+        client->status_bits = NULL;
+    }
+    if (client->task_stack_buffer) {
+        heap_caps_free(client->task_stack_buffer);
+        client->task_stack_buffer = NULL;
+    }
+    if (client->task_buffer) {
+        heap_caps_free(client->task_buffer);
+        client->task_buffer = NULL;
     }
     free(client);
     client = NULL;
@@ -514,11 +526,18 @@ static esp_err_t stop_wait_task(esp_websocket_client_handle_t client)
     client->run = false;
     xEventGroupSetBits(client->status_bits, REQUESTED_STOP_BIT);
     xEventGroupWaitBits(client->status_bits, STOPPED_BIT, false, true, portMAX_DELAY);
+    
+    if (client->task_handle) {
+        while (eTaskGetState(client->task_handle) != eSuspended) {
+             vTaskDelay(1);
+        }
+        vTaskDelete(client->task_handle);
+        client->task_handle = NULL;
+    }
+
     client->state = WEBSOCKET_STATE_UNKNOW;
     return ESP_OK;
 }
-
-#if WS_TRANSPORT_HEADER_CALLBACK_SUPPORT
 static void websocket_header_hook(void * client, const char * line, int line_len)
 {
     ESP_LOGD(TAG, "%s header:%.*s", __func__, line_len, line);
@@ -747,7 +766,11 @@ static int esp_websocket_client_send_with_exact_opcode(esp_websocket_client_hand
 
 esp_websocket_client_handle_t esp_websocket_client_init(const esp_websocket_client_config_t *config)
 {
+    #if CONFIG_ESP_WS_CLIENT_ALLOC_IN_EXT_RAM
+    esp_websocket_client_handle_t client = heap_caps_calloc(1, sizeof(struct esp_websocket_client), MALLOC_CAP_SPIRAM);
+#else
     esp_websocket_client_handle_t client = calloc(1, sizeof(struct esp_websocket_client));
+#endif
     ESP_WS_CLIENT_MEM_CHECK(TAG, client, return NULL);
 
     esp_event_loop_args_t event_args = {
@@ -782,7 +805,11 @@ esp_websocket_client_handle_t esp_websocket_client_init(const esp_websocket_clie
     ESP_WS_CLIENT_MEM_CHECK(TAG, client->tx_lock, goto _websocket_init_fail);
 #endif
 
+    #if CONFIG_ESP_WS_CLIENT_ALLOC_IN_EXT_RAM
+    client->config = heap_caps_calloc(1, sizeof(websocket_config_storage_t), MALLOC_CAP_SPIRAM);
+#else
     client->config = calloc(1, sizeof(websocket_config_storage_t));
+#endif
     ESP_WS_CLIENT_MEM_CHECK(TAG, client->config, goto _websocket_init_fail);
 
     if (config->transport == WEBSOCKET_TRANSPORT_OVER_TCP) {
@@ -886,6 +913,29 @@ esp_err_t esp_websocket_client_destroy(esp_websocket_client_handle_t client)
         return ESP_ERR_INVALID_ARG;
     }
 
+    if (client->status_bits) {
+        xSemaphoreTakeRecursive(client->lock, portMAX_DELAY);
+        EventBits_t bits = xEventGroupGetBits(client->status_bits);
+        if (bits & DESTRUCTION_IN_PROGRESS_BIT) {
+            xSemaphoreGiveRecursive(client->lock);
+            ESP_LOGD(TAG, "Destruction already in progress, skipping");
+            return ESP_OK;
+        }
+
+        if ((bits & STOPPED_BIT) == 0) {
+            if (xTaskGetCurrentTaskHandle() == client->task_handle) {
+                ESP_LOGI(TAG, "esp_websocket_client_destroy called from task, deferring...");
+                client->run = false;
+                client->selected_for_destroying = true;
+                xEventGroupSetBits(client->status_bits, DESTRUCTION_IN_PROGRESS_BIT);
+                xSemaphoreGiveRecursive(client->lock);
+                return ESP_OK;
+            }
+        }
+        xEventGroupSetBits(client->status_bits, DESTRUCTION_IN_PROGRESS_BIT);
+        xSemaphoreGiveRecursive(client->lock);
+    }
+
     if (client->status_bits && (STOPPED_BIT & xEventGroupGetBits(client->status_bits)) == 0) {
         stop_wait_task(client);
     }
@@ -1118,6 +1168,23 @@ static esp_err_t esp_websocket_client_recv(esp_websocket_client_handle_t client)
 
 static int esp_websocket_client_send_close(esp_websocket_client_handle_t client, int code, const char *additional_data, int total_len, TickType_t timeout);
 
+static void esp_websocket_client_destroy_task(void *pv)
+{
+    esp_websocket_client_handle_t client = (esp_websocket_client_handle_t)pv;
+    ESP_LOGI(TAG, "Deferred destruction of websocket client");
+    
+    if (client->task_handle) {
+        while (eTaskGetState(client->task_handle) != eSuspended) {
+             vTaskDelay(1);
+        }
+        vTaskDelete(client->task_handle);
+        client->task_handle = NULL;
+    }
+
+    destroy_and_free_resources(client);
+    vTaskDelete(NULL);
+}
+
 static void esp_websocket_client_task(void *pv)
 {
     const int lock_timeout = portMAX_DELAY;
@@ -1377,9 +1444,23 @@ static void esp_websocket_client_task(void *pv)
     xEventGroupSetBits(client->status_bits, STOPPED_BIT);
     client->state = WEBSOCKET_STATE_UNKNOW;
     if (client->selected_for_destroying == true) {
-        destroy_and_free_resources(client);
+        bool already_destroying = false;
+        xSemaphoreTakeRecursive(client->lock, portMAX_DELAY);
+        if (xEventGroupGetBits(client->status_bits) & DESTRUCTION_IN_PROGRESS_BIT) {
+            already_destroying = true;
+        } else {
+            xEventGroupSetBits(client->status_bits, DESTRUCTION_IN_PROGRESS_BIT);
+        }
+        xSemaphoreGiveRecursive(client->lock);
+
+        if (!already_destroying) {
+            if (xTaskCreate(esp_websocket_client_destroy_task, "ws_destroy", 4096, client, client->config->task_prio, NULL) != pdPASS) {
+                ESP_LOGE(TAG, "Failed to create destroy task, memory will leak");
+                xEventGroupClearBits(client->status_bits, DESTRUCTION_IN_PROGRESS_BIT);
+            }
+        }
     }
-    vTaskDelete(NULL);
+    vTaskSuspend(NULL);
 }
 
 esp_err_t esp_websocket_client_start(esp_websocket_client_handle_t client)
@@ -1400,8 +1481,58 @@ esp_err_t esp_websocket_client_start(esp_websocket_client_handle_t client)
         }
     }
 
-    if (xTaskCreate(esp_websocket_client_task, client->config->task_name ? client->config->task_name : "websocket_task",
-                    client->config->task_stack, client, client->config->task_prio, &client->task_handle) != pdTRUE) {
+    client->task_handle = NULL;
+    BaseType_t res = pdPASS;
+#if CONFIG_ESP_WS_CLIENT_TASK_STACK_IN_EXT_RAM
+    if (client->config->task_stack > 0) {
+        if (client->task_stack_buffer == NULL) {
+            client->task_stack_buffer = (StackType_t *)heap_caps_calloc(1, client->config->task_stack, MALLOC_CAP_SPIRAM);
+        }
+        // TCB must be in internal RAM for xTaskCreateStaticPinnedToCore
+        if (client->task_buffer == NULL) {
+            client->task_buffer = (StaticTask_t *)heap_caps_calloc(1, sizeof(StaticTask_t), MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
+        }
+
+        if (client->task_stack_buffer && client->task_buffer) {
+            ESP_LOGI(TAG, "Allocated %d bytes stack in PSRAM for WebSocket task", client->config->task_stack);
+            client->task_handle = xTaskCreateStaticPinnedToCore(
+                esp_websocket_client_task,
+                client->config->task_name ? client->config->task_name : "websocket_task",
+                client->config->task_stack / sizeof(StackType_t),
+                client,
+                client->config->task_prio,
+                client->task_stack_buffer,
+                client->task_buffer,
+                tskNO_AFFINITY
+            );
+            if (client->task_handle == NULL) {
+                res = pdFAIL;
+            }
+        } else {
+            ESP_LOGW(TAG, "Failed to allocate PSRAM stack, falling back to internal RAM");
+            if (client->task_stack_buffer) {
+                heap_caps_free(client->task_stack_buffer);
+                client->task_stack_buffer = NULL;
+            }
+            if (client->task_buffer) {
+                heap_caps_free(client->task_buffer);
+                client->task_buffer = NULL;
+            }
+
+            res = xTaskCreatePinnedToCore(esp_websocket_client_task, client->config->task_name ? client->config->task_name : "websocket_task",
+                                        client->config->task_stack, client, client->config->task_prio, &client->task_handle, tskNO_AFFINITY);
+        }
+    } else {
+        res = xTaskCreatePinnedToCore(esp_websocket_client_task, client->config->task_name ? client->config->task_name : "websocket_task",
+                                    client->config->task_stack, client, client->config->task_prio, &client->task_handle, tskNO_AFFINITY);
+    }
+#else
+    res = xTaskCreatePinnedToCore(esp_websocket_client_task, client->config->task_name ? client->config->task_name : "websocket_task",
+                                client->config->task_stack, client, client->config->task_prio, &client->task_handle, tskNO_AFFINITY);
+#endif
+
+    if (res != pdPASS || client->task_handle == NULL) {
+        client->task_handle = NULL;
         ESP_LOGE(TAG, "Error create websocket task");
         return ESP_FAIL;
     }
