<commit_before>%% -------------------------------------------------------------------
%%
%% riak_kv_bitcask_backend: Bitcask Driver for Riak
%%
%% Copyright (c) 2007-2010 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_bitcask_backend).
-behavior(riak_kv_backend).

%% KV Backend API
-export([api_version/0,
         capabilities/1,
         capabilities/2,
         start/2,
         stop/1,
         get/3,
         put/5,
         delete/4,
         drop/1,
         fold_buckets/4,
         fold_keys/4,
         fold_objects/4,
         is_empty/1,
         status/1,
         callback/3]).

-export([data_size/1]).

%% Helper API
-export([key_counts/0,
         key_counts/1]).

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-endif.

-include_lib("bitcask/include/bitcask.hrl").

-define(MERGE_CHECK_INTERVAL, timer:minutes(3)).
-define(API_VERSION, 1).
-define(CAPABILITIES, [async_fold,size]).

%% must not be 131, otherwise will match t2b in error
%% yes, I know that this is horrible.
-define(VERSION_1, 1).
-define(VERSION_BYTE, ?VERSION_1).
-define(CURRENT_KEY_TRANS, fun key_transform_to_1/1).

-record(state, {ref :: reference(),
                data_dir :: string(),
                opts :: [{atom(), term()}],
                partition :: integer(),
                root :: string(),
                key_vsn :: integer()}).

-type state() :: #state{}.
-type config() :: [{atom(), term()}].

%% ===================================================================
%% Public API
%% ===================================================================

%% @doc Return the major version of the
%% current API.
-spec api_version() -> {ok, integer()}.
api_version() ->
    {ok, ?API_VERSION}.

%% @doc Return the capabilities of the backend.
-spec capabilities(state()) -> {ok, [atom()]}.
capabilities(_) ->
    {ok, ?CAPABILITIES}.

%% @doc Return the capabilities of the backend.
-spec capabilities(riak_object:bucket(), state()) -> {ok, [atom()]}.
capabilities(_, _) ->
    {ok, ?CAPABILITIES}.

%% @doc Transformation functions for the keys coming off the disk.
key_transform_to_1(<<?VERSION_1:8,_Rest/binary>> = Key) ->
    Key;
key_transform_to_1(<<131:8,_Rest/binary>> = Key0) ->
    {Bucket, Key} = binary_to_term(Key0),
    BucketSz = size(Bucket),
    <<?VERSION_BYTE, BucketSz:16/integer, 
     Bucket/binary, Key/binary>>.

key_transform_to_0(<<?VERSION_1:8,_Rest/binary>> = Key0) ->
    <<?VERSION_BYTE, BucketSz:16/integer, 
     Bucket:BucketSz/bytes, Key/binary>> = Key0,
    term_to_binary({Bucket, Key});
key_transform_to_0(<<131:8,_Rest/binary>> = Key) ->
    Key.

bk_to_tuple(<<?VERSION_1:8,Sz:16/integer,
             Bucket:Sz/bytes, Key/binary>>) ->
    {Bucket, Key};
bk_to_tuple(<<131:8,_Rest/binary>> = BK) ->
    binary_to_term(BK).

make_bk(0, Bucket, Key) ->
    term_to_binary({Bucket, Key});
make_bk(1, Bucket, Key) ->
    BucketSz = size(Bucket),
    <<?VERSION_BYTE:8, BucketSz:16/integer,
     Bucket/binary, Key/binary>>.

%% @doc Start the bitcask backend
-spec start(integer(), config()) -> {ok, state()} | {error, term()}.
start(Partition, Config0) ->
    {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,
        
    %% Get the data root directory  
    case app_helper:get_prop_or_env(data_root, Config, bitcask) of
        undefined ->
            lager:error("Failed to create bitcask dir: data_root is not set"),
            {error, data_root_unset};
        DataRoot ->
            %% Check if a directory exists for the partition
            PartitionStr = integer_to_list(Partition),
            case get_data_dir(DataRoot, PartitionStr) of
                {ok, DataDir} ->
                    BitcaskOpts = set_mode(read_write, Config),
                    case bitcask:open(filename:join(DataRoot, DataDir), BitcaskOpts) of <commit_msg>Modify the version 1 format to support bucket types<commit_after>%% -------------------------------------------------------------------
%%
%% riak_kv_bitcask_backend: Bitcask Driver for Riak
%%
%% Copyright (c) 2007-2010 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_bitcask_backend).
-behavior(riak_kv_backend).

%% KV Backend API
-export([api_version/0,
         capabilities/1,
         capabilities/2,
         start/2,
         stop/1,
         get/3,
         put/5,
         delete/4,
         drop/1,
         fold_buckets/4,
         fold_keys/4,
         fold_objects/4,
         is_empty/1,
         status/1,
         callback/3]).

-export([data_size/1]).

%% Helper API
-export([key_counts/0,
         key_counts/1]).

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-endif.

-include_lib("bitcask/include/bitcask.hrl").

-define(MERGE_CHECK_INTERVAL, timer:minutes(3)).
-define(API_VERSION, 1).
-define(CAPABILITIES, [async_fold,size]).

%% must not be 131, otherwise will match t2b in error
%% yes, I know that this is horrible.
-define(VERSION_1, 1).
-define(VERSION_BYTE, ?VERSION_1).
-define(CURRENT_KEY_TRANS, fun key_transform_to_1/1).

-record(state, {ref :: reference(),
                data_dir :: string(),
                opts :: [{atom(), term()}],
                partition :: integer(),
                root :: string(),
                key_vsn :: integer()}).

-type state() :: #state{}.
-type config() :: [{atom(), term()}].

%% ===================================================================
%% Public API
%% ===================================================================

%% @doc Return the major version of the
%% current API.
-spec api_version() -> {ok, integer()}.
api_version() ->
    {ok, ?API_VERSION}.

%% @doc Return the capabilities of the backend.
-spec capabilities(state()) -> {ok, [atom()]}.
capabilities(_) ->
    {ok, ?CAPABILITIES}.

%% @doc Return the capabilities of the backend.
-spec capabilities(riak_object:bucket(), state()) -> {ok, [atom()]}.
capabilities(_, _) ->
    {ok, ?CAPABILITIES}.

%% @doc Transformation functions for the keys coming off the disk.
key_transform_to_1(<<?VERSION_1:7, _:1, _Rest/binary>> = Key) ->
    Key;
key_transform_to_1(<<131:8,_Rest/bits>> = Key0) ->
    {Bucket, Key} = binary_to_term(Key0),
    make_bk(?VERSION_BYTE, Bucket, Key).

key_transform_to_0(<<?VERSION_1:7,_Rest/bits>> = Key0) ->
    term_to_binary(bk_to_tuple(Key0));
key_transform_to_0(<<131:8,_Rest/binary>> = Key) ->
    Key.

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) ->
    term_to_binary({Bucket, Key});
make_bk(1, {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_bk(1, 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) ->
    {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,
        
    %% Get the data root directory  
    case app_helper:get_prop_or_env(data_root, Config, bitcask) of
        undefined ->
            lager:error("Failed to create bitcask dir: data_root is not set"),
            {error, data_root_unset};
        DataRoot ->
            %% Check if a directory exists for the partition
            PartitionStr = integer_to_list(Partition),
            case get_data_dir(DataRoot, PartitionStr) of
                {ok, DataDir} ->
                    BitcaskOpts = set_mode(read_write, Config),
                    case bitcask:open(filename:join(DataRoot, DataDir), BitcaskOpts) of 
{'fraction_non_alphanumeric': 0.18978890525282277, 'fraction_numerical': 0.01953853706431026, 'mean_word_length': 2.5878830574145826, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 84, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1859728', 'n_tokens_mistral': 3316, 'n_tokens_neox': 3172, 'n_words': 789}