bugfix: to stop riak_kv_vnode crashing on a node shutdown

rebar.config

@@ -30,7 +30,7 @@
         {sidejob, ".*", {git, "git://github.com/basho/sidejob.git", {tag, "2.0.2"}}},
         {erlang_js, ".*", {git, "git://github.com/basho/erlang_js.git", {tag, "1.3.0"}}},
         {eper, ".*", {git, "git://github.com/basho/eper.git", {tag, "0.97.6p1"}}},
-        {bitcask, ".*", {git, "git://github.com/bet365/bitcask.git", {tag, "2.2.7c"}}},
+        {bitcask, ".*", {git, "git://github.com/bet365/bitcask.git", {tag, "2.2.9"}}},
         {sext, ".*", {git, "git://github.com/basho/sext.git", {tag, "1.1p3"}}},
         {riak_pipe, ".*", {git, "git://github.com/bet365/riak_pipe.git", {tag, "2.2.7"}}},
         {riak_dt, ".*", {git, "git://github.com/basho/riak_dt.git", {tag, "2.1.3-225"}}},

src/riak_kv_bitcask_backend.erl

@@ -42,8 +42,8 @@
          callback/3]).
 
 -export([data_size/1,
-         key_transform_to_1/1,
-         key_transform_to_2/1]).
+         encode_disk_key/2,
+         decode_disk_key/1]).
 
 -ifdef(TEST).
 -include_lib("eunit/include/eunit.hrl").
@@ -63,10 +63,15 @@
 
 %% must not be 131, otherwise will match t2b in error
 %% yes, I know that this is horrible.
+-define(VERSION_0, 131).
 -define(VERSION_1, 1).
 -define(VERSION_2, 2).
--define(VERSION_BYTE, ?VERSION_2).
--define(CURRENT_KEY_TRANS, fun key_transform_to_2/1).
+
+-define(CURRENT_VERSION, ?VERSION_2).
+-define(ENCODE_DISK_KEY_FUN, fun encode_disk_key/2).
+-define(DECODE_DISK_KEY_FUN, fun decode_disk_key/1).
+
+
 
 -record(state, {ref :: reference(),
                 data_dir :: string(),
@@ -100,126 +105,19 @@ capabilities(_) ->
 capabilities(_, _) ->
     {ok, ?CAPABILITIES}.
 
-%% @doc Transformation functions for the keys coming off the disk.
-key_transform_to_2(<<?VERSION_BYTE:7, Type:1, TstampExpire:32/integer, Rest/binary>>) ->
-    {<<?VERSION_BYTE:7, Type:1, Rest/binary>>, #keymeta{tstamp_expire = TstampExpire}};
-key_transform_to_2(<<?VERSION_1:7, 0:1, BucketSz:16/integer, Bucket:BucketSz/binary, Key/binary>>) ->
-    {make_kd(?VERSION_BYTE, Bucket, Key), #keymeta{tstamp_expire = ?DEFAULT_TSTAMP_EXPIRE}};
-key_transform_to_2(<<?VERSION_1:7, 1:1, TypeSz:16/integer, Type:TypeSz/binary, BucketSz:16/integer, Bucket:BucketSz/binary, Key/binary>>) ->
-    {make_kd(?VERSION_BYTE, {Type, Bucket}, Key), #keymeta{tstamp_expire = ?DEFAULT_TSTAMP_EXPIRE}};
-key_transform_to_2(<<131:8,_Rest/bits>> = Key0) ->
-    {Bucket, Key} = binary_to_term(Key0),
-    {make_kd(?VERSION_BYTE, Bucket, Key), #keymeta{tstamp_expire = ?DEFAULT_TSTAMP_EXPIRE}}.
-
-
-
-key_transform_to_1(<<?VERSION_2:7, Type:1/integer, _:32, Rest/binary>>) ->
-    {<<?VERSION_1:7, Type:1, Rest/binary>>, #keymeta{}};
-key_transform_to_1(<<?VERSION_1:7, _:1, _Rest/binary>> = Key) ->
-    {Key, #keymeta{}};
-key_transform_to_1(<<131:8,_Rest/bits>> = Key0) ->
-    {Bucket, Key} = binary_to_term(Key0),
-    {make_kd(?VERSION_1, Bucket, Key), #keymeta{}}.
-
-key_transform_to_0(<<?VERSION_2:7,_Rest/bits>> = Key0) ->
-    {term_to_binary(bk_to_tuple_0(Key0)), #keymeta{}};
-key_transform_to_0(<<?VERSION_1:7,_Rest/bits>> = Key0) ->
-    {term_to_binary(bk_to_tuple(Key0)), #keymeta{}};
-key_transform_to_0(<<131:8,_Rest/binary>> = Key) ->
-    {Key, #keymeta{}}.
-
-bk_to_tuple_0(<<?VERSION_2:7, HasType:1, _:32, Sz:16/integer,
-             TypeOrBucket:Sz/bytes, Rest/binary>>) ->
-    case HasType of
-        0 ->
-            %% no type, first field is bucket
-            {TypeOrBucket, Rest};
-        1 ->
-            %% has a tyoe, extract bucket as well
-            <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
-            {{TypeOrBucket, Bucket}, Key}
-    end.
-
-bk_to_tuple(<<?VERSION_2:7, HasType:1, Sz:16/integer,
-             TypeOrBucket:Sz/bytes, Rest/binary>>) ->
-    case HasType of
-        0 ->
-            %% no type, first field is bucket
-            {TypeOrBucket, Rest};
-        1 ->
-            %% has a tyoe, extract bucket as well
-            <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
-            {{TypeOrBucket, Bucket}, Key}
-    end;
-bk_to_tuple(<<?VERSION_1:7, HasType:1, Sz:16/integer,
-             TypeOrBucket:Sz/bytes, Rest/binary>>) ->
-    case HasType of
-        0 ->
-            %% no type, first field is bucket
-            {TypeOrBucket, Rest};
-        1 ->
-            %% has a tyoe, extract bucket as well
-            <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
-            {{TypeOrBucket, Bucket}, Key}
-    end;
-bk_to_tuple(<<131:8,_Rest/binary>> = BK) ->
-    binary_to_term(BK).
-
-make_bk(0, Bucket, Key, _TstampExpire) ->
-    term_to_binary({Bucket, Key});
-make_bk(1, {Type, Bucket}, Key, _TstampExpire) ->
-    TypeSz = size(Type),
-    BucketSz = size(Bucket),
-    <<?VERSION_1:7, 1:1, TypeSz:16/integer, Type/binary,
-      BucketSz:16/integer, Bucket/binary, Key/binary>>;
-make_bk(1, Bucket, Key, _TstampExpire) ->
-    BucketSz = size(Bucket),
-    <<?VERSION_1:7, 0:1, BucketSz:16/integer,
-     Bucket/binary, Key/binary>>;
-make_bk(2, {Type, Bucket}, Key, TstampExpire) ->
-    TypeSz = size(Type),
-    BucketSz = size(Bucket),
-    <<?VERSION_BYTE:7, 1:1, TstampExpire:32/integer, TypeSz:16/integer, Type/binary, 
-      BucketSz:16/integer, Bucket/binary, Key/binary>>;
-make_bk(2, Bucket, Key, TstampExpire) ->
-    BucketSz = size(Bucket),
-    <<?VERSION_BYTE:7, 0:1, TstampExpire:32/integer, BucketSz:16/integer,
-     Bucket/binary, Key/binary>>.
-
-make_kd(0, Bucket, Key) ->
-    term_to_binary({Bucket, Key});
-make_kd(1, {Type, Bucket}, Key) ->
-    TypeSz = size(Type),
-    BucketSz = size(Bucket),
-    <<?VERSION_1:7, 1:1, TypeSz:16/integer, Type/binary,
-      BucketSz:16/integer, Bucket/binary, Key/binary>>;
-make_kd(1, Bucket, Key) ->
-    BucketSz = size(Bucket),
-    <<?VERSION_1:7, 0:1, BucketSz:16/integer,
-     Bucket/binary, Key/binary>>;
-make_kd(2, {Type, Bucket}, Key) ->
-    TypeSz = size(Type),
-    BucketSz = size(Bucket),
-    <<?VERSION_BYTE:7, 1:1, TypeSz:16/integer, Type/binary, BucketSz:16/integer, Bucket/binary, Key/binary>>;
-make_kd(2, Bucket, Key) ->
-    BucketSz = size(Bucket),
-    <<?VERSION_BYTE:7, 0:1, BucketSz:16/integer, Bucket/binary, Key/binary>>.
-
 %% @doc Start the bitcask backend
 -spec start(integer(), config()) -> {ok, state()} | {error, term()}.
 start(Partition, Config0) ->
     random:seed(erlang:now()),
-    {Config, KeyVsn} =
-        case app_helper:get_prop_or_env(small_keys, Config0, bitcask) of
-            false ->
-                C0 = proplists:delete(small_keys, Config0),
-                C1 = C0 ++ [{key_transform, fun key_transform_to_0/1}],
-                {C1, 0};
-            _ ->
-                C0 = proplists:delete(small_keys, Config0),
-                C1 = C0 ++ [{key_transform, ?CURRENT_KEY_TRANS}],
-                {C1, ?VERSION_BYTE}
-        end,
+
+    BaseConfig = proplists:delete(small_keys, Config0),
+    KeyOpts =
+        [
+            {encode_disk_key_fun, ?ENCODE_DISK_KEY_FUN},
+            {decode_disk_key_fun, ?DECODE_DISK_KEY_FUN}
+        ],
+    Config = BaseConfig ++ KeyOpts,
+    KeyVsn = ?CURRENT_VERSION,
 
     %% Get the data root directory
     case app_helper:get_prop_or_env(data_root, Config, bitcask) of
@@ -274,7 +172,7 @@ stop(#state{ref=Ref}) ->
                  {error, not_found, state()} |
                  {error, term(), state()}.
 get(Bucket, Key, #state{ref=Ref, key_vsn=KVers}=State) ->
-    BitcaskKey = make_kd(KVers, Bucket, Key),
+    BitcaskKey = make_bitcask_key(KVers, Bucket, Key),
     case bitcask:get(Ref, BitcaskKey) of
         {ok, Value} ->
             {ok, Value, State};
@@ -298,11 +196,10 @@ get(Bucket, Key, #state{ref=Ref, key_vsn=KVers}=State) ->
 -spec put(riak_object:bucket(), riak_object:key(), [index_spec()], binary(), integer(), state()) ->
                  {ok, state()} |
                  {error, term(), state()}.
-put(Bucket, PrimaryKey, _IndexSpecs, Val, TstampExpire,
-    #state{ref=Ref, key_vsn=KeyVsn}=State) ->
-    BitcaskKey = make_bk(KeyVsn, Bucket, PrimaryKey, TstampExpire),
-    {KeyDirKey, _} = ?CURRENT_KEY_TRANS(BitcaskKey),
-    case bitcask:put(Ref, KeyDirKey, BitcaskKey, Val, TstampExpire) of
+put(Bucket, PrimaryKey, _IndexSpecs, Val, TstampExpire, #state{ref=Ref, key_vsn=KeyVsn}=State) ->
+    BitcaskKey = make_bitcask_key(KeyVsn, Bucket, PrimaryKey),
+    Opts = [{?TSTAMP_EXPIRE_KEY, TstampExpire}],
+    case bitcask:put(Ref, BitcaskKey, Val, Opts) of
         ok ->
             {ok, State};
         {error, Reason} ->
@@ -312,11 +209,9 @@ put(Bucket, PrimaryKey, _IndexSpecs, Val, TstampExpire,
 -spec put(riak_object:bucket(), riak_object:key(), [index_spec()], binary(), state()) ->
                  {ok, state()} |
                  {error, term(), state()}.
-put(Bucket, PrimaryKey, _IndexSpecs, Val,
-    #state{ref=Ref, key_vsn=KeyVsn}=State) ->
-    BitcaskKey = make_bk(KeyVsn, Bucket, PrimaryKey, ?DEFAULT_TSTAMP_EXPIRE),
-    {KeyDirKey, _} = ?CURRENT_KEY_TRANS(BitcaskKey),
-    case bitcask:put(Ref, KeyDirKey, BitcaskKey, Val, ?DEFAULT_TSTAMP_EXPIRE) of
+put(Bucket, PrimaryKey, _IndexSpecs, Val, #state{ref=Ref, key_vsn=KeyVsn}=State) ->
+    BitcaskKey = make_bitcask_key(KeyVsn, Bucket, PrimaryKey),
+    case bitcask:put(Ref, BitcaskKey, Val) of
         ok ->
             {ok, State};
         {error, Reason} ->
@@ -329,11 +224,9 @@ put(Bucket, PrimaryKey, _IndexSpecs, Val,
 %% is ignored.
 -spec delete(riak_object:bucket(), riak_object:key(), [index_spec()], state()) ->
                     {ok, state()}.
-delete(Bucket, Key, _IndexSpecs,
-       #state{ref=Ref, key_vsn=KeyVsn}=State) ->
-    BitcaskKey = make_bk(KeyVsn, Bucket, Key, ?DEFAULT_TSTAMP_EXPIRE),
-    {KeyDirKey, _} = ?CURRENT_KEY_TRANS(BitcaskKey),
-    ok = bitcask:delete(Ref, KeyDirKey, BitcaskKey),
+delete(Bucket, Key, _IndexSpecs, #state{ref=Ref, key_vsn=KeyVsn}=State) ->
+    BitcaskKey = make_bitcask_key(KeyVsn, Bucket, Key),
+    ok = bitcask:delete(Ref, BitcaskKey),
     {ok, State}.
 
 %% @doc Fold over all the buckets.
@@ -549,6 +442,107 @@ callback(Ref,
 callback(_Ref, _Msg, State) ->
     {ok, State}.
 
+%% ===================================================================
+%% Key Encode and Decode Functions
+%% ===================================================================
+
+%% Make the bitcask keydir key, based on the version number we are using and the bucket, key from riak
+make_bitcask_key(0, Bucket, Key) ->
+    term_to_binary({Bucket, Key});
+
+make_bitcask_key(1, {Type, Bucket}, Key) ->
+    TypeSz = size(Type),
+    BucketSz = size(Bucket),
+    <<?VERSION_1:7, 1:1, TypeSz:16/integer, Type/binary, BucketSz:16/integer, Bucket/binary, Key/binary>>;
+
+make_bitcask_key(1, Bucket, Key) ->
+    BucketSz = size(Bucket),
+    <<?VERSION_1:7, 0:1, BucketSz:16/integer, Bucket/binary, Key/binary>>;
+
+make_bitcask_key(2, {Type, Bucket}, Key) ->
+    TypeSz = size(Type),
+    BucketSz = size(Bucket),
+    <<?VERSION_2:7, 1:1, TypeSz:16/integer, Type/binary, BucketSz:16/integer, Bucket/binary, Key/binary>>;
+
+make_bitcask_key(2, Bucket, Key) ->
+    BucketSz = size(Bucket),
+    <<?VERSION_2:7, 0:1, BucketSz:16/integer, Bucket/binary, Key/binary>>.
+
+
+
+
+%% Make the riak key from the bitcask key, based on the version number
+make_riak_key(<<?VERSION_2:7, HasType:1, Sz:16/integer,
+    TypeOrBucket:Sz/bytes, Rest/binary>>) ->
+    case HasType of
+        0 ->
+            %% no type, first field is bucket
+            {TypeOrBucket, Rest};
+        1 ->
+            %% has a tyoe, extract bucket as well
+            <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
+            {{TypeOrBucket, Bucket}, Key}
+    end;
+make_riak_key(<<?VERSION_1:7, HasType:1, Sz:16/integer,
+    TypeOrBucket:Sz/bytes, Rest/binary>>) ->
+    case HasType of
+        0 ->
+            %% no type, first field is bucket
+            {TypeOrBucket, Rest};
+        1 ->
+            %% has a tyoe, extract bucket as well
+            <<BucketSz:16/integer, Bucket:BucketSz/bytes, Key/binary>> = Rest,
+            {{TypeOrBucket, Bucket}, Key}
+    end;
+make_riak_key(<<?VERSION_0:8,_Rest/binary>> = BK) ->
+    binary_to_term(BK).
+
+
+
+
+%% Take the bitcask key stored down to disk, and decode it to match the keydir stored in memory
+%% return the keyinfo record, to find additional information that was stored down in the encoding
+decode_disk_key(<<?VERSION_2:7, Type:1, TstampExpire:32/integer, Rest/binary>>) ->
+    #keyinfo{
+        key = <<?CURRENT_VERSION:7, Type:1, Rest/binary>>,
+        tstamp_expire = TstampExpire
+    };
+
+decode_disk_key(<<?VERSION_1:7, 0:1, BucketSz:16/integer, Bucket:BucketSz/binary, Key/binary>>) ->
+    #keyinfo{
+        key = make_bitcask_key(?CURRENT_VERSION, Bucket, Key)
+    };
+
+decode_disk_key(<<?VERSION_1:7, 1:1, TypeSz:16/integer, Type:TypeSz/binary, BucketSz:16/integer, Bucket:BucketSz/binary, Key/binary>>) ->
+    #keyinfo{
+        key = make_bitcask_key(?CURRENT_VERSION, {Type, Bucket}, Key)
+    };
+
+decode_disk_key(<<?VERSION_0:8,_Rest/bits>> = Key0) ->
+    {Bucket, Key} = binary_to_term(Key0),
+    #keyinfo{
+        key = make_bitcask_key(?CURRENT_VERSION, Bucket, Key)
+    }.
+
+
+
+
+%% Take the bitcask key generated by make_bitask_key, and encode it with additional information to be stored down to
+%% disk, but not stored in the keydir.
+encode_disk_key(<<?VERSION_2:7, Type:1, Rest/bits>>, Opts) ->
+    TstampExpire = proplists:get_value(?TSTAMP_EXPIRE_KEY, Opts, ?DEFAULT_TSTAMP_EXPIRE),
+    <<?VERSION_2:7, Type:1, TstampExpire:32/integer, Rest/bits>>;
+
+encode_disk_key(<<?VERSION_1:7, _Rest/bits>> = BitcaskKey, _Opts) ->
+    BitcaskKey;
+
+encode_disk_key(<<?VERSION_0:8,_Rest/bits>> = BitcaskKey, _Opts) ->
+    BitcaskKey.
+
+
+
+
+
 %% ===================================================================
 %% Internal functions
 %% ===================================================================
@@ -593,7 +587,7 @@ check_fcntl() ->
 %% Return a function to fold over the buckets on this backend
 fold_buckets_fun(FoldBucketsFun) ->
     fun(#bitcask_entry{key=BK}, {Acc, BucketSet}) ->
-            {Bucket, _} = bk_to_tuple(BK),
+            {Bucket, _} = make_riak_key(BK),
             case sets:is_element(Bucket, BucketSet) of
                 true ->
                     {Acc, BucketSet};
@@ -607,12 +601,12 @@ fold_buckets_fun(FoldBucketsFun) ->
 %% Return a function to fold over keys on this backend
 fold_keys_fun(FoldKeysFun, undefined) ->
     fun(#bitcask_entry{key=BK}, Acc) ->
-            {Bucket, Key} = bk_to_tuple(BK),
+            {Bucket, Key} = make_riak_key(BK),
             FoldKeysFun(Bucket, Key, Acc)
     end;
 fold_keys_fun(FoldKeysFun, Bucket) ->
     fun(#bitcask_entry{key=BK}, Acc) ->
-            {B, Key} = bk_to_tuple(BK),
+            {B, Key} = make_riak_key(BK),
             case B =:= Bucket of
                 true ->
                     FoldKeysFun(B, Key, Acc);
@@ -625,12 +619,12 @@ fold_keys_fun(FoldKeysFun, Bucket) ->
 %% Return a function to fold over keys on this backend
 fold_objects_fun(FoldObjectsFun, undefined) ->
     fun(BK, Value, Acc) ->
-            {Bucket, Key} = bk_to_tuple(BK),
+            {Bucket, Key} = make_riak_key(BK),
             FoldObjectsFun(Bucket, Key, Value, Acc)
     end;
 fold_objects_fun(FoldObjectsFun, Bucket) ->
     fun(BK, Value, Acc) ->
-            {B, Key} = bk_to_tuple(BK),
+            {B, Key} = make_riak_key(BK),
             case B =:= Bucket of
                 true ->
                     FoldObjectsFun(B, Key, Value, Acc);