17
17
*
18
18
*/
19
19
use actix_web:: rt:: spawn;
20
- use datafusion:: arrow;
21
20
use datafusion:: arrow:: datatypes:: Schema ;
22
21
use datafusion:: arrow:: error:: ArrowError ;
23
22
use datafusion:: arrow:: ipc:: writer:: StreamWriter ;
@@ -36,11 +35,10 @@ use std::sync::RwLock;
36
35
37
36
use crate :: metadata;
38
37
use crate :: metadata:: LOCK_EXPECT ;
39
- use crate :: option:: CONFIG ;
40
38
use crate :: s3;
41
- use crate :: storage:: ObjectStorage ;
39
+ use crate :: storage:: { ObjectStorage , StorageDir } ;
42
40
43
- use self :: error:: EventError ;
41
+ use self :: error:: { EventError , StreamWriterError } ;
44
42
45
43
type LocalWriter = Mutex < Option < StreamWriter < std:: fs:: File > > > ;
46
44
type LocalWriterGuard < ' a > = MutexGuard < ' a , Option < StreamWriter < std:: fs:: File > > > ;
@@ -91,18 +89,7 @@ impl STREAM_WRITERS {
91
89
. write ( )
92
90
. map_err ( |_| StreamWriterError :: RwPoisoned ) ?;
93
91
94
- let file = OpenOptions :: new ( )
95
- . append ( true )
96
- . create_new ( true )
97
- . open ( data_file_path ( & stream) )
98
- . map_err ( StreamWriterError :: Io ) ?;
99
-
100
- let mut stream_writer = StreamWriter :: try_new ( file, & record. schema ( ) )
101
- . expect ( "File and RecordBatch both are checked" ) ;
102
-
103
- stream_writer
104
- . write ( record)
105
- . map_err ( StreamWriterError :: Writer ) ?;
92
+ let stream_writer = init_new_stream_writer_file ( & stream, record) ?;
106
93
107
94
hashmap_guard. insert ( stream, Mutex :: new ( Some ( stream_writer) ) ) ;
108
95
@@ -125,63 +112,51 @@ impl STREAM_WRITERS {
125
112
stream : & str ,
126
113
record : & RecordBatch ,
127
114
) -> Result < ( ) , StreamWriterError > {
128
- let file = OpenOptions :: new ( )
129
- . append ( true )
130
- . create_new ( true )
131
- . open ( data_file_path ( stream) )
132
- . map_err ( StreamWriterError :: Io ) ?;
133
-
134
- let mut stream_writer = StreamWriter :: try_new ( file, & record. schema ( ) )
135
- . expect ( "File and RecordBatch both are checked" ) ;
136
-
137
- stream_writer
138
- . write ( record)
139
- . map_err ( StreamWriterError :: Writer ) ?;
115
+ let stream_writer = init_new_stream_writer_file ( stream, record) ?;
140
116
141
117
writer_guard. replace ( stream_writer) ; // replace the stream writer behind this mutex
142
118
143
119
Ok ( ( ) )
144
120
}
145
121
146
- // Unset the entry so that
147
- pub fn unset_entry ( stream : & str ) -> Result < ( ) , StreamWriterError > {
148
- let guard = STREAM_WRITERS
122
+ pub fn unset_all ( ) -> Result < ( ) , StreamWriterError > {
123
+ let map = STREAM_WRITERS
149
124
. read ( )
150
125
. map_err ( |_| StreamWriterError :: RwPoisoned ) ?;
151
- let stream_writer = match guard. get ( stream) {
152
- Some ( writer) => writer,
153
- None => return Ok ( ( ) ) ,
154
- } ;
155
- stream_writer
156
- . lock ( )
157
- . map_err ( |_| StreamWriterError :: MutexPoisoned ) ?
158
- . take ( ) ;
126
+
127
+ for writer in map. values ( ) {
128
+ if let Some ( mut streamwriter) = writer
129
+ . lock ( )
130
+ . map_err ( |_| StreamWriterError :: MutexPoisoned ) ?
131
+ . take ( )
132
+ {
133
+ let _ = streamwriter. finish ( ) ;
134
+ }
135
+ }
159
136
160
137
Ok ( ( ) )
161
138
}
162
139
}
163
140
164
- #[ derive( Debug , thiserror:: Error ) ]
165
- pub enum StreamWriterError {
166
- #[ error( "Arrow writer failed: {0}" ) ]
167
- Writer ( arrow:: error:: ArrowError ) ,
168
- #[ error( "Io Error when creating new file: {0}" ) ]
169
- Io ( std:: io:: Error ) ,
170
- #[ error( "RwLock was poisoned" ) ]
171
- RwPoisoned ,
172
- #[ error( "Mutex was poisoned" ) ]
173
- MutexPoisoned ,
174
- }
141
+ fn init_new_stream_writer_file (
142
+ stream_name : & str ,
143
+ record : & RecordBatch ,
144
+ ) -> Result < StreamWriter < std:: fs:: File > , StreamWriterError > {
145
+ let dir = StorageDir :: new ( stream_name) ;
146
+ let path = dir. path_by_current_time ( ) ;
147
+
148
+ std:: fs:: create_dir_all ( dir. data_path ) ?;
149
+
150
+ let file = OpenOptions :: new ( ) . create ( true ) . append ( true ) . open ( path) ?;
151
+
152
+ let mut stream_writer = StreamWriter :: try_new ( file, & record. schema ( ) )
153
+ . expect ( "File and RecordBatch both are checked" ) ;
154
+
155
+ stream_writer
156
+ . write ( record)
157
+ . map_err ( StreamWriterError :: Writer ) ?;
175
158
176
- fn data_file_path ( stream_name : & str ) -> String {
177
- format ! (
178
- "{}/{}" ,
179
- CONFIG
180
- . parseable
181
- . local_stream_data_path( stream_name)
182
- . to_string_lossy( ) ,
183
- "data.records"
184
- )
159
+ Ok ( stream_writer)
185
160
}
186
161
187
162
#[ derive( Clone ) ]
@@ -309,7 +284,7 @@ impl Event {
309
284
infer_json_schema ( & mut buf_reader, None )
310
285
}
311
286
312
- fn get_reader ( & self , arrow_schema : arrow :: datatypes :: Schema ) -> json:: Reader < & [ u8 ] > {
287
+ fn get_reader ( & self , arrow_schema : Schema ) -> json:: Reader < & [ u8 ] > {
313
288
json:: Reader :: new (
314
289
self . body . as_bytes ( ) ,
315
290
Arc :: new ( arrow_schema) ,
@@ -348,8 +323,6 @@ pub mod error {
348
323
use crate :: storage:: ObjectStorageError ;
349
324
use datafusion:: arrow:: error:: ArrowError ;
350
325
351
- use super :: StreamWriterError ;
352
-
353
326
#[ derive( Debug , thiserror:: Error ) ]
354
327
pub enum EventError {
355
328
#[ error( "Missing Record from event body" ) ]
@@ -365,4 +338,16 @@ pub mod error {
365
338
#[ error( "Schema Mismatch: {0}" ) ]
366
339
ObjectStorage ( #[ from] ObjectStorageError ) ,
367
340
}
341
+
342
+ #[ derive( Debug , thiserror:: Error ) ]
343
+ pub enum StreamWriterError {
344
+ #[ error( "Arrow writer failed: {0}" ) ]
345
+ Writer ( #[ from] ArrowError ) ,
346
+ #[ error( "Io Error when creating new file: {0}" ) ]
347
+ Io ( #[ from] std:: io:: Error ) ,
348
+ #[ error( "RwLock was poisoned" ) ]
349
+ RwPoisoned ,
350
+ #[ error( "Mutex was poisoned" ) ]
351
+ MutexPoisoned ,
352
+ }
368
353
}
0 commit comments