1+ use std:: collections:: HashMap ;
2+ use std:: error:: Error ;
3+ use std:: ffi:: OsStr ;
4+ use std:: fs;
5+ use std:: ops:: Deref ;
6+ use std:: path:: { Path , PathBuf } ;
7+ use std:: sync:: { Arc , Mutex } ;
8+ use log:: { debug, error} ;
9+ use crate :: MessageOffset ;
10+
11+ pub struct MessageOffsetHolder {
12+ journal_dir : String ,
13+ inner : Arc < Mutex < HashMap < OffsetKey , Offset > > >
14+ }
15+
16+ type Offset = i64 ;
17+
18+ #[ derive( Clone , Eq , PartialEq , Hash , Debug ) ]
19+ pub struct OffsetKey ( pub String , pub i32 ) ;
20+
21+ impl MessageOffsetHolder {
22+ pub fn with_offsets_in ( journal_dir : String ) -> Result < MessageOffsetHolder , Box < dyn Error > > {
23+ let offsets = read_offsets_from ( & journal_dir) ?;
24+ Ok ( MessageOffsetHolder {
25+ journal_dir,
26+ inner : Arc :: new ( Mutex :: new ( offsets) ) ,
27+ } )
28+ }
29+
30+ pub fn update ( & self , offset : MessageOffset ) {
31+ let mut guard = self . inner . lock ( ) . unwrap ( ) ;
32+ ( * guard) . insert ( OffsetKey ( offset. topic , offset. partition ) , offset. offset ) ;
33+ }
34+
35+ pub fn flush ( & self ) {
36+ let offsets = {
37+ let guard = self . inner . lock ( ) . unwrap ( ) ;
38+ guard. clone ( )
39+ } ;
40+
41+ save_offsets_in ( offsets, & self . journal_dir ) ;
42+ }
43+ }
44+
45+ fn read_offsets_from < P : AsRef < Path > > ( directory : P ) -> Result < HashMap < OffsetKey , Offset > , Box < dyn Error > > {
46+ ensure_dir_exists ( & directory) ?;
47+
48+ let map = fs:: read_dir ( directory. as_ref ( ) ) ?
49+ . filter_map ( |file| {
50+ file. map_err ( |e| {
51+ error ! ( "Error reading file: {e}" )
52+ } ) . ok ( )
53+ } )
54+ . filter_map ( |file| {
55+ let path = file. path ( ) ;
56+
57+ fs:: read_to_string ( & path)
58+ . map_err ( |e| {
59+ error ! ( "Error reading file {path:?}! Reason: {e}" )
60+ } )
61+ . ok ( )
62+ . and_then ( |content| {
63+ let offset_key = file_name_to_offset_key ( path. file_name ( ) , path. file_stem ( ) ) ?;
64+ Some ( ( offset_key, content. parse ( ) . unwrap ( ) ) )
65+ } )
66+ } )
67+ . collect ( ) ;
68+
69+ Ok ( map)
70+ }
71+
72+ /// Parses file_name as "$topic.$partition" (eg. sampletopic.1) and returns OffsetKey if successful
73+ fn file_name_to_offset_key ( file_name : Option < & OsStr > , file_stem : Option < & OsStr > ) -> Option < OffsetKey > {
74+ let file_name = file_name?. to_str ( ) ?;
75+ let file_stem = file_stem?. to_str ( ) ?;
76+
77+ let topic = file_name. strip_suffix ( & file_stem) ?. to_string ( ) ;
78+
79+ Some ( OffsetKey ( topic, file_stem. parse ( ) . ok ( ) ?) )
80+ }
81+
82+ fn save_offsets_in < P : AsRef < Path > > ( offsets : HashMap < OffsetKey , Offset > , directory : P ) {
83+ if let Err ( e) = ensure_dir_exists ( & directory) {
84+ error ! ( "Cannot save offsets! {e}" ) ;
85+ return ;
86+ }
87+
88+ offsets. into_iter ( )
89+ . for_each ( |( key, offset) | {
90+ save_offset ( & directory, key, offset) ;
91+ } ) ;
92+ }
93+
94+ fn save_offset < P : AsRef < Path > > ( base_path : P , offset_key : OffsetKey , offset : Offset ) {
95+ let base_path = PathBuf :: from ( base_path. as_ref ( ) ) ;
96+ let file_path = base_path. join ( format ! ( "{}.{}" , & offset_key. 0 , & offset_key. 1 ) ) ;
97+
98+ debug ! ( "Saving offset [Topic: {}] [Partition: {}] [Offset: {}] to {:?}" , & offset_key. 0 , & offset_key. 1 , offset, & file_path) ;
99+
100+ // save offset to file with name "$topic.$partition" (eg. sampletopic.1)
101+ if let Err ( e) = fs:: write ( & file_path, format ! ( "{}" , offset) ) {
102+ error ! ( "Failed to write journal to file {file_path:?}. Topic: [{}], Partition: [{}], Offset: [{}]. Reason: {e}" , offset_key. 0 , offset_key. 1 , offset) ;
103+ }
104+ }
105+
106+ fn ensure_dir_exists < P : AsRef < Path > > ( directory : P ) -> Result < ( ) , String > {
107+ let dir = directory. as_ref ( ) ;
108+ if !dir. exists ( ) {
109+ if let Err ( e) = fs:: create_dir_all ( dir) {
110+ return Err ( format ! ( "Cannot create directory {dir:?}! Journal will not be saved. Reason: {e}" ) ) ;
111+ }
112+ }
113+ if dir. is_file ( ) {
114+ return Err ( format ! ( "Directory is a file: {dir:?}! Cannot use as journal." ) ) ;
115+ }
116+
117+ Ok ( ( ) )
118+ }
119+
120+ impl Deref for MessageOffsetHolder {
121+ type Target = Mutex < HashMap < OffsetKey , Offset > > ;
122+
123+ fn deref ( & self ) -> & Self :: Target {
124+ self . inner . deref ( )
125+ }
126+ }
127+
128+ impl Drop for MessageOffsetHolder {
129+ fn drop ( & mut self ) {
130+ // when json-processor finishes, we should save current offsets
131+ self . flush ( )
132+ }
133+ }
0 commit comments