Remove peerState if there are no pending requests with that peer

Author: martinsaposnic

lightning-liquidity/src/lsps5/client.rs

@@ -56,6 +56,12 @@ impl PeerState {
 			pending_remove_webhook_requests: BoundedMap::new(MAX_PENDING_REQUESTS),
 		}
 	}
+
+	fn is_empty(&self) -> bool {
+		self.pending_set_webhook_requests.is_empty()
+			&& self.pending_list_webhooks_requests.is_empty()
+			&& self.pending_remove_webhook_requests.is_empty()
+	}
 }
 
 /// Client-side handler for the LSPS5 (bLIP-55) webhook registration protocol.
@@ -345,8 +351,26 @@ where
 			}
 		};
 		self.with_peer_state(*counterparty_node_id, handle_response);
+
+		self.check_and_remove_empty_peer_state(counterparty_node_id);
+
 		result
 	}
+
+	fn check_and_remove_empty_peer_state(&self, counterparty_node_id: &PublicKey) {
+		let mut outer_state_lock = self.per_peer_state.write().unwrap();
+		let should_remove =
+			if let Some(peer_state_mutex) = outer_state_lock.get(counterparty_node_id) {
+				let peer_state = peer_state_mutex.lock().unwrap();
+				peer_state.is_empty()
+			} else {
+				false
+			};
+
+		if should_remove {
+			outer_state_lock.remove(counterparty_node_id);
+		}
+	}
 }
 
 impl<ES: Deref> LSPSProtocolMessageHandler for LSPS5ClientHandler<ES>
@@ -461,36 +485,6 @@ mod tests {
 		}
 	}
 
-	#[test]
-	fn test_handle_response_clears_pending_state() {
-		let (client, _, _, peer, _) = setup_test_client();
-
-		let req_id = client
-			.set_webhook(peer, "test-app".to_string(), "https://example.com/hook".to_string())
-			.unwrap();
-
-		let response = LSPS5Response::SetWebhook(SetWebhookResponse {
-			num_webhooks: 1,
-			max_webhooks: 5,
-			no_change: false,
-		});
-		let response_msg = LSPS5Message::Response(req_id.clone(), response);
-
-		{
-			let outer_state_lock = client.per_peer_state.read().unwrap();
-			let peer_state = outer_state_lock.get(&peer).unwrap().lock().unwrap();
-			assert!(peer_state.pending_set_webhook_requests.contains_key(&req_id));
-		}
-
-		client.handle_message(response_msg, &peer).unwrap();
-
-		{
-			let outer_state_lock = client.per_peer_state.read().unwrap();
-			let peer_state = outer_state_lock.get(&peer).unwrap().lock().unwrap();
-			assert!(!peer_state.pending_set_webhook_requests.contains_key(&req_id));
-		}
-	}
-
 	#[test]
 	fn test_unknown_request_id_handling() {
 		let (client, _message_queue, _, peer, _) = setup_test_client();
@@ -552,4 +546,68 @@ mod tests {
 			assert!(peer_state.pending_set_webhook_requests.contains_key(&new_req_id));
 		}
 	}
+
+	#[test]
+	fn test_peer_state_cleanup_and_recreation() {
+		let (client, _, _, peer, _) = setup_test_client();
+
+		let set_webhook_req_id = client
+			.set_webhook(peer, "test-app".to_string(), "https://example.com/hook".to_string())
+			.unwrap();
+
+		let list_webhooks_req_id = client.list_webhooks(peer);
+
+		{
+			let state = client.per_peer_state.read().unwrap();
+			assert!(state.contains_key(&peer));
+			let peer_state = state.get(&peer).unwrap().lock().unwrap();
+			assert!(peer_state.pending_set_webhook_requests.contains_key(&set_webhook_req_id));
+			assert!(peer_state.pending_list_webhooks_requests.contains_key(&list_webhooks_req_id));
+		}
+
+		let set_webhook_response = LSPS5Response::SetWebhook(SetWebhookResponse {
+			num_webhooks: 1,
+			max_webhooks: 5,
+			no_change: false,
+		});
+		let response_msg = LSPS5Message::Response(set_webhook_req_id.clone(), set_webhook_response);
+		// trigger cleanup but there is still a pending request
+		// so the peer state should not be removed
+		client.handle_message(response_msg, &peer).unwrap();
+
+		{
+			let state = client.per_peer_state.read().unwrap();
+			assert!(state.contains_key(&peer));
+			let peer_state = state.get(&peer).unwrap().lock().unwrap();
+			assert!(!peer_state.pending_set_webhook_requests.contains_key(&set_webhook_req_id));
+			assert!(peer_state.pending_list_webhooks_requests.contains_key(&list_webhooks_req_id));
+		}
+
+		let list_webhooks_response =
+			LSPS5Response::ListWebhooks(crate::lsps5::msgs::ListWebhooksResponse {
+				app_names: vec![],
+				max_webhooks: 5,
+			});
+		let response_msg = LSPS5Message::Response(list_webhooks_req_id, list_webhooks_response);
+
+		// now the pending request is handled, so the peer state should be removed
+		client.handle_message(response_msg, &peer).unwrap();
+
+		{
+			let state = client.per_peer_state.read().unwrap();
+			assert!(!state.contains_key(&peer));
+		}
+
+		// check that it's possible to recreate the peer state by sending a new request
+		let new_req_id = client
+			.set_webhook(peer, "test-app-2".to_string(), "https://example.com/hook2".to_string())
+			.unwrap();
+
+		{
+			let state = client.per_peer_state.read().unwrap();
+			assert!(state.contains_key(&peer));
+			let peer_state = state.get(&peer).unwrap().lock().unwrap();
+			assert!(peer_state.pending_set_webhook_requests.contains_key(&new_req_id));
+		}
+	}
 }

lightning-liquidity/src/utils/bounded_map.rs

@@ -74,6 +74,11 @@ where
 	pub(crate) fn contains_key(&self, key: &K) -> bool {
 		self.map.contains_key(key)
 	}
+
+	/// Check whether the map is empty.
+	pub(crate) fn is_empty(&self) -> bool {
+		self.map.is_empty()
+	}
 }
 
 #[cfg(test)]
@@ -87,21 +92,36 @@ mod tests {
 	fn new_empty() {
 		let m: BoundedMap<&str, i32> = BoundedMap::new(3);
 		assert_eq!(m.len(), 0);
+		assert!(m.is_empty());
 		assert!(m.get(&"key").is_none());
 	}
 
 	#[test]
 	fn insert_within_capacity() {
 		let mut m = BoundedMap::new(2);
+		assert!(m.is_empty());
 		m.insert("a", 1);
 		m.insert("b", 2);
 
+		assert!(!m.is_empty());
 		assert_eq!(m.len(), 2);
 		assert!(m.contains_key(&"a"));
 		assert_eq!(m.get(&"a"), Some(&1));
 		assert_eq!(m.get(&"b"), Some(&2));
 	}
 
+	#[test]
+	fn is_empty() {
+		let mut m: BoundedMap<&str, i32> = BoundedMap::new(3);
+		assert!(m.is_empty());
+
+		m.insert("a", 1);
+		assert!(!m.is_empty());
+
+		m.remove(&"a");
+		assert!(m.is_empty());
+	}
+
 	#[test]
 	fn eviction_fifo() {
 		let mut m = BoundedMap::new(2);
@@ -110,6 +130,7 @@ mod tests {
 		m.insert("third", 3); // evicts "first"
 
 		assert_eq!(m.len(), 2);
+		assert!(!m.is_empty());
 		assert!(m.get(&"first").is_none());
 		assert_eq!(m.get(&"second"), Some(&2));
 		assert_eq!(m.get(&"third"), Some(&3));
@@ -123,6 +144,7 @@ mod tests {
 		m.insert("old", 10); // update moves "old" to back
 		m.insert("newer", 3); // should evict "new", not "old"
 
+		assert!(!m.is_empty());
 		assert_eq!(m.get(&"old"), Some(&10));
 		assert!(m.get(&"new").is_none());
 		assert_eq!(m.get(&"newer"), Some(&3));
@@ -133,38 +155,50 @@ mod tests {
 		let mut m = BoundedMap::new(2);
 		m.insert("keep", 1);
 		m.insert("remove", 2);
+		assert!(!m.is_empty());
 
 		assert_eq!(m.remove(&"remove"), Some(2));
 		assert_eq!(m.len(), 1);
+		assert!(!m.is_empty());
 		assert!(m.get(&"remove").is_none());
 		assert_eq!(m.get(&"keep"), Some(&1));
+
+		m.remove(&"keep");
+		assert!(m.is_empty());
 	}
 
 	#[test]
 	fn remove_nonexistent() {
 		let mut m = BoundedMap::new(2);
 		m.insert("exists", 1);
+		assert!(!m.is_empty());
 
 		assert_eq!(m.remove(&"missing"), None);
 		assert_eq!(m.len(), 1);
+		assert!(!m.is_empty());
 	}
 
 	#[test]
 	fn zero_capacity() {
 		let mut m = BoundedMap::new(0);
+		assert!(m.is_empty());
 		m.insert("any", 42);
 
 		assert_eq!(m.len(), 0);
+		assert!(m.is_empty());
 		assert!(m.get(&"any").is_none());
 	}
 
 	#[test]
 	fn capacity_one() {
 		let mut m = BoundedMap::new(1);
+		assert!(m.is_empty());
 		m.insert("first", 1);
+		assert!(!m.is_empty());
 		assert_eq!(m.get(&"first"), Some(&1));
 
 		m.insert("second", 2);
+		assert!(!m.is_empty());
 		assert!(m.get(&"first").is_none());
 		assert_eq!(m.get(&"second"), Some(&2));
 	}