Only lock Session.WriteBuffer and wipe cache on lock timeout

Author: zoldar

lib/plausible/event/write_buffer.ex

@@ -35,12 +35,4 @@ defmodule Plausible.Event.WriteBuffer do
   def flush do
     Plausible.Ingestion.WriteBuffer.flush(__MODULE__)
   end
-
-  def lock(timeout) do
-    Plausible.Ingestion.WriteBuffer.lock(__MODULE__, timeout)
-  end
-
-  def unlock() do
-    Plausible.Ingestion.WriteBuffer.unlock(__MODULE__)
-  end
 end

lib/plausible/ingestion/write_buffer.ex

@@ -5,9 +5,6 @@ defmodule Plausible.Ingestion.WriteBuffer do
 
   alias Plausible.IngestRepo
 
-  @lock_timeout :timer.seconds(3)
-  @lock_interval 100
-
   def start_link(opts) do
     GenServer.start_link(__MODULE__, opts, name: Keyword.fetch!(opts, :name))
   end
@@ -20,54 +17,39 @@ defmodule Plausible.Ingestion.WriteBuffer do
     GenServer.call(server, :flush, :infinity)
   end
 
+  def flush_async(server) do
+    GenServer.cast(server, :flush)
+  end
+
   @impl true
   def init(opts) do
     name = Keyword.fetch!(opts, :name)
     buffer = opts[:buffer] || []
     max_buffer_size = opts[:max_buffer_size] || default_max_buffer_size()
     flush_interval_ms = opts[:flush_interval_ms] || default_flush_interval_ms()
-    lock_timeout_ms = opts[:lock_timeouts_ms] || @lock_timeout
-    lock_interval_ms = opts[:lock_interval_ms] || @lock_interval
+    on_init = Keyword.get(opts, :on_init, fn _opts -> %{} end)
 
     Process.flag(:trap_exit, true)
     timer = Process.send_after(self(), :tick, flush_interval_ms)
 
-    ^name = :ets.new(name, [:named_table, :set, :public])
+    extra_state = on_init.(opts)
 
     {:ok,
-     %{
-       buffer: buffer,
-       timer: timer,
-       name: name,
-       insert_sql: Keyword.fetch!(opts, :insert_sql),
-       insert_opts: Keyword.fetch!(opts, :insert_opts),
-       header: Keyword.fetch!(opts, :header),
-       buffer_size: IO.iodata_length(buffer),
-       max_buffer_size: max_buffer_size,
-       flush_interval_ms: flush_interval_ms,
-       lock_timeout_ms: lock_timeout_ms,
-       lock_interval_ms: lock_interval_ms
-     }}
-  end
-
-  @spec lock(atom(), pid(), pos_integer()) :: :ok | {:error, :timeout}
-  def lock(name, locker \\ self(), timeout) do
-    true = :ets.insert(name, {:state, %{locker: locker}})
-
-    flush(name)
-
-    receive do
-      {:locked, ^name} -> :ok
-    after
-      timeout -> {:error, :timeout}
-    end
-  end
-
-  @spec unlock(atom()) :: :ok
-  def unlock(name) do
-    true = :ets.insert(name, {:state, %{locker: nil}})
-
-    :ok
+     Map.merge(
+       %{
+         buffer: buffer,
+         timer: timer,
+         name: name,
+         insert_sql: Keyword.fetch!(opts, :insert_sql),
+         insert_opts: Keyword.fetch!(opts, :insert_opts),
+         on_flush: Keyword.get(opts, :on_flush, fn _result, state -> state end),
+         header: Keyword.fetch!(opts, :header),
+         buffer_size: IO.iodata_length(buffer),
+         max_buffer_size: max_buffer_size,
+         flush_interval_ms: flush_interval_ms
+       },
+       extra_state
+     )}
   end
 
   @impl true
@@ -89,6 +71,14 @@ defmodule Plausible.Ingestion.WriteBuffer do
     end
   end
 
+  def handle_cast(:flush, state) do
+    %{timer: timer, flush_interval_ms: flush_interval_ms} = state
+    Process.cancel_timer(timer)
+    do_flush(state)
+    new_timer = Process.send_after(self(), :tick, flush_interval_ms)
+    {:noreply, %{state | buffer: [], buffer_size: 0, timer: new_timer}}
+  end
+
   @impl true
   def handle_info(:tick, state) do
     do_flush(state)
@@ -118,44 +108,18 @@ defmodule Plausible.Ingestion.WriteBuffer do
       insert_opts: insert_opts,
       insert_sql: insert_sql,
       header: header,
-      name: name
+      name: name,
+      on_flush: on_flush
     } = state
 
     case buffer do
       [] ->
-        nil
+        on_flush.(:empty, state)
 
       _not_empty ->
         Logger.notice("Flushing #{buffer_size} byte(s) RowBinary from #{name}")
         IngestRepo.query!(insert_sql, [header | buffer], insert_opts)
-    end
-
-    case :ets.lookup(state.name, :state) do
-      [state: %{locker: pid}] when is_pid(pid) ->
-        send(pid, {:locked, state.name})
-        now = System.monotonic_time()
-        lock_loop(state.name, now, state.lock_timeout_ms, state.lock_interval_ms)
-
-      _ ->
-        :ignore
-    end
-  end
-
-  defp lock_loop(name, start, lock_timeout, lock_interval) do
-    now = System.monotonic_time()
-
-    if now - start <= lock_timeout do
-      Process.sleep(lock_interval)
-
-      case :ets.lookup(name, :state) do
-        [state: %{locker: pid}] when is_pid(pid) ->
-          lock_loop(name, start, lock_timeout, lock_interval)
-
-        _ ->
-          :pass
-      end
-    else
-      unlock(name)
+        on_flush.(:success, state)
     end
   end
 

lib/plausible/session/transfer.ex

@@ -120,8 +120,7 @@ defmodule Plausible.Session.Transfer do
       @cmd_takeover_done ->
         # Wipe the cache after transfer is complete to avoid making the transferred cache stale
         Cache.Adapter.wipe(:sessions)
-        # Unblock ingest after trasnsfer is done
-        Plausible.Event.WriteBuffer.unlock()
+        # Unblock ingest after transfer is done
         Plausible.Session.WriteBuffer.unlock()
         :counters.add(given_counter, 1, 1)
     end
@@ -130,19 +129,12 @@ defmodule Plausible.Session.Transfer do
   defp list_cache_names(session_version, parent) do
     if session_version == session_version() and attempted?(parent) do
       # Blocking ingest before transfer
-      locks_acquired? =
-        [
-          Task.async(fn -> Plausible.Event.WriteBuffer.lock(@lock_acquire_timeout) end),
-          Task.async(fn -> Plausible.Session.WriteBuffer.lock(@lock_acquire_timeout) end)
-        ]
-        # timeouts are managed by WriteBuffer.lock/1 calls instead
-        |> Task.await_many(:infinity)
-        |> Enum.all?(&(&1 == :ok))
-
-      if locks_acquired? do
-        Cache.Adapter.get_names(:sessions)
-      else
-        []
+      case Plausible.Session.WriteBuffer.lock(@lock_acquire_timeout) do
+        :ok ->
+          Cache.Adapter.get_names(:sessions)
+
+        _ ->
+          []
       end
     else
       []

lib/plausible/session/write_buffer.ex

@@ -1,6 +1,9 @@
 defmodule Plausible.Session.WriteBuffer do
   @moduledoc false
 
+  @lock_timeout :timer.seconds(3)
+  @lock_interval 100
+
   %{
     header: header,
     insert_sql: insert_sql,
@@ -16,7 +19,9 @@ defmodule Plausible.Session.WriteBuffer do
         name: __MODULE__,
         header: unquote(header),
         insert_sql: unquote(insert_sql),
-        insert_opts: unquote(insert_opts)
+        insert_opts: unquote(insert_opts),
+        on_init: &Plausible.Session.WriteBuffer.on_init/1,
+        on_flush: &Plausible.Session.WriteBuffer.on_flush/2
       )
 
     Plausible.Ingestion.WriteBuffer.child_spec(opts)
@@ -41,11 +46,71 @@ defmodule Plausible.Session.WriteBuffer do
     Plausible.Ingestion.WriteBuffer.flush(__MODULE__)
   end
 
-  def lock(timeout) do
-    Plausible.Ingestion.WriteBuffer.lock(__MODULE__, timeout)
+  @doc false
+  def on_init(opts) do
+    name = Keyword.fetch!(opts, :name)
+
+    ^name = :ets.new(name, [:named_table, :set, :public])
+
+    %{
+      lock_timeout_ms: opts[:lock_timeouts_ms] || @lock_timeout,
+      lock_interval_ms: opts[:lock_interval_ms] || @lock_interval
+    }
+  end
+
+  @doc false
+  def on_flush(_, state) do
+    case :ets.lookup(state.name, :state) do
+      [state: %{locker: pid}] when is_pid(pid) ->
+        send(pid, {:locked, state.name})
+        now = System.monotonic_time()
+        lock_loop(state.name, now, state.lock_timeout_ms, state.lock_interval_ms)
+
+      _ ->
+        :ignore
+    end
+
+    state
+  end
+
+  def lock(locker \\ self(), timeout) do
+    name = __MODULE__
+
+    true = :ets.insert(name, {:state, %{locker: locker}})
+    Plausible.Ingestion.WriteBuffer.flush_async(name)
+
+    receive do
+      {:locked, ^name} -> :ok
+    after
+      timeout -> {:error, :timeout}
+    end
   end
 
   def unlock() do
-    Plausible.Ingestion.WriteBuffer.unlock(__MODULE__)
+    name = __MODULE__
+    true = :ets.insert(name, {:state, %{locker: nil}})
+
+    :ok
+  end
+
+  defp lock_loop(name, start, lock_timeout, lock_interval) do
+    now = System.monotonic_time()
+
+    if now - start <= lock_timeout do
+      Process.sleep(lock_interval)
+
+      case :ets.lookup(name, :state) do
+        [state: %{locker: pid}] when is_pid(pid) ->
+          lock_loop(name, start, lock_timeout, lock_interval)
+
+        _ ->
+          :pass
+      end
+    else
+      # Wipe the cache before unlocking to prevent stale session in case
+      # transfer actually occurs, either partially or completely
+      Plausible.Cache.Adapter.wipe(:sessions)
+      unlock()
+    end
   end
 end