@@ -80,6 +80,7 @@ fn get_k_colouring_locks() -> &'static DashMap<KColouringLockKey, Arc<RwLock<()>
8080/// A lock is considered unused if its Arc strong count is 1 (only the map holds a reference).
8181fn cleanup_k_colouring_locks ( ) {
8282 if let Some ( locks) = K_COLOURING_LOCKS . get ( ) {
83+ // TODO[DK]: Track average lock map size as part of metrics for later
8384 locks. retain ( |_, v| Arc :: strong_count ( v) > 1 ) ;
8485 }
8586}
@@ -422,8 +423,32 @@ impl<C: DagknightStore + DagknightStoreReader, O: HeaderStoreReader, D: Relation
422423 // subgroup = the current subgroup
423424 // tips = all tips in this conflict. part of which is the subgroup
424425 //
426+ // `next_chain_ancestor`:
427+ // - If `free_search` is true, this must be `None` (asserted).
428+ // - If `free_search` is false, this must be `Some(nca)` where `nca` is a block whose selected
429+ // parent is `root` (asserted). This bounds the zone fill to the subgroup's region.
430+ //
425431 // Returns the conflict zone manager which gives access to the coloring data of the conflict zone
426- pub fn fill_zone_data ( & self , tips : & [ Hash ] ) -> BlockHashSet {
432+ pub fn fill_zone_data ( & self , tips : & [ Hash ] , next_chain_ancestor : Option < Hash > ) -> BlockHashSet {
433+ // Construct lock key and validate parameters
434+ let lock_key = if self . free_search {
435+ assert ! ( next_chain_ancestor. is_none( ) , "free_search expects None for next_chain_ancestor" ) ;
436+ KColouringLockKey :: free_search_key ( self . root , self . k )
437+ } else {
438+ let nca = next_chain_ancestor. expect ( "committed_search expects Some(next_chain_ancestor)" ) ;
439+ // Assert NCA's selected parent is self.root (conflict_genesis)
440+ assert ! (
441+ self . reachability_service. is_chain_ancestor_of( self . root, next_chain_ancestor. unwrap( ) ) ,
442+ "conflict_genesis must be a chain ancestor of next_chain_ancestor"
443+ ) ;
444+ KColouringLockKey :: committed_search_key ( self . root , self . k , nca)
445+ } ;
446+
447+ // Acquire lock for this zone fill
448+ let locks = get_k_colouring_locks ( ) ;
449+ let lock_arc = locks. entry ( lock_key) . or_insert_with ( || Arc :: new ( RwLock :: new ( ( ) ) ) ) . clone ( ) ;
450+ let _guard = lock_arc. write ( ) ;
451+
427452 self . init_root ( ) ;
428453
429454 let ( last_known_tips, visited_subdag) = self . find_last_known_tips ( tips) ;
@@ -459,7 +484,7 @@ impl<C: DagknightStore + DagknightStoreReader, O: HeaderStoreReader, D: Relation
459484 let selected_parent = if self . free_search {
460485 self . find_selected_parent ( parents. iter ( ) . copied ( ) )
461486 } else {
462- let next_chain_ancestor_of_current = self . reachability_service . get_next_chain_ancestor ( current_hash , self . root ) ;
487+ let next_chain_ancestor_of_current = next_chain_ancestor . unwrap ( ) ;
463488 let agreeing_parents = parents
464489 . iter ( )
465490 . copied ( )
@@ -487,7 +512,7 @@ impl<C: DagknightStore + DagknightStoreReader, O: HeaderStoreReader, D: Relation
487512 let is_in_zone = if self . free_search {
488513 self . reachability_service . try_is_dag_ancestor_of ( self . root , child) . unwrap_or ( false )
489514 } else {
490- self . reachability_service . try_is_chain_ancestor_of ( self . root , child) . unwrap_or ( false )
515+ self . reachability_service . try_is_chain_ancestor_of ( next_chain_ancestor . unwrap ( ) , child) . unwrap_or ( false )
491516 } ;
492517 if !is_in_zone {
493518 continue ;
@@ -497,6 +522,9 @@ impl<C: DagknightStore + DagknightStoreReader, O: HeaderStoreReader, D: Relation
497522 }
498523 }
499524
525+ // Opportunistically cleanup unused locks after zone fill
526+ cleanup_k_colouring_locks ( ) ;
527+
500528 visited_subdag
501529 }
502530 // END Copied from GD Manager
@@ -850,8 +878,10 @@ mod tests {
850878
851879 // Now fill zone data
852880 let tips = vec ! [ hash_x, hash_d] ;
853- manager_committed. fill_zone_data ( & tips) ;
854- manager_free. fill_zone_data ( & tips) ;
881+ // For committed search, NCA is hash_z (whose selected parent is hash_a, the conflict genesis)
882+ manager_committed. fill_zone_data ( & tips, Some ( hash_z) ) ;
883+ // For free search, NCA is None
884+ manager_free. fill_zone_data ( & tips, None ) ;
855885
856886 // Get X's selected parent from both managers
857887 let committed_sp = manager_committed. get_selected_parent ( hash_x) . unwrap ( ) ;
0 commit comments