diff --git a/config/test.exs b/config/test.exs index 05d6b225d..1222dad44 100644 --- a/config/test.exs +++ b/config/test.exs @@ -14,7 +14,14 @@ config :logger, level: :warn # Disable authentication mode during test config :livebook, :authentication_mode, :disabled -config :livebook, :data_path, Path.expand("tmp/livebook_data/test") +data_path = Path.expand("tmp/livebook_data/test") + +# Clear data path for tests +if File.exists?(data_path) do + File.rm_rf!(data_path) +end + +config :livebook, :data_path, data_path # Use the embedded runtime in tests by default, so they # are cheaper to run. Other runtimes can be tested by starting diff --git a/lib/livebook/storage/ets.ex b/lib/livebook/storage/ets.ex index 7234ae70c..85dd6110c 100644 --- a/lib/livebook/storage/ets.ex +++ b/lib/livebook/storage/ets.ex @@ -10,13 +10,11 @@ defmodule Livebook.Storage.Ets do """ @behaviour Livebook.Storage - @table_name __MODULE__ - use GenServer @impl Livebook.Storage def all(namespace) do - @table_name + table_name() |> :ets.match({{namespace, :"$1"}, :"$2", :"$3", :_}) |> Enum.group_by( fn [entity_id, _attr, _val] -> entity_id end, @@ -31,7 +29,7 @@ defmodule Livebook.Storage.Ets do @impl Livebook.Storage def fetch(namespace, entity_id) do - @table_name + table_name() |> :ets.lookup({namespace, entity_id}) |> case do [] -> @@ -48,7 +46,7 @@ defmodule Livebook.Storage.Ets do @impl Livebook.Storage def fetch_key(namespace, entity_id, key) do - @table_name + table_name() |> :ets.match({{namespace, entity_id}, key, :"$1", :_}) |> case do [[value]] -> {:ok, value} @@ -56,6 +54,11 @@ defmodule Livebook.Storage.Ets do end end + @spec config_file_path() :: Path.t() + def config_file_path() do + Path.join([Livebook.Config.data_path(), "storage.ets"]) + end + @impl Livebook.Storage def insert(namespace, entity_id, attributes) do GenServer.call(__MODULE__, {:insert, namespace, entity_id, attributes}) @@ -73,7 +76,12 @@ defmodule Livebook.Storage.Ets do @impl GenServer def init(_opts) do - table = :ets.new(@table_name, [:named_table, :protected, :duplicate_bag]) + # Make sure that this process does not terminate abruptly + # in case it is persisting to disk. terminate/2 is still a no-op. + Process.flag(:trap_exit, true) + + table = load_or_create_table() + :persistent_term.put(__MODULE__, table) {:ok, %{table: table}} end @@ -98,6 +106,8 @@ defmodule Livebook.Storage.Ets do :ets.insert(table, attributes) + :ok = save_to_file(state) + {:reply, :ok, state} end @@ -105,6 +115,29 @@ defmodule Livebook.Storage.Ets do def handle_call({:delete, namespace, entity_id}, _from, %{table: table} = state) do :ets.delete(table, {namespace, entity_id}) + :ok = save_to_file(state) + {:reply, :ok, state} end + + defp table_name(), do: :persistent_term.get(__MODULE__) + + defp load_or_create_table() do + config_file_path() + |> String.to_charlist() + |> :ets.file2tab() + |> case do + {:ok, tab} -> + tab + + {:error, _reason} -> + :ets.new(__MODULE__, [:protected, :duplicate_bag]) + end + end + + defp save_to_file(%{table: table}) do + file_path = String.to_charlist(config_file_path()) + + :ets.tab2file(table, file_path) + end end diff --git a/test/livebook/storage/ets_test.exs b/test/livebook/storage/ets_test.exs index 81e395bd6..cfab5edf3 100644 --- a/test/livebook/storage/ets_test.exs +++ b/test/livebook/storage/ets_test.exs @@ -87,4 +87,35 @@ defmodule Livebook.Storage.EtsTest do assert [] = Ets.all(:unknown_namespace) end end + + describe "persistence" do + defp read_table_and_lookup(path, entity) do + {:ok, tab} = + path + |> String.to_charlist() + |> :ets.file2tab() + + :ets.lookup(tab, {:persistence, entity}) + end + + test "insert triggers saving to file" do + :ok = Ets.insert(:persistence, "insert", key: "val") + + path = Ets.config_file_path() + assert File.exists?(path) + + assert [_test] = read_table_and_lookup(path, "insert") + end + + test "delete triggers saving to file" do + :ok = Ets.insert(:persistence, "delete", key: "val") + + path = Ets.config_file_path() + assert File.exists?(path) + + :ok = Ets.delete(:persistence, "delete") + + assert [] = read_table_and_lookup(path, "delete") + end + end end