Fix spans in tokio::spawn-ed tasks

tokio::spawn is a span boundary, the spawned future has no parent span.

For short futures, we simply inherit the current span with
`.in_current_span()`.

For long running futures containing a sleeping infinite loop, we don't
actually want a span on the entire task or even the entire loop body,
both would result in very long spans. Instead, we put the outermost span
(created using #[tracing::instrument] or .instrument()) around the
actual work happening after the sleep, which results in a new root span
being created after every sleep.
This commit is contained in:
Lambda 2024-05-19 21:39:13 +00:00
parent 5e9e5b76bc
commit ac42e0bfff
3 changed files with 109 additions and 70 deletions

View file

@ -218,71 +218,84 @@ impl Service {
pub(crate) fn start_handler(self: &Arc<Self>) {
let self2 = Arc::clone(self);
tokio::spawn(async move {
self2.handler().await;
});
}
let mut receiver = self2.receiver.lock().await;
let grapevine_user = UserId::parse(format!(
"@{}:{}",
if services().globals.config.conduit_compat {
"conduit"
} else {
"grapevine"
},
services().globals.server_name()
))
.expect("admin bot username should be valid");
async fn handler(&self) {
let mut receiver = self.receiver.lock().await;
// TODO: Use futures when we have long admin commands
let Ok(Some(grapevine_room)) = services().admin.get_admin_room()
else {
return;
};
let grapevine_user = UserId::parse(format!(
"@{}:{}",
if services().globals.config.conduit_compat {
"conduit"
} else {
"grapevine"
},
services().globals.server_name()
))
.expect("admin bot username should be valid");
if let Ok(Some(grapevine_room)) = services().admin.get_admin_room() {
loop {
let event = receiver
.recv()
.await
.expect("admin command channel has been closed");
let message_content = match event {
AdminRoomEvent::SendMessage(content) => content,
AdminRoomEvent::ProcessMessage(room_message) => {
self.process_admin_message(room_message).await
}
};
let mutex_state = Arc::clone(
services()
.globals
.roomid_mutex_state
.write()
.await
.entry(grapevine_room.clone())
.or_default(),
);
let state_lock = mutex_state.lock().await;
services()
.rooms
.timeline
.build_and_append_pdu(
PduBuilder {
event_type: TimelineEventType::RoomMessage,
content: to_raw_value(&message_content)
.expect("event is valid, we just created it"),
unsigned: None,
state_key: None,
redacts: None,
},
&grapevine_user,
&grapevine_room,
&state_lock,
)
.await
.unwrap();
Self::handle_event(
&self2,
event,
&grapevine_room,
&grapevine_user,
)
.await;
}
}
});
}
#[tracing::instrument(skip(self, grapevine_room, grapevine_user))]
async fn handle_event(
&self,
event: AdminRoomEvent,
grapevine_room: &OwnedRoomId,
grapevine_user: &ruma::OwnedUserId,
) {
let message_content = match event {
AdminRoomEvent::SendMessage(content) => content,
AdminRoomEvent::ProcessMessage(room_message) => {
self.process_admin_message(room_message).await
}
};
let mutex_state = Arc::clone(
services()
.globals
.roomid_mutex_state
.write()
.await
.entry(grapevine_room.clone())
.or_default(),
);
let state_lock = mutex_state.lock().await;
services()
.rooms
.timeline
.build_and_append_pdu(
PduBuilder {
event_type: TimelineEventType::RoomMessage,
content: to_raw_value(&message_content)
.expect("event is valid, we just created it"),
unsigned: None,
state_key: None,
redacts: None,
},
grapevine_user,
grapevine_room,
&state_lock,
)
.await
.unwrap();
}
pub(crate) fn process_message(&self, room_message: String) {