1% Licensed under the Apache License, Version 2.0 (the "License"); you may not
2% use this file except in compliance with the License. You may obtain a copy of
3% the License at
4%
5%   http://www.apache.org/licenses/LICENSE-2.0
6%
7% Unless required by applicable law or agreed to in writing, software
8% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10% License for the specific language governing permissions and limitations under
11% the License.
12
13-module(couch_btree_tests).
14
15-include_lib("couch/include/couch_eunit.hrl").
16-include_lib("couch/include/couch_db.hrl").
17
18-define(ROWS, 1000).
19-define(TIMEOUT, 60). % seconds
20
21
22setup() ->
23    {ok, Fd} = couch_file:open(?tempfile(), [create, overwrite]),
24    {ok, Btree} = couch_btree:open(nil, Fd, [{compression, none},
25                                             {reduce, fun reduce_fun/2}]),
26    {Fd, Btree}.
27
28setup_kvs(_) ->
29    setup().
30
31setup_red() ->
32    {_, EvenOddKVs} = lists:foldl(
33        fun(Idx, {Key, Acc}) ->
34            case Key of
35                "even" -> {"odd", [{{Key, Idx}, 1} | Acc]};
36                _ -> {"even", [{{Key, Idx}, 1} | Acc]}
37            end
38        end, {"odd", []}, lists:seq(1, ?ROWS)),
39    {Fd, Btree} = setup(),
40    {ok, Btree1} = couch_btree:add_remove(Btree, EvenOddKVs, []),
41    {Fd, Btree1}.
42setup_red(_) ->
43    setup_red().
44
45teardown(Fd) when is_pid(Fd) ->
46    ok = couch_file:close(Fd);
47teardown({Fd, _}) ->
48    teardown(Fd).
49teardown(_, {Fd, _}) ->
50    teardown(Fd).
51
52
53kvs_test_funs() ->
54    [
55        fun should_set_fd_correctly/2,
56        fun should_set_root_correctly/2,
57        fun should_create_zero_sized_btree/2,
58        fun should_set_reduce_option/2,
59        fun should_fold_over_empty_btree/2,
60        fun should_add_all_keys/2,
61        fun should_continuously_add_new_kv/2,
62        fun should_continuously_remove_keys/2,
63        fun should_insert_keys_in_reversed_order/2,
64        fun should_add_every_odd_key_remove_every_even/2,
65        fun should_add_every_even_key_remove_every_old/2
66    ].
67
68red_test_funs() ->
69    [
70        fun should_reduce_whole_range/2,
71        fun should_reduce_first_half/2,
72        fun should_reduce_second_half/2
73    ].
74
75
76btree_open_test_() ->
77    {ok, Fd} = couch_file:open(?tempfile(), [create, overwrite]),
78    {ok, Btree} = couch_btree:open(nil, Fd, [{compression, none}]),
79    {
80        "Ensure that created btree is really a btree record",
81        ?_assert(is_record(Btree, btree))
82    }.
83
84sorted_kvs_test_() ->
85    Funs = kvs_test_funs(),
86    Sorted = [{Seq, couch_rand:uniform()} || Seq <- lists:seq(1, ?ROWS)],
87    {
88        "BTree with sorted keys",
89        {
90            setup,
91            fun() -> test_util:start(?MODULE, [ioq]) end, fun test_util:stop/1,
92            {
93                foreachx,
94                fun setup_kvs/1, fun teardown/2,
95                [{Sorted, Fun} || Fun <- Funs]
96            }
97        }
98    }.
99
100rsorted_kvs_test_() ->
101    Sorted = [{Seq, couch_rand:uniform()} || Seq <- lists:seq(1, ?ROWS)],
102    Funs = kvs_test_funs(),
103    Reversed = Sorted,
104    {
105        "BTree with backward sorted keys",
106        {
107            setup,
108            fun() -> test_util:start(?MODULE, [ioq]) end, fun test_util:stop/1,
109            {
110                foreachx,
111                fun setup_kvs/1, fun teardown/2,
112                [{Reversed, Fun} || Fun <- Funs]
113            }
114        }
115    }.
116
117shuffled_kvs_test_() ->
118    Funs = kvs_test_funs(),
119    Sorted = [{Seq, couch_rand:uniform()} || Seq <- lists:seq(1, ?ROWS)],
120    Shuffled = shuffle(Sorted),
121    {
122        "BTree with shuffled keys",
123        {
124            setup,
125            fun() -> test_util:start(?MODULE, [ioq]) end, fun test_util:stop/1,
126            {
127                foreachx,
128                fun setup_kvs/1, fun teardown/2,
129                [{Shuffled, Fun} || Fun <- Funs]
130            }
131        }
132    }.
133
134reductions_test_() ->
135    {
136        "BTree reductions",
137        {
138            setup,
139            fun() -> test_util:start(?MODULE, [ioq]) end, fun test_util:stop/1,
140            [
141                {
142                    "Common tests",
143                    {
144                        foreach,
145                        fun setup_red/0, fun teardown/1,
146                        [
147                            fun should_reduce_without_specified_direction/1,
148                            fun should_reduce_forward/1,
149                            fun should_reduce_backward/1
150                        ]
151                    }
152                },
153                {
154                    "Range requests",
155                    [
156                        {
157                            "Forward direction",
158                            {
159                                foreachx,
160                                fun setup_red/1, fun teardown/2,
161                                [{fwd, F} || F <- red_test_funs()]
162                            }
163                        },
164                        {
165                            "Backward direction",
166                            {
167                                foreachx,
168                                fun setup_red/1, fun teardown/2,
169                                [{rev, F} || F <- red_test_funs()]
170                            }
171                        }
172                    ]
173                }
174            ]
175        }
176    }.
177
178
179should_set_fd_correctly(_, {Fd, Btree}) ->
180    ?_assertMatch(Fd, Btree#btree.fd).
181
182should_set_root_correctly(_, {_, Btree}) ->
183    ?_assertMatch(nil, Btree#btree.root).
184
185should_create_zero_sized_btree(_, {_, Btree}) ->
186    ?_assertMatch(0, couch_btree:size(Btree)).
187
188should_set_reduce_option(_, {_, Btree}) ->
189    ReduceFun = fun reduce_fun/2,
190    Btree1 = couch_btree:set_options(Btree, [{reduce, ReduceFun}]),
191    ?_assertMatch(ReduceFun, Btree1#btree.reduce).
192
193should_fold_over_empty_btree(_, {_, Btree}) ->
194    {ok, _, EmptyRes} = couch_btree:foldl(Btree, fun(_, X) -> {ok, X+1} end, 0),
195    ?_assertEqual(EmptyRes, 0).
196
197should_add_all_keys(KeyValues, {Fd, Btree}) ->
198    {ok, Btree1} = couch_btree:add_remove(Btree, KeyValues, []),
199    [
200        should_return_complete_btree_on_adding_all_keys(KeyValues, Btree1),
201        should_have_non_zero_size(Btree1),
202        should_have_lesser_size_than_file(Fd, Btree1),
203        should_keep_root_pointer_to_kp_node(Fd, Btree1),
204        should_remove_all_keys(KeyValues, Btree1)
205    ].
206
207should_return_complete_btree_on_adding_all_keys(KeyValues, Btree) ->
208    ?_assert(test_btree(Btree, KeyValues)).
209
210should_have_non_zero_size(Btree) ->
211    ?_assert(couch_btree:size(Btree) > 0).
212
213should_have_lesser_size_than_file(Fd, Btree) ->
214    ?_assert((couch_btree:size(Btree) =< couch_file:bytes(Fd))).
215
216should_keep_root_pointer_to_kp_node(Fd, Btree) ->
217    ?_assertMatch({ok, {kp_node, _}},
218                  couch_file:pread_term(Fd, element(1, Btree#btree.root))).
219
220should_remove_all_keys(KeyValues, Btree) ->
221    Keys = keys(KeyValues),
222    {ok, Btree1} = couch_btree:add_remove(Btree, [], Keys),
223    {
224        "Should remove all the keys",
225        [
226            should_produce_valid_btree(Btree1, []),
227            should_be_empty(Btree1)
228        ]
229    }.
230
231should_continuously_add_new_kv(KeyValues, {_, Btree}) ->
232    {Btree1, _} = lists:foldl(
233        fun(KV, {BtAcc, PrevSize}) ->
234            {ok, BtAcc2} = couch_btree:add_remove(BtAcc, [KV], []),
235            ?assert(couch_btree:size(BtAcc2) > PrevSize),
236            {BtAcc2, couch_btree:size(BtAcc2)}
237        end, {Btree, couch_btree:size(Btree)}, KeyValues),
238    {
239        "Should continuously add key-values to btree",
240        [
241            should_produce_valid_btree(Btree1, KeyValues),
242            should_not_be_empty(Btree1)
243        ]
244    }.
245
246should_continuously_remove_keys(KeyValues, {_, Btree}) ->
247    {ok, Btree1} = couch_btree:add_remove(Btree, KeyValues, []),
248    {Btree2, _} = lists:foldl(
249        fun({K, _}, {BtAcc, PrevSize}) ->
250            {ok, BtAcc2} = couch_btree:add_remove(BtAcc, [], [K]),
251            ?assert(couch_btree:size(BtAcc2) < PrevSize),
252            {BtAcc2, couch_btree:size(BtAcc2)}
253        end, {Btree1, couch_btree:size(Btree1)}, KeyValues),
254    {
255        "Should continuously remove keys from btree",
256        [
257            should_produce_valid_btree(Btree2, []),
258            should_be_empty(Btree2)
259        ]
260    }.
261
262should_insert_keys_in_reversed_order(KeyValues, {_, Btree}) ->
263    KeyValuesRev = lists:reverse(KeyValues),
264    {Btree1, _} = lists:foldl(
265        fun(KV, {BtAcc, PrevSize}) ->
266            {ok, BtAcc2} = couch_btree:add_remove(BtAcc, [KV], []),
267            ?assert(couch_btree:size(BtAcc2) > PrevSize),
268            {BtAcc2, couch_btree:size(BtAcc2)}
269        end, {Btree, couch_btree:size(Btree)}, KeyValuesRev),
270    should_produce_valid_btree(Btree1, KeyValues).
271
272should_add_every_odd_key_remove_every_even(KeyValues, {_, Btree}) ->
273    {ok, Btree1} = couch_btree:add_remove(Btree, KeyValues, []),
274    {_, Rem2Keys0, Rem2Keys1} = lists:foldl(fun(X, {Count, Left, Right}) ->
275        case Count rem 2 == 0 of
276            true -> {Count + 1, [X | Left], Right};
277            false -> {Count + 1, Left, [X | Right]}
278        end
279                                            end, {0, [], []}, KeyValues),
280    {timeout, ?TIMEOUT,
281        ?_assert(test_add_remove(Btree1, Rem2Keys0, Rem2Keys1))
282    }.
283
284should_add_every_even_key_remove_every_old(KeyValues, {_, Btree}) ->
285    {ok, Btree1} = couch_btree:add_remove(Btree, KeyValues, []),
286    {_, Rem2Keys0, Rem2Keys1} = lists:foldl(fun(X, {Count, Left, Right}) ->
287        case Count rem 2 == 0 of
288            true -> {Count + 1, [X | Left], Right};
289            false -> {Count + 1, Left, [X | Right]}
290        end
291                                            end, {0, [], []}, KeyValues),
292    {timeout, ?TIMEOUT,
293        ?_assert(test_add_remove(Btree1, Rem2Keys1, Rem2Keys0))
294    }.
295
296
297should_reduce_without_specified_direction({_, Btree}) ->
298    ?_assertMatch(
299        {ok, [{{"odd", _}, ?ROWS div 2}, {{"even", _}, ?ROWS div 2}]},
300        fold_reduce(Btree, [])).
301
302should_reduce_forward({_, Btree}) ->
303    ?_assertMatch(
304        {ok, [{{"odd", _}, ?ROWS div 2}, {{"even", _}, ?ROWS div 2}]},
305        fold_reduce(Btree, [{dir, fwd}])).
306
307should_reduce_backward({_, Btree}) ->
308    ?_assertMatch(
309        {ok, [{{"even", _}, ?ROWS div 2}, {{"odd", _}, ?ROWS div 2}]},
310        fold_reduce(Btree, [{dir, rev}])).
311
312should_reduce_whole_range(fwd, {_, Btree}) ->
313    {SK, EK} = {{"even", 0}, {"odd", ?ROWS - 1}},
314    [
315        {
316            "include endkey",
317            ?_assertMatch(
318                {ok, [{{"odd", 1}, ?ROWS div 2},
319                      {{"even", 2}, ?ROWS div 2}]},
320                fold_reduce(Btree, [{dir, fwd},
321                                    {start_key, SK},
322                                    {end_key, EK}]))
323        },
324        {
325            "exclude endkey",
326            ?_assertMatch(
327                {ok, [{{"odd", 1}, (?ROWS div 2) - 1},
328                      {{"even", 2}, ?ROWS div 2}]},
329                fold_reduce(Btree, [{dir, fwd},
330                                    {start_key, SK},
331                                    {end_key_gt, EK}]))
332        }
333    ];
334should_reduce_whole_range(rev, {_, Btree}) ->
335    {SK, EK} = {{"odd", ?ROWS - 1}, {"even", 2}},
336    [
337        {
338            "include endkey",
339            ?_assertMatch(
340                {ok, [{{"even", ?ROWS}, ?ROWS div 2},
341                      {{"odd", ?ROWS - 1}, ?ROWS div 2}]},
342                fold_reduce(Btree, [{dir, rev},
343                                    {start_key, SK},
344                                    {end_key, EK}]))
345        },
346        {
347            "exclude endkey",
348            ?_assertMatch(
349                {ok, [{{"even", ?ROWS}, (?ROWS div 2) - 1},
350                      {{"odd", ?ROWS - 1}, ?ROWS div 2}]},
351                fold_reduce(Btree, [{dir, rev},
352                                    {start_key, SK},
353                                    {end_key_gt, EK}]))
354        }
355    ].
356
357should_reduce_first_half(fwd, {_, Btree}) ->
358    {SK, EK} = {{"even", 0}, {"odd", (?ROWS div 2) - 1}},
359    [
360        {
361            "include endkey",
362            ?_assertMatch(
363                {ok, [{{"odd", 1}, ?ROWS div 4},
364                      {{"even", 2}, ?ROWS div 2}]},
365                fold_reduce(Btree, [{dir, fwd},
366                                    {start_key, SK}, {end_key, EK}]))
367        },
368        {
369            "exclude endkey",
370            ?_assertMatch(
371                {ok, [{{"odd", 1}, (?ROWS div 4) - 1},
372                      {{"even", 2}, ?ROWS div 2}]},
373                fold_reduce(Btree, [{dir, fwd},
374                                    {start_key, SK},
375                                    {end_key_gt, EK}]))
376        }
377    ];
378should_reduce_first_half(rev, {_, Btree}) ->
379    {SK, EK} = {{"odd", ?ROWS - 1}, {"even", ?ROWS div 2}},
380    [
381        {
382            "include endkey",
383            ?_assertMatch(
384                {ok, [{{"even", ?ROWS}, (?ROWS div 4) + 1},
385                      {{"odd", ?ROWS - 1}, ?ROWS div 2}]},
386                fold_reduce(Btree, [{dir, rev},
387                                    {start_key, SK},
388                                    {end_key, EK}]))
389        },
390        {
391            "exclude endkey",
392            ?_assertMatch(
393                {ok, [{{"even", ?ROWS}, ?ROWS div 4},
394                      {{"odd", ?ROWS - 1}, ?ROWS div 2}]},
395                fold_reduce(Btree, [{dir, rev},
396                                    {start_key, SK},
397                                    {end_key_gt, EK}]))
398        }
399    ].
400
401should_reduce_second_half(fwd, {_, Btree}) ->
402    {SK, EK} = {{"even", ?ROWS div 2}, {"odd", ?ROWS - 1}},
403    [
404        {
405            "include endkey",
406            ?_assertMatch(
407                {ok, [{{"odd", 1}, ?ROWS div 2},
408                      {{"even", ?ROWS div 2}, (?ROWS div 4) + 1}]},
409                fold_reduce(Btree, [{dir, fwd},
410                                    {start_key, SK},
411                                    {end_key, EK}]))
412        },
413        {
414            "exclude endkey",
415            ?_assertMatch(
416                {ok, [{{"odd", 1}, (?ROWS div 2) - 1},
417                      {{"even", ?ROWS div 2}, (?ROWS div 4) + 1}]},
418                fold_reduce(Btree, [{dir, fwd},
419                                    {start_key, SK},
420                                    {end_key_gt, EK}]))
421        }
422    ];
423should_reduce_second_half(rev, {_, Btree}) ->
424    {SK, EK} = {{"odd", (?ROWS div 2) + 1}, {"even", 2}},
425    [
426        {
427            "include endkey",
428            ?_assertMatch(
429                {ok, [{{"even", ?ROWS}, ?ROWS div 2},
430                      {{"odd", (?ROWS div 2) + 1}, (?ROWS div 4) + 1}]},
431                fold_reduce(Btree, [{dir, rev},
432                                    {start_key, SK},
433                                    {end_key, EK}]))
434        },
435        {
436            "exclude endkey",
437            ?_assertMatch(
438                {ok, [{{"even", ?ROWS}, (?ROWS div 2) - 1},
439                      {{"odd", (?ROWS div 2) + 1}, (?ROWS div 4) + 1}]},
440                fold_reduce(Btree, [{dir, rev},
441                                    {start_key, SK},
442                                    {end_key_gt, EK}]))
443        }
444    ].
445
446should_produce_valid_btree(Btree, KeyValues) ->
447    ?_assert(test_btree(Btree, KeyValues)).
448
449should_be_empty(Btree) ->
450    ?_assertEqual(couch_btree:size(Btree), 0).
451
452should_not_be_empty(Btree) ->
453    ?_assert(couch_btree:size(Btree) > 0).
454
455fold_reduce(Btree, Opts) ->
456    GroupFun = fun({K1, _}, {K2, _}) ->
457        K1 == K2
458    end,
459    FoldFun = fun(GroupedKey, Unreduced, Acc) ->
460        {ok, [{GroupedKey, couch_btree:final_reduce(Btree, Unreduced)} | Acc]}
461    end,
462    couch_btree:fold_reduce(Btree, FoldFun, [],
463                            [{key_group_fun, GroupFun}] ++ Opts).
464
465
466keys(KVs) ->
467    [K || {K, _} <- KVs].
468
469reduce_fun(reduce, KVs) ->
470    length(KVs);
471reduce_fun(rereduce, Reds) ->
472    lists:sum(Reds).
473
474
475shuffle(List) ->
476    randomize(round(math:log(length(List)) + 0.5), List).
477
478randomize(1, List) ->
479    randomize(List);
480randomize(T, List) ->
481    lists:foldl(
482        fun(_E, Acc) ->
483            randomize(Acc)
484        end, randomize(List), lists:seq(1, (T - 1))).
485
486randomize(List) ->
487    D = lists:map(fun(A) -> {couch_rand:uniform(), A} end, List),
488    {_, D1} = lists:unzip(lists:keysort(1, D)),
489    D1.
490
491test_btree(Btree, KeyValues) ->
492    ok = test_key_access(Btree, KeyValues),
493    ok = test_lookup_access(Btree, KeyValues),
494    ok = test_final_reductions(Btree, KeyValues),
495    ok = test_traversal_callbacks(Btree, KeyValues),
496    true.
497
498test_add_remove(Btree, OutKeyValues, RemainingKeyValues) ->
499    Btree2 = lists:foldl(
500        fun({K, _}, BtAcc) ->
501            {ok, BtAcc2} = couch_btree:add_remove(BtAcc, [], [K]),
502            BtAcc2
503        end, Btree, OutKeyValues),
504    true = test_btree(Btree2, RemainingKeyValues),
505
506    Btree3 = lists:foldl(
507        fun(KV, BtAcc) ->
508            {ok, BtAcc2} = couch_btree:add_remove(BtAcc, [KV], []),
509            BtAcc2
510        end, Btree2, OutKeyValues),
511    true = test_btree(Btree3, OutKeyValues ++ RemainingKeyValues).
512
513test_key_access(Btree, List) ->
514    FoldFun = fun(Element, {[HAcc|TAcc], Count}) ->
515        case Element == HAcc of
516            true -> {ok, {TAcc, Count + 1}};
517            _ -> {ok, {TAcc, Count + 1}}
518        end
519    end,
520    Length = length(List),
521    Sorted = lists:sort(List),
522    {ok, _, {[], Length}} = couch_btree:foldl(Btree, FoldFun, {Sorted, 0}),
523    {ok, _, {[], Length}} = couch_btree:fold(Btree, FoldFun,
524                                             {Sorted, 0}, [{dir, rev}]),
525    ok.
526
527test_lookup_access(Btree, KeyValues) ->
528    FoldFun = fun({Key, Value}, {Key, Value}) -> {stop, true} end,
529    lists:foreach(
530        fun({Key, Value}) ->
531            [{ok, {Key, Value}}] = couch_btree:lookup(Btree, [Key]),
532            {ok, _, true} = couch_btree:foldl(Btree, FoldFun,
533                                              {Key, Value}, [{start_key, Key}])
534        end, KeyValues).
535
536test_final_reductions(Btree, KeyValues) ->
537    KVLen = length(KeyValues),
538    FoldLFun = fun(_X, LeadingReds, Acc) ->
539        CountToStart = KVLen div 3 + Acc,
540        CountToStart = couch_btree:final_reduce(Btree, LeadingReds),
541        {ok, Acc + 1}
542    end,
543    FoldRFun = fun(_X, LeadingReds, Acc) ->
544        CountToEnd = KVLen - KVLen div 3 + Acc,
545        CountToEnd = couch_btree:final_reduce(Btree, LeadingReds),
546        {ok, Acc + 1}
547    end,
548    {LStartKey, _} = case KVLen of
549        0 -> {nil, nil};
550        _ -> lists:nth(KVLen div 3 + 1, lists:sort(KeyValues))
551    end,
552    {RStartKey, _} = case KVLen of
553        0 -> {nil, nil};
554        _ -> lists:nth(KVLen div 3, lists:sort(KeyValues))
555    end,
556    {ok, _, FoldLRed} = couch_btree:foldl(Btree, FoldLFun, 0,
557                                          [{start_key, LStartKey}]),
558    {ok, _, FoldRRed} = couch_btree:fold(Btree, FoldRFun, 0,
559                                         [{dir, rev}, {start_key, RStartKey}]),
560    KVLen = FoldLRed + FoldRRed,
561    ok.
562
563test_traversal_callbacks(Btree, _KeyValues) ->
564    FoldFun = fun
565        (visit, _GroupedKey, _Unreduced, Acc) ->
566            {ok, Acc andalso false};
567        (traverse, _LK, _Red, Acc) ->
568            {skip, Acc andalso true}
569    end,
570    % With 250 items the root is a kp. Always skipping should reduce to true.
571    {ok, _, true} = couch_btree:fold(Btree, FoldFun, true, [{dir, fwd}]),
572    ok.
573