@@ -195,6 +195,43 @@ impl GeneralBucket {
195195 }
196196}
197197
198+ struct BucketResources {
199+ slots_allocated : u16 ,
200+ slots_used : u16 ,
201+ liquidity_allocated : u64 ,
202+ liquidity_used : u64 ,
203+ }
204+
205+ impl BucketResources {
206+ fn new ( slots_allocated : u16 , liquidity_allocated : u64 ) -> Self {
207+ BucketResources { slots_allocated, slots_used : 0 , liquidity_allocated, liquidity_used : 0 }
208+ }
209+
210+ fn resources_available ( & self , htlc_amount_msat : u64 ) -> bool {
211+ return ( self . liquidity_used + htlc_amount_msat <= self . liquidity_allocated )
212+ && ( self . slots_used < self . slots_allocated ) ;
213+ }
214+
215+ fn add_htlc ( & mut self , htlc_amount_msat : u64 ) -> Result < ( ) , ( ) > {
216+ if !self . resources_available ( htlc_amount_msat) {
217+ return Err ( ( ) ) ;
218+ }
219+
220+ self . slots_used += 1 ;
221+ self . liquidity_used += htlc_amount_msat;
222+ Ok ( ( ) )
223+ }
224+
225+ fn remove_htlc ( & mut self , htlc_amount_msat : u64 ) -> Result < ( ) , ( ) > {
226+ if self . slots_used == 0 || self . liquidity_used < htlc_amount_msat {
227+ return Err ( ( ) ) ;
228+ }
229+ self . slots_used -= 1 ;
230+ self . liquidity_used -= htlc_amount_msat;
231+ Ok ( ( ) )
232+ }
233+ }
234+
198235/// A weighted average that decays over a specified window.
199236///
200237/// It enables tracking of historical behavior without storing individual data points.
@@ -289,7 +326,9 @@ mod tests {
289326
290327 use crate :: {
291328 crypto:: chacha20:: ChaCha20 ,
292- ln:: resource_manager:: { AggregatedWindowAverage , DecayingAverage , GeneralBucket } ,
329+ ln:: resource_manager:: {
330+ AggregatedWindowAverage , BucketResources , DecayingAverage , GeneralBucket ,
331+ } ,
293332 util:: test_utils:: TestKeysInterface ,
294333 } ;
295334 use bitcoin:: Network ;
@@ -450,6 +489,60 @@ mod tests {
450489 assert ! ( general_bucket. slots_occupied[ slot_occupied as usize ] . is_none( ) ) ;
451490 }
452491
492+ fn test_bucket_resources ( ) -> BucketResources {
493+ BucketResources {
494+ slots_allocated : 10 ,
495+ slots_used : 0 ,
496+ liquidity_allocated : 100_000 ,
497+ liquidity_used : 0 ,
498+ }
499+ }
500+
501+ #[ test]
502+ fn test_bucket_resources_add_htlc ( ) {
503+ let mut bucket_resources = test_bucket_resources ( ) ;
504+ let available_liquidity = bucket_resources. liquidity_allocated ;
505+ assert ! ( bucket_resources. add_htlc( available_liquidity + 1000 ) . is_err( ) ) ;
506+
507+ assert ! ( bucket_resources. add_htlc( 21_000 ) . is_ok( ) ) ;
508+ assert ! ( bucket_resources. add_htlc( 42_000 ) . is_ok( ) ) ;
509+ assert_eq ! ( bucket_resources. slots_used, 2 ) ;
510+ assert_eq ! ( bucket_resources. liquidity_used, 63_000 ) ;
511+ }
512+
513+ #[ test]
514+ fn test_bucket_resources_add_htlc_over_resources_available ( ) {
515+ // Test trying to go over slot limit
516+ let mut bucket_resources = test_bucket_resources ( ) ;
517+ let slots_available = bucket_resources. slots_allocated ;
518+ for _ in 0 ..slots_available {
519+ assert ! ( bucket_resources. add_htlc( 10 ) . is_ok( ) ) ;
520+ }
521+ assert_eq ! ( bucket_resources. slots_used, slots_available) ;
522+ assert ! ( bucket_resources. add_htlc( 10 ) . is_err( ) ) ;
523+
524+ // Test trying to go over liquidity limit
525+ let mut bucket = test_bucket_resources ( ) ;
526+ assert ! ( bucket. add_htlc( bucket. liquidity_allocated - 1000 ) . is_ok( ) ) ;
527+ assert ! ( bucket. add_htlc( 2000 ) . is_err( ) ) ;
528+ }
529+
530+ #[ test]
531+ fn test_bucket_resources_remove_htlc ( ) {
532+ let mut bucket_resources = test_bucket_resources ( ) ;
533+
534+ // If no resources have been used, removing HTLC should fail
535+ assert ! ( bucket_resources. remove_htlc( 100 ) . is_err( ) ) ;
536+
537+ bucket_resources. add_htlc ( 1000 ) . unwrap ( ) ;
538+ // Test failure if it tries to remove amount over what is currently in use.
539+ assert ! ( bucket_resources. remove_htlc( 1001 ) . is_err( ) ) ;
540+
541+ assert ! ( bucket_resources. remove_htlc( 1000 ) . is_ok( ) ) ;
542+ assert_eq ! ( bucket_resources. slots_used, 0 ) ;
543+ assert_eq ! ( bucket_resources. liquidity_used, 0 ) ;
544+ }
545+
453546 #[ test]
454547 fn test_decaying_average_error ( ) {
455548 let timestamp = 1000 ;
0 commit comments