sending.rs: try to preserve requester spans

Add a `follows_from` relationship to the dispatched worker's span
pointing to the caller of `send_*()`, if known.
This commit is contained in:
Lambda 2024-05-23 20:11:01 +00:00
parent 18992b4d1d
commit fb52ded222

View file

@ -35,7 +35,7 @@ use tokio::{
select,
sync::{mpsc, Mutex, Semaphore},
};
use tracing::{debug, error, warn};
use tracing::{debug, error, warn, Span};
use crate::{
api::{appservice_server, server_server},
@ -105,6 +105,8 @@ pub(crate) struct RequestData {
destination: Destination,
event_type: SendingEventType,
key: RequestKey,
/// Span of the original `send_*()` method call
requester_span: Span,
}
pub(crate) struct Service {
@ -128,17 +130,27 @@ enum TransactionStatus {
struct HandlerInputs {
destination: Destination,
events: Vec<SendingEventType>,
/// Span of the original `send_*()` method call, if known (gets lost when
/// event is persisted to database)
requester_span: Option<Span>,
}
type HandlerResponse = Result<Destination, (Destination, Error)>;
fn destination_from_response(response: &HandlerResponse) -> &Destination {
match response {
Ok(kind) | Err((kind, _)) => kind,
}
#[derive(Debug)]
struct HandlerResponse {
destination: Destination,
result: Result<()>,
/// The span of the just-completed handler, for follows-from relationships.
handler_span: Span,
}
type TransactionStatusMap = HashMap<Destination, TransactionStatus>;
enum SelectedEvents {
None,
Retries(Vec<SendingEventType>),
New(Vec<SendingEventType>),
}
impl Service {
pub(crate) fn build(db: &'static dyn Data, config: &Config) -> Arc<Self> {
let (sender, receiver) = mpsc::unbounded_channel();
@ -191,30 +203,28 @@ impl Service {
for (destination, events) in initial_transactions {
current_transaction_status
.insert(destination.clone(), TransactionStatus::Running);
futures.push(handle_events(destination.clone(), events));
futures.push(handle_events(HandlerInputs {
destination: destination.clone(),
events,
requester_span: None,
}));
}
loop {
select! {
Some(response) = futures.next() => {
if let Some(HandlerInputs {
destination,
events,
}) = self.handle_response(
if let Some(inputs) = self.handle_response(
response,
&mut current_transaction_status,
)? {
futures.push(handle_events(destination, events));
futures.push(handle_events(inputs));
}
}
Some(data) = receiver.recv() => {
if let Some(HandlerInputs {
destination,
events,
}) = self
.handle_receiver(data, &mut current_transaction_status)
{
futures.push(handle_events(destination, events));
if let Some(inputs) = self.handle_receiver(
data, &mut current_transaction_status
) {
futures.push(handle_events(inputs));
}
}
}
@ -222,20 +232,27 @@ impl Service {
}
#[tracing::instrument(
skip(self, current_transaction_status),
skip(self, handler_span, current_transaction_status),
fields(
current_status = ?current_transaction_status.get(
destination_from_response(&response)
&destination
),
),
)]
fn handle_response(
&self,
response: HandlerResponse,
HandlerResponse {
destination,
result,
handler_span,
}: HandlerResponse,
current_transaction_status: &mut TransactionStatusMap,
) -> Result<Option<HandlerInputs>> {
match response {
Ok(destination) => {
// clone() is required for the relationship to show up in jaeger
Span::current().follows_from(handler_span.clone());
match result {
Ok(()) => {
self.db.delete_all_active_requests_for(&destination)?;
// Find events that have been added since starting the
@ -260,10 +277,12 @@ impl Service {
.into_iter()
.map(|(event, _)| event)
.collect(),
requester_span: None,
}))
}
}
Err((destination, _)) => {
Err(_err) => {
warn!("Marking transaction as failed");
current_transaction_status.entry(destination).and_modify(|e| {
*e = match e {
TransactionStatus::Running => {
@ -286,7 +305,7 @@ impl Service {
}
#[tracing::instrument(
skip(self, event_type, key, current_transaction_status),
skip(self, event_type, key, requester_span, current_transaction_status),
fields(
current_status = ?current_transaction_status.get(&destination),
),
@ -297,21 +316,43 @@ impl Service {
destination,
event_type,
key,
requester_span,
}: RequestData,
current_transaction_status: &mut TransactionStatusMap,
) -> Option<HandlerInputs> {
if let Ok(Some(events)) = self.select_events(
// clone() is required for the relationship to show up in jaeger
Span::current().follows_from(requester_span.clone());
match self.select_events(
&destination,
vec![(event_type, key)],
current_transaction_status,
) {
Ok(SelectedEvents::Retries(events)) => {
debug!("retrying old events");
Some(HandlerInputs {
destination,
events,
requester_span: None,
})
} else {
}
Ok(SelectedEvents::New(events)) => {
debug!("sending new event");
Some(HandlerInputs {
destination,
events,
requester_span: Some(requester_span),
})
}
Ok(SelectedEvents::None) => {
debug!("holding off from sending any events");
None
}
Err(error) => {
error!(%error, "Failed to select events to send");
None
}
}
}
#[tracing::instrument(
@ -330,7 +371,7 @@ impl Service {
Destination,
TransactionStatus,
>,
) -> Result<Option<Vec<SendingEventType>>> {
) -> Result<SelectedEvents> {
let mut retry = false;
let mut allow = true;
@ -364,19 +405,22 @@ impl Service {
.or_insert(TransactionStatus::Running);
if !allow {
return Ok(None);
return Ok(SelectedEvents::None);
}
let mut events = Vec::new();
if retry {
// We retry the previous transaction
for (_, e) in
self.db.active_requests_for(destination).filter_map(Result::ok)
{
events.push(e);
}
let events = self
.db
.active_requests_for(destination)
.filter_map(Result::ok)
.map(|(_, e)| e)
.collect();
Ok(SelectedEvents::Retries(events))
} else {
let mut events = Vec::new();
self.db.mark_as_active(&new_events)?;
for (e, _) in new_events {
events.push(e);
@ -393,9 +437,9 @@ impl Service {
self.db.set_latest_educount(server_name, last_count)?;
}
}
}
Ok(Some(events))
Ok(SelectedEvents::New(events))
}
}
#[tracing::instrument(skip(self))]
@ -537,6 +581,7 @@ impl Service {
destination,
event_type,
key: keys.into_iter().next().unwrap(),
requester_span: Span::current(),
})
.unwrap();
@ -567,6 +612,7 @@ impl Service {
destination: destination.clone(),
event_type,
key,
requester_span: Span::current(),
})
.unwrap();
}
@ -590,6 +636,7 @@ impl Service {
destination,
event_type,
key: keys.into_iter().next().unwrap(),
requester_span: Span::current(),
})
.unwrap();
@ -611,6 +658,7 @@ impl Service {
destination,
event_type,
key: keys.into_iter().next().unwrap(),
requester_span: Span::current(),
})
.unwrap();
@ -883,10 +931,18 @@ async fn handle_federation_event(
#[tracing::instrument(skip_all)]
async fn handle_events(
destination: Destination,
events: Vec<SendingEventType>,
HandlerInputs {
destination,
events,
requester_span,
}: HandlerInputs,
) -> HandlerResponse {
let ret = match &destination {
if let Some(span) = requester_span {
// clone() is required for the relationship to show up in jaeger
Span::current().follows_from(span.clone());
}
let result = match &destination {
Destination::Appservice(id) => {
handle_appservice_event(id, events).await
}
@ -898,8 +954,9 @@ async fn handle_events(
}
};
match ret {
Ok(()) => Ok(destination),
Err(e) => Err((destination, e)),
HandlerResponse {
destination,
result,
handler_span: Span::current(),
}
}