@@ -117,7 +117,7 @@ use crate::{
117
117
BaoFileStorage , BaoFileStorageSubscriber , CompleteStorage , DataReader ,
118
118
OutboardReader ,
119
119
} ,
120
- util:: entity_manager,
120
+ util:: entity_manager:: { self , EntityState } ,
121
121
} ,
122
122
util:: { BaoTreeSender , FixedSize , MemOrFile , ValueOrPoisioned } ,
123
123
Hash , IROH_BLOCK_SIZE ,
@@ -211,12 +211,12 @@ impl TaskContext {
211
211
}
212
212
}
213
213
214
- impl entity_manager:: Params for HashContext {
215
- type EntityId = Hash ;
214
+ impl entity_manager:: EntityState for HashContext {
215
+ type Id = Hash ;
216
216
217
217
type GlobalState = Arc < TaskContext > ;
218
218
219
- fn id ( & self ) -> & Self :: EntityId {
219
+ fn id ( & self ) -> & Self :: Id {
220
220
& self . id
221
221
}
222
222
@@ -228,24 +228,41 @@ impl entity_manager::Params for HashContext {
228
228
self . state . sender_count ( ) + self . state . receiver_count ( )
229
229
}
230
230
231
- fn new ( id : & Self :: EntityId , global : & Self :: GlobalState ) -> Self {
231
+ fn new ( id : & Self :: Id , global : & Self :: GlobalState ) -> Self {
232
232
Self {
233
233
id : * id,
234
234
global : global. clone ( ) ,
235
235
state : BaoFileHandle :: default ( ) ,
236
236
}
237
237
}
238
238
239
- fn reset ( & mut self , id : & Self :: EntityId , global : & Self :: GlobalState ) {
239
+ fn reset ( & mut self , id : & Self :: Id , global : & Self :: GlobalState ) {
240
240
self . id = * id;
241
241
self . global = global. clone ( ) ;
242
242
// this is identical to self.state = BaoFileHandle::default(),
243
243
// but does not allocate a new handle.
244
244
self . state . send_replace ( BaoFileStorage :: Initial ) ;
245
245
}
246
246
247
+ #[ instrument( skip_all, fields( hash = %self . id. fmt_short( ) ) ) ]
247
248
async fn on_shutdown ( & self , _cause : entity_manager:: ShutdownCause ) {
248
- self . persist ( ) . await ;
249
+ self . state . send_if_modified ( |guard| {
250
+ let hash = & self . id ;
251
+ let BaoFileStorage :: Partial ( fs) = guard. take ( ) else {
252
+ return false ;
253
+ } ;
254
+ let path = self . global . options . path . bitfield_path ( hash) ;
255
+ trace ! ( "writing bitfield for hash {} to {}" , hash, path. display( ) ) ;
256
+ if let Err ( cause) = fs. sync_all ( & path) {
257
+ error ! (
258
+ "failed to write bitfield for {} at {}: {:?}" ,
259
+ hash,
260
+ path. display( ) ,
261
+ cause
262
+ ) ;
263
+ }
264
+ false
265
+ } ) ;
249
266
}
250
267
}
251
268
@@ -878,7 +895,7 @@ async fn handle_batch_impl(cmd: BatchMsg, id: Scope, scope: &Arc<TempTagScope>)
878
895
}
879
896
880
897
/// The minimal API you need to implement for an entity for a store to work.
881
- trait EntityApi {
898
+ trait EntityApi : EntityState {
882
899
/// Import from a stream of n0 bao encoded data.
883
900
async fn import_bao ( & self , cmd : ImportBaoMsg ) ;
884
901
/// Finish an import from a local file or memory.
@@ -891,8 +908,6 @@ trait EntityApi {
891
908
async fn export_bao ( & self , cmd : ExportBaoMsg ) ;
892
909
/// Export the entry to a local file.
893
910
async fn export_path ( & self , cmd : ExportPathMsg ) ;
894
- /// Persist the entry at the end of its lifecycle.
895
- async fn persist ( & self ) ;
896
911
}
897
912
898
913
/// A more opinionated API that can be used as a helper to save implementation
@@ -1000,27 +1015,6 @@ impl EntityApi for HashContext {
1000
1015
} ;
1001
1016
cmd. tx . send ( res) . await . ok ( ) ;
1002
1017
}
1003
-
1004
- #[ instrument( skip_all, fields( hash = %self . id. fmt_short( ) ) ) ]
1005
- async fn persist ( & self ) {
1006
- self . state . send_if_modified ( |guard| {
1007
- let hash = & self . id ;
1008
- let BaoFileStorage :: Partial ( fs) = guard. take ( ) else {
1009
- return false ;
1010
- } ;
1011
- let path = self . global . options . path . bitfield_path ( hash) ;
1012
- trace ! ( "writing bitfield for hash {} to {}" , hash, path. display( ) ) ;
1013
- if let Err ( cause) = fs. sync_all ( & path) {
1014
- error ! (
1015
- "failed to write bitfield for {} at {}: {:?}" ,
1016
- hash,
1017
- path. display( ) ,
1018
- cause
1019
- ) ;
1020
- }
1021
- false
1022
- } ) ;
1023
- }
1024
1018
}
1025
1019
1026
1020
async fn finish_import_impl ( ctx : & HashContext , import_data : ImportEntry ) -> io:: Result < ( ) > {
0 commit comments