Просмотр исходного кода

ensure_reg() and ensure_reg_other()

Ulf Wiger 8 лет назад
Родитель
Сommit
7c0d81abed
4 измененных файлов с 274 добавлено и 77 удалено
  1. 147 58
      src/gproc.erl
  2. 27 14
      src/gproc_dist.erl
  3. 45 4
      test/gproc_dist_tests.erl
  4. 55 1
      test/gproc_tests.erl

+ 147 - 58
src/gproc.erl

@@ -51,6 +51,8 @@
 	 reg_or_locate/1, reg_or_locate/2, reg_or_locate/3,
 	 reg_shared/1, reg_shared/2, reg_shared/3, unreg_shared/1,
 	 set_attributes_shared/2, set_value_shared/2,
+         ensure_reg/1, ensure_reg/2, ensure_reg/3,
+         ensure_reg_other/2, ensure_reg_other/3, ensure_reg_other/4,
          mreg/3,
          munreg/3,
          set_value/2,
@@ -148,7 +150,9 @@
 -type pidpat()  :: pid() | sel_var().
 -type headpat() :: {keypat(), pidpat(), any()}.
 -type key()     :: {type(), scope(), any()}.
-
+-type value()   :: any().
+-type attr()    :: {atom(), any()}.
+-type attrs()   :: [attr()].
 -type sel_pattern() :: [{headpat(), list(), list()}].
 
 -type reg_id() :: {type(), scope(), any()}.
@@ -202,7 +206,7 @@ start_link() ->
 %% @end
 %%
 add_local_name(Name)  ->
-    ?CATCH_GPROC_ERROR(reg1({n,l,Name}, undefined, []), [Name]).
+    ?CATCH_GPROC_ERROR(reg1({n,l,Name}, undefined, [], reg), [Name]).
 
 
 %% spec(Name::any()) -> true
@@ -211,7 +215,7 @@ add_local_name(Name)  ->
 %% @end
 %%
 add_global_name(Name) ->
-    ?CATCH_GPROC_ERROR(reg1({n,g,Name}, undefined, []), [Name]).
+    ?CATCH_GPROC_ERROR(reg1({n,g,Name}, undefined, [], reg), [Name]).
 
 
 %% spec(Name::any(), Value::any()) -> true
@@ -220,7 +224,7 @@ add_global_name(Name) ->
 %% @end
 %%
 add_local_property(Name , Value) ->
-    ?CATCH_GPROC_ERROR(reg1({p,l,Name}, Value, []), [Name, Value]).
+    ?CATCH_GPROC_ERROR(reg1({p,l,Name}, Value, [], reg), [Name, Value]).
 
 %% spec(Name::any(), Value::any()) -> true
 %%
@@ -228,7 +232,7 @@ add_local_property(Name , Value) ->
 %% @end
 %%
 add_global_property(Name, Value) ->
-    ?CATCH_GPROC_ERROR(reg1({p,g,Name}, Value, []), [Name, Value]).
+    ?CATCH_GPROC_ERROR(reg1({p,g,Name}, Value, [], reg), [Name, Value]).
 
 %% spec(Name::any(), Initial::integer()) -> true
 %%
@@ -236,7 +240,7 @@ add_global_property(Name, Value) ->
 %% @end
 %%
 add_local_counter(Name, Initial) when is_integer(Initial) ->
-    ?CATCH_GPROC_ERROR(reg1({c,l,Name}, Initial, []), [Name, Initial]).
+    ?CATCH_GPROC_ERROR(reg1({c,l,Name}, Initial, [], reg), [Name, Initial]).
 
 
 %% spec(Name::any(), Initial::integer()) -> true
@@ -255,7 +259,7 @@ add_shared_local_counter(Name, Initial) when is_integer(Initial) ->
 %% @end
 %%
 add_global_counter(Name, Initial) when is_integer(Initial) ->
-    ?CATCH_GPROC_ERROR(reg1({c,g,Name}, Initial, []), [Name, Initial]).
+    ?CATCH_GPROC_ERROR(reg1({c,g,Name}, Initial, [], reg), [Name, Initial]).
 
 %% spec(Name::any()) -> true
 %%
@@ -271,7 +275,8 @@ add_local_aggr_counter(Name)  -> ?CATCH_GPROC_ERROR(reg1({a,l,Name}), [Name]).
 %% @equiv reg({a,g,Name})
 %% @end
 %%
-add_global_aggr_counter(Name) -> ?CATCH_GPROC_ERROR(reg1({a,g,Name}), [Name]).
+add_global_aggr_counter(Name) ->
+    ?CATCH_GPROC_ERROR(reg1({a,g,Name}), [Name]).
 
 
 %% @spec (Name::any()) -> pid()
@@ -545,7 +550,7 @@ lookup_env(Scope, App, Key, P) ->
 
 cache_env(Scope, App, Key, Value) ->
     ?CATCH_GPROC_ERROR(
-       reg1({p, Scope, {gproc_env, App, Key}}, Value, []),
+       reg1({p, Scope, {gproc_env, App, Key}}, Value, [], reg),
        [Scope,App,Key,Value]).
 
 update_cached_env(Scope, App, Key, Value) ->
@@ -617,7 +622,7 @@ reg(Key) ->
     ?CATCH_GPROC_ERROR(reg1(Key), [Key]).
 
 reg1(Key) ->
-    reg1(Key, default(Key), []).
+    reg1(Key, default(Key), [], reg).
 
 %% @spec reg_or_locate(Key::key()) -> {pid(), NewValue}
 %%
@@ -948,15 +953,15 @@ demonitor1({T,l,_} = Key, Ref) when T==n; T==a ->
 demonitor1(_, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
-%% @spec reg(Key::key(), Value) -> true
+%% @spec reg(Key::key(), Value::value()) -> true
 %%
 %% @doc Register a name or property for the current process
 %%
 %%
 reg(Key, Value) ->
-    ?CATCH_GPROC_ERROR(reg1(Key, Value, []), [Key, Value]).
+    ?CATCH_GPROC_ERROR(reg1(Key, Value, [], reg), [Key, Value]).
 
-%% @spec reg(Key::key(), Value, Attrs::[{atom(),any()}]) -> true
+%% @spec reg(Key::key(), Value::value(), Attrs::attrs()) -> true
 %%
 %% @doc Register a name or property for the current process
 %% `Attrs' (default: `[]') can be inspected using {@link get_attribute/2}.
@@ -994,34 +999,64 @@ reg(Key, Value) ->
 %%  `{Type, Context, Name}'
 %% @end
 reg(Key, Value, Attrs) ->
-    ?CATCH_GPROC_ERROR(reg1(Key, Value, Attrs), [Key, Value, Attrs]).
+    ?CATCH_GPROC_ERROR(reg1(Key, Value, Attrs, reg), [Key, Value, Attrs]).
+
+%% @equiv ensure_reg(Key, default(Key), [])
+ensure_reg(Key) ->
+    ?CATCH_GPROC_ERROR(reg1(Key, ensure), [Key]).
+
+%% @equiv ensure_reg(Key, Value, [])
+-spec ensure_reg(key(), value()) -> new | updated.
+ensure_reg(Key, Value) ->
+    ?CATCH_GPROC_ERROR(reg1(Key, Value, ensure), [Key, Value]).
+
+%% @spec ensure_reg(Key::key(), Value::value(), Attrs::attrs()) ->
+%%          new | updated
+%%
+%% @doc Registers a new name or property unless such and entry (by key) has
+%% already been registered by the current process. If `Key' already exists,
+%% the entry will be updated with the given `Value' and `Attrs'.
+%%
+%% This function allows the caller to efficiently register an entry without
+%% first checking whether it has already been registered. An exception is
+%% raised if the name or property is already registered by someone else.
+%% @end
+-spec ensure_reg(key(), value(), attrs()) -> new | updated.
+ensure_reg(Key, Value, Attrs) ->
+    ?CATCH_GPROC_ERROR(reg1(Key, Value, Attrs, ensure), [Key, Value, Attrs]).
+
+reg1(Key, Op) ->
+    reg1(Key, default(Key), [], Op).
 
-reg1({T,g,_} = Key, Value, As) when T==p; T==a; T==c; T==n; T==r; T==rc ->
+reg1(Key, Value, Op) ->
+    reg1(Key, Value, [], Op).
+
+reg1({T,g,_} = Key, Value, As, Op) when T==p; T==a; T==c; T==n; T==r; T==rc ->
     %% anything global
     ?CHK_DIST,
-    gproc_dist:reg(Key, Value, As);
-reg1({p,l,_} = Key, Value, As) ->
-    local_reg(Key, Value, As);
-reg1({a,l,_} = Key, undefined, As) ->
-    call({reg, Key, undefined, As});
-reg1({c,l,_} = Key, Value, As) when is_integer(Value) ->
-    call({reg, Key, Value, As});
-reg1({n,l,_} = Key, Value, As) ->
-    call({reg, Key, Value, As});
-reg1({r,l,_} = Key, Value, As) ->
-    call({reg, Key, Value, As});
-reg1({rc,l,_} = Key, Value, As) ->
-    call({reg, Key, Value, As});
-reg1(_, _, _) ->
+    gproc_dist:reg(Key, Value, As, Op);
+reg1({p,l,_} = Key, Value, As, Op) ->
+    local_reg(Key, Value, As, Op);
+reg1({a,l,_} = Key, undefined, As, Op) ->
+    call({reg, Key, undefined, As, Op});
+reg1({c,l,_} = Key, Value, As, Op) when is_integer(Value) ->
+    call({reg, Key, Value, As, Op});
+reg1({n,l,_} = Key, Value, As, Op) ->
+    call({reg, Key, Value, As, Op});
+reg1({r,l,_} = Key, Value, As, Op) ->
+    call({reg, Key, Value, As, Op});
+reg1({rc,l,_} = Key, Value, As, Op) ->
+    call({reg, Key, Value, As, Op});
+reg1(_, _, _, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
 %% @equiv reg_other(Key, Pid, default(Key), [])
 reg_other(Key, Pid) ->
-    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid), [Key, Pid]).
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, reg), [Key, Pid]).
 
 %% @equiv reg_other(Key, Pid, Value, [])
 reg_other(Key, Pid, Value) ->
-    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, []), [Key, Pid, Value]).
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, [], reg), [Key, Pid, Value]).
 
 %% @spec reg_other(Key, Pid, Value, Attrs) -> true
 %% @doc Register name or property to another process.
@@ -1041,18 +1076,39 @@ reg_other(Key, Pid, Value) ->
 %% * `rc' - resource counters
 %% @end
 reg_other(Key, Pid, Value, Attrs) ->
-    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, Attrs),
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, Attrs, reg),
                        [Key, Pid, Value, Attrs]).
 
-reg_other1(Key, Pid) ->
-    reg_other1(Key, Pid, default(Key), []).
+ensure_reg_other(Key, Pid) ->
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, ensure), [Key, Pid]).
+
+%% @equiv ensure_reg_other(Key, Pid, Value, [])
+ensure_reg_other(Key, Pid, Value) ->
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, [], ensure),
+                       [Key, Pid, Value]).
 
-reg_other1({_,g,_} = Key, Pid, Value, As) when is_pid(Pid) ->
+%% @spec ensure_reg_other(Key::key(), Pid::pid(),
+%%                        Value::value(), Attrs::attrs()) ->
+%%          new | updated
+%%
+%% @doc Register or update name or property to another process.
+%%
+%% Equivalent to {@link reg_other/3}, but allows for registration of another
+%% process instead of the current process. Also see {@link ensure_reg/3}.
+%% @end
+ensure_reg_other(Key, Pid, Value, Attrs) ->
+    ?CATCH_GPROC_ERROR(reg_other1(Key, Pid, Value, Attrs, ensure),
+                       [Key, Pid, Value, Attrs]).
+
+reg_other1(Key, Pid, Op) ->
+    reg_other1(Key, Pid, default(Key), [], Op).
+
+reg_other1({_,g,_} = Key, Pid, Value, As, Op) when is_pid(Pid) ->
     ?CHK_DIST,
-    gproc_dist:reg_other(Key, Pid, Value, As);
-reg_other1({T,l,_} = Key, Pid, Value, As) when is_pid(Pid) ->
+    gproc_dist:reg_other(Key, Pid, Value, As, Op);
+reg_other1({T,l,_} = Key, Pid, Value, As, Op) when is_pid(Pid) ->
     if T==n; T==a; T==r; T==rc ->
-            call({reg_other, Key, Pid, Value, As});
+            call({reg_other, Key, Pid, Value, As, Op});
        true ->
             ?THROW_GPROC_ERROR(badarg)
     end.
@@ -1132,13 +1188,13 @@ reg_shared1({_,g,_} = Key, Value, As) ->
     ?CHK_DIST,
     gproc_dist:reg_shared(Key, Value, As);
 reg_shared1({a,l,_} = Key, undefined, As) ->
-    call({reg_shared, Key, undefined, As});
+    call({reg_shared, Key, undefined, As, reg});
 reg_shared1({c,l,_} = Key, Value, As) when is_integer(Value) ->
-    call({reg_shared, Key, Value, As});
+    call({reg_shared, Key, Value, As, reg});
 reg_shared1({p,l,_} = Key, Value, As) ->
-    call({reg_shared, Key, Value, As});
+    call({reg_shared, Key, Value, As, reg});
 reg_shared1({rc,l,_} = Key, undefined, As) ->
-    call({reg_shared, Key, undefined, As});
+    call({reg_shared, Key, undefined, As, reg});
 reg_shared1(_, _, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
@@ -1406,17 +1462,30 @@ select_count(Context, Pat) ->
 %%% Local properties can be registered in the local process, since
 %%% no other process can interfere.
 %%%
-local_reg({_,Scope,_} = Key, Value, As) ->
+local_reg({_,Scope,_} = Key, Value, As, Op) ->
     case gproc_lib:insert_reg(Key, Value, self(), l) of
-        false -> ?THROW_GPROC_ERROR(badarg);
+        false ->
+            case ets:member(?TAB, {Key, self()}) of
+                true when Op == ensure ->
+                    gproc_lib:do_set_value(Key, Value, self()),
+                    set_attrs(As, Key, self()),
+                    updated;
+                _ ->
+                    ?THROW_GPROC_ERROR(badarg)
+            end;
         true  ->
             monitor_me(),
             if As =/= [] ->
-                    gproc_lib:insert_attr(Key, As, self(), Scope);
-               true -> true
+                    gproc_lib:insert_attr(Key, As, self(), Scope),
+                    regged_new(Op);
+               true ->
+                    regged_new(Op)
             end
     end.
 
+regged_new(reg   ) -> true;
+regged_new(ensure) -> new.
+
 local_mreg(_, []) -> true;
 local_mreg(T, [_|_] = KVL) ->
     case gproc_lib:insert_many(T, l, KVL, self()) of
@@ -2185,10 +2254,10 @@ handle_cast({cancel_wait_or_monitor, Pid, {T,_,_} = Key}, S) ->
     {noreply, S}.
 
 %% @hidden
-handle_call({reg, {_T,l,_} = Key, Val, Attrs}, {Pid,_}, S) ->
-    handle_reg_call(Key, Pid, Val, Attrs, S);
-handle_call({reg_other, {_T,l,_} = Key, Pid, Val, Attrs}, _, S) ->
-    handle_reg_call(Key, Pid, Val, Attrs, S);
+handle_call({reg, {_T,l,_} = Key, Val, Attrs, Op}, {Pid,_}, S) ->
+    handle_reg_call(Key, Pid, Val, Attrs, Op, S);
+handle_call({reg_other, {_T,l,_} = Key, Pid, Val, Attrs, Op}, _, S) ->
+    handle_reg_call(Key, Pid, Val, Attrs, Op, S);
 handle_call({set_attributes, {_,l,_} = Key, Attrs}, {Pid,_}, S) ->
     case gproc_lib:insert_attr(Key, Attrs, Pid, l) of
 	false -> {reply, badarg, S};
@@ -2275,7 +2344,7 @@ handle_call({demonitor, {T,l,_} = Key, Ref, Pid}, _From, S)
 		end
 	end,
     {reply, ok, S};
-handle_call({reg_shared, {_T,l,_} = Key, Val, Attrs}, _From, S) ->
+handle_call({reg_shared, {_T,l,_} = Key, Val, Attrs, Op}, _From, S) ->
     case try_insert_reg(Key, Val, shared) of
 	true ->
             _ = if Attrs =/= [] ->
@@ -2283,7 +2352,14 @@ handle_call({reg_shared, {_T,l,_} = Key, Val, Attrs}, _From, S) ->
                    true -> true
                 end,
 	    {reply, true, S};
-	false ->
+        already_registered when Op == ensure ->
+            case gproc_lib:do_set_value(Key, Val, shared) of
+                true ->
+                    {reply, updated, S};
+                false ->
+                    {reply, badarg, S}
+            end;
+	_ ->
 	    {reply, badarg, S}
     end;
 handle_call({unreg, {_,l,_} = Key}, {Pid,_}, S) ->
@@ -2375,19 +2451,30 @@ terminate(_Reason, _S) ->
 
 %% handle_call body common to reg and reg_other.
 %%
-handle_reg_call(Key, Pid, Val, Attrs, S) ->
+handle_reg_call(Key, Pid, Val, Attrs, Op, S) ->
     case try_insert_reg(Key, Val, Pid) of
         true ->
             _ = gproc_lib:ensure_monitor(Pid,l),
-            _ = if Attrs =/= [] ->
-                        gproc_lib:insert_attr(Key, Attrs, Pid, l);
-                   true -> true
-                end,
-            {reply, true, S};
+            _ = set_attrs(Attrs, Key, Pid),
+            {reply, regged_new(Op), S};
+        already_registered when Op == ensure ->
+            case gproc_lib:do_set_value(Key, Val, Pid) of
+                true ->
+                    _ = set_attrs(Attrs, Key, Pid),
+                    {reply, updated, S};
+                false ->
+                    %% actually pretty bad, if it ever happens
+                    {reply, badarg, S}
+            end;
         false ->
             {reply, badarg, S}
     end.
 
+set_attrs([], _, _) ->
+    true;
+set_attrs([_|_] = Attrs, Key, Pid) ->
+    gproc_lib:insert_attr(Key, Attrs, Pid, l).
+
 handle_unreg_call(Key, Pid, S) ->
     case ets:lookup(?TAB, {Pid,Key}) of
         [{_, r}] ->
@@ -2440,6 +2527,8 @@ try_insert_reg({T,l,_} = Key, Val, Pid) ->
                 %% In this particular case, the lookup cannot result in
                 %% [{_, Waiters}], since the insert_reg/4 function would
                 %% have succeeded then.
+                [{_, Pid, _}] ->
+                    already_registered;
                 [{_, OtherPid, _}] ->
                     case is_process_alive(OtherPid) of
                         true ->

+ 27 - 14
src/gproc_dist.erl

@@ -23,8 +23,8 @@
 -behaviour(gen_leader).
 
 -export([start_link/0, start_link/1,
-         reg/1, reg/3, unreg/1,
-         reg_other/4, unreg_other/2,
+         reg/1, reg/4, unreg/1,
+         reg_other/5, unreg_other/2,
 	 reg_or_locate/3,
 	 reg_shared/3, unreg_shared/1,
          monitor/2,
@@ -93,7 +93,7 @@ start_link({Nodes, Opts}) ->
 %% {@see gproc:reg/1}
 %%
 reg(Key) ->
-    reg(Key, gproc:default(Key), []).
+    reg(Key, gproc:default(Key), [], reg).
 
 %% {@see gproc:reg_or_locate/2}
 %%
@@ -120,13 +120,13 @@ reg_or_locate(_, _, _) ->
 %%%          | r  - resource property
 %%%          | rc - resource counter
 %%% @end
-reg({_,g,_} = Key, Value, Attrs) ->
+reg({_,g,_} = Key, Value, Attrs, Op) ->
     %% anything global
-    leader_call({reg, Key, Value, self(), Attrs});
-reg(_, _, _) ->
+    leader_call({reg, Key, Value, self(), Attrs, Op});
+reg(_, _, _, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
-%% @spec ({Class,g,Key}, pid(), Value, Attrs) -> true
+%% @spec ({Class,g,Key}, pid(), Value, Attrs, Op::reg | unreg) -> true
 %% @doc
 %%    Class = n  - unique name
 %%          | a  - aggregated counter
@@ -135,13 +135,13 @@ reg(_, _, _) ->
 %%    Value = term()
 %%    Attrs = [{Key, Value}]
 %% @end
-reg_other({T,g,_} = Key, Pid, Value, Attrs) when is_pid(Pid) ->
+reg_other({T,g,_} = Key, Pid, Value, Attrs, Op) when is_pid(Pid) ->
     if T==n; T==a; T==r; T==rc ->
-            leader_call({reg_other, Key, Value, Pid, Attrs});
+            leader_call({reg_other, Key, Value, Pid, Attrs, Op});
        true ->
             ?THROW_GPROC_ERROR(badarg)
     end;
-reg_other(_, _, _, _) ->
+reg_other(_, _, _, _, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
 unreg_other({T,g,_} = Key, Pid) when is_pid(Pid) ->
@@ -154,7 +154,7 @@ unreg_other(_, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
 reg_shared({_,g,_} = Key, Value, Attrs) ->
-    leader_call({reg, Key, Value, shared, Attrs});
+    leader_call({reg, Key, Value, shared, Attrs, reg});
 reg_shared(_, _, _) ->
     ?THROW_GPROC_ERROR(badarg).
 
@@ -364,11 +364,21 @@ handle_leader_call(sync, From, #state{sync_requests = SReqs} = S, E) ->
             GenLeader:broadcast({from_leader, {sync, From}}, Alive, E),
             {noreply, S#state{sync_requests = [{From, Alive}|SReqs]}}
     end;
-handle_leader_call({Reg, {_C,g,_Name} = K, Value, Pid, As}, _From, S, _E)
+handle_leader_call({Reg, {_C,g,_Name} = K, Value, Pid, As, Op}, _From, S, _E)
   when Reg==reg; Reg==reg_other ->
     case gproc_lib:insert_reg(K, Value, Pid, g) of
-        false ->
+        false when Op == reg ->
             {reply, badarg, S};
+        false when Op == ensure ->
+            case ets:lookup(?TAB, ets_key(K, Pid)) of
+                [{_, Pid, _}] ->
+                    gproc_lib:do_set_value(K, Value, Pid),
+                    gproc_lib:insert_attr(K, As, Pid, g),
+                    Vals = mk_broadcast_insert_vals([{K, Pid, Value}]),
+                    {reply, updated, [{insert, Vals}], S};
+                _ ->
+                    {reply, badarg, [], S}
+            end;
         true ->
             _ = gproc_lib:ensure_monitor(Pid,g),
             _ = if As =/= [] ->
@@ -376,7 +386,7 @@ handle_leader_call({Reg, {_C,g,_Name} = K, Value, Pid, As}, _From, S, _E)
                    true -> []
                 end,
 	    Vals = mk_broadcast_insert_vals([{K, Pid, Value}]),
-            {reply, true, [{insert, Vals}], S}
+            {reply, regged_new(Op), [{insert, Vals}], S}
     end;
 handle_leader_call({monitor, {T,g,_} = K, MPid, Type}, _From, S, _E) when T==n;
                                                                           T==a ->
@@ -1057,3 +1067,6 @@ add_follow_to_waiters(Waiters, {T,_,_} = K, Pid, Ref, S) ->
             {reply, Ref, [{insert, [Obj, Rev]},
                           {notify, [{Pid, Msg}]}], S}
     end.
+
+regged_new(reg   ) -> true;
+regged_new(ensure) -> new.

+ 45 - 4
test/gproc_dist_tests.erl

@@ -43,21 +43,30 @@ dist_test_() ->
                 [{inorder, basic_tests(Ns)}]
         end,
         fun(Ns) ->
-                [?f(t_sync_cand_dies(Ns))]
+                tests(Ns, [?f(t_sync_cand_dies(Ns))])
         end,
         fun(Ns) ->
-                [?f(t_fail_node(Ns))]
+                tests(Ns, [?f(t_fail_node(Ns))])
         end,
         fun(Ns) ->
-                [?f(t_master_dies(Ns))]
+                tests(Ns, [?f(t_master_dies(Ns))])
         end
        ]}
      ]}.
 
+tests(skip, _) ->
+    [];
+tests(_, L) ->
+    L.
+
+basic_tests(skip) ->
+    [];
 basic_tests(Ns) ->
     [
      ?f(t_simple_reg(Ns)),
      ?f(t_simple_reg_other(Ns)),
+     ?f(t_simple_ensure(Ns)),
+     ?f(t_simple_ensure_other(Ns)),
      ?f(t_simple_reg_or_locate(Ns)),
      ?f(t_simple_counter(Ns)),
      ?f(t_aggr_counter(Ns)),
@@ -95,6 +104,8 @@ dist_setup() ->
             skip
     end.
 
+dist_cleanup(skip) ->
+    ok;
 dist_cleanup(Ns) ->
     [slave:stop(N) || N <- Ns],
     ok.
@@ -112,7 +123,7 @@ run_dist_tests() ->
 	    end
     end.
 
--define(T_NAME, {n, g, {?MODULE, ?LINE, erlang:now()}}).
+-define(T_NAME, {n, g, {?MODULE, ?LINE, os:timestamp()}}).
 -define(T_KVL, [{foo, "foo"}, {bar, "bar"}]).
 -define(T_COUNTER, {c, g, {?MODULE, ?LINE}}).
 -define(T_RESOURCE, {r, g, {?MODULE, ?LINE}}).
@@ -137,6 +148,36 @@ t_simple_reg_other([A, B|_] = Ns) ->
     ?assertMatch(ok, t_call(P1, die)),
     ?assertMatch(ok, t_call(P2, die)).
 
+t_simple_ensure([H|_] = Ns) ->
+    Name = ?T_NAME,
+    P = t_spawn_reg(H, Name),
+    ?assertMatch(ok, t_lookup_everywhere(Name, Ns, P)),
+    ?assertMatch(
+       updated, t_call(
+                  P, {apply, gproc, ensure_reg, [Name, new_val, [{a,1}]]})),
+    ?assertMatch(
+       [{a,1}], t_call(
+                  P, {apply, gproc, get_attributes, [Name]})),
+    ?assertMatch(ok, t_read_everywhere(Name, P, Ns, new_val)),
+    ?assertMatch(true, t_call(P, {apply, gproc, unreg, [Name]})),
+    ?assertMatch(ok, t_lookup_everywhere(Name, Ns, undefined)),
+    ?assertMatch(ok, t_call(P, die)).
+
+t_simple_ensure_other([A, B|_] = Ns) ->
+    Name = ?T_NAME,
+    P1 = t_spawn(A),
+    P2 = t_spawn(B),
+    ?assertMatch(true, t_call(P1, {apply, gproc, reg_other, [Name, P2]})),
+    ?assertMatch(ok, t_lookup_everywhere(Name, Ns, P2)),
+    ?assertMatch(
+       updated, t_call(
+                  P1, {apply, gproc, ensure_reg_other, [Name, P2, new_val]})),
+    ?assertMatch(ok, t_read_everywhere(Name, P2, Ns, new_val)),
+    ?assertMatch(true, t_call(P1, {apply, gproc, unreg_other, [Name, P2]})),
+    ?assertMatch(ok, t_lookup_everywhere(Name, Ns, undefined)),
+    ?assertMatch(ok, t_call(P1, die)),
+    ?assertMatch(ok, t_call(P2, die)).
+
 t_simple_reg_or_locate([A,B|_] = _Ns) ->
     Name = ?T_NAME,
     P1 = t_spawn(A),

+ 55 - 1
test/gproc_tests.erl

@@ -76,6 +76,12 @@ reg_test_() ->
       , ?_test(t_is_clean())
       , {spawn, ?_test(?debugVal(t_simple_reg_other()))}
       , ?_test(t_is_clean())
+      , ?_test(?debugVal(t_simple_ensure()))
+      , ?_test(t_is_clean())
+      , ?_test(?debugVal(t_simple_ensure_other()))
+      , ?_test(t_is_clean())
+      , ?_test(?debugVal(t_simple_ensure_prop()))
+      , ?_test(t_is_clean())
       , {spawn, ?_test(?debugVal(t_simple_reg_or_locate()))}
       , ?_test(t_is_clean())
       , {spawn, ?_test(?debugVal(t_reg_or_locate2()))}
@@ -178,6 +184,52 @@ t_simple_reg_other() ->
             ok
     end.
 
+t_simple_ensure() ->
+    P = self(),
+    Key = {n,l,name},
+    ?assert(gproc:reg(Key) =:= true),
+    ?assert(gproc:where(Key) =:= P),
+    ?assert(gproc:ensure_reg(Key, new_val, [{a,1}]) =:= updated),
+    ?assert(gproc:get_attributes(Key) =:= [{a,1}]),
+    ?assert(gproc:unreg(Key) =:= true),
+    ?assert(gproc:where(Key) =:= undefined).
+
+t_simple_ensure_other() ->
+    P = self(),
+    Key = {n,l,name},
+    P1 = spawn_link(fun() ->
+                            receive
+                                {P, goodbye} -> ok
+                            end
+                    end),
+    Ref = erlang:monitor(process, P1),
+    ?assert(gproc:reg_other(Key, P1) =:= true),
+    ?assert(gproc:where(Key) =:= P1),
+    ?assert(gproc:ensure_reg_other(Key, P1, new_val, [{a,1}]) =:= updated),
+    ?assert(gproc:get_value(Key, P1) =:= new_val),
+    ?assert(gproc:get_attributes(Key, P1) =:= [{a,1}]),
+    ?assert(gproc:unreg_other(Key, P1) =:= true),
+    ?assert(gproc:where({n,l,name}) =:= undefined),
+    P1 ! {self(), goodbye},
+    receive
+        {'DOWN', Ref, _, _, _} ->
+            ok
+    end.
+
+t_simple_ensure_prop() ->
+    Key = {p,l,?LINE},
+    P = self(),
+    Select = fun() ->
+                     gproc:select({l,p}, [{ {Key,'_','_'},[],['$_'] }])
+             end,
+    ?assertMatch(new, gproc:ensure_reg(Key, first_val, [])),
+    ?assertMatch([{Key,P,first_val}], Select()),
+    ?assertMatch(updated, gproc:ensure_reg(Key, new_val, [{a,1}])),
+    ?assertMatch([{Key,P,new_val}], Select()),
+    ?assertMatch([{a,1}], gproc:get_attributes(Key)),
+    ?assertMatch(true, gproc:unreg(Key)),
+    ?assertMatch([], Select()).
+
 t_simple_reg_or_locate() ->
     P = self(),
     ?assertMatch({P, undefined}, gproc:reg_or_locate({n,l,name})),
@@ -412,8 +464,10 @@ t_is_clean() ->
     sys:get_status(gproc_monitor),
     T = ets:tab2list(gproc),
     Tm = ets:tab2list(gproc_monitor),
+    ?debugFmt("self() = ~p~n", [self()]),
     ?assertMatch([], Tm),
-    ?assertMatch([], T -- [{{whereis(gproc_monitor), l}}]).
+    ?assertMatch([], T -- [{{whereis(gproc_monitor), l}},
+                           {{self(), l}}]).
 
 t_simple_mreg() ->
     P = self(),