13
13
//!
14
14
//! [`ChannelManager`]: crate::ln::channelmanager::ChannelManager
15
15
16
+ use alloc:: sync:: Arc ;
16
17
use core:: mem;
17
18
use core:: time:: Duration ;
18
19
use sync:: { Condvar , Mutex } ;
19
20
20
21
#[ cfg( any( test, feature = "std" ) ) ]
21
22
use std:: time:: Instant ;
22
23
24
+ #[ cfg( feature = "std" ) ]
25
+ use std:: future:: Future as StdFuture ;
26
+ #[ cfg( feature = "std" ) ]
27
+ use std:: task:: { Context , Poll } ;
28
+ #[ cfg( feature = "std" ) ]
29
+ use std:: pin:: Pin ;
30
+
23
31
/// Used to signal to one of many waiters that the condition they're waiting on has happened.
24
32
pub ( crate ) struct Notifier {
25
- /// Users won't access the lock directly, but rather wait on its bool using
26
- /// `wait_timeout` and `wait`.
27
- lock : ( Mutex < bool > , Condvar ) ,
33
+ notify_pending : Mutex < ( bool , Option < Arc < Mutex < FutureState > > > ) > ,
34
+ condvar : Condvar ,
28
35
}
29
36
30
37
impl Notifier {
31
38
pub ( crate ) fn new ( ) -> Self {
32
39
Self {
33
- lock : ( Mutex :: new ( false ) , Condvar :: new ( ) ) ,
40
+ notify_pending : Mutex :: new ( ( false , None ) ) ,
41
+ condvar : Condvar :: new ( ) ,
34
42
}
35
43
}
36
44
37
45
pub ( crate ) fn wait ( & self ) {
38
46
loop {
39
- let & ( ref mtx, ref cvar) = & self . lock ;
40
- let mut guard = mtx. lock ( ) . unwrap ( ) ;
41
- if * guard {
42
- * guard = false ;
47
+ let mut guard = self . notify_pending . lock ( ) . unwrap ( ) ;
48
+ if guard. 0 {
49
+ guard. 0 = false ;
43
50
return ;
44
51
}
45
- guard = cvar . wait ( guard) . unwrap ( ) ;
46
- let result = * guard;
52
+ guard = self . condvar . wait ( guard) . unwrap ( ) ;
53
+ let result = guard. 0 ;
47
54
if result {
48
- * guard = false ;
55
+ guard. 0 = false ;
49
56
return
50
57
}
51
58
}
@@ -55,22 +62,21 @@ impl Notifier {
55
62
pub ( crate ) fn wait_timeout ( & self , max_wait : Duration ) -> bool {
56
63
let current_time = Instant :: now ( ) ;
57
64
loop {
58
- let & ( ref mtx, ref cvar) = & self . lock ;
59
- let mut guard = mtx. lock ( ) . unwrap ( ) ;
60
- if * guard {
61
- * guard = false ;
65
+ let mut guard = self . notify_pending . lock ( ) . unwrap ( ) ;
66
+ if guard. 0 {
67
+ guard. 0 = false ;
62
68
return true ;
63
69
}
64
- guard = cvar . wait_timeout ( guard, max_wait) . unwrap ( ) . 0 ;
70
+ guard = self . condvar . wait_timeout ( guard, max_wait) . unwrap ( ) . 0 ;
65
71
// Due to spurious wakeups that can happen on `wait_timeout`, here we need to check if the
66
72
// desired wait time has actually passed, and if not then restart the loop with a reduced wait
67
73
// time. Note that this logic can be highly simplified through the use of
68
74
// `Condvar::wait_while` and `Condvar::wait_timeout_while`, if and when our MSRV is raised to
69
75
// 1.42.0.
70
76
let elapsed = current_time. elapsed ( ) ;
71
- let result = * guard;
77
+ let result = guard. 0 ;
72
78
if result || elapsed >= max_wait {
73
- * guard = false ;
79
+ guard. 0 = false ;
74
80
return result;
75
81
}
76
82
match max_wait. checked_sub ( elapsed) {
@@ -82,29 +88,130 @@ impl Notifier {
82
88
83
89
/// Wake waiters, tracking that wake needs to occur even if there are currently no waiters.
84
90
pub ( crate ) fn notify ( & self ) {
85
- let & ( ref persist_mtx, ref cnd) = & self . lock ;
86
- let mut lock = persist_mtx. lock ( ) . unwrap ( ) ;
87
- * lock = true ;
91
+ let mut lock = self . notify_pending . lock ( ) . unwrap ( ) ;
92
+ lock. 0 = true ;
93
+ if let Some ( future_state) = lock. 1 . take ( ) {
94
+ future_state. lock ( ) . unwrap ( ) . complete ( ) ;
95
+ }
88
96
mem:: drop ( lock) ;
89
- cnd. notify_all ( ) ;
97
+ self . condvar . notify_all ( ) ;
98
+ }
99
+
100
+ /// Gets a [`Future`] that will get woken up with any waiters
101
+ pub ( crate ) fn get_future ( & self ) -> Future {
102
+ let mut lock = self . notify_pending . lock ( ) . unwrap ( ) ;
103
+ if lock. 0 {
104
+ Future {
105
+ state : Arc :: new ( Mutex :: new ( FutureState {
106
+ callbacks : Vec :: new ( ) ,
107
+ complete : false ,
108
+ } ) )
109
+ }
110
+ } else if let Some ( existing_state) = & lock. 1 {
111
+ Future { state : Arc :: clone ( & existing_state) }
112
+ } else {
113
+ let state = Arc :: new ( Mutex :: new ( FutureState {
114
+ callbacks : Vec :: new ( ) ,
115
+ complete : false ,
116
+ } ) ) ;
117
+ lock. 1 = Some ( Arc :: clone ( & state) ) ;
118
+ Future { state }
119
+ }
90
120
}
91
121
92
122
#[ cfg( any( test, feature = "_test_utils" ) ) ]
93
123
pub fn notify_pending ( & self ) -> bool {
94
- let & ( ref mtx, _) = & self . lock ;
95
- let guard = mtx. lock ( ) . unwrap ( ) ;
96
- * guard
124
+ self . notify_pending . lock ( ) . unwrap ( ) . 0
125
+ }
126
+ }
127
+
128
+ /// A callback which is called when a [`Future`] completes.
129
+ ///
130
+ /// Note that this MUST NOT call back into LDK directly, it must instead schedule actions to be
131
+ /// taken later. Rust users should use the [`std::future::Future`] implementation for [`Future`]
132
+ /// instead.
133
+ ///
134
+ /// Note that the [`std::future::Future`] implementation may only work for runtimes which schedule
135
+ /// futures when they receive a wake, rather than immediately executing them.
136
+ pub trait FutureCallback : Send {
137
+ /// The method which is called.
138
+ fn call ( & self ) ;
139
+ }
140
+
141
+ impl < F : Fn ( ) + Send > FutureCallback for F {
142
+ fn call ( & self ) { ( self ) ( ) ; }
143
+ }
144
+
145
+ pub ( crate ) struct FutureState {
146
+ callbacks : Vec < Box < dyn FutureCallback > > ,
147
+ complete : bool ,
148
+ }
149
+
150
+ impl FutureState {
151
+ fn complete ( & mut self ) {
152
+ for callback in self . callbacks . drain ( ..) {
153
+ callback. call ( ) ;
154
+ }
155
+ self . complete = true ;
156
+ }
157
+ }
158
+
159
+ /// A simple future which can complete once, and calls some callback(s) when it does so.
160
+ pub struct Future {
161
+ state : Arc < Mutex < FutureState > > ,
162
+ }
163
+
164
+ impl Future {
165
+ /// Registers a callback to be called upon completion of this future. If the future has already
166
+ /// completed, the callback will be called immediately.
167
+ pub fn register_callback ( & self , callback : Box < dyn FutureCallback > ) {
168
+ let mut state = self . state . lock ( ) . unwrap ( ) ;
169
+ if state. complete {
170
+ mem:: drop ( state) ;
171
+ callback. call ( ) ;
172
+ } else {
173
+ state. callbacks . push ( callback) ;
174
+ }
175
+ }
176
+ }
177
+
178
+ #[ cfg( feature = "std" ) ]
179
+ mod std_future {
180
+ use std:: task:: Waker ;
181
+ pub struct StdWaker ( pub Waker ) ;
182
+ impl super :: FutureCallback for StdWaker {
183
+ fn call ( & self ) { self . 0 . wake_by_ref ( ) }
184
+ }
185
+ }
186
+
187
+ #[ cfg( feature = "std" ) ]
188
+ /// (C-not exported) as Rust Futures aren't usable in language bindings.
189
+ impl < ' a > StdFuture for Future {
190
+ type Output = ( ) ;
191
+
192
+ fn poll ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Self :: Output > {
193
+ let mut state = self . state . lock ( ) . unwrap ( ) ;
194
+ if state. complete {
195
+ Poll :: Ready ( ( ) )
196
+ } else {
197
+ let waker = cx. waker ( ) . clone ( ) ;
198
+ state. callbacks . push ( Box :: new ( std_future:: StdWaker ( waker) ) ) ;
199
+ Poll :: Pending
200
+ }
97
201
}
98
202
}
99
203
100
204
#[ cfg( test) ]
101
205
mod tests {
206
+ use super :: * ;
207
+ use core:: sync:: atomic:: { AtomicBool , Ordering } ;
208
+ use core:: future:: Future as FutureTrait ;
209
+ use core:: task:: { Context , Poll , RawWaker , RawWakerVTable , Waker } ;
210
+
102
211
#[ cfg( feature = "std" ) ]
103
212
#[ test]
104
213
fn test_wait_timeout ( ) {
105
- use super :: * ;
106
214
use sync:: Arc ;
107
- use core:: sync:: atomic:: { AtomicBool , Ordering } ;
108
215
use std:: thread;
109
216
110
217
let persistence_notifier = Arc :: new ( Notifier :: new ( ) ) ;
@@ -114,10 +221,9 @@ mod tests {
114
221
let exit_thread_clone = exit_thread. clone ( ) ;
115
222
thread:: spawn ( move || {
116
223
loop {
117
- let & ( ref persist_mtx, ref cnd) = & thread_notifier. lock ;
118
- let mut lock = persist_mtx. lock ( ) . unwrap ( ) ;
119
- * lock = true ;
120
- cnd. notify_all ( ) ;
224
+ let mut lock = thread_notifier. notify_pending . lock ( ) . unwrap ( ) ;
225
+ lock. 0 = true ;
226
+ thread_notifier. condvar . notify_all ( ) ;
121
227
122
228
if exit_thread_clone. load ( Ordering :: SeqCst ) {
123
229
break
@@ -146,4 +252,84 @@ mod tests {
146
252
}
147
253
}
148
254
}
255
+
256
+ #[ test]
257
+ fn test_future_callbacks ( ) {
258
+ let future = Future {
259
+ state : Arc :: new ( Mutex :: new ( FutureState {
260
+ callbacks : Vec :: new ( ) ,
261
+ complete : false ,
262
+ } ) )
263
+ } ;
264
+ let callback = Arc :: new ( AtomicBool :: new ( false ) ) ;
265
+ let callback_ref = Arc :: clone ( & callback) ;
266
+ future. register_callback ( Box :: new ( move || assert ! ( !callback_ref. fetch_or( true , Ordering :: SeqCst ) ) ) ) ;
267
+
268
+ assert ! ( !callback. load( Ordering :: SeqCst ) ) ;
269
+ future. state . lock ( ) . unwrap ( ) . complete ( ) ;
270
+ assert ! ( callback. load( Ordering :: SeqCst ) ) ;
271
+ future. state . lock ( ) . unwrap ( ) . complete ( ) ;
272
+ }
273
+
274
+ #[ test]
275
+ fn test_pre_completed_future_callbacks ( ) {
276
+ let future = Future {
277
+ state : Arc :: new ( Mutex :: new ( FutureState {
278
+ callbacks : Vec :: new ( ) ,
279
+ complete : false ,
280
+ } ) )
281
+ } ;
282
+ future. state . lock ( ) . unwrap ( ) . complete ( ) ;
283
+
284
+ let callback = Arc :: new ( AtomicBool :: new ( false ) ) ;
285
+ let callback_ref = Arc :: clone ( & callback) ;
286
+ future. register_callback ( Box :: new ( move || assert ! ( !callback_ref. fetch_or( true , Ordering :: SeqCst ) ) ) ) ;
287
+
288
+ assert ! ( callback. load( Ordering :: SeqCst ) ) ;
289
+ assert ! ( future. state. lock( ) . unwrap( ) . callbacks. is_empty( ) ) ;
290
+ }
291
+
292
+ // Rather annoyingly, there's no safe way in Rust std to construct a Waker despite it being
293
+ // totally possible to construct from a trait implementation (though somewhat less effecient
294
+ // compared to a raw VTable). Instead, we have to write out a lot of boilerplate to build a
295
+ // waker, which we do here with a trivial Arc<AtomicBool> data element to track woke-ness.
296
+ const WAKER_V_TABLE : RawWakerVTable = RawWakerVTable :: new ( waker_clone, wake, wake_by_ref, drop) ;
297
+ unsafe fn wake_by_ref ( ptr : * const ( ) ) { let p = ptr as * const Arc < AtomicBool > ; assert ! ( !( * p) . fetch_or( true , Ordering :: SeqCst ) ) ; }
298
+ unsafe fn drop ( ptr : * const ( ) ) { let p = ptr as * mut Arc < AtomicBool > ; Box :: from_raw ( p) ; }
299
+ unsafe fn wake ( ptr : * const ( ) ) { wake_by_ref ( ptr) ; drop ( ptr) ; }
300
+ unsafe fn waker_clone ( ptr : * const ( ) ) -> RawWaker {
301
+ let p = ptr as * const Arc < AtomicBool > ;
302
+ RawWaker :: new ( Box :: into_raw ( Box :: new ( Arc :: clone ( & * p) ) ) as * const ( ) , & WAKER_V_TABLE )
303
+ }
304
+
305
+ fn create_waker ( ) -> ( Arc < AtomicBool > , Waker ) {
306
+ let a = Arc :: new ( AtomicBool :: new ( false ) ) ;
307
+ let waker = unsafe { Waker :: from_raw ( waker_clone ( ( & a as * const Arc < AtomicBool > ) as * const ( ) ) ) } ;
308
+ ( a, waker)
309
+ }
310
+
311
+ #[ test]
312
+ fn test_future ( ) {
313
+ let mut future = Future {
314
+ state : Arc :: new ( Mutex :: new ( FutureState {
315
+ callbacks : Vec :: new ( ) ,
316
+ complete : false ,
317
+ } ) )
318
+ } ;
319
+ let mut second_future = Future { state : Arc :: clone ( & future. state ) } ;
320
+
321
+ let ( woken, waker) = create_waker ( ) ;
322
+ assert_eq ! ( Pin :: new( & mut future) . poll( & mut Context :: from_waker( & waker) ) , Poll :: Pending ) ;
323
+ assert ! ( !woken. load( Ordering :: SeqCst ) ) ;
324
+
325
+ let ( second_woken, second_waker) = create_waker ( ) ;
326
+ assert_eq ! ( Pin :: new( & mut second_future) . poll( & mut Context :: from_waker( & second_waker) ) , Poll :: Pending ) ;
327
+ assert ! ( !second_woken. load( Ordering :: SeqCst ) ) ;
328
+
329
+ future. state . lock ( ) . unwrap ( ) . complete ( ) ;
330
+ assert ! ( woken. load( Ordering :: SeqCst ) ) ;
331
+ assert ! ( second_woken. load( Ordering :: SeqCst ) ) ;
332
+ assert_eq ! ( Pin :: new( & mut future) . poll( & mut Context :: from_waker( & waker) ) , Poll :: Ready ( ( ) ) ) ;
333
+ assert_eq ! ( Pin :: new( & mut second_future) . poll( & mut Context :: from_waker( & second_waker) ) , Poll :: Ready ( ( ) ) ) ;
334
+ }
149
335
}
0 commit comments