Streams with parallel processing, lazy filtering and random sampling

.gitignore

@@ -17,6 +17,7 @@ _build
 *.iml
 rebar3.crashdump
 data/
+!src/data/
 .DS_Store
 src/pb/
 src/grpc/autogen

src/blockchain_rocks.erl

@@ -0,0 +1,116 @@
+-module(blockchain_rocks).
+
+-export([
+    fold/4,
+    fold/5,
+    foreach/3,
+    foreach/4,
+    stream/2,
+    stream/3,
+    sample/3,
+    sample/4
+]).
+
+%% API ========================================================================
+
+-spec fold(
+    rocksdb:db_handle(),
+    rocksdb:read_options(),
+    Acc,
+    fun(({K :: binary(), V :: binary()}, Acc) -> Acc)
+) ->
+    Acc.
+fold(DB, Opts, Acc, F) ->
+    data_stream:fold(stream(DB, Opts), Acc, F).
+
+-spec fold(
+    rocksdb:db_handle(),
+    rocksdb:cf_handle(),
+    rocksdb:read_options(),
+    Acc,
+    fun(({K :: binary(), V :: binary()}, Acc) -> Acc)
+) ->
+    Acc.
+fold(DB, CF, Opts, Acc, F) ->
+    data_stream:fold(stream(DB, CF, Opts), Acc, F).
+
+-spec foreach(
+    rocksdb:db_handle(),
+    rocksdb:read_options(),
+    fun(({K :: binary(), V :: binary()}) -> ok)
+) ->
+    ok.
+foreach(DB, Opts, F) ->
+    data_stream:foreach(stream(DB, Opts), F).
+
+-spec foreach(
+    rocksdb:db_handle(),
+    rocksdb:cf_handle(),
+    rocksdb:read_options(),
+    fun(({K :: binary(), V :: binary()}) -> ok)
+) ->
+    ok.
+foreach(DB, CF, Opts, F) ->
+    data_stream:foreach(stream(DB, CF, Opts), F).
+
+-spec stream(rocksdb:db_handle(), rocksdb:read_options()) ->
+    data_stream:t({K :: binary(), V :: binary()}).
+stream(DB, Opts) ->
+    stream_(fun () -> rocksdb:iterator(DB, Opts) end).
+
+-spec stream(
+    rocksdb:db_handle(),
+    rocksdb:cf_handle(),
+    rocksdb:read_options()
+) ->
+    data_stream:t({K :: binary(), V :: binary()}).
+stream(DB, CF, Opts) ->
+    stream_(fun () -> rocksdb:iterator(DB, CF, Opts) end).
+
+%% @doc Select K random records from database.
+-spec sample(rocksdb:db_handle(), rocksdb:read_options(), pos_integer()) ->
+    [{K :: binary(), V :: binary()}].
+sample(DB, Opts, K) ->
+    Stream = stream(DB, Opts),
+    data_stream:sample(Stream, K).
+
+%% @doc Select K random records from CF.
+-spec sample(
+    rocksdb:db_handle(),
+    rocksdb:cf_handle(),
+    rocksdb:read_options(),
+    pos_integer()
+) ->
+    [{K :: binary(), V :: binary()}].
+sample(DB, CF, Opts, K) ->
+    Stream = stream(DB, CF, Opts),
+    data_stream:sample(Stream, K).
+
+%% Internal ===================================================================
+
+-spec stream_(fun(() -> {ok, rocksdb:itr_handle()} | {error, term()})) ->
+    data_stream:t({K :: binary(), V :: binary()}).
+stream_(IterOpen) ->
+    case IterOpen() of
+        {error, Reason} ->
+            error({blockchain_rocks_iter_make, Reason});
+        {ok, Iter} ->
+            Move =
+                fun Move_ (Target) ->
+                    fun () ->
+                        case rocksdb:iterator_move(Iter, Target) of
+                            {ok, K, V} ->
+                                {some, {{K, V}, Move_(next)}};
+                            {error, invalid_iterator} ->
+                                ok = rocksdb:iterator_close(Iter),
+                                none;
+                            Error ->
+                                error({blockchain_rocks_iter_move, Target, Error})
+                        end
+                    end
+                end,
+            data_stream:from_fun(Move(first))
+    end.
+
+%% Test =======================================================================
+%% See test/blockchain_rocks_SUITE.erl

src/blockchain_term.erl

@@ -12,7 +12,6 @@
 -export_type([
     t/0,
     result/0,
-    stream/1, % TODO Find stream def a better home module than this one.
     error/0,
     frame/0,
     unsound/0,
@@ -105,69 +104,70 @@
 
 -include("blockchain_term.hrl").
 
-%% TODO Maybe use a map?
+%% TODO Switch to a map or a record?
 -type file_handle() ::
     {
         file:fd(),
         Pos :: non_neg_integer(),
         Len :: pos_integer()
     }.
 
--type stream(A) :: fun(() -> none | {some, {A, stream(A)}}).
-
 -spec from_bin(binary()) -> result().
 from_bin(<<Bin/binary>>) ->
     envelope(Bin).
 
 %% Tries to stream a list of binaries from file.
-%% TODO Generalize.
 -spec from_file_stream_bin_list(file_handle()) ->
-    stream({ok, binary()} | {error, term()}).
+    data_stream:t({ok, binary()} | {error, term()}).
 from_file_stream_bin_list({Fd, Pos, Len}) ->
     {ok, Pos} = file:position(Fd, {bof, Pos}),
-    case file:read(Fd, 6) of
-        {ok, <<?ETF_VERSION, ?ETF_TAG_LIST_EXT, N:32/integer-unsigned-big>>} ->
-            stream_bin_list_elements(N, {Fd, Pos + 6, Len});
-        {ok, <<V/binary>>} ->
-            fun () -> {some, {{error, {bad_etf_version_and_tag_and_len, V}}, stream_end()}} end;
-        {error, _}=Err ->
-            fun () -> {some, {Err, stream_end()}} end
-    end.
+    Next =
+        case file:read(Fd, 6) of
+            {ok, <<?ETF_VERSION, ?ETF_TAG_LIST_EXT, N:32/integer-unsigned-big>>} ->
+                next_bin_list_elements(N, {Fd, Pos + 6, Len});
+            {ok, <<V/binary>>} ->
+                fun () -> {some, {{error, {bad_etf_version_and_tag_and_len, V}}, next_end()}} end;
+            {error, _}=Err ->
+                fun () -> {some, {Err, next_end()}} end
+        end,
+    data_stream:from_fun(Next).
 
--spec stream_bin_list_elements(non_neg_integer(), file_handle()) ->
-    stream({ok, binary()} | {error, term()}).
-stream_bin_list_elements(0, {Fd, Pos, _}) ->
+-spec next_bin_list_elements(non_neg_integer(), file_handle()) ->
+    data_stream:next({ok, binary()} | {error, term()}).
+next_bin_list_elements(0, {Fd, Pos, _}) ->
     fun () ->
         {ok, Pos} = file:position(Fd, {bof, Pos}),
         case file:read(Fd, 1) of
             {ok, <<?ETF_TAG_NIL_EXT>>} ->
                 none;
             {ok, <<_/binary>>} ->
-                {some, {{error, bad_bin_list_nil_tag}, stream_end()}};
+                {some, {{error, bad_bin_list_nil_tag}, next_end()}};
             {error, _}=Err ->
-                {some, {Err, stream_end()}}
+                {some, {Err, next_end()}}
         end
     end;
-stream_bin_list_elements(N, {Fd, Pos0, L}) ->
+next_bin_list_elements(N, {Fd, Pos0, L}) ->
     fun () ->
         {ok, Pos1} = file:position(Fd, {bof, Pos0}),
         case file:read(Fd, 5) of
             {ok, <<?ETF_TAG_BINARY_EXT, Len:32/integer-unsigned-big>>} ->
                 {ok, Pos2} = file:position(Fd, {bof, Pos1 + 5}),
                 case file:read(Fd, Len) of
                     {ok, <<Bin/binary>>} ->
-                        {some, {Bin, stream_bin_list_elements(N - 1, {Fd, Pos2 + Len, L})}};
+                        {some, {Bin, next_bin_list_elements(N - 1, {Fd, Pos2 + Len, L})}};
                     {error, _}=Err ->
-                        {some, {Err, stream_end()}}
+                        {some, {Err, next_end()}}
                 end;
             {ok, <<_/binary>>} ->
-                {some, {{error, bad_bin_list_element}, stream_end()}};
+                {some, {{error, bad_bin_list_element}, next_end()}};
             {error, _}=Err ->
-                {some, {Err, stream_end()}}
+                {some, {Err, next_end()}}
         end
     end.
 
-stream_end() ->
+-spec next_end() ->
+    data_stream:next({ok, binary()} | {error, term()}).
+next_end() ->
     fun () -> none end.
 
 %% TODO -spec from_bin_with_contract(binary(), blockchain_contract:t()) ->

src/blockchain_utils.erl

@@ -10,6 +10,7 @@
 -include("blockchain_vars.hrl").
 
 -export([
+    cpus/0,
     shuffle_from_hash/2,
     shuffle/1, shuffle/2,
     rand_from_hash/1, rand_state/1,
@@ -302,6 +303,7 @@ validation_width() ->
             N
     end.
 
+-spec cpus() -> non_neg_integer().
 cpus() ->
     Ct = erlang:system_info(schedulers_online),
     max(2, ceil(Ct/2) + 1).