fix: sending code got stuck sometimes
This commit is contained in:
		
							parent
							
								
									566b8ebabb
								
							
						
					
					
						commit
						f7713fdf2e
					
				
					 2 changed files with 59 additions and 23 deletions
				
			
		|  | @ -8,7 +8,7 @@ use std::{ | ||||||
| 
 | 
 | ||||||
| use crate::{appservice_server, server_server, utils, Error, PduEvent, Result}; | use crate::{appservice_server, server_server, utils, Error, PduEvent, Result}; | ||||||
| use federation::transactions::send_transaction_message; | use federation::transactions::send_transaction_message; | ||||||
| use log::info; | use log::{info, warn}; | ||||||
| use ring::digest; | use ring::digest; | ||||||
| use rocket::futures::stream::{FuturesUnordered, StreamExt}; | use rocket::futures::stream::{FuturesUnordered, StreamExt}; | ||||||
| use ruma::{ | use ruma::{ | ||||||
|  | @ -36,6 +36,7 @@ impl Sending { | ||||||
|     ) { |     ) { | ||||||
|         let servernamepduids = self.servernamepduids.clone(); |         let servernamepduids = self.servernamepduids.clone(); | ||||||
|         let servercurrentpdus = self.servercurrentpdus.clone(); |         let servercurrentpdus = self.servercurrentpdus.clone(); | ||||||
|  |         let maximum_requests = self.maximum_requests.clone(); | ||||||
|         let rooms = rooms.clone(); |         let rooms = rooms.clone(); | ||||||
|         let globals = globals.clone(); |         let globals = globals.clone(); | ||||||
|         let appservice = appservice.clone(); |         let appservice = appservice.clone(); | ||||||
|  | @ -44,23 +45,43 @@ impl Sending { | ||||||
|             let mut futures = FuturesUnordered::new(); |             let mut futures = FuturesUnordered::new(); | ||||||
| 
 | 
 | ||||||
|             // Retry requests we could not finish yet
 |             // Retry requests we could not finish yet
 | ||||||
|             let mut current_transactions = HashMap::new(); |             let mut current_transactions = HashMap::<(Box<ServerName>, bool), Vec<IVec>>::new(); | ||||||
| 
 | 
 | ||||||
|             for (server, pdu, is_appservice) in servercurrentpdus |             for (key, server, pdu, is_appservice) in servercurrentpdus | ||||||
|                 .iter() |                 .iter() | ||||||
|                 .filter_map(|r| r.ok()) |                 .filter_map(|r| r.ok()) | ||||||
|                 .filter_map(|(key, _)| Self::parse_servercurrentpdus(key).ok()) |                 .filter_map(|(key, _)| Self::parse_servercurrentpdus(key).ok()) | ||||||
|                 .filter(|(_, pdu, _)| !pdu.is_empty()) // Skip reservation key
 |  | ||||||
|                 .take(50) |  | ||||||
|             // This should not contain more than 50 anyway
 |  | ||||||
|             { |             { | ||||||
|                 current_transactions |                 if pdu.is_empty() { | ||||||
|  |                     // Remove old reservation key
 | ||||||
|  |                     servercurrentpdus.remove(key).unwrap(); | ||||||
|  |                     continue; | ||||||
|  |                 } | ||||||
|  | 
 | ||||||
|  |                 let entry = current_transactions | ||||||
|                     .entry((server, is_appservice)) |                     .entry((server, is_appservice)) | ||||||
|                     .or_insert_with(Vec::new) |                     .or_insert_with(Vec::new); | ||||||
|                     .push(pdu); | 
 | ||||||
|  |                 if entry.len() > 30 { | ||||||
|  |                     warn!("Dropping some current pdus because too many were queued. This should not happen."); | ||||||
|  |                     servercurrentpdus.remove(key).unwrap(); | ||||||
|  |                     continue; | ||||||
|  |                 } | ||||||
|  | 
 | ||||||
|  |                 entry.push(pdu); | ||||||
|             } |             } | ||||||
| 
 | 
 | ||||||
|             for ((server, is_appservice), pdus) in current_transactions { |             for ((server, is_appservice), pdus) in current_transactions { | ||||||
|  |                 // Create new reservation
 | ||||||
|  |                 let mut prefix = if is_appservice { | ||||||
|  |                     "+".as_bytes().to_vec() | ||||||
|  |                 } else { | ||||||
|  |                     Vec::new() | ||||||
|  |                 }; | ||||||
|  |                 prefix.extend_from_slice(server.as_bytes()); | ||||||
|  |                 prefix.push(0xff); | ||||||
|  |                 servercurrentpdus.insert(prefix, &[]).unwrap(); | ||||||
|  | 
 | ||||||
|                 futures.push(Self::handle_event( |                 futures.push(Self::handle_event( | ||||||
|                     server, |                     server, | ||||||
|                     is_appservice, |                     is_appservice, | ||||||
|  | @ -68,6 +89,7 @@ impl Sending { | ||||||
|                     &globals, |                     &globals, | ||||||
|                     &rooms, |                     &rooms, | ||||||
|                     &appservice, |                     &appservice, | ||||||
|  |                     &maximum_requests, | ||||||
|                 )); |                 )); | ||||||
|             } |             } | ||||||
| 
 | 
 | ||||||
|  | @ -106,7 +128,7 @@ impl Sending { | ||||||
|                                     .map(|k| { |                                     .map(|k| { | ||||||
|                                         k.subslice(prefix.len(), k.len() - prefix.len()) |                                         k.subslice(prefix.len(), k.len() - prefix.len()) | ||||||
|                                     }) |                                     }) | ||||||
|                                     .take(50) |                                     .take(30) | ||||||
|                                     .collect::<Vec<_>>(); |                                     .collect::<Vec<_>>(); | ||||||
| 
 | 
 | ||||||
|                                 if !new_pdus.is_empty() { |                                 if !new_pdus.is_empty() { | ||||||
|  | @ -117,7 +139,7 @@ impl Sending { | ||||||
|                                         servernamepduids.remove(¤t_key).unwrap(); |                                         servernamepduids.remove(¤t_key).unwrap(); | ||||||
|                                     } |                                     } | ||||||
| 
 | 
 | ||||||
|                                     futures.push(Self::handle_event(server, is_appservice, new_pdus, &globals, &rooms, &appservice)); |                                     futures.push(Self::handle_event(server, is_appservice, new_pdus, &globals, &rooms, &appservice, &maximum_requests)); | ||||||
|                                 } else { |                                 } else { | ||||||
|                                     servercurrentpdus.remove(&prefix).unwrap(); |                                     servercurrentpdus.remove(&prefix).unwrap(); | ||||||
|                                     // servercurrentpdus with the prefix should be empty now
 |                                     // servercurrentpdus with the prefix should be empty now
 | ||||||
|  | @ -194,15 +216,17 @@ impl Sending { | ||||||
|                                     prefix.extend_from_slice(server.as_bytes()); |                                     prefix.extend_from_slice(server.as_bytes()); | ||||||
|                                     prefix.push(0xff); |                                     prefix.push(0xff); | ||||||
| 
 | 
 | ||||||
|                                     servercurrentpdus |                                     if servercurrentpdus | ||||||
|                                         .compare_and_swap(prefix, Option::<&[u8]>::None, Some(&[])) // Try to reserve
 |                                         .compare_and_swap(prefix, Option::<&[u8]>::None, Some(&[])) // Try to reserve
 | ||||||
|                                         == Ok(Ok(())) |                                         == Ok(Ok(())) { true } else { | ||||||
|  |                                             false | ||||||
|  |                                     } | ||||||
|                                 }) |                                 }) | ||||||
|                             { |                             { | ||||||
|                                 servercurrentpdus.insert(&key, &[]).unwrap(); |                                 servercurrentpdus.insert(&key, &[]).unwrap(); | ||||||
|                                 servernamepduids.remove(&key).unwrap(); |                                 servernamepduids.remove(&key).unwrap(); | ||||||
| 
 | 
 | ||||||
|                                 futures.push(Self::handle_event(server, is_appservice, vec![pdu_id.into()], &globals, &rooms, &appservice)); |                                 futures.push(Self::handle_event(server, is_appservice, vec![pdu_id.into()], &globals, &rooms, &appservice, &maximum_requests)); | ||||||
|                             } |                             } | ||||||
|                         } |                         } | ||||||
|                     } |                     } | ||||||
|  | @ -244,6 +268,7 @@ impl Sending { | ||||||
|         globals: &super::globals::Globals, |         globals: &super::globals::Globals, | ||||||
|         rooms: &super::rooms::Rooms, |         rooms: &super::rooms::Rooms, | ||||||
|         appservice: &super::appservice::Appservice, |         appservice: &super::appservice::Appservice, | ||||||
|  |         maximum_requests: &Semaphore, | ||||||
|     ) -> std::result::Result<(Box<ServerName>, bool), (Box<ServerName>, bool, Error)> { |     ) -> std::result::Result<(Box<ServerName>, bool), (Box<ServerName>, bool, Error)> { | ||||||
|         if is_appservice { |         if is_appservice { | ||||||
|             let pdu_jsons = pdu_ids |             let pdu_jsons = pdu_ids | ||||||
|  | @ -266,7 +291,9 @@ impl Sending { | ||||||
|                 }) |                 }) | ||||||
|                 .filter_map(|r| r.ok()) |                 .filter_map(|r| r.ok()) | ||||||
|                 .collect::<Vec<_>>(); |                 .collect::<Vec<_>>(); | ||||||
|             appservice_server::send_request( | 
 | ||||||
|  |             let permit = maximum_requests.acquire().await; | ||||||
|  |             let response = appservice_server::send_request( | ||||||
|                 &globals, |                 &globals, | ||||||
|                 appservice |                 appservice | ||||||
|                     .get_registration(server.as_str()) |                     .get_registration(server.as_str()) | ||||||
|  | @ -282,7 +309,11 @@ impl Sending { | ||||||
|             ) |             ) | ||||||
|             .await |             .await | ||||||
|             .map(|_response| (server.clone(), is_appservice)) |             .map(|_response| (server.clone(), is_appservice)) | ||||||
|             .map_err(|e| (server, is_appservice, e)) |             .map_err(|e| (server, is_appservice, e)); | ||||||
|  | 
 | ||||||
|  |             drop(permit); | ||||||
|  | 
 | ||||||
|  |             response | ||||||
|         } else { |         } else { | ||||||
|             let pdu_jsons = pdu_ids |             let pdu_jsons = pdu_ids | ||||||
|                 .iter() |                 .iter() | ||||||
|  | @ -312,7 +343,8 @@ impl Sending { | ||||||
|                 .filter_map(|r| r.ok()) |                 .filter_map(|r| r.ok()) | ||||||
|                 .collect::<Vec<_>>(); |                 .collect::<Vec<_>>(); | ||||||
| 
 | 
 | ||||||
|             server_server::send_request( |             let permit = maximum_requests.acquire().await; | ||||||
|  |             let response = server_server::send_request( | ||||||
|                 &globals, |                 &globals, | ||||||
|                 server.clone(), |                 server.clone(), | ||||||
|                 send_transaction_message::v1::Request { |                 send_transaction_message::v1::Request { | ||||||
|  | @ -328,12 +360,17 @@ impl Sending { | ||||||
|             ) |             ) | ||||||
|             .await |             .await | ||||||
|             .map(|_response| (server.clone(), is_appservice)) |             .map(|_response| (server.clone(), is_appservice)) | ||||||
|             .map_err(|e| (server, is_appservice, e)) |             .map_err(|e| (server, is_appservice, e)); | ||||||
|  | 
 | ||||||
|  |             drop(permit); | ||||||
|  | 
 | ||||||
|  |             response | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
| 
 | 
 | ||||||
|     fn parse_servercurrentpdus(key: IVec) -> Result<(Box<ServerName>, IVec, bool)> { |     fn parse_servercurrentpdus(key: IVec) -> Result<(IVec, Box<ServerName>, IVec, bool)> { | ||||||
|         let mut parts = key.splitn(2, |&b| b == 0xff); |         let key2 = key.clone(); | ||||||
|  |         let mut parts = key2.splitn(2, |&b| b == 0xff); | ||||||
|         let server = parts.next().expect("splitn always returns one element"); |         let server = parts.next().expect("splitn always returns one element"); | ||||||
|         let pdu = parts |         let pdu = parts | ||||||
|             .next() |             .next() | ||||||
|  | @ -351,6 +388,7 @@ impl Sending { | ||||||
|         }; |         }; | ||||||
| 
 | 
 | ||||||
|         Ok::<_, Error>(( |         Ok::<_, Error>(( | ||||||
|  |             key, | ||||||
|             Box::<ServerName>::try_from(server).map_err(|_| { |             Box::<ServerName>::try_from(server).map_err(|_| { | ||||||
|                 Error::bad_database("Invalid server string in server_currenttransaction") |                 Error::bad_database("Invalid server string in server_currenttransaction") | ||||||
|             })?, |             })?, | ||||||
|  |  | ||||||
|  | @ -82,9 +82,7 @@ where | ||||||
|                         registration |                         registration | ||||||
|                             .get("as_token") |                             .get("as_token") | ||||||
|                             .and_then(|as_token| as_token.as_str()) |                             .and_then(|as_token| as_token.as_str()) | ||||||
|                             .map_or(false, |as_token| { |                             .map_or(false, |as_token| token.as_deref() == Some(as_token)) | ||||||
|                                 dbg!(token.as_deref()) == dbg!(Some(as_token)) |  | ||||||
|                             }) |  | ||||||
|                     }) { |                     }) { | ||||||
|                 match T::METADATA.authentication { |                 match T::METADATA.authentication { | ||||||
|                     AuthScheme::AccessToken | AuthScheme::QueryOnlyAccessToken => { |                     AuthScheme::AccessToken | AuthScheme::QueryOnlyAccessToken => { | ||||||
|  |  | ||||||
		Loading…
	
		Reference in a new issue