18
18
*/
19
19
20
20
mod file_writer;
21
- mod mem_writer;
22
21
23
22
use std:: {
24
23
collections:: HashMap ,
25
24
sync:: { Mutex , RwLock } ,
26
25
} ;
27
26
28
- use crate :: {
29
- option:: CONFIG ,
30
- storage:: staging:: { self , ReadBuf } ,
31
- } ;
32
-
33
27
use self :: { errors:: StreamWriterError , file_writer:: FileWriter } ;
34
28
use arrow_array:: RecordBatch ;
35
- use chrono:: { NaiveDateTime , Utc } ;
36
29
use derive_more:: { Deref , DerefMut } ;
37
- use mem_writer:: MemWriter ;
38
30
use once_cell:: sync:: Lazy ;
39
31
40
- type InMemWriter = MemWriter < 8192 > ;
41
-
42
32
pub static STREAM_WRITERS : Lazy < WriterTable > = Lazy :: new ( WriterTable :: default) ;
43
33
44
- pub enum StreamWriter {
45
- Mem ( InMemWriter ) ,
46
- Disk ( FileWriter , InMemWriter ) ,
47
- }
48
-
49
- impl StreamWriter {
50
- pub fn push (
51
- & mut self ,
52
- stream_name : & str ,
53
- schema_key : & str ,
54
- rb : RecordBatch ,
55
- ) -> Result < ( ) , StreamWriterError > {
56
- match self {
57
- StreamWriter :: Mem ( mem) => {
58
- mem. push ( schema_key, rb) ;
59
- }
60
- StreamWriter :: Disk ( disk, mem) => {
61
- disk. push ( stream_name, schema_key, & rb) ?;
62
- mem. push ( schema_key, rb) ;
63
- }
64
- }
65
- Ok ( ( ) )
66
- }
67
- }
68
-
69
- // Each entry in writer table is initialized with some context
70
- // This is helpful for generating prefix when writer is finalized
71
- pub struct WriterContext {
72
- stream_name : String ,
73
- time : NaiveDateTime ,
74
- }
75
-
76
34
#[ derive( Deref , DerefMut , Default ) ]
77
- pub struct WriterTable ( RwLock < HashMap < String , ( Mutex < StreamWriter > , WriterContext ) > > ) ;
35
+ pub struct WriterTable ( RwLock < HashMap < String , Mutex < FileWriter > > > ) ;
78
36
79
37
impl WriterTable {
80
38
// append to a existing stream
@@ -87,36 +45,26 @@ impl WriterTable {
87
45
let hashmap_guard = self . read ( ) . unwrap ( ) ;
88
46
89
47
match hashmap_guard. get ( stream_name) {
90
- Some ( ( stream_writer, _ ) ) => {
48
+ Some ( stream_writer) => {
91
49
stream_writer
92
50
. lock ( )
93
51
. unwrap ( )
94
- . push ( stream_name, schema_key, record) ?;
52
+ . push ( stream_name, schema_key, & record) ?;
95
53
}
96
54
None => {
97
55
drop ( hashmap_guard) ;
98
56
let mut map = self . write ( ) . unwrap ( ) ;
99
57
// check for race condition
100
58
// if map contains entry then just
101
- if let Some ( ( writer, _ ) ) = map. get ( stream_name) {
59
+ if let Some ( writer) = map. get ( stream_name) {
102
60
writer
103
61
. lock ( )
104
62
. unwrap ( )
105
- . push ( stream_name, schema_key, record) ?;
63
+ . push ( stream_name, schema_key, & record) ?;
106
64
} else {
107
- // there is no entry so this can be inserted safely
108
- let context = WriterContext {
109
- stream_name : stream_name. to_owned ( ) ,
110
- time : Utc :: now ( ) . naive_utc ( ) ,
111
- } ;
112
- let mut writer = if CONFIG . parseable . in_mem_ingestion {
113
- StreamWriter :: Mem ( InMemWriter :: default ( ) )
114
- } else {
115
- StreamWriter :: Disk ( FileWriter :: default ( ) , InMemWriter :: default ( ) )
116
- } ;
117
-
118
- writer. push ( stream_name, schema_key, record) ?;
119
- map. insert ( stream_name. to_owned ( ) , ( Mutex :: new ( writer) , context) ) ;
65
+ let mut writer = FileWriter :: default ( ) ;
66
+ writer. push ( stream_name, schema_key, & record) ?;
67
+ map. insert ( stream_name. to_owned ( ) , Mutex :: new ( writer) ) ;
120
68
}
121
69
}
122
70
} ;
@@ -131,40 +79,11 @@ impl WriterTable {
131
79
let mut table = self . write ( ) . unwrap ( ) ;
132
80
let map = std:: mem:: take ( & mut * table) ;
133
81
drop ( table) ;
134
- for ( writer, context ) in map. into_values ( ) {
82
+ for writer in map. into_values ( ) {
135
83
let writer = writer. into_inner ( ) . unwrap ( ) ;
136
- match writer {
137
- StreamWriter :: Mem ( mem) => {
138
- let rb = mem. finalize ( ) ;
139
- let mut read_bufs = staging:: MEMORY_READ_BUFFERS . write ( ) . unwrap ( ) ;
140
-
141
- read_bufs
142
- . entry ( context. stream_name )
143
- . or_insert ( Vec :: default ( ) )
144
- . push ( ReadBuf {
145
- time : context. time ,
146
- buf : rb,
147
- } ) ;
148
- }
149
- StreamWriter :: Disk ( disk, _) => disk. close_all ( ) ,
150
- }
84
+ writer. close_all ( ) ;
151
85
}
152
86
}
153
-
154
- pub fn clone_read_buf ( & self , stream_name : & str ) -> Option < ReadBuf > {
155
- let hashmap_guard = self . read ( ) . unwrap ( ) ;
156
- let ( writer, context) = hashmap_guard. get ( stream_name) ?;
157
- let writer = writer. lock ( ) . unwrap ( ) ;
158
- let mem = match & * writer {
159
- StreamWriter :: Mem ( mem) => mem,
160
- StreamWriter :: Disk ( _, mem) => mem,
161
- } ;
162
-
163
- Some ( ReadBuf {
164
- time : context. time ,
165
- buf : mem. recordbatch_cloned ( ) ,
166
- } )
167
- }
168
87
}
169
88
170
89
pub mod errors {
0 commit comments