Add support for conditional postcommit hooks

src/riak_kv_bucket.erl

@@ -40,6 +40,8 @@
 -type props() :: [prop()].
 -type errors() :: [error()].
 
+-export_type([props/0]).
+
 %% @doc called by riak_core in a few places to ensure bucket
 %%  properties are sane. The arguments combinations have the following
 %%  meanings:-

src/riak_kv_hooks.erl

@@ -0,0 +1,118 @@
+%% -------------------------------------------------------------------
+%%
+%% Copyright (c) 2012-2013 Basho Technologies, Inc.  All Rights Reserved.
+%%
+%% This file is provided to you 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(riak_kv_hooks).
+
+%% API
+-export([add_conditional_postcommit/1,
+         del_conditional_postcommit/1,
+         get_conditional_postcommit/2]).
+
+%% Exported for internal use by `riak_kv_sup'
+-export([create_table/0]).
+
+%% Types
+-type hook()      :: {module(), atom()}.
+-type hook_type() :: conditional_postcommit.
+-type bucket() :: riak_object:bucket().
+-type key()    :: riak_object:key().
+-type bucket_props() :: riak_kv_bucket:props().
+
+%%%===================================================================
+
+%% @doc
+%% Called by {@link riak_kv_sup} to create the public ETS table used to
+%% track registered hooks. Having `riak_kv_sup' own the table ensures
+%% that the table exists aslong as riak_kv is running.
+-spec create_table() -> ok.
+create_table() ->
+    ?MODULE = ets:new(?MODULE, [named_table, public, bag,
+                                {write_concurrency, true},
+                                {read_concurrency, true}]),
+    restore_state(),
+    ok.
+
+%% @doc
+%% Add a global conditional postcommit hook that is called for each
+%% PUT operation. The hook is of the form `{Module, Fun}'. The specified
+%% function is called with the relevent bucket, key, and bucket properties
+%% at the time of the PUT operation and is expected to return `false' or
+%% a normal postcommit hook specification that should be invoked.
+-spec add_conditional_postcommit(hook()) -> ok.
+add_conditional_postcommit(Hook) ->
+    add_hook(conditional_postcommit, Hook).
+
+%% @doc Remove a previously registered conditional postcommit hook
+-spec del_conditional_postcommit(hook()) -> ok.
+del_conditional_postcommit(Hook) ->
+    del_hook(conditional_postcommit, Hook).
+
+%% @doc
+%% This function invokes each registered conditional postcommit
+%% hook. Each hook will return either `false' or a list of active
+%% hooks. This function then returns the combined list of active hooks.
+-spec get_conditional_postcommit({bucket(), key()}, bucket_props()) -> [any()].
+get_conditional_postcommit({{BucketType, Bucket}, _Key}, BucketProps) ->
+    Hooks = get_hooks(conditional_postcommit),
+    ActiveHooks =
+        [ActualHook || {Mod, Fun} <- Hooks,
+                       ActualHook <- [Mod:Fun(BucketType, Bucket, BucketProps)],
+                       ActualHook =/= false],
+    lists:flatten(ActiveHooks);
+get_conditional_postcommit(_BKey, _BucketProps) ->
+    %% For now, we only support typed buckets.
+    [].
+
+%%%===================================================================
+
+-spec add_hook(hook_type(), hook()) -> ok.
+add_hook(Type, Hook) ->
+    ets:insert(?MODULE, {Type, Hook}),
+    save_state(),
+    ok.
+
+-spec del_hook(hook_type(), hook()) -> ok.
+del_hook(Type, Hook) ->
+    ets:delete_object(?MODULE, {Type, Hook}),
+    save_state(),
+    ok.
+
+-spec get_hooks(hook_type()) -> [hook()].
+get_hooks(Type) ->
+    [Hook || {_, Hook} <- ets:lookup(?MODULE, Type)].
+
+%% Backup the current ETS state to the application environment just in case
+%% riak_kv_sup dies and the ETS table is lost.
+-spec save_state() -> ok.
+save_state() ->
+    Hooks = ets:tab2list(?MODULE),
+    ok = application:set_env(riak_kv, riak_kv_hooks, Hooks, infinity),
+    ok.
+
+%% Restore registered hooks in the unlikely case that riak_kv_sup died and
+%% the ETS table was lost/recreated.
+-spec restore_state() -> ok.
+restore_state() ->
+    case application:get_env(riak_kv, riak_kv_hooks) of
+        undefined ->
+            ok;
+        {ok, Hooks} ->
+            true = ets:insert_new(?MODULE, Hooks),
+            ok
+    end.

src/riak_kv_put_fsm.erl

@@ -415,7 +415,7 @@ validate(timeout, StateData0 = #state{from = {raw, ReqId, _Pid},
                 end,
             Postcommit =
                 if Disable -> [];
-                   true -> get_hooks(postcommit, BucketProps)
+                   true -> get_hooks(postcommit, BucketProps, StateData0)
                 end,
             StateData1 = StateData0#state{n=N,
                                           w=W,
@@ -893,6 +893,11 @@ get_hooks(HookType, BucketProps) ->
             Hooks
     end.
 
+get_hooks(postcommit, BucketProps, #state{bkey=BKey}) ->
+    BaseHooks = get_hooks(postcommit, BucketProps),
+    CondHooks = riak_kv_hooks:get_conditional_postcommit(BKey, BucketProps),
+    BaseHooks ++ (CondHooks -- BaseHooks).
+
 get_option(Name, Options) ->
     get_option(Name, Options, undefined).
 

src/riak_kv_sup.erl

@@ -44,6 +44,7 @@ init([]) ->
     catch dtrace:init(),                   % NIF load trigger (R14B04)
     catch dyntrace:p(),                    % NIF load trigger (R15B01+)
     riak_kv_entropy_info:create_table(),
+    riak_kv_hooks:create_table(),
     VMaster = {riak_kv_vnode_master,
                {riak_core_vnode_master, start_link,
                 [riak_kv_vnode, riak_kv_legacy_vnode, riak_kv]},