kvs_rocks.erl 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  1. -module(kvs_rocks).
  2. -include("backend.hrl").
  3. -include("kvs.hrl").
  4. -include("metainfo.hrl").
  5. -include_lib("stdlib/include/qlc.hrl").
  6. -export(?BACKEND).
  7. -export([ref/0,ref/1,bt/1,key/2,key/1,fd/1,tb/1,estimate/0,estimate/1]).
  8. -export([seek_it/1, seek_it/2, move_it/3, move_it/4, take_it/4, take_it/5]).
  9. e(X,Y) -> element(X,Y).
  10. bt([]) -> [];
  11. bt(X) -> binary_to_term(X).
  12. tb([]) -> <<>>;
  13. tb(T) when is_list(T) -> unicode:characters_to_nfkc_binary(T);
  14. tb(T) when is_atom(T) -> erlang:atom_to_binary(T, utf8);
  15. tb(T) when is_binary(T) -> T;
  16. tb(T) -> term_to_binary(T).
  17. sz([]) -> 0;
  18. sz(B) -> byte_size(B).
  19. key(R) when is_tuple(R) andalso tuple_size(R) > 1 -> key(e(1,R), e(2,R));
  20. key(R) -> key(R,[]).
  21. key(writer,R) -> % allow old writers
  22. iolist_to_binary([lists:join(<<"/">>, lists:flatten([<<>>, erlang:atom_to_binary(writer, utf8), tb(R)]))]);
  23. key(Tab,R) -> Fd = case Tab of [] -> []; _ -> tb(Tab) end,
  24. iolist_to_binary([lists:join(<<"/">>, lists:flatten([<<>>, Fd, fmt(R)]))]).
  25. fmt([]) -> [];
  26. fmt(K) -> Key = tb(K),
  27. End = sz(Key),
  28. {S,E} = case binary:matches(Key, [<<"/">>], []) of
  29. [{0,1}] -> {1, End-1};
  30. [{0,1},{1,1}] -> {2, End-2};
  31. [{0,1},{1,1}|_] -> {2, End-2};
  32. [{0,1}|_] -> {1, End-1};
  33. _ -> {0, End}
  34. end,
  35. binary:part(Key,{S,E}).
  36. fd(K) -> Key = tb(K),
  37. End = sz(Key),
  38. {S,_} = case binary:matches(Key,[<<"/">>],[]) of
  39. [{0,1}] -> {End,End};
  40. [{0,1},{1,1}] -> {End,End};
  41. [{0,1},{1,1}|T] -> hd(lists:reverse(T));
  42. [{0,1}|T] -> hd(lists:reverse(T));
  43. _ -> {End,End}
  44. end,
  45. binary:part(Key,{0,S}).
  46. run(<<>>,SK,_,_,_) -> {ok,SK,[],[]};
  47. run(Key, % key
  48. SK, % sup-key
  49. Dir, % direction next/prev
  50. Compiled_Operations,
  51. Db) ->
  52. % H is iterator reference
  53. S = sz(SK),
  54. Initial_Object = {ref(Db), []},
  55. Run = fun (F,K,H,V,Acc) when binary_part(K,{0,S}) == SK -> {F(H,Dir),H,[V|Acc]}; % continue +------------+
  56. (_,K,H,V,Acc) -> stop_it(H), % fail-safe closing |
  57. throw({ok, fd(K), bt(V), [bt(A1) || A1 <- Acc]}) end, % acc unfold |
  58. % |
  59. Range_Check = fun(F,K,H,V) -> case F(H,prev) of % backward prefetch |
  60. {ok,K1,V1} when binary_part(K,{0,S}) == SK -> {{ok,K1,V1},H,[V]}; % return (range-check error) |
  61. {ok,K1,V1} -> Run(F,K1,H,V1,[]); % run prev-take chain | loop
  62. _ -> stop_it(H), % fail-safe closing |
  63. throw({ok, fd(K), bt(V), [bt(V)]}) % acc unfold |
  64. end end, % |
  65. % |
  66. State_Machine = fun % |
  67. (F,{ok,H}) -> {F(H,{seek,Key}),H}; % first move (seek) |
  68. (F,{{ok,K,V},H}) when Dir =:= prev -> Range_Check(F,K,H,V); % first chained prev-take |
  69. (F,{{ok,K,V},H}) -> Run(F,K,H,V,[]); % first chained next-take |
  70. (F,{{ok,K,V},H,A}) -> Run(F,K,H,V,A); % chained CPS-take continuator +---+
  71. (_,{{error,E},H,Acc}) -> {{error,E},H,Acc}; % error effects
  72. (F,{I,O}) -> F(I,O) % chain constructor from initial object
  73. end,
  74. catch case lists:foldl(State_Machine, Initial_Object, Compiled_Operations) of
  75. {{ok,K,Bin},H,A} -> stop_it(H), {ok, fd(K), bt(Bin), [bt(A1) || A1 <- A]};
  76. {{ok,K,Bin},H} -> stop_it(H), {ok, fd(K), bt(Bin), []};
  77. {{error,_},H,Acc} -> stop_it(H), {ok, fd(SK), bt(shd(Acc)), [bt(A1) || A1 <- Acc]};
  78. {{error,_},H} -> stop_it(H), {ok, fd(SK), [], []}
  79. end.
  80. initialize() -> [ kvs:initialize(kvs_rocks,Module) || Module <- kvs:modules() ].
  81. index(_,_,_) -> [].
  82. ref_env(Db) -> list_to_atom("rocks_ref_" ++ Db).
  83. db() -> application:get_env(kvs,rocks_name,"rocksdb").
  84. start() -> ok.
  85. stop() -> ok.
  86. destroy() -> destroy(db()).
  87. destroy(Db) -> rocksdb:destroy(Db, []).
  88. version() -> {version,"KVS ROCKSDB"}.
  89. dir() -> [].
  90. ref() -> ref(db()).
  91. ref(Db) -> application:get_env(kvs,ref_env(Db),[]).
  92. leave() -> leave(db()).
  93. leave(Db) -> case ref(Db) of [] -> skip; X -> rocksdb:close(X), application:set_env(kvs,ref_env(Db),[]), ok end.
  94. join(_,Db) ->
  95. application:start(rocksdb),
  96. leave(Db), {ok, Ref} = rocksdb:open(Db, [{create_if_missing, true}]),
  97. initialize(),
  98. application:set_env(kvs,ref_env(Db),Ref).
  99. compile(seek) -> [fun rocksdb:iterator/2,fun rocksdb:iterator_move/2];
  100. compile(move) -> [fun rocksdb:iterator_move/2];
  101. compile(close) -> [fun rocksdb:iterator_close/1].
  102. compile(take,N) -> lists:map(fun(_) -> fun rocksdb:iterator_move/2 end, lists:seq(1, N)).
  103. stop_it(H) -> try begin [F]=compile(close), F(H) end catch error:badarg -> ok end.
  104. seek_it(K) -> seek_it(K,db()).
  105. seek_it(K,Db) -> run(K,K,ok,compile(seek),Db).
  106. move_it(Key,SK,Dir) -> move_it(Key,SK,Dir,db()).
  107. move_it(Key,SK,Dir,Db) -> run(Key,SK,Dir,compile(seek) ++ compile(move),Db).
  108. take_it(Key,SK,Dir,N) -> take_it(Key,SK,Dir,N,db()).
  109. take_it(Key,SK,Dir,N,Db) when is_integer(N) andalso N >= 0 -> run(Key,SK,Dir,compile(seek) ++ compile(take,N),Db);
  110. take_it(Key,SK,Dir,_,Db) -> take_it(Key,SK,Dir,0,Db).
  111. all(R,Db) -> kvs_st:feed(R,Db).
  112. get(Tab, {step,N,[208|_]=Key}, Db) -> get(Tab, {step,N,list_to_binary(Key)},Db);
  113. get(Tab, [208|_]=Key, Db) -> get(Tab, list_to_binary(Key), Db);
  114. get(Tab, Key, Db) ->
  115. case rocksdb:get(ref(Db), key(Tab,Key), []) of
  116. not_found -> {error,not_found};
  117. {ok,Bin} -> {ok,bt(Bin)} end.
  118. put(Record) -> put(Record,db()).
  119. put(Records,Db) when is_list(Records) -> lists:map(fun(Record) -> put(Record,Db) end, Records);
  120. put(Record,Db) -> rocksdb:put(ref(Db), key(Record), term_to_binary(Record), [{sync,true}]).
  121. delete(Feed, Id, Db) -> rocksdb:delete(ref(Db), key(Feed,Id), []).
  122. count(_) -> 0.
  123. estimate() -> estimate(db()).
  124. estimate(Db) -> case rocksdb:get_property(ref(Db), <<"rocksdb.estimate-num-keys">>) of
  125. {ok, Est} when is_binary(Est) -> binary_to_integer(Est);
  126. {ok, Est} when is_list(Est) -> list_to_integer(Est);
  127. {ok, Est} when is_integer(Est) -> Est;
  128. _ -> 0
  129. end.
  130. shd([]) -> [];
  131. shd(X) -> hd(X).
  132. create_table(_,_) -> [].
  133. add_table_index(_, _) -> ok.
  134. dump() -> ok.
  135. seq(_,_) ->
  136. Val =
  137. case os:type() of
  138. {win32,nt} -> {Mega,Sec,Micro} = erlang:timestamp(), integer_to_list((Mega*1000000+Sec)*1000000+Micro);
  139. _ -> erlang:integer_to_list(element(2,hd(lists:reverse(erlang:system_info(os_monotonic_time_source)))))
  140. end,
  141. case 20 - length(Val) > 0 of true -> string:copies("0", 20 - length(Val)); _ -> "" end ++ Val.