From e4e6982355ac8886fdb164f25abbdcfb96a83984 Mon Sep 17 00:00:00 2001
From: Josh Kalderimis <josh.kalderimis@gmail.com>
Date: Wed, 2 Oct 2024 16:13:50 +1300
Subject: [PATCH] Allow the socket drainer to be configured at runtime

---
 config/runtime.exs                           |  6 ++++++
 lib/nerves_hub_web/channels/device_socket.ex | 10 ++++++++++
 lib/nerves_hub_web/device_endpoint.ex        |  6 ++++--
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/config/runtime.exs b/config/runtime.exs
index 4ec868beb..4023f5301 100644
--- a/config/runtime.exs
+++ b/config/runtime.exs
@@ -39,6 +39,12 @@ config :nerves_hub,
   mapbox_access_token: System.get_env("MAPBOX_ACCESS_TOKEN"),
   dashboard_enabled: System.get_env("DASHBOARD_ENABLED", "false") == "true"
 
+config :nerves_hub, :device_socket_drainer,
+  batch_size: String.to_integer(System.get_env("DEVICE_SOCKET_DRAINER_BATCH_SIZE", "1000")),
+  batch_interval:
+    String.to_integer(System.get_env("DEVICE_SOCKET_DRAINER_BATCH_INTERVAL", "4000")),
+  shutdown: String.to_integer(System.get_env("DEVICE_SOCKET_DRAINER_SHUTDOWN", "30000"))
+
 # only set this in :prod as not to override the :dev config
 if config_env() == :prod do
   config :nerves_hub,
diff --git a/lib/nerves_hub_web/channels/device_socket.ex b/lib/nerves_hub_web/channels/device_socket.ex
index a34bcf7cd..43a2dad31 100644
--- a/lib/nerves_hub_web/channels/device_socket.ex
+++ b/lib/nerves_hub_web/channels/device_socket.ex
@@ -116,6 +116,16 @@ defmodule NervesHubWeb.DeviceSocket do
   def id(%{assigns: %{device: device}}), do: "device_socket:#{device.id}"
   def id(_socket), do: nil
 
+  def drainer_configuration() do
+    config = Application.get_env(:nerves_hub, :device_socket_drainer)
+
+    [
+      batch_size: config[:batch_size],
+      batch_interval: config[:batch_interval],
+      shutdown: config[:shutdown]
+    ]
+  end
+
   defp decode_from_headers(%{"x-nh-alg" => "NH1-HMAC-" <> alg} = headers) do
     with [digest_str, iter_str, klen_str] <- String.split(alg, "-"),
          digest <- String.to_existing_atom(String.downcase(digest_str)),
diff --git a/lib/nerves_hub_web/device_endpoint.ex b/lib/nerves_hub_web/device_endpoint.ex
index 5172c1b06..0768153ad 100644
--- a/lib/nerves_hub_web/device_endpoint.ex
+++ b/lib/nerves_hub_web/device_endpoint.ex
@@ -16,7 +16,8 @@ defmodule NervesHubWeb.DeviceEndpoint do
       timeout: 180_000,
       fullsweep_after: 0,
       error_handler: {WebsocketConnectionError, :handle_error, []}
-    ]
+    ],
+    drainer: {NervesHubWeb.DeviceSocket, :drainer_configuration, []}
   )
 
   socket(
@@ -28,7 +29,8 @@ defmodule NervesHubWeb.DeviceEndpoint do
       timeout: 180_000,
       fullsweep_after: 0,
       error_handler: {WebsocketConnectionError, :handle_error, []}
-    ]
+    ],
+    drainer: {NervesHubWeb.DeviceSocket, :drainer_configuration, []}
   )
 
   plug(NervesHubWeb.Plugs.ImAlive)