Populate zone attribute from COUCHDB_ZONE environment variable

src/couch/src/couch_util.erl

@@ -228,8 +228,14 @@ get_value(Key, List, Default) ->
             Default
     end.
 
-set_value(Key, List, Value) ->
-    lists:keyreplace(Key, 1, List, {Key, Value}).
+% insert or update a {Key, Value} tuple in a list of tuples
+-spec set_value(Key, TupleList1, Value) -> TupleList2 when
+    Key :: term(),
+    TupleList1 :: [tuple()],
+    Value :: term(),
+    TupleList2 :: [tuple()].
+set_value(Key, TupleList1, Value) ->
+    lists:keystore(Key, 1, TupleList1, {Key, Value}).
 
 get_nested_json_value({Props}, [Key | Keys]) ->
     case couch_util:get_value(Key, Props, nil) of

src/docs/src/cluster/databases.rst

@@ -66,6 +66,9 @@ Add a key value pair of the form:
 
     "zone": "metro-dc-a"
 
+Alternatively, you can set the ``COUCHDB_ZONE`` environment variable
+on each node and CouchDB will configure this document for you on startup.
+
 Do this for all of the nodes in your cluster.
 
 In your config file (``local.ini`` or ``default.ini``) on each node, define a

src/docs/src/cluster/sharding.rst

@@ -650,6 +650,9 @@ Do this for all of the nodes in your cluster. For example:
             "zone": "{zone-name}"
             }'
 
+Alternatively, you can set the ``COUCHDB_ZONE`` environment variable
+on each node and CouchDB will configure this document for you on startup.
+
 In the local config file (``local.ini``) of each node, define a
 consistent cluster-wide setting like:
 
@@ -669,7 +672,7 @@ when the database is created, using the same syntax as the ini file:
 
 .. code-block:: bash
 
-    curl -X PUT $COUCH_URL:5984/{db}?zone={zone}
+    curl -X PUT $COUCH_URL:5984/{db}?placement={zone}
 
 The ``placement`` argument may also be specified. Note that this *will*
 override the logic that determines the number of created replicas!

src/mem3/src/mem3_nodes.erl

@@ -66,7 +66,7 @@ handle_call({get_node_info, Node, Key}, _From, State) ->
     {reply, Resp, State};
 handle_call({add_node, Node, NodeInfo}, _From, State) ->
     gen_event:notify(mem3_events, {add_node, Node}),
-    ets:insert(?MODULE, {Node, NodeInfo}),
+    update_ets(Node, NodeInfo),
     {reply, ok, State};
 handle_call({remove_node, Node}, _From, State) ->
     gen_event:notify(mem3_events, {remove_node, Node}),
@@ -95,12 +95,26 @@ handle_info(_Info, State) ->
     {noreply, State}.
 
 %% internal functions
-
 initialize_nodelist() ->
     DbName = mem3_sync:nodes_db(),
     {ok, Db} = mem3_util:ensure_exists(DbName),
     {ok, _} = couch_db:fold_docs(Db, fun first_fold/2, Db, []),
+
     insert_if_missing(Db, [config:node_name() | mem3_seeds:get_seeds()]),
+
+    % when creating the document for the local node, populate
+    % the placement zone as defined by the COUCHDB_ZONE environment
+    % variable. This is an additional update on top of the first,
+    % empty document so that we don't create conflicting revisions
+    % between different nodes in the cluster when using a seedlist.
+    case os:getenv("COUCHDB_ZONE") of
+        false ->
+            % do not support unsetting a zone.
+            ok;
+        Zone ->
+            set_zone(DbName, config:node_name(), ?l2b(Zone))
+    end,
+
     Seq = couch_db:get_update_seq(Db),
     couch_db:close(Db),
     Seq.
@@ -111,7 +125,7 @@ first_fold(#full_doc_info{deleted = true}, Acc) ->
     {ok, Acc};
 first_fold(#full_doc_info{id = Id} = DocInfo, Db) ->
     {ok, #doc{body = {Props}}} = couch_db:open_doc(Db, DocInfo, [ejson_body]),
-    ets:insert(?MODULE, {mem3_util:to_atom(Id), Props}),
+    update_ets(mem3_util:to_atom(Id), Props),
     {ok, Db}.
 
 listen_for_changes(Since) ->
@@ -156,7 +170,7 @@ insert_if_missing(Db, Nodes) ->
                 [_] ->
                     Acc;
                 [] ->
-                    ets:insert(?MODULE, {Node, []}),
+                    update_ets(Node, []),
                     [#doc{id = couch_util:to_binary(Node)} | Acc]
             end
         end,
@@ -169,3 +183,42 @@ insert_if_missing(Db, Nodes) ->
         true ->
             {ok, []}
     end.
+
+-spec update_ets(Node :: term(), NodeInfo :: [tuple()]) -> true.
+update_ets(Node, NodeInfo) ->
+    ets:insert(?MODULE, {Node, NodeInfo}).
+
+% sets the placement zone for the given node document.
+-spec set_zone(DbName :: binary(), Node :: string() | binary(), Zone :: binary()) -> ok.
+set_zone(DbName, Node, Zone) ->
+    {ok, Db} = couch_db:open(DbName, [sys_db, ?ADMIN_CTX]),
+    Props = get_from_db(Db, Node),
+    CurrentZone = couch_util:get_value(<<"zone">>, Props),
+    case CurrentZone of
+        Zone ->
+            ok;
+        _ ->
+            couch_log:info("Setting node zone attribute to ~s~n", [Zone]),
+            Props1 = couch_util:set_value(<<"zone">>, Props, Zone),
+            save_to_db(Db, Node, Props1)
+    end,
+    couch_db:close(Db),
+    ok.
+
+% get a node document from the system nodes db as a property list
+-spec get_from_db(Db :: any(), Node :: string() | binary()) -> [tuple()].
+get_from_db(Db, Node) ->
+    Id = couch_util:to_binary(Node),
+    {ok, Doc} = couch_db:open_doc(Db, Id, [ejson_body]),
+    {Props} = couch_doc:to_json_obj(Doc, []),
+    Props.
+
+% save a node document (represented as a property list)
+% to the system nodes db and update the ETS cache.
+-spec save_to_db(Db :: any(), Node :: string() | binary(), Props :: [tuple()]) -> ok.
+save_to_db(Db, Node, Props) ->
+    Doc = couch_doc:from_json_obj({Props}),
+    #doc{body = {NodeInfo}} = Doc,
+    {ok, _} = couch_db:update_doc(Db, Doc, []),
+    update_ets(Node, NodeInfo),
+    ok.

src/mem3/test/eunit/mem3_zone_test.erl

@@ -0,0 +1,77 @@
+% Licensed under the Apache License, Version 2.0 (the "License"); you may not
+% use this file except in compliance with the License. You may obtain a copy of
+% the License at
+%
+%   http://www.apache.org/licenses/LICENSE-2.0
+%
+% Unless required by applicable law or agreed to in writing, software
+% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+% License for the specific language governing permissions and limitations under
+% the License.
+
+-module(mem3_zone_test).
+
+-include_lib("couch/include/couch_eunit.hrl").
+-include_lib("eunit/include/eunit.hrl").
+
+mem3_zone_test_() ->
+    {
+        foreach,
+        fun setup/0,
+        fun teardown/1,
+        [
+            ?TDEF_FE(t_empty_zone),
+            ?TDEF_FE(t_set_zone_from_env),
+            ?TDEF_FE(t_set_zone_when_node_in_seedlist),
+            ?TDEF_FE(t_zone_already_set)
+        ]
+    }.
+
+assertZoneEqual(Expected) ->
+    [Node | _] = mem3:nodes(),
+    Actual = mem3:node_info(Node, <<"zone">>),
+    ?assertEqual(Expected, Actual).
+
+t_empty_zone(_) ->
+    ok = application:start(mem3),
+    assertZoneEqual(undefined).
+
+t_set_zone_from_env(_) ->
+    Zone = "zone1",
+    os:putenv("COUCHDB_ZONE", Zone),
+    ok = application:start(mem3),
+    assertZoneEqual(iolist_to_binary(Zone)).
+
+t_set_zone_when_node_in_seedlist(_) ->
+    CfgSeeds = "nonode@nohost",
+    config:set("cluster", "seedlist", CfgSeeds, false),
+    Zone = "zone1",
+    os:putenv("COUCHDB_ZONE", Zone),
+    ok = application:start(mem3),
+    assertZoneEqual(iolist_to_binary(Zone)).
+
+t_zone_already_set(_) ->
+    Zone = "zone1",
+    os:putenv("COUCHDB_ZONE", Zone),
+    ok = application:start(mem3),
+    application:stop(mem3),
+    ok = application:start(mem3),
+    assertZoneEqual(iolist_to_binary(Zone)).
+
+setup() ->
+    meck:new(mem3_seeds, [passthrough]),
+    meck:new(mem3_rpc, [passthrough]),
+    test_util:start_couch([rexi]).
+
+teardown(Ctx) ->
+    catch application:stop(mem3),
+    os:unsetenv("COUCHDB_ZONE"),
+    Filename = config:get("mem3", "nodes_db", "_nodes") ++ ".couch",
+    file:delete(filename:join([?BUILDDIR(), "tmp", "data", Filename])),
+    case config:get("couch_httpd_auth", "authentication_db") of
+        undefined -> ok;
+        DbName -> couch_server:delete(list_to_binary(DbName), [])
+    end,
+    meck:unload(),
+    test_util:stop_couch(Ctx).