From ef646d78101529e0813439f1ce1cc4921b5d4fbf Mon Sep 17 00:00:00 2001
From: Bastien CHAMAGNE <bastien@chamagne.fr>
Date: Wed, 22 Mar 2023 15:37:24 +0100
Subject: [PATCH] Add a timeout for the hydrating function

---
 .../oracle_chain/services/hydrating_cache.ex  | 15 ++++++++++++++-
 .../services/hydrating_cache_test.exs         | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/lib/archethic/oracle_chain/services/hydrating_cache.ex b/lib/archethic/oracle_chain/services/hydrating_cache.ex
index f8b4fba1f..cb05f1fa0 100644
--- a/lib/archethic/oracle_chain/services/hydrating_cache.ex
+++ b/lib/archethic/oracle_chain/services/hydrating_cache.ex
@@ -20,7 +20,8 @@ defmodule Archethic.OracleChain.Services.HydratingCache do
       :refresh_interval,
       :value,
       :hydrating_task,
-      :hydrating_timer
+      :hydrating_timer,
+      :hydrating_function_timeout
     ])
   end
 
@@ -44,6 +45,7 @@ defmodule Archethic.OracleChain.Services.HydratingCache do
     refresh_interval = Keyword.fetch!(options, :refresh_interval)
     mfa = Keyword.fetch!(options, :mfa)
     ttl = Keyword.get(options, :ttl, :infinity)
+    hydrating_function_timeout = Keyword.get(options, :hydrating_function_timeout, 5000)
 
     # start hydrating as soon as init is done
     hydrating_timer = Process.send_after(self(), :hydrate, 0)
@@ -53,6 +55,7 @@ defmodule Archethic.OracleChain.Services.HydratingCache do
      %State{
        mfa: mfa,
        ttl: ttl,
+       hydrating_function_timeout: hydrating_function_timeout,
        refresh_interval: refresh_interval,
        hydrating_timer: hydrating_timer
      }}
@@ -69,6 +72,7 @@ defmodule Archethic.OracleChain.Services.HydratingCache do
   def handle_info(
         :hydrate,
         state = %State{
+          hydrating_function_timeout: hydrating_function_timeout,
           mfa: {m, f, a}
         }
       ) do
@@ -82,9 +86,18 @@ defmodule Archethic.OracleChain.Services.HydratingCache do
         end
       end)
 
+    # we make sure that our hydrating function does not hang
+    Process.send_after(self(), {:kill_hydrating_task, hydrating_task}, hydrating_function_timeout)
+
     {:noreply, %State{state | hydrating_task: hydrating_task}}
   end
 
+  def handle_info({:kill_hydrating_task, task}, state) do
+    Task.shutdown(task, :brutal_kill)
+
+    {:noreply, state}
+  end
+
   def handle_info(
         {ref, result},
         state = %State{
diff --git a/test/archethic/oracle_chain/services/hydrating_cache_test.exs b/test/archethic/oracle_chain/services/hydrating_cache_test.exs
index 3d3e33f08..efe80b19a 100644
--- a/test/archethic/oracle_chain/services/hydrating_cache_test.exs
+++ b/test/archethic/oracle_chain/services/hydrating_cache_test.exs
@@ -126,4 +126,23 @@ defmodule ArchethicCache.OracleChain.Services.HydratingCacheTest do
 
     assert :error = HydratingCache.get(pid, 1)
   end
+
+  test "should kill the task if hydrating function takes longer than timeout" do
+    {:ok, pid} =
+      HydratingCache.start_link(
+        mfa: {Process, :sleep, [100]},
+        refresh_interval: 100,
+        hydrating_function_timeout: 10,
+        ttl: :infinity
+      )
+
+    :erlang.trace(pid, true, [:receive])
+    assert_receive {:trace, ^pid, :receive, :hydrate}
+    assert_receive {:trace, ^pid, :receive, {:kill_hydrating_task, _}}
+
+    refute Process.alive?(:sys.get_state(pid).hydrating_task.pid)
+    assert Process.alive?(pid)
+
+    assert :error = HydratingCache.get(pid)
+  end
 end