unaligned subblock patch from cliff

beerriot · Jul 9, 2010 · 401e825 · 401e825
1 parent a1e7594
commit 401e825
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 6 deletions.
diff --git a/include/luwak.hrl b/include/luwak.hrl
@@ -8,6 +8,8 @@
 
 -define(TIMEOUT_DEFAULT, 60000).
 
+-define(fmt(S, As), lists:flatten(io_lib:format(S, As))).
+
 -record(split, {head=[], midhead=[], middle=[], midtail=[], tail=[]}).
 
 -ifndef(EUNIT_HRL).

diff --git a/src/luwak_io.erl b/src/luwak_io.erl
@@ -127,9 +127,10 @@ write_blocks(Riak, File, undefined, Start, Data, BlockSize, Written)
     {ok, Block} = luwak_block:create(Riak, Slice),
     write_blocks(Riak, File, undefined, Start+BlockSize, Tail, BlockSize,
                  [{luwak_block:name(Block),BlockSize}|Written]);
-%% we are doing a sub-block write
-write_blocks(Riak, _File, PartialStartBlock, Start, Data, BlockSize, Written)
-  when is_list(Written), byte_size(Data) < BlockSize ->
+%% we are doing a sub-block write on a full block
+write_blocks(Riak, File, PartialStartBlock, Start, Data, BlockSize, Written)
+  when is_list(Written), byte_size(Data) < BlockSize,
+       byte_size(PartialStartBlock) == BlockSize ->
     ?debugFmt("D write_blocks(Riak, File, ~p, ~p, ~p, ~p, ~p) ~n",
               [PartialStartBlock, Start, Data, BlockSize, Written]),
     DataSize = byte_size(Data),
@@ -141,6 +142,30 @@ write_blocks(Riak, _File, PartialStartBlock, Start, Data, BlockSize, Written)
     {ok, lists:reverse([{luwak_block:name(Block),
                          byte_size(BlockData)}|
                         Written])};
+%% sub block write on a partial write
+write_blocks(Riak, File, PartialStartBlock, Start, Data, BlockSize, Written)
+  when is_list(Written), byte_size(Data) < BlockSize ->
+  PartialStartData = luwak_block:data(PartialStartBlock),
+  DataSize = byte_size(Data),
+  PartialStart = Start rem BlockSize,
+  BlockData = case byte_size(PartialStartData) of
+    S when S < PartialStart ->
+      Difference = (PartialStart - S) * 8,
+      <<PartialStartData/binary, 0:Difference, Data/binary>>;
+    S when S == PartialStart ->
+      <<PartialStartData/binary, Data/binary>>;
+    S when S > PartialStart ->
+      <<Head:PartialStart/binary, Left/binary>> = PartialStartData,
+      if
+        byte_size(Left) > DataSize ->
+          <<_:DataSize/binary, Tail>> = Left,
+          <<Head/binary, Data/binary, Tail/binary>>;
+        true ->
+          <<Head/binary, Data/binary>>
+      end
+  end,
+  {ok, Block} = luwak_block:create(Riak, BlockData),
+  {ok, lists:reverse([{luwak_block:name(Block),byte_size(BlockData)}|Written])};
 %% we are starting with a sub-block write
 write_blocks(Riak, File, PartialStartBlock, Start, Data, BlockSize, Written)
   when Start > BlockSize ->

diff --git a/src/luwak_put_stream.erl b/src/luwak_put_stream.erl
@@ -172,11 +172,12 @@ flush(Riak, State=#state{offset=Offset,file=File,buffer=Buffer,written=Written})
                                            buffer=[],
                                            buffersize=0,
                                            written=Written++Written1}));
-flush(Riak, State=#state{offset=Offset,file=File,written=Written})
+flush(Riak, State=#state{offset=Offset,file=File,written=Written,blocksize=BlockSize})
   when length(Written) > 0 ->
     ?debugFmt("B flush(Riak, ~p)~n", [State]),
     OriginalOffset = Offset - luwak_tree_utils:blocklist_length(Written),
-    {ok, NewFile} = luwak_tree:update(Riak, File, OriginalOffset, Written),
+    OriginalOffsetAligned = OriginalOffset - OriginalOffset rem BlockSize,  
+    {ok, NewFile} = luwak_tree:update(Riak, File, OriginalOffsetAligned, Written),
     update_checksum(Riak, State#state{file=NewFile});
 flush(_Riak, State) ->
     State.

diff --git a/src/luwak_tree.erl b/src/luwak_tree.erl
@@ -20,7 +20,7 @@ update(Riak, File, StartingPos, Blocks) ->
     BlockSize = luwak_file:get_property(File, block_size),
     if
         StartingPos rem BlockSize =/= 0 ->
-            throw({error, "StartingPos must be a multiple of blocksize"});
+            throw({error, ?fmt("StartingPos (~p) must be a multiple of blocksize", [StartingPos])});
         true ->
             ok
     end,

diff --git a/test/luwak_io_tests.erl b/test/luwak_io_tests.erl
@@ -62,3 +62,23 @@ truncate_test() ->
       Blocks = luwak_io:get_range(Riak, File2, 0, 7),
       ?assertEqual(<<"wontyou">>, iolist_to_binary(Blocks))
     end).
+
+sub_block_partial_write_test() ->
+  test_helper:riak_test(fun(Riak) ->
+      {ok, File} = luwak_file:create(Riak, <<"basho">>, [{block_size,1000}], dict:new()),
+      S = luwak_put_stream:start_link(Riak, File, 0, 1000),
+      B1 = crypto:rand_bytes(33),
+      luwak_put_stream:send(S, B1),
+      luwak_put_stream:close(S),
+      timer:sleep(100),
+      {ok, File2} = luwak_file:get(Riak, <<"basho">>),
+      S1 = luwak_put_stream:start_link(Riak, File2, 33, 1000),
+      B2 = crypto:rand_bytes(33),
+      luwak_put_stream:send(S1, B2),
+      luwak_put_stream:close(S1),
+      timer:sleep(100),
+      {ok, File3} = luwak_file:get(Riak, <<"basho">>),
+      Blocks = luwak_io:get_range(Riak, File3, 0, 66),
+      ?debugFmt("blocks ~p~n", [Blocks]),
+      ?assertEqual(<<B1/binary, B2/binary>>, iolist_to_binary(Blocks))
+    end).