esphome/esphome/core/lock_free_queue.h

#pragma once

#include "esphome/core/defines.h"

#include <atomic>
#include <cstddef>

#ifdef USE_ESP32
#include <freertos/FreeRTOS.h>
#include <freertos/task.h>
#endif

/*
 * Lock-free queue for single-producer single-consumer scenarios.
 * This allows one thread to push items and another to pop them without
 * blocking each other.
 *
 * This is a Single-Producer Single-Consumer (SPSC) lock-free ring buffer.
 * Available on multi-threaded platforms (ESP32, LibreTiny) where another task
 * produces or consumes, and on single-threaded platforms (RP2) where the
 * producer runs in interrupt context.
 *
 * Common use cases:
 * - BLE events: BLE task produces, main loop consumes
 * - MQTT messages: main task produces, MQTT thread consumes
 *
 * @tparam T The type of elements stored in the queue (must be a pointer type)
 * @tparam SIZE The maximum number of elements (1-255, limited by uint8_t indices)
 */

namespace esphome {

namespace lockfree_internal {
#if defined(ESPHOME_THREAD_MULTI_NO_ATOMICS) || defined(ESPHOME_THREAD_SINGLE)
// Platforms where std::atomic RMW operations are unavailable or unnecessary:
//  - ESPHOME_THREAD_MULTI_NO_ATOMICS: cores lacking atomic read-modify-write
//    instructions (currently the ARMv5TE BK72xx SoCs — no LDREX/STREX, no
//    libatomic; other LibreTiny chips such as LN882x/RTL87xx are ARMv7-M and
//    keep std::atomic).
//  - ESPHOME_THREAD_SINGLE: every platform on this model (ESP8266, RP2,
//    nRF52) runs everything on one core (the chip may have more — RP2 is
//    dual-core, but ESPHome and its interrupt producers stay on core 0), so
//    the only possible concurrency is same-core interrupt preemption (on RP2
//    the BTstack packet handler runs in the CYW43 async-context low-priority
//    IRQ on the core that initialized it, core 0). Using plain accesses here
//    also avoids __atomic_* library calls on RP2040 (Cortex-M0+, no
//    LDREX/STREX).
// For this queue's SPSC contract RMW atomics are not needed: aligned 8/16-bit
// loads and stores are single instructions on these cores, so torn reads
// cannot occur, and on a single in-order core a compiler barrier supplies all
// the acquire/release ordering the algorithm requires. Each index has exactly
// one writer (head_: consumer, tail_: producer). The dropped counter's
// increment/exchange pair is not atomic here — a concurrent reset can lose
// counts — which is acceptable for a diagnostic drop counter.
#define ESPHOME_LFQ_COMPILER_BARRIER() __asm__ __volatile__("" ::: "memory")
template<typename T> class PlainAtomic {
 public:
  PlainAtomic() = default;
  constexpr PlainAtomic(T value) : value_(value) {}
  T load(std::memory_order order = std::memory_order_seq_cst) const {
    T value = value_;
    if (order != std::memory_order_relaxed)
      ESPHOME_LFQ_COMPILER_BARRIER();  // acquire: later reads may not hoist above this load
    return value;
  }
  void store(T value, std::memory_order order = std::memory_order_seq_cst) {
    if (order != std::memory_order_relaxed)
      ESPHOME_LFQ_COMPILER_BARRIER();  // release: earlier writes may not sink below this store
    value_ = value;
  }
  T fetch_add(T amount, std::memory_order /*order*/ = std::memory_order_seq_cst) {
    T value = value_;
    value_ = value + amount;
    return value;
  }
  T exchange(T desired, std::memory_order /*order*/ = std::memory_order_seq_cst) {
    T value = value_;
    value_ = desired;
    return value;
  }

 private:
  volatile T value_{0};
};
template<typename T> using AtomicIndex = PlainAtomic<T>;
#else
template<typename T> using AtomicIndex = std::atomic<T>;
#endif
}  // namespace lockfree_internal

// Base lock-free queue without task notification
template<class T, uint8_t SIZE> class LockFreeQueue {
 public:
  LockFreeQueue() : dropped_count_(0), head_(0), tail_(0) {}

  bool push(T *element) {
    bool was_empty;
    uint8_t old_tail;
    return push_internal_(element, was_empty, old_tail);
  }

 protected:
  // Advance ring buffer index by one, wrapping at SIZE.
  // Power-of-2 sizes use modulo (compiler emits single mask instruction).
  // Non-power-of-2 sizes use comparison to avoid expensive multiply-shift sequences.
  static constexpr uint8_t next_index(uint8_t index) {
    if constexpr ((SIZE & (SIZE - 1)) == 0) {
      return (index + 1) % SIZE;
    } else {
      uint8_t next = index + 1;
      if (next >= SIZE) [[unlikely]]
        next = 0;
      return next;
    }
  }

  // Internal push that reports queue state - for use by derived classes
  bool push_internal_(T *element, bool &was_empty, uint8_t &old_tail) {
    if (element == nullptr)
      return false;

    uint8_t current_tail = tail_.load(std::memory_order_relaxed);
    uint8_t next_tail = next_index(current_tail);

    // Read head before incrementing tail
    uint8_t head_before = head_.load(std::memory_order_acquire);

    if (next_tail == head_before) {
      // Buffer full
      dropped_count_.fetch_add(1, std::memory_order_relaxed);
      return false;
    }

    was_empty = (current_tail == head_before);
    old_tail = current_tail;

    buffer_[current_tail] = element;
    tail_.store(next_tail, std::memory_order_release);

    return true;
  }

 public:
  T *pop() {
    uint8_t current_head = head_.load(std::memory_order_relaxed);

    if (current_head == tail_.load(std::memory_order_acquire)) {
      return nullptr;  // Empty
    }

    T *element = buffer_[current_head];
    head_.store(next_index(current_head), std::memory_order_release);
    return element;
  }

  size_t size() const {
    uint8_t tail = tail_.load(std::memory_order_acquire);
    uint8_t head = head_.load(std::memory_order_acquire);
    if constexpr ((SIZE & (SIZE - 1)) == 0) {
      return (tail - head + SIZE) % SIZE;
    } else {
      int diff = static_cast<int>(tail) - static_cast<int>(head);
      if (diff < 0)
        diff += SIZE;
      return static_cast<size_t>(diff);
    }
  }

  uint16_t get_and_reset_dropped_count() {
    // Fast path: relaxed load is a single instruction on all platforms.
    // The atomic exchange (especially for uint16_t on Xtensa) compiles to
    // an expensive sub-word CAS retry loop (~25 instructions + memory barriers).
    // Since drops are rare, avoid the exchange in the common case.
    if (dropped_count_.load(std::memory_order_relaxed) == 0)
      return 0;
    return dropped_count_.exchange(0, std::memory_order_relaxed);
  }

  void increment_dropped_count() { dropped_count_.fetch_add(1, std::memory_order_relaxed); }

  bool empty() const { return head_.load(std::memory_order_acquire) == tail_.load(std::memory_order_acquire); }

  bool full() const {
    uint8_t next_tail = next_index(tail_.load(std::memory_order_relaxed));
    return next_tail == head_.load(std::memory_order_acquire);
  }

 protected:
  T *buffer_[SIZE]{};
  // Atomic: written by producer (push/increment), read+reset by consumer (get_and_reset)
  lockfree_internal::AtomicIndex<uint16_t> dropped_count_;  // 65535 max - more than enough for drop tracking
  // Atomic: written by consumer (pop), read by producer (push) to check if full
  // Using uint8_t limits queue size to 255 elements but saves memory and ensures
  // atomic operations are efficient on all platforms
  lockfree_internal::AtomicIndex<uint8_t> head_;
  // Atomic: written by producer (push), read by consumer (pop) to check if empty
  lockfree_internal::AtomicIndex<uint8_t> tail_;
};

#ifdef USE_ESP32
// Extended queue with task notification support
template<class T, uint8_t SIZE> class NotifyingLockFreeQueue : public LockFreeQueue<T, SIZE> {
 public:
  NotifyingLockFreeQueue() : LockFreeQueue<T, SIZE>(), task_to_notify_(nullptr) {}

  bool push(T *element) {
    bool was_empty;
    uint8_t old_tail;
    bool result = this->push_internal_(element, was_empty, old_tail);

    // Notify optimization: only notify if we need to
    if (result && task_to_notify_ != nullptr &&
        (was_empty || this->head_.load(std::memory_order_acquire) == old_tail)) {
      // Notify in two cases:
      // 1. Queue was empty - consumer might be going to sleep
      // 2. Consumer just caught up to where tail was - might go to sleep
      // Note: There's a benign race in case 2 - between reading head and calling
      // xTaskNotifyGive(), the consumer could advance further. This would result
      // in an unnecessary wake-up, but is harmless and extremely rare in practice.
      xTaskNotifyGive(task_to_notify_);
    }
    // Otherwise: consumer is still behind, no need to notify

    return result;
  }

  // Set the FreeRTOS task handle to notify when items are pushed to the queue
  // This enables efficient wake-up of a consumer task that's waiting for data
  // @param task The FreeRTOS task handle to notify, or nullptr to disable notifications
  void set_task_to_notify(TaskHandle_t task) { task_to_notify_ = task; }

 private:
  TaskHandle_t task_to_notify_;
};
#endif

}  // namespace esphome