MB-61292: Add cb_atomic_persistent_term

The idea is to modify encryption keys in persistent_term atomically
and avoid scenarios like:

1. proc1 reads keys from disk
2. proc3 changes keys are changed on disk
3. proc2 reads keys from disk
3. proc2 writes keys to persistent_term based on #3
4. proc1 writes keys to persistent_term based on #1 (overwrites #2)

Change-Id: I8d08717170e7b9c920778b7918fc74877d06bbe8
Reviewed-on: https://review.couchbase.org/c/ns_server/+/213279
Tested-by: Timofey Barmin <[email protected]>
Reviewed-by: Navdeep S Boparai <[email protected]>
Well-Formed: Build Bot <[email protected]>

Author: timofey-barmin

apps/ns_server/src/cb_atomic_persistent_term.erl

@@ -0,0 +1,61 @@
+%% @author Couchbase <[email protected]>
+%% @copyright 2024-Present Couchbase, Inc.
+%%
+%% Use of this software is governed by the Business Source License included in
+%% the file licenses/BSL-Couchbase.txt.  As of the Change Date specified in that
+%% file, in accordance with the Business Source License, use of this software
+%% will be governed by the Apache License, Version 2.0, included in the file
+%% licenses/APL2.txt.
+%%
+-module(cb_atomic_persistent_term).
+
+%% API
+-export([start_link/0, set/2, get_or_set_if_undefined/2]).
+
+%%%===================================================================
+%%% API
+%%%===================================================================
+
+start_link() ->
+    work_queue:start_link(?MODULE).
+
+get_or_set_if_undefined(Name, ValueFun) ->
+    maybe
+        undefined ?= persistent_term:get(Name, undefined),
+        work_queue:submit_sync_work(
+          ?MODULE,
+          fun () ->
+              case persistent_term:get(Name, undefined) of
+                  undefined ->
+                      case ValueFun() of
+                          {ok, Value} ->
+                              persistent_term:put(Name, Value),
+                              {ok, Value};
+                          {error, _} = Error ->
+                              Error
+                      end;
+                  Value ->
+                      {ok, Value}
+              end
+          end)
+    else
+        V -> {ok, V}
+    end.
+
+set(Name, SetFun) ->
+    work_queue:submit_sync_work(
+      ?MODULE,
+      fun () ->
+          PrevValue = persistent_term:get(Name, undefined),
+          case SetFun(PrevValue) of
+              {set, Value, Return} ->
+                  persistent_term:put(Name, Value),
+                  Return;
+              {ignore, Return} ->
+                  Return
+          end
+      end).
+
+%%%===================================================================
+%%% Internal functions
+%%%===================================================================

apps/ns_server/src/ns_server_cluster_sup.erl

@@ -34,7 +34,9 @@ start_link() ->
 
 init([]) ->
     {ok, {{one_for_one, 10, 1},
-          [{local_tasks, {local_tasks, start_link, []},
+          [{atomic_persistent_term, {cb_atomic_persistent_term, start_link, []},
+            permanent, 5000, worker, [cb_atomic_persistent_term]},
+           {local_tasks, {local_tasks, start_link, []},
             permanent, brutal_kill, worker, [local_tasks]},
            {log_os_info, {log_os_info, start_link, []},
             transient, 1000, worker, [log_os_info]},