Add optional (on by default) shocker Keep-Alive (#125)

* Add optional (on by default) shocker Keep-Alive!

* Move Keep-Alive manager into it's own task

* Allow Idle Task to alongside Keep-Alive Task, plus iterator logic fix

* Revert RFTransmitter Changes

* CommandHandler Keep-Alive

* Keep-Alive working in Command Handler

Co-Authored-By: hhvrc <[email protected]>

* Switch to using a queue kill message to avoid panic when deleting a mid-listening queue

* Gateway Keep-Alive wording :)

* Some cleanup

* Mini CommandHandler cleanup

* Add keep-alive toggling to Serial commands

* Fix USB Serial on OpenShock Core

* Revert keep-alive duration used during testing

---------

Co-authored-by: hhvrc <[email protected]>

.env

@@ -1,4 +1,4 @@
 OPENSHOCK_API_DOMAIN=api.shocklink.net
 OPENSHOCK_FW_VERSION=0.0.0-unknown
 OPENSHOCK_FW_HOSTNAME=OpenShock
-OPENSHOCK_FW_AP_PREFIX=OpenShock-
+OPENSHOCK_FW_AP_PREFIX=OpenShock-

WebUI/src/lib/_fbs/open-shock/serialization/configuration/rfconfig.ts

@@ -15,12 +15,17 @@ txPin():number {
   return this.bb!.readUint8(this.bb_pos);
 }
 
+keepaliveEnabled():boolean {
+  return !!this.bb!.readInt8(this.bb_pos + 1);
+}
+
 static sizeOf():number {
-  return 1;
+  return 2;
 }
 
-static createRFConfig(builder:flatbuffers.Builder, tx_pin: number):flatbuffers.Offset {
-  builder.prep(1, 1);
+static createRFConfig(builder:flatbuffers.Builder, tx_pin: number, keepalive_enabled: boolean):flatbuffers.Offset {
+  builder.prep(1, 2);
+  builder.writeInt8(Number(Boolean(keepalive_enabled)));
   builder.writeInt8(tx_pin);
   return builder.offset();
 }

include/CommandHandler.h

@@ -15,5 +15,8 @@ namespace OpenShock::CommandHandler {
   SetRfPinResultCode SetRfTxPin(std::uint8_t txPin);
   std::uint8_t GetRfTxPin();
 
+  bool SetKeepAliveEnabled(bool enabled);
+  bool SetKeepAlivePaused(bool paused);
+
   bool HandleCommand(ShockerModelType shockerModel, std::uint16_t shockerId, ShockerCommandType type, std::uint8_t intensity, std::uint16_t durationMs);
 }  // namespace OpenShock::CommandHandler

include/Config.h

@@ -10,6 +10,7 @@ namespace OpenShock::Config {
   // This is a copy of the flatbuffers schema defined in schemas/ConfigFile.fbs
   struct RFConfig {
     std::uint8_t txPin;
+    bool keepAliveEnabled;
   };
   struct WiFiCredentials {
     std::uint8_t id;
@@ -51,6 +52,7 @@ namespace OpenShock::Config {
   bool SetBackendConfig(const BackendConfig& config);
 
   bool SetRFConfigTxPin(std::uint8_t txPin);
+  bool SetRFConfigKeepAliveEnabled(bool enabled);
 
   std::uint8_t AddWiFiCredentials(const std::string& ssid, const std::uint8_t (&bssid)[6], const std::string& password);
   bool TryGetWiFiCredentialsByID(std::uint8_t id, WiFiCredentials& out);

include/radio/RFTransmitter.h

@@ -22,7 +22,7 @@ namespace OpenShock {
 
     inline bool ok() const { return m_rmtHandle != nullptr && m_queueHandle != nullptr && m_taskHandle != nullptr; }
 
-    bool SendCommand(ShockerModelType model, std::uint16_t shockerId, ShockerCommandType type, std::uint8_t intensity, std::uint16_t durationMs);
+    bool SendCommand(ShockerModelType model, std::uint16_t shockerId, ShockerCommandType type, std::uint8_t intensity, std::uint16_t durationMs, bool overwriteExisting = true);
     void ClearPendingCommands();
 
   private:

include/serialization/_fbs/ConfigFile_generated.h

@@ -63,23 +63,29 @@ struct BSSID::Traits {
 FLATBUFFERS_MANUALLY_ALIGNED_STRUCT(1) RFConfig FLATBUFFERS_FINAL_CLASS {
  private:
   uint8_t tx_pin_;
+  uint8_t keepalive_enabled_;
 
  public:
   struct Traits;
   static FLATBUFFERS_CONSTEXPR_CPP11 const char *GetFullyQualifiedName() {
     return "OpenShock.Serialization.Configuration.RFConfig";
   }
   RFConfig()
-      : tx_pin_(0) {
+      : tx_pin_(0),
+        keepalive_enabled_(0) {
   }
-  RFConfig(uint8_t _tx_pin)
-      : tx_pin_(::flatbuffers::EndianScalar(_tx_pin)) {
+  RFConfig(uint8_t _tx_pin, bool _keepalive_enabled)
+      : tx_pin_(::flatbuffers::EndianScalar(_tx_pin)),
+        keepalive_enabled_(::flatbuffers::EndianScalar(static_cast<uint8_t>(_keepalive_enabled))) {
   }
   uint8_t tx_pin() const {
     return ::flatbuffers::EndianScalar(tx_pin_);
   }
+  bool keepalive_enabled() const {
+    return ::flatbuffers::EndianScalar(keepalive_enabled_) != 0;
+  }
 };
-FLATBUFFERS_STRUCT_END(RFConfig, 1);
+FLATBUFFERS_STRUCT_END(RFConfig, 2);
 
 struct RFConfig::Traits {
   using type = RFConfig;

platformio.ini

@@ -98,6 +98,7 @@ build_flags =
 	-DOPENSHOCK_LED_GPIO=35
 	-DOPENSHOCK_TX_PIN=15
 	-DOPENSHOCK_ESTOP_PIN=13
+	-DARDUINO_USB_CDC_ON_BOOT=1
 
 ; TODO:
 ; https://docs.platformio.org/en/latest/boards/espressif32/upesy_wroom.html;upesy-esp32-wroom-devkit

schemas/ConfigFile.fbs

@@ -6,6 +6,7 @@ struct BSSID {
 
 struct RFConfig {
   tx_pin:uint8;
+  keepalive_enabled:bool;
 }
 
 table WiFiCredentials {

src/CommandHandler.cpp

@@ -5,30 +5,182 @@
 #include "Constants.h"
 #include "Logging.h"
 #include "radio/RFTransmitter.h"
+#include "Time.h"
+#include "util/TaskUtils.h"
+
+#include <freertos/FreeRTOS.h>
+#include <freertos/queue.h>
+#include <freertos/semphr.h>
 
 #include <memory>
+#include <unordered_map>
 
 const char* const TAG = "CommandHandler";
 
+const std::int64_t KEEP_ALIVE_INTERVAL  = 60'000;
+const std::uint16_t KEEP_ALIVE_DURATION = 300;
+
 using namespace OpenShock;
 
+template<typename T>
+constexpr T saturate(T value, T min, T max) {
+  return std::min(std::max(value, min), max);
+}
+std::uint32_t calculateEepyTime(std::int64_t timeToKeepAlive) {
+  std::int64_t now = OpenShock::millis();
+  return static_cast<std::uint32_t>(saturate<std::int64_t>(timeToKeepAlive - now, 0LL, KEEP_ALIVE_INTERVAL));
+}
+
+struct KnownShocker {
+  ShockerModelType model;
+  std::uint16_t shockerId;
+  std::int64_t lastActivityTimestamp;
+  bool killTask;
+};
+
+static SemaphoreHandle_t s_rfTransmitterSemaphore     = nullptr;
 static std::unique_ptr<RFTransmitter> s_rfTransmitter = nullptr;
 
+static SemaphoreHandle_t s_keepAliveSemaphore = nullptr;
+static QueueHandle_t s_keepAliveQueue         = nullptr;
+static TaskHandle_t s_keepAliveTaskHandle     = nullptr;
+
+void _keepAliveTask(void* arg) {
+  std::int64_t timeToKeepAlive = KEEP_ALIVE_INTERVAL;
+
+  // Map of shocker IDs to time of next keep alive
+  std::unordered_map<std::uint16_t, KnownShocker> activityMap;
+
+  while (true) {
+    // Calculate eepyTime based on the timeToKeepAlive
+    std::uint32_t eepyTime = calculateEepyTime(timeToKeepAlive);
+
+    KnownShocker cmd;
+    while (xQueueReceive(s_keepAliveQueue, &cmd, pdMS_TO_TICKS(eepyTime)) == pdTRUE) {
+      if (cmd.killTask) {
+        ESP_LOGW(TAG, "Received kill command, exiting keep alive task");
+        vTaskDelete(nullptr);
+        break;  // This should never be reached
+      }
+
+      activityMap[cmd.shockerId] = cmd;
+
+      eepyTime = calculateEepyTime(std::min(timeToKeepAlive, cmd.lastActivityTimestamp + KEEP_ALIVE_INTERVAL));
+    }
+
+    // Update the time to now
+    std::int64_t now = OpenShock::millis();
+
+    // Keep track of the minimum activity time, so we know when to wake up
+    timeToKeepAlive = now + KEEP_ALIVE_INTERVAL;
+
+    // For every entry that has a keep alive time less than now, send a keep alive
+    for (auto it = activityMap.begin(); it != activityMap.end(); ++it) {
+      auto& cmd = it->second;
+
+      if (cmd.lastActivityTimestamp + KEEP_ALIVE_INTERVAL < now) {
+        ESP_LOGV(TAG, "Sending keep alive for shocker %u", cmd.shockerId);
+
+        if (s_rfTransmitter == nullptr) {
+          ESP_LOGW(TAG, "RF Transmitter is not initialized, ignoring keep alive");
+          break;
+        }
+
+        if (!s_rfTransmitter->SendCommand(cmd.model, cmd.shockerId, ShockerCommandType::Vibrate, 0, KEEP_ALIVE_DURATION, false)) {
+          ESP_LOGW(TAG, "Failed to send keep alive for shocker %u", cmd.shockerId);
+        }
+
+        cmd.lastActivityTimestamp = now;
+      }
+
+      timeToKeepAlive = std::min(timeToKeepAlive, cmd.lastActivityTimestamp + KEEP_ALIVE_INTERVAL);
+    }
+  }
+}
+
+bool _internalSetKeepAliveEnabled(bool enabled) {
+  bool wasEnabled = s_keepAliveQueue != nullptr && s_keepAliveTaskHandle != nullptr;
+
+  if (enabled == wasEnabled) {
+    ESP_LOGV(TAG, "Keep alive task is already %s", enabled ? "enabled" : "disabled");
+    return true;
+  }
+
+  xSemaphoreTake(s_keepAliveSemaphore, portMAX_DELAY);
+
+  if (enabled) {
+    ESP_LOGV(TAG, "Enabling keep alive task");
+
+    s_keepAliveQueue = xQueueCreate(32, sizeof(KnownShocker));
+    if (s_keepAliveQueue == nullptr) {
+      ESP_LOGE(TAG, "Failed to create keep-alive task");
+
+      xSemaphoreGive(s_keepAliveSemaphore);
+      return false;
+    }
+
+    if (TaskUtils::TaskCreateExpensive(_keepAliveTask, "KeepAliveTask", 4096, nullptr, 1, &s_keepAliveTaskHandle) != pdPASS) {
+      ESP_LOGE(TAG, "Failed to create keep-alive task");
+
+      vQueueDelete(s_keepAliveQueue);
+      s_keepAliveQueue = nullptr;
+
+      xSemaphoreGive(s_keepAliveSemaphore);
+      return false;
+    }
+  } else {
+    ESP_LOGV(TAG, "Disabling keep alive task");
+    if (s_keepAliveTaskHandle != nullptr && s_keepAliveQueue != nullptr) {
+      // Wait for the task to stop
+      KnownShocker cmd {.killTask = true};
+      while (eTaskGetState(s_keepAliveTaskHandle) != eDeleted) {
+        vTaskDelay(pdMS_TO_TICKS(10));
+
+        // Send nullptr to stop the task gracefully
+        xQueueSend(s_keepAliveQueue, &cmd, pdMS_TO_TICKS(10));
+      }
+      vQueueDelete(s_keepAliveQueue);
+      s_keepAliveQueue = nullptr;
+    } else {
+      ESP_LOGW(TAG, "Keep alive task is already disabled? Something might be wrong.");
+    }
+  }
+
+  xSemaphoreGive(s_keepAliveSemaphore);
+
+  return true;
+}
+
 bool CommandHandler::Init() {
-  std::uint8_t txPin = Config::GetRFConfig().txPin;
+  if (s_rfTransmitterSemaphore != nullptr) {
+    ESP_LOGW(TAG, "RF Transmitter is already initialized");
+    return true;
+  }
+
+  auto& rfConfig = Config::GetRFConfig();
+
+  std::uint8_t txPin = rfConfig.txPin;
   if (!OpenShock::IsValidOutputPin(txPin)) {
     ESP_LOGW(TAG, "Clearing invalid RF TX pin");
     Config::SetRFConfigTxPin(Constants::GPIO_INVALID);
     return false;
   }
 
+  s_rfTransmitterSemaphore = xSemaphoreCreateBinary();
+  xSemaphoreGive(s_rfTransmitterSemaphore);
   s_rfTransmitter = std::make_unique<RFTransmitter>(txPin, 32);
   if (!s_rfTransmitter->ok()) {
-    ESP_LOGE(TAG, "Failed to initialize RF transmitter");
+    ESP_LOGE(TAG, "Failed to initialize RF Transmitter");
     s_rfTransmitter = nullptr;
     return false;
   }
 
+  s_keepAliveSemaphore = xSemaphoreCreateBinary();
+  xSemaphoreGive(s_keepAliveSemaphore);
+  if (rfConfig.keepAliveEnabled) {
+    _internalSetKeepAliveEnabled(true);
+  }
+
   return true;
 }
 
@@ -41,6 +193,8 @@ SetRfPinResultCode CommandHandler::SetRfTxPin(std::uint8_t txPin) {
     return SetRfPinResultCode::InvalidPin;
   }
 
+  xSemaphoreTake(s_rfTransmitterSemaphore, portMAX_DELAY);
+
   if (s_rfTransmitter != nullptr) {
     ESP_LOGV(TAG, "Destroying existing RF transmitter");
     s_rfTransmitter = nullptr;
@@ -50,26 +204,56 @@ SetRfPinResultCode CommandHandler::SetRfTxPin(std::uint8_t txPin) {
   auto rfxmit = std::make_unique<RFTransmitter>(txPin, 32);
   if (!rfxmit->ok()) {
     ESP_LOGE(TAG, "Failed to initialize RF transmitter");
+
+    xSemaphoreGive(s_rfTransmitterSemaphore);
     return SetRfPinResultCode::InternalError;
   }
 
   if (!Config::SetRFConfigTxPin(txPin)) {
     ESP_LOGE(TAG, "Failed to set RF TX pin in config");
+
+    xSemaphoreGive(s_rfTransmitterSemaphore);
     return SetRfPinResultCode::InternalError;
   }
 
   s_rfTransmitter = std::move(rfxmit);
 
+  xSemaphoreGive(s_rfTransmitterSemaphore);
   return SetRfPinResultCode::Success;
 }
 
+bool CommandHandler::SetKeepAliveEnabled(bool enabled) {
+  if (!_internalSetKeepAliveEnabled(enabled)) {
+    return false;
+  }
+
+  if (!Config::SetRFConfigKeepAliveEnabled(enabled)) {
+    ESP_LOGE(TAG, "Failed to set keep alive enabled in config");
+    return false;
+  }
+
+  return true;
+}
+
+bool CommandHandler::SetKeepAlivePaused(bool paused) {
+  if (!_internalSetKeepAliveEnabled(!paused)) {
+    return false;
+  }
+
+  return true;
+}
+
 std::uint8_t CommandHandler::GetRfTxPin() {
   return Config::GetRFConfig().txPin;
 }
 
 bool CommandHandler::HandleCommand(ShockerModelType model, std::uint16_t shockerId, ShockerCommandType type, std::uint8_t intensity, std::uint16_t durationMs) {
+  xSemaphoreTake(s_rfTransmitterSemaphore, portMAX_DELAY);
+
   if (s_rfTransmitter == nullptr) {
     ESP_LOGW(TAG, "RF Transmitter is not initialized, ignoring command");
+
+    xSemaphoreGive(s_rfTransmitterSemaphore);
     return false;
   }
 
@@ -86,5 +270,19 @@ bool CommandHandler::HandleCommand(ShockerModelType model, std::uint16_t shocker
     ESP_LOGV(TAG, "Command received: %u %u %u %u", model, shockerId, type, intensity);
   }
 
-  return s_rfTransmitter->SendCommand(model, shockerId, type, intensity, durationMs);
+  bool ok = s_rfTransmitter->SendCommand(model, shockerId, type, intensity, durationMs);
+
+  xSemaphoreGive(s_rfTransmitterSemaphore);
+  xSemaphoreTake(s_keepAliveSemaphore, portMAX_DELAY);
+
+  if (ok && s_keepAliveQueue != nullptr) {
+    KnownShocker cmd {.model = model, .shockerId = shockerId, .lastActivityTimestamp = OpenShock::millis() + durationMs};
+    if (xQueueSend(s_keepAliveQueue, &cmd, pdMS_TO_TICKS(10)) != pdTRUE) {
+      ESP_LOGE(TAG, "Failed to send keep-alive command to queue");
+    }
+  }
+
+  xSemaphoreGive(s_keepAliveSemaphore);
+
+  return ok;
 }