f add callbacker

Author: TheBlueMatt

src/onion_message_resolver.rs

@@ -7,7 +7,7 @@ use std::future::Future;
 use std::ops::Deref;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, RwLock};
 use std::task::{Context, Poll, Waker};
 use std::vec::Vec;
 
@@ -82,10 +82,13 @@ fn channel() -> (ChannelSend, ChannelRecv) {
 /// A [`HrnResolver`] which uses lightning onion messages and DNSSEC proofs to request DNS
 /// resolution directly from untrusted lightning nodes, providing privacy through onion routing.
 ///
-/// This implements LDK's [`DNSResolverMessageHandler`], which it uses to send onion messages (you
-/// should make sure to call LDK's [`PeerManager::process_events`] after a query begins) and
+/// This implements LDK's [`DNSResolverMessageHandler`], which it uses to send onion messages and
 /// process response messages.
 ///
+/// Note that after a query begines, [`PeerManager::process_events`] should be called to ensure the
+/// query message goes out in a timely manner. You can call [`Self::register_post_queue_action`] to
+/// have this happen automatically.
+///
 /// [`PeerManager::process_events`]: lightning::ln::peer_handler::PeerManager::process_events
 pub struct LDKOnionMessageDNSSECHrnResolver<N: Deref<Target = NetworkGraph<L>>, L: Deref>
 where
@@ -96,6 +99,7 @@ where
 	next_id: AtomicUsize,
 	pending_resolutions: Mutex<HashMap<HumanReadableName, Vec<(PaymentId, ChannelSend)>>>,
 	message_queue: Mutex<Vec<(DNSResolverMessage, MessageSendInstructions)>>,
+	pm_event_poker: RwLock<Option<Box<dyn Fn() + Send + Sync>>>,
 }
 
 impl<N: Deref<Target = NetworkGraph<L>>, L: Deref> LDKOnionMessageDNSSECHrnResolver<N, L>
@@ -113,9 +117,18 @@ where
 			resolver: OMNameResolver::new(0, 0),
 			pending_resolutions: Mutex::new(HashMap::new()),
 			message_queue: Mutex::new(Vec::new()),
+			pm_event_poker: RwLock::new(None),
 		}
 	}
 
+	/// Sets a callback which is called any time a new resolution begins and a message is available
+	/// to be sent. This should generally call [`PeerManager::process_events`].
+	///
+	/// [`PeerManager::process_events`]: lightning::ln::peer_handler::PeerManager::process_events
+	pub fn register_post_queue_action(&self, callback: Box<dyn Fn() + Send + Sync>) {
+		*self.pm_event_poker.write().unwrap() = Some(callback);
+	}
+
 	fn init_resolve_hrn<'a>(
 		&'a self, hrn: &HumanReadableName,
 	) -> Result<ChannelRecv, &'static str> {
@@ -168,31 +181,40 @@ where
 			self.resolver.resolve_name(payment_id, hrn.clone(), &OsRng).map_err(|_| err)?;
 		let context = MessageContext::DNSResolver(dns_context);
 
-		let mut queue = self.message_queue.lock().unwrap();
-		for destination in dns_resolvers {
-			let instructions =
-				MessageSendInstructions::WithReplyPath { destination, context: context.clone() };
-			queue.push((DNSResolverMessage::DNSSECQuery(query.clone()), instructions));
-		}
-
 		let (send, recv) = channel();
-		let mut pending_resolutions = self.pending_resolutions.lock().unwrap();
-		let senders = pending_resolutions.entry(hrn.clone()).or_insert_with(Vec::new);
-		senders.push((payment_id, send));
-
-		// If we're running in no-std, we won't expire lookups with the time updates above, so walk
-		// the pending resolution list and expire them here.
-		pending_resolutions.retain(|_name, resolutions| {
-			resolutions.retain(|(_payment_id, resolution)| {
-				let has_receiver = resolution.receiver_alive();
-				if !has_receiver {
-					// TODO: Once LDK 0.2 ships, expire the pending resolution in the resolver:
-					// self.resolver.expire_pending_resolution(name, payment_id);
-				}
-				has_receiver
+		{
+			let mut pending_resolutions = self.pending_resolutions.lock().unwrap();
+			let senders = pending_resolutions.entry(hrn.clone()).or_insert_with(Vec::new);
+			senders.push((payment_id, send));
+
+			// If we're running in no-std, we won't expire lookups with the time updates above, so walk
+			// the pending resolution list and expire them here.
+			pending_resolutions.retain(|_name, resolutions| {
+				resolutions.retain(|(_payment_id, resolution)| {
+					let has_receiver = resolution.receiver_alive();
+					if !has_receiver {
+						// TODO: Once LDK 0.2 ships, expire the pending resolution in the resolver:
+						// self.resolver.expire_pending_resolution(name, payment_id);
+					}
+					has_receiver
+				});
+				!resolutions.is_empty()
 			});
-			!resolutions.is_empty()
-		});
+		}
+
+		{
+			let mut queue = self.message_queue.lock().unwrap();
+			for destination in dns_resolvers {
+				let instructions =
+					MessageSendInstructions::WithReplyPath { destination, context: context.clone() };
+				queue.push((DNSResolverMessage::DNSSECQuery(query.clone()), instructions));
+			}
+		}
+
+		let callback = self.pm_event_poker.read().unwrap();
+		if let Some(callback) = &*callback {
+			callback();
+		}
 
 		Ok(recv)
 	}