1- use std:: io:: { Cursor , Seek , SeekFrom , Write } ;
2- use std:: path:: PathBuf ;
3- use std:: sync:: { Arc , Mutex } ;
1+ use std:: io:: { Seek , SeekFrom , Write } ;
2+ use std:: path:: { Path , PathBuf } ;
3+ use std:: pin:: Pin ;
4+ use std:: task:: { Context , Poll } ;
45
6+ use tokio:: io:: AsyncWrite ;
7+
8+ use crate :: CasClientError ;
59use crate :: error:: Result ;
610
7- /// Enum of different output formats to write reconstructed files.
11+ /// type that represents all acceptable sequential output mechanisms
12+ /// To convert something that is Write rather than AsyncWrite uses the AsyncWriteFromWrite adapter
13+ pub type SequentialOutput = Box < dyn AsyncWrite + Send + Unpin > ;
14+
15+ pub fn sequential_output_from_filepath ( filename : impl AsRef < Path > ) -> Result < SequentialOutput > {
16+ let file = std:: fs:: OpenOptions :: new ( )
17+ . write ( true )
18+ . truncate ( false )
19+ . create ( true )
20+ . open ( & filename) ?;
21+ Ok ( Box :: new ( AsyncWriteFromWrite ( Some ( Box :: new ( file) ) ) ) )
22+ }
23+
24+ /// Enum of different output formats to write reconstructed files
25+ /// where the result writer can be set at a specific position and new handles can be created
826#[ derive( Debug , Clone ) ]
9- pub enum OutputProvider {
27+ pub enum SeekingOutputProvider {
1028 File ( FileProvider ) ,
1129 #[ cfg( test) ]
12- Buffer ( BufferProvider ) ,
30+ Buffer ( buffer_provider :: BufferProvider ) ,
1331}
1432
15- impl OutputProvider {
33+ impl SeekingOutputProvider {
34+ // shortcut to create a new FileProvider variant from filename
35+ pub fn new_file_provider ( filename : PathBuf ) -> Self {
36+ Self :: File ( FileProvider :: new ( filename) )
37+ }
38+
1639 /// Create a new writer to start writing at the indicated start location.
1740pub ( crate ) fn get_writer_at ( & self , start : u64 ) -> Result < Box < dyn Write + Send > > {
1841 match self {
19- OutputProvider :: File ( fp) => fp. get_writer_at ( start) ,
42+ SeekingOutputProvider :: File ( fp) => fp. get_writer_at ( start) ,
2043 #[ cfg( test) ]
21- OutputProvider :: Buffer ( bp) => bp. get_writer_at ( start) ,
44+ SeekingOutputProvider :: Buffer ( bp) => bp. get_writer_at ( start) ,
2245 }
2346 }
2447}
2548
49+ // Adapter used to create an AsyncWrite from a Writer.
50+ struct AsyncWriteFromWrite ( Option < Box < dyn Write + Send > > ) ;
51+
52+ impl AsyncWrite for AsyncWriteFromWrite {
53+ fn poll_write ( mut self : Pin < & mut Self > , _cx : & mut Context < ' _ > , buf : & [ u8 ] ) -> Poll < std:: io:: Result < usize > > {
54+ let Some ( inner) = self . 0 . as_mut ( ) else {
55+ return Poll :: Ready ( Ok ( 0 ) ) ;
56+ } ;
57+ Poll :: Ready ( inner. write ( buf) )
58+ }
59+
60+ fn poll_flush ( mut self : Pin < & mut Self > , _cx : & mut Context < ' _ > ) -> Poll < std:: io:: Result < ( ) > > {
61+ let Some ( inner) = self . 0 . as_mut ( ) else {
62+ return Poll :: Ready ( Err ( std:: io:: Error :: new (
63+ std:: io:: ErrorKind :: BrokenPipe ,
64+ "writer closed, already dropped" ,
65+ ) ) ) ;
66+ } ;
67+ Poll :: Ready ( inner. flush ( ) )
68+ }
69+
70+ fn poll_shutdown ( mut self : Pin < & mut Self > , _cx : & mut Context < ' _ > ) -> Poll < std:: io:: Result < ( ) > > {
71+ let _ = self . 0 . take ( ) ;
72+ Poll :: Ready ( Ok ( ( ) ) )
73+ }
74+ }
75+
76+ impl TryFrom < SeekingOutputProvider > for SequentialOutput {
77+ type Error = CasClientError ;
78+
79+ fn try_from ( value : SeekingOutputProvider ) -> std:: result:: Result < Self , Self :: Error > {
80+ let w = value. get_writer_at ( 0 ) ?;
81+ Ok ( Box :: new ( AsyncWriteFromWrite ( Some ( w) ) ) )
82+ }
83+ }
84+
2685/// Provides new Writers to a file located at a particular location
2786#[ derive( Debug , Clone ) ]
2887pub struct FileProvider {
@@ -45,44 +104,69 @@ impl FileProvider {
45104 }
46105}
47106
48- #[ derive( Debug , Default , Clone ) ]
49- pub struct BufferProvider {
50- pub buf : ThreadSafeBuffer ,
51- }
107+ #[ cfg( test) ]
108+ pub ( crate ) mod buffer_provider {
109+ use std:: io:: { Cursor , Write } ;
110+ use std:: sync:: { Arc , Mutex } ;
111+
112+ use crate :: error:: Result ;
113+ use crate :: output_provider:: AsyncWriteFromWrite ;
114+ use crate :: { SeekingOutputProvider , SequentialOutput } ;
52115
53- impl BufferProvider {
54- pub fn get_writer_at ( & self , start : u64 ) -> crate :: error :: Result < Box < dyn std :: io :: Write + Send > > {
55- let mut buffer = self . buf . clone ( ) ;
56- buffer . idx = start ;
57- Ok ( Box :: new ( buffer ) )
116+ /// BufferProvider may be Seeking or Sequential
117+ /// only used in testing
118+ # [ derive ( Debug , Clone ) ]
119+ pub struct BufferProvider {
120+ pub buf : ThreadSafeBuffer ,
58121 }
59- }
60122
61- #[ derive( Debug , Default , Clone ) ]
62- /// Thread-safe in-memory buffer that implements [Write](Write) trait at some position
63- /// within an underlying buffer and allows access to inner buffer.
64- /// Thread-safe in-memory buffer that implements [Write](Write) trait and allows
65- /// access to inner buffer
66- pub struct ThreadSafeBuffer {
67- idx : u64 ,
68- inner : Arc < Mutex < Cursor < Vec < u8 > > > > ,
69- }
70- impl ThreadSafeBuffer {
71- pub fn value ( & self ) -> Vec < u8 > {
72- self . inner . lock ( ) . unwrap ( ) . get_ref ( ) . clone ( )
123+ impl BufferProvider {
124+ pub fn get_writer_at ( & self , start : u64 ) -> Result < Box < dyn Write + Send > > {
125+ let mut buffer = self . buf . clone ( ) ;
126+ buffer. idx = start;
127+ Ok ( Box :: new ( buffer) )
128+ }
129+ }
130+
131+ #[ derive( Debug , Default , Clone ) ]
132+ /// Thread-safe in-memory buffer that implements [Write](Write) trait at some position
133+ /// within an underlying buffer and allows access to the inner buffer.
134+ /// Thread-safe in-memory buffer that implements [Write](Write) trait and allows
135+ /// access to the inner buffer
136+ pub struct ThreadSafeBuffer {
137+ idx : u64 ,
138+ inner : Arc < Mutex < Cursor < Vec < u8 > > > > ,
139+ }
140+
141+ impl ThreadSafeBuffer {
142+ pub fn value ( & self ) -> Vec < u8 > {
143+ self . inner . lock ( ) . unwrap ( ) . get_ref ( ) . clone ( )
144+ }
73145 }
74- }
75146
76- impl std:: io:: Write for ThreadSafeBuffer {
77- fn write ( & mut self , buf : & [ u8 ] ) -> std:: io:: Result < usize > {
78- let mut guard = self . inner . lock ( ) . map_err ( |e| std:: io:: Error :: other ( format ! ( "{e}" ) ) ) ?;
79- guard. set_position ( self . idx ) ;
80- let num_written = guard. write ( buf) ?;
81- self . idx = guard. position ( ) ;
82- Ok ( num_written)
147+ impl Write for ThreadSafeBuffer {
148+ fn write ( & mut self , buf : & [ u8 ] ) -> std:: io:: Result < usize > {
149+ let mut guard = self . inner . lock ( ) . map_err ( |e| std:: io:: Error :: other ( format ! ( "{e}" ) ) ) ?;
150+ guard. set_position ( self . idx ) ;
151+ let num_written = Write :: write ( guard. get_mut ( ) , buf) ?;
152+ self . idx = guard. position ( ) ;
153+ Ok ( num_written)
154+ }
155+
156+ fn flush ( & mut self ) -> std:: io:: Result < ( ) > {
157+ Ok ( ( ) )
158+ }
83159 }
84160
85- fn flush ( & mut self ) -> std:: io:: Result < ( ) > {
86- Ok ( ( ) )
161+ impl From < ThreadSafeBuffer > for SequentialOutput {
162+ fn from ( value : ThreadSafeBuffer ) -> Self {
163+ Box :: new ( AsyncWriteFromWrite ( Some ( Box :: new ( value) ) ) )
164+ }
165+ }
166+
167+ impl From < ThreadSafeBuffer > for SeekingOutputProvider {
168+ fn from ( value : ThreadSafeBuffer ) -> Self {
169+ SeekingOutputProvider :: Buffer ( BufferProvider { buf : value } )
170+ }
87171 }
88172}
0 commit comments