/src/ra_log_snapshot.erl

https://github.com/rabbitmq/ra · Erlang · 224 lines · 179 code · 22 blank · 23 comment · 0 complexity · 0ba019dbfc9778d9886884991a5ed08e MD5 · raw file

  1. %% This Source Code Form is subject to the terms of the Mozilla Public
  2. %% License, v. 2.0. If a copy of the MPL was not distributed with this
  3. %% file, You can obtain one at https://mozilla.org/MPL/2.0/.
  4. %%
  5. %% Copyright (c) 2017-2020 VMware, Inc. or its affiliates. All rights reserved.
  6. %%
  7. %% @hidden
  8. -module(ra_log_snapshot).
  9. -behaviour(ra_snapshot).
  10. -export([
  11. prepare/2,
  12. write/3,
  13. begin_accept/2,
  14. accept_chunk/2,
  15. complete_accept/2,
  16. begin_read/1,
  17. read_chunk/3,
  18. recover/1,
  19. validate/1,
  20. read_meta/1
  21. ]).
  22. -include("ra.hrl").
  23. -define(MAGIC, "RASN").
  24. -define(VERSION, 1).
  25. -type file_err() :: ra_snapshot:file_err().
  26. -type meta() :: ra_snapshot:meta().
  27. %% DO nothing. There is no preparation for snapshotting
  28. prepare(_Index, State) -> State.
  29. %% @doc
  30. %% Snapshot file format:
  31. %% "RASN"
  32. %% Version (byte)
  33. %% Checksum (unsigned 32)
  34. %% MetaData Len (unsigned 32)
  35. %% MetaData (binary)
  36. %% Snapshot Data (binary)
  37. %% @end
  38. -spec write(file:filename(), meta(), term()) ->
  39. ok | {error, file_err()}.
  40. write(Dir, Meta, MacState) ->
  41. %% no compression on meta data to make sure reading it is as fast
  42. %% as possible
  43. MetaBin = term_to_binary(Meta),
  44. %% the data can however be compressed
  45. Bin = term_to_binary(MacState, [{compressed, 9}]),
  46. Data = [<<(size(MetaBin)):32/unsigned>>, MetaBin, Bin],
  47. Checksum = erlang:crc32(Data),
  48. File = filename(Dir),
  49. ra_lib:write_file(File, [<<?MAGIC,
  50. ?VERSION:8/unsigned,
  51. Checksum:32/integer>>,
  52. Data]).
  53. begin_accept(SnapDir, Meta) ->
  54. File = filename(SnapDir),
  55. {ok, Fd} = file:open(File, [write, binary, raw]),
  56. MetaBin = term_to_binary(Meta),
  57. Data = [<<(size(MetaBin)):32/unsigned>>, MetaBin],
  58. PartialCrc = erlang:crc32(Data),
  59. ok = file:write(Fd, [<<?MAGIC,
  60. ?VERSION:8/unsigned,
  61. 0:32/integer>>,
  62. Data]),
  63. {ok, {PartialCrc, Fd}}.
  64. accept_chunk(Chunk, {PartialCrc, Fd}) ->
  65. <<Crc:32/integer, Rest/binary>> = Chunk,
  66. accept_chunk(Rest, {PartialCrc, Crc, Fd});
  67. accept_chunk(Chunk, {PartialCrc0, Crc, Fd}) ->
  68. ok = file:write(Fd, Chunk),
  69. PartialCrc = erlang:crc32(PartialCrc0, Chunk),
  70. {ok, {PartialCrc, Crc, Fd}}.
  71. complete_accept(Chunk, {PartialCrc, Fd}) ->
  72. <<Crc:32/integer, Rest/binary>> = Chunk,
  73. complete_accept(Rest, {PartialCrc, Crc, Fd});
  74. complete_accept(Chunk, {PartialCrc0, Crc, Fd}) ->
  75. ok = file:write(Fd, Chunk),
  76. {ok, 5} = file:position(Fd, 5),
  77. ok = file:write(Fd, <<Crc:32/integer>>),
  78. Crc = erlang:crc32(PartialCrc0, Chunk),
  79. ok = file:sync(Fd),
  80. ok = file:close(Fd),
  81. ok.
  82. begin_read(Dir) ->
  83. File = filename(Dir),
  84. case file:open(File, [read, binary, raw]) of
  85. {ok, Fd} ->
  86. case read_meta_internal(Fd) of
  87. {ok, Meta, Crc} ->
  88. {ok, DataStart} = file:position(Fd, cur),
  89. {ok, Eof} = file:position(Fd, eof),
  90. {ok, Meta, {Crc, {DataStart, Eof, Fd}}};
  91. {error, _} = Err ->
  92. _ = file:close(Fd),
  93. Err
  94. end;
  95. Err ->
  96. Err
  97. end.
  98. read_chunk({Crc, ReadState}, Size, Dir) when is_integer(Crc) ->
  99. case read_chunk(ReadState, Size - 4, Dir) of
  100. {ok, Data, ReadState1} ->
  101. {ok, <<Crc:32/integer, Data/binary>>, ReadState1};
  102. {error, _} = Err ->
  103. Err
  104. end;
  105. read_chunk({Pos, Eof, Fd}, Size, _Dir) ->
  106. {ok, _} = file:position(Fd, Pos),
  107. case file:read(Fd, Size) of
  108. {ok, Data} ->
  109. case Pos + Size >= Eof of
  110. true ->
  111. _ = file:close(Fd),
  112. {ok, Data, last};
  113. false ->
  114. {ok, Data, {next, {Pos + Size, Eof, Fd}}}
  115. end;
  116. {error, _} = Err ->
  117. Err;
  118. eof ->
  119. {error, unexpected_eof}
  120. end.
  121. -spec recover(file:filename()) ->
  122. {ok, meta(), term()} |
  123. {error, invalid_format |
  124. {invalid_version, integer()} |
  125. checksum_error |
  126. file_err()}.
  127. recover(Dir) ->
  128. File = filename(Dir),
  129. case file:read_file(File) of
  130. {ok, <<?MAGIC, ?VERSION:8/unsigned, Crc:32/integer, Data/binary>>} ->
  131. validate(Crc, Data);
  132. {ok, <<?MAGIC, Version:8/unsigned, _:32/integer, _/binary>>} ->
  133. {error, {invalid_version, Version}};
  134. {ok, _} ->
  135. {error, invalid_format};
  136. {error, _} = Err ->
  137. Err
  138. end.
  139. validate(Dir) ->
  140. case recover(Dir) of
  141. {ok, _, _} -> ok;
  142. Err -> Err
  143. end.
  144. %% @doc reads the index and term from the snapshot file without reading the
  145. %% entire binary body. NB: this does not do checksum validation.
  146. -spec read_meta(file:filename()) ->
  147. {ok, meta()} | {error, invalid_format |
  148. {invalid_version, integer()} |
  149. checksum_error |
  150. file_err()}.
  151. read_meta(Dir) ->
  152. File = filename(Dir),
  153. case file:open(File, [read, binary, raw]) of
  154. {ok, Fd} ->
  155. case read_meta_internal(Fd) of
  156. {ok, Meta, _} ->
  157. _ = file:close(Fd),
  158. {ok, Meta};
  159. {error, _} = Err ->
  160. _ = file:close(Fd),
  161. Err
  162. end;
  163. Err ->
  164. Err
  165. end.
  166. %% Internal
  167. read_meta_internal(Fd) ->
  168. HeaderSize = 9 + 4,
  169. case file:read(Fd, HeaderSize) of
  170. {ok, <<?MAGIC, ?VERSION:8/unsigned, Crc:32/integer,
  171. MetaSize:32/unsigned>>} ->
  172. case file:read(Fd, MetaSize) of
  173. {ok, MetaBin} ->
  174. {ok, binary_to_term(MetaBin), Crc};
  175. Err ->
  176. Err
  177. end;
  178. {ok, <<?MAGIC, Version:8/unsigned, _:32/integer, _/binary>>} ->
  179. {error, {invalid_version, Version}};
  180. {ok, _} ->
  181. {error, invalid_format};
  182. eof ->
  183. {error, unexpected_eof_when_parsing_header};
  184. Err ->
  185. Err
  186. end.
  187. validate(Crc, Data) ->
  188. case erlang:crc32(Data) of
  189. Crc ->
  190. parse_snapshot(Data);
  191. _ ->
  192. {error, checksum_error}
  193. end.
  194. parse_snapshot(<<MetaSize:32/unsigned, MetaBin:MetaSize/binary,
  195. Rest/binary>>) ->
  196. Meta = binary_to_term(MetaBin),
  197. {ok, Meta, binary_to_term(Rest)}.
  198. filename(Dir) ->
  199. filename:join(Dir, "snapshot.dat").
  200. -ifdef(TEST).
  201. -include_lib("eunit/include/eunit.hrl").
  202. -endif.