1 files changed, 26 insertions, 4 deletions
diff --git a/lib/pleroma/web/mastodon_api/websocket_handler.ex b/lib/pleroma/web/mastodon_api/websocket_handler.ex
index e2ffd02d0..393d093e5 100644
--- a/lib/pleroma/web/mastodon_api/websocket_handler.ex
+++ b/lib/pleroma/web/mastodon_api/websocket_handler.ex
@@ -12,8 +12,10 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
 
   @behaviour :cowboy_websocket
 
+  # Client ping period.
+  @tick :timer.seconds(30)
   # Cowboy timeout period.
-  @timeout :timer.seconds(30)
+  @timeout :timer.seconds(60)
   # Hibernate every X messages
   @hibernate_every 100
 
@@ -44,7 +46,8 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
           req
         end
 
-      {:cowboy_websocket, req, %{user: user, topic: topic, count: 0}, %{idle_timeout: @timeout}}
+      {:cowboy_websocket, req, %{user: user, topic: topic, count: 0, timer: nil},
+       %{idle_timeout: @timeout}}
     else
       {:error, code} ->
         Logger.debug("#{__MODULE__} denied connection: #{inspect(code)} - #{inspect(req)}")
@@ -66,11 +69,18 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
     )
 
     Streamer.add_socket(state.topic, state.user)
-    {:ok, state}
+    {:ok, %{state | timer: timer()}}
+  end
+
+  # Client's Pong frame.
+  def websocket_handle(:pong, state) do
+    if state.timer, do: Process.cancel_timer(state.timer)
+    {:ok, %{state | timer: timer()}}
   end
 
   # We never receive messages.
-  def websocket_handle(_frame, state) do
+  def websocket_handle(frame, state) do
+    Logger.error("#{__MODULE__} received frame: #{inspect(frame)}")
     {:ok, state}
   end
 
@@ -94,6 +104,14 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
     end
   end
 
+  # Ping tick. We don't re-queue a timer there, it is instead queued when :pong is received.
+  # As we hibernate there, reset the count to 0.
+  # If the client misses :pong, Cowboy will automatically timeout the connection after
+  # `@idle_timeout`.
+  def websocket_info(:tick, state) do
+    {:reply, :ping, %{state | timer: nil, count: 0}, :hibernate}
+  end
+
   def terminate(reason, _req, state) do
     Logger.debug(
       "#{__MODULE__} terminating websocket connection for user #{
@@ -149,4 +167,8 @@ defmodule Pleroma.Web.MastodonAPI.WebsocketHandler do
   end
 
   defp expand_topic(topic, _), do: topic
+
+  defp timer do
+    Process.send_after(self(), :tick, @tick)
+  end
 end