1- use crate :: notification_listener:: { NotificationListener , SafeChannelName } ;
1+ use futures:: sync:: mpsc:: { channel, Sender } ;
2+ use std:: collections:: HashMap ;
3+ use std:: sync:: { Arc , RwLock } ;
4+ use uuid:: Uuid ;
5+
26use graph:: prelude:: { ChainHeadUpdateListener as ChainHeadUpdateListenerTrait , * } ;
37use graph:: serde_json;
48
9+ use crate :: notification_listener:: { NotificationListener , SafeChannelName } ;
10+
11+ type ChainHeadUpdateSubscribers = Arc < RwLock < HashMap < String , Sender < ChainHeadUpdate > > > > ;
12+
513pub struct ChainHeadUpdateListener {
6- notification_listener : NotificationListener ,
7- network_name : String ,
14+ logger : Logger ,
15+ subscribers : ChainHeadUpdateSubscribers ,
16+ _listener : NotificationListener ,
817}
918
1019impl ChainHeadUpdateListener {
1120 pub fn new ( logger : & Logger , postgres_url : String , network_name : String ) -> Self {
21+ let logger = logger. new ( o ! ( "component" => "ChainHeadUpdateListener" ) ) ;
22+ let subscribers = Arc :: new ( RwLock :: new ( HashMap :: new ( ) ) ) ;
23+
24+ // Create a Postgres notification listener for chain head updates
25+ let mut listener = NotificationListener :: new (
26+ & logger,
27+ postgres_url,
28+ SafeChannelName :: i_promise_this_is_safe ( "chain_head_updates" ) ,
29+ ) ;
30+
31+ Self :: listen ( & logger, & mut listener, network_name, subscribers. clone ( ) ) ;
32+
1233 ChainHeadUpdateListener {
13- notification_listener : NotificationListener :: new (
14- logger,
15- postgres_url,
16- SafeChannelName :: i_promise_this_is_safe ( "chain_head_updates" ) ,
17- ) ,
18- network_name,
19- }
20- }
21- }
34+ logger,
35+ subscribers,
2236
23- impl ChainHeadUpdateListenerTrait for ChainHeadUpdateListener {
24- fn start ( & mut self ) {
25- self . notification_listener . start ( )
37+ // We keep the listener around to tie its stream's lifetime to
38+ // that of the chain head update listener and prevent it from
39+ // terminating early
40+ _listener : listener,
41+ }
2642 }
27- }
2843
29- impl EventProducer < ChainHeadUpdate > for ChainHeadUpdateListener {
30- fn take_event_stream (
31- & mut self ,
32- ) -> Option < Box < Stream < Item = ChainHeadUpdate , Error = ( ) > + Send > > {
33- let network_name = self . network_name . clone ( ) ;
44+ fn listen (
45+ logger : & Logger ,
46+ listener : & mut NotificationListener ,
47+ network_name : String ,
48+ subscribers : ChainHeadUpdateSubscribers ,
49+ ) {
50+ let logger = logger. clone ( ) ;
3451
35- self . notification_listener . take_event_stream ( ) . map (
36- move |stream| -> Box < Stream < Item = _ , Error = _ > + Send > {
37- Box :: new ( stream. filter_map ( move |notification| {
52+ // Process chain head updates in a dedicated task
53+ tokio:: spawn (
54+ listener
55+ . take_event_stream ( )
56+ . unwrap ( )
57+ . filter_map ( move |notification| {
3858 // Create ChainHeadUpdate from JSON
3959 let update: ChainHeadUpdate =
4060 serde_json:: from_value ( notification. payload . clone ( ) ) . unwrap_or_else ( |_| {
@@ -44,14 +64,65 @@ impl EventProducer<ChainHeadUpdate> for ChainHeadUpdateListener {
4464 )
4565 } ) ;
4666
47- // Only include update if about the right network
67+ // Only include update if it is for the network we're interested in
4868 if update. network_name == network_name {
4969 Some ( update)
5070 } else {
5171 None
5272 }
53- } ) )
54- } ,
55- )
73+ } )
74+ . for_each ( move |update| {
75+ let logger = logger. clone ( ) ;
76+ let senders = subscribers. read ( ) . unwrap ( ) . clone ( ) ;
77+ let subscribers = subscribers. clone ( ) ;
78+
79+ debug ! (
80+ logger,
81+ "Received chain head update" ;
82+ "network" => & update. network_name,
83+ "head_block_hash" => format!( "{}" , update. head_block_hash) ,
84+ "head_block_number" => & update. head_block_number,
85+ ) ;
86+
87+ // Forward update to all susbcribers
88+ stream:: iter_ok :: < _ , ( ) > ( senders) . for_each ( move |( id, sender) | {
89+ let logger = logger. clone ( ) ;
90+ let subscribers = subscribers. clone ( ) ;
91+
92+ sender. send ( update. clone ( ) ) . then ( move |result| {
93+ if result. is_err ( ) {
94+ // If sending to a subscriber fails, we'll assume that
95+ // the receiving end has been dropped. In this case we
96+ // remove the subscriber
97+ debug ! ( logger, "Unsubscribe" ; "id" => & id) ;
98+ subscribers. write ( ) . unwrap ( ) . remove ( & id) ;
99+ }
100+
101+ // Move on to the next subscriber
102+ Ok ( ( ) )
103+ } )
104+ } )
105+ } ) ,
106+ ) ;
107+
108+ // We're ready, start listening to chain head updaates
109+ listener. start ( ) ;
110+ }
111+ }
112+
113+ impl ChainHeadUpdateListenerTrait for ChainHeadUpdateListener {
114+ fn subscribe ( & self ) -> ChainHeadUpdateStream {
115+ // Generate a new (unique) UUID; we're looping just to be sure we avoid collisions
116+ let mut id = Uuid :: new_v4 ( ) . to_string ( ) ;
117+ while self . subscribers . read ( ) . unwrap ( ) . contains_key ( & id) {
118+ id = Uuid :: new_v4 ( ) . to_string ( ) ;
119+ }
120+
121+ debug ! ( self . logger, "Subscribe" ; "id" => & id) ;
122+
123+ // Create a subscriber and return the receiving end
124+ let ( sender, receiver) = channel ( 100 ) ;
125+ self . subscribers . write ( ) . unwrap ( ) . insert ( id, sender) ;
126+ Box :: new ( receiver)
56127 }
57128}
0 commit comments