1use std::collections::hash_map::Entry;
14use std::collections::{HashMap, HashSet, VecDeque};
15use std::sync::{Arc, Mutex as StdMutex};
16use std::time::{Duration, Instant};
17
18use lsp_types::{LogMessageParams, PublishDiagnosticsParams, ShowMessageParams, Uri};
19use tokio::sync::{Notify, mpsc};
20use tracing::{debug, warn};
21
22use crate::bridge::{BoundedDiagnostics, DiagnosticsKey, MAX_DIAGNOSTIC_ENTRIES};
23use crate::util::{MAX_LOG_STRING_BYTES, WarnLimiter, escape_bounded, lock_std};
24
25pub const MAX_PENDING_PUBLISHES: usize = MAX_DIAGNOSTIC_ENTRIES;
28
29pub const MAILBOX_MAX_BYTES: usize = 64 * 1024 * 1024;
31
32const MAX_LOST_FILES: usize = MAX_DIAGNOSTIC_ENTRIES;
34
35pub const DROP_WARN_PERIOD: Duration = WarnLimiter::DEFAULT_PERIOD;
37
38#[derive(Debug, Clone, Copy)]
40pub struct MailboxLimits {
41 pub max_pending: usize,
43 pub max_bytes: usize,
45}
46
47impl Default for MailboxLimits {
48 fn default() -> Self {
49 Self {
50 max_pending: MAX_PENDING_PUBLISHES,
51 max_bytes: MAILBOX_MAX_BYTES,
52 }
53 }
54}
55
56#[derive(Debug, Default)]
59pub struct DropCounter {
60 warn: WarnLimiter,
61 suppressed: usize,
62}
63
64impl DropCounter {
65 pub const fn new() -> Self {
67 Self {
68 warn: WarnLimiter::new(),
69 suppressed: 0,
70 }
71 }
72
73 pub fn record(&mut self, now: Instant) -> Option<usize> {
77 if self.warn.due(now, DROP_WARN_PERIOD) {
78 Some(std::mem::take(&mut self.suppressed))
79 } else {
80 self.suppressed = self.suppressed.saturating_add(1);
81 None
82 }
83 }
84}
85
86#[derive(Debug)]
89pub struct BoundedPublish {
90 uri: Uri,
91 version: Option<i32>,
92 diagnostics: BoundedDiagnostics,
93}
94
95impl BoundedPublish {
96 fn new(params: PublishDiagnosticsParams) -> (Self, usize) {
97 let PublishDiagnosticsParams {
98 uri,
99 diagnostics,
100 version,
101 } = params;
102 let diagnostics = BoundedDiagnostics::new(&uri, diagnostics);
103 let bytes = diagnostics
104 .buffered_bytes()
105 .saturating_add(AsRef::<str>::as_ref(&uri).len());
106 (
107 Self {
108 uri,
109 version,
110 diagnostics,
111 },
112 bytes,
113 )
114 }
115
116 #[must_use]
118 pub const fn uri(&self) -> &Uri {
119 &self.uri
120 }
121
122 #[must_use]
124 pub const fn version(&self) -> Option<i32> {
125 self.version
126 }
127
128 #[must_use]
130 pub const fn diagnostics(&self) -> &BoundedDiagnostics {
131 &self.diagnostics
132 }
133
134 #[must_use]
136 pub fn into_parts(self) -> (Uri, Option<i32>, BoundedDiagnostics) {
137 (self.uri, self.version, self.diagnostics)
138 }
139}
140
141#[derive(Debug, PartialEq, Eq)]
144pub struct LostFiles(Vec<Uri>);
145
146impl LostFiles {
147 fn new(uris: Vec<Uri>) -> Option<Self> {
148 (!uris.is_empty()).then_some(Self(uris))
149 }
150
151 #[must_use]
153 pub fn uris(&self) -> &[Uri] {
154 &self.0
155 }
156}
157
158#[derive(Debug)]
160pub enum PublishDelivery {
161 Publish(BoundedPublish),
163 Lost(LostFiles),
165 LostUnnamed,
168}
169
170#[derive(Debug)]
172pub enum ServerMessage {
173 Log(LogMessageParams),
175 Show(ShowMessageParams),
177}
178
179#[derive(Debug)]
180struct Pending {
181 publish: BoundedPublish,
182 bytes: usize,
183 seq: u64,
186}
187
188#[derive(Debug, Default)]
189struct State {
190 pending: HashMap<DiagnosticsKey, Pending>,
191 order: VecDeque<(DiagnosticsKey, u64)>,
195 next_seq: u64,
196 bytes: usize,
197 lost: Vec<Uri>,
198 lost_keys: HashSet<DiagnosticsKey>,
199 lost_overflowed: bool,
200 drops: DropCounter,
201 writer_gone: bool,
202 reader_gone: bool,
203}
204
205#[derive(Debug)]
206struct Shared {
207 state: StdMutex<State>,
208 wake: Notify,
209 limits: MailboxLimits,
210}
211
212fn is_current(pending: &HashMap<DiagnosticsKey, Pending>, key: &DiagnosticsKey, seq: u64) -> bool {
214 pending.get(key).is_some_and(|pending| pending.seq == seq)
215}
216
217impl State {
218 fn enqueue(&mut self, key: DiagnosticsKey) -> u64 {
220 let seq = self.next_seq;
221 self.next_seq = seq.saturating_add(1);
222 self.order.push_back((key, seq));
223 seq
224 }
225
226 fn withdraw(&mut self, key: &DiagnosticsKey) {
227 if let Some(pending) = self.pending.remove(key) {
228 self.bytes = self.bytes.saturating_sub(pending.bytes);
229 }
230 }
231
232 fn compact_order(&mut self, max_pending: usize) {
236 if self.order.len() > max_pending.saturating_mul(2).max(16) {
237 let pending = &self.pending;
238 self.order
239 .retain(|(key, seq)| is_current(pending, key, *seq));
240 }
241 }
242
243 fn next_ready(&mut self) -> Option<PublishDelivery> {
246 if let Some(lost) = self.take_lost_files() {
247 return Some(PublishDelivery::Lost(lost));
248 }
249 while let Some((key, seq)) = self.order.pop_front() {
250 let Entry::Occupied(entry) = self.pending.entry(key) else {
251 continue;
252 };
253 if entry.get().seq != seq {
254 continue;
255 }
256 let pending = entry.remove();
257 self.bytes = self.bytes.saturating_sub(pending.bytes);
258 return Some(PublishDelivery::Publish(pending.publish));
259 }
260 std::mem::take(&mut self.lost_overflowed).then_some(PublishDelivery::LostUnnamed)
261 }
262
263 fn mark_lost(&mut self, key: DiagnosticsKey, uri: Uri) {
264 if self.lost_keys.contains(&key) {
265 return;
266 }
267 if self.lost.len() >= MAX_LOST_FILES {
268 self.lost_overflowed = true;
269 return;
270 }
271 self.lost_keys.insert(key);
272 self.lost.push(uri);
273 }
274
275 fn take_lost_files(&mut self) -> Option<LostFiles> {
276 let lost = LostFiles::new(std::mem::take(&mut self.lost))?;
277 self.lost_keys.clear();
278 Some(lost)
279 }
280
281 fn count_drop(&mut self) -> Option<usize> {
283 self.drops.record(Instant::now())
284 }
285}
286
287#[derive(Debug)]
290pub struct PublishWriter {
291 shared: Arc<Shared>,
292}
293
294#[derive(Debug)]
297pub struct PublishReader {
298 shared: Arc<Shared>,
299}
300
301pub fn mailbox(limits: MailboxLimits) -> (PublishWriter, PublishReader) {
304 let shared = Arc::new(Shared {
305 state: StdMutex::new(State::default()),
306 wake: Notify::new(),
307 limits,
308 });
309 (
310 PublishWriter {
311 shared: Arc::clone(&shared),
312 },
313 PublishReader { shared },
314 )
315}
316
317impl PublishWriter {
318 pub fn publish(&self, params: PublishDiagnosticsParams) {
326 let (publish, bytes) = BoundedPublish::new(params);
327 let key = DiagnosticsKey::of(publish.uri());
328 let limits = self.shared.limits;
329 let mut state = lock_std(&self.shared.state);
330 if state.reader_gone {
331 debug!(
332 "dropping a publish for {}: nobody reads the mailbox",
333 publish.uri().as_ref()
334 );
335 return;
336 }
337 let replaced = state
338 .pending
339 .get(&key)
340 .map(|pending| (pending.bytes, pending.seq));
341 let is_new = replaced.is_none();
342 let fits = state
343 .bytes
344 .saturating_sub(replaced.map_or(0, |(bytes, _)| bytes))
345 .saturating_add(bytes)
346 <= limits.max_bytes
347 && (!is_new || state.pending.len() < limits.max_pending);
348 state.compact_order(limits.max_pending);
349 state.withdraw(&key);
350 let warn_due = if fits {
351 state.bytes = state.bytes.saturating_add(bytes);
352 let seq = match replaced {
353 Some((_, seq)) => seq,
354 None => state.enqueue(key.clone()),
355 };
356 state.pending.insert(
357 key,
358 Pending {
359 publish,
360 bytes,
361 seq,
362 },
363 );
364 None
365 } else {
366 let (uri, ..) = publish.into_parts();
367 state.mark_lost(key, uri);
368 state.count_drop()
369 };
370 drop(state);
371 if let Some(suppressed) = warn_due {
372 warn!(
373 "publishDiagnostics mailbox is full: dropped a publish ({suppressed} others \
374 since the last warning); the files read as evicted until the server publishes \
375 again"
376 );
377 }
378 self.shared.wake.notify_one();
379 }
380}
381
382impl Drop for PublishWriter {
383 fn drop(&mut self) {
384 lock_std(&self.shared.state).writer_gone = true;
385 self.shared.wake.notify_one();
386 }
387}
388
389impl PublishReader {
390 #[must_use]
392 pub fn closed() -> Self {
393 mailbox(MailboxLimits::default()).1
394 }
395
396 pub fn try_recv(&mut self) -> Option<PublishDelivery> {
398 lock_std(&self.shared.state).next_ready()
399 }
400
401 pub async fn recv(&mut self) -> Option<PublishDelivery> {
404 loop {
405 {
406 let mut state = lock_std(&self.shared.state);
407 if let Some(delivery) = state.next_ready() {
408 return Some(delivery);
409 }
410 if state.writer_gone {
411 return None;
412 }
413 }
414 self.shared.wake.notified().await;
415 }
416 }
417}
418
419impl Drop for PublishReader {
420 fn drop(&mut self) {
421 let mut state = lock_std(&self.shared.state);
422 state.reader_gone = true;
423 state.order.clear();
424 state.bytes = 0;
425 let pending = std::mem::take(&mut state.pending);
426 drop(state);
427 drop(pending);
428 }
429}
430
431#[derive(Debug)]
436pub struct NotificationSink {
437 messages: mpsc::Sender<ServerMessage>,
438 publishes: PublishWriter,
439}
440
441impl NotificationSink {
442 pub fn new(messages: mpsc::Sender<ServerMessage>) -> (Self, PublishReader) {
444 Self::with_limits(messages, MailboxLimits::default())
445 }
446
447 pub fn with_limits(
449 messages: mpsc::Sender<ServerMessage>,
450 limits: MailboxLimits,
451 ) -> (Self, PublishReader) {
452 let (publishes, reader) = mailbox(limits);
453 (
454 Self {
455 messages,
456 publishes,
457 },
458 reader,
459 )
460 }
461
462 pub const fn messages(&self) -> &mpsc::Sender<ServerMessage> {
464 &self.messages
465 }
466
467 pub const fn publishes(&self) -> &PublishWriter {
469 &self.publishes
470 }
471}
472
473#[derive(Debug)]
476pub struct NotificationInbox {
477 pub messages: mpsc::Receiver<ServerMessage>,
479 pub publishes: PublishReader,
481}
482
483impl NotificationInbox {
484 #[must_use]
486 pub const fn new(messages: mpsc::Receiver<ServerMessage>, publishes: PublishReader) -> Self {
487 Self {
488 messages,
489 publishes,
490 }
491 }
492}
493
494#[cfg(test)]
495impl From<mpsc::Receiver<ServerMessage>> for NotificationInbox {
496 fn from(messages: mpsc::Receiver<ServerMessage>) -> Self {
497 Self::new(messages, PublishReader::closed())
498 }
499}
500
501#[derive(Debug, Clone, Copy, PartialEq, Eq)]
503pub enum Lane {
504 Notification,
506 Lifecycle,
508}
509
510impl Lane {
511 const fn name(self) -> &'static str {
512 match self {
513 Self::Notification => "notification",
514 Self::Lifecycle => "lifecycle",
515 }
516 }
517}
518
519#[derive(Debug)]
522pub struct DropLog {
523 lane: Lane,
524 drops: StdMutex<DropCounter>,
525}
526
527impl DropLog {
528 const fn new(lane: Lane) -> Self {
529 Self {
530 lane,
531 drops: StdMutex::new(DropCounter::new()),
532 }
533 }
534
535 pub fn of(lane: Lane) -> &'static Self {
538 static NOTIFICATION: DropLog = DropLog::new(Lane::Notification);
539 static LIFECYCLE: DropLog = DropLog::new(Lane::Lifecycle);
540 match lane {
541 Lane::Notification => &NOTIFICATION,
542 Lane::Lifecycle => &LIFECYCLE,
543 }
544 }
545
546 pub fn record(&self, method: &str) {
549 let lane = self.lane.name();
550 let method = escape_bounded(method, MAX_LOG_STRING_BYTES);
551 debug!("dropping notification: lane={lane}, method={method}");
552 let warn_due = lock_std(&self.drops).record(Instant::now());
553 if let Some(suppressed) = warn_due {
554 warn!(
555 "dropped a notification on the {lane} lane ({suppressed} others since the last \
556 warning; channel full or closed)"
557 );
558 }
559 }
560}
561
562#[cfg(test)]
563mod tests {
564 use std::assert_matches;
565
566 use lsp_types::Diagnostic;
567
568 use super::*;
569
570 fn uri(n: usize) -> Uri {
571 Uri::from(format!("file:///ws/f{n}.rs"))
572 }
573
574 fn publish_of(n: usize, message: &str) -> PublishDiagnosticsParams {
575 PublishDiagnosticsParams {
576 uri: uri(n),
577 diagnostics: vec![Diagnostic {
578 message: message.to_owned().into(),
579 ..Diagnostic::default()
580 }],
581 version: None,
582 }
583 }
584
585 fn bytes_of(n: usize, message: &str) -> usize {
586 BoundedPublish::new(publish_of(n, message)).1
587 }
588
589 #[test]
591 fn drop_log_escapes_and_bounds_the_method() {
592 use tracing_subscriber::layer::SubscriberExt as _;
593
594 let captured = crate::test_lsp::CapturedLogs::default();
595 let guard =
596 tracing::subscriber::set_default(tracing_subscriber::registry().with(captured.clone()));
597
598 DropLog::of(Lane::Notification).record(&format!("evil\nFORGED{}", "x".repeat(10_000)));
599 drop(guard);
600
601 let line = captured
602 .messages()
603 .into_iter()
604 .find(|message| message.contains("dropping notification"))
605 .expect("the drop is logged at debug");
606 assert!(!line.contains('\n'), "{line:.80}");
607 assert!(line.contains("evil\\nFORGED"), "{line:.80}");
608 assert!(line.len() < 2 * MAX_LOG_STRING_BYTES, "{}", line.len());
609 }
610
611 #[test]
613 fn publish_bytes_include_the_uri() {
614 let params = publish_of(1, "x");
615 let diagnostics = BoundedDiagnostics::new(¶ms.uri, params.diagnostics.clone());
616
617 assert_eq!(
618 bytes_of(1, "x"),
619 diagnostics.buffered_bytes() + "file:///ws/f1.rs".len()
620 );
621 }
622
623 #[tokio::test]
626 async fn republished_file_follows_arrival_order_after_a_withdrawal() {
627 let one = bytes_of(1, "x");
628 let (writer, mut reader) = mailbox(small(10, one * 2 + one / 2));
629 writer.publish(publish_of(1, "x"));
630 writer.publish(publish_of(1, &"y".repeat(one * 4)));
631 writer.publish(publish_of(2, "x"));
632 writer.publish(publish_of(1, "x"));
633 drop(writer);
634
635 let Some(PublishDelivery::Lost(lost)) = reader.recv().await else {
636 panic!("expected the lost file first");
637 };
638 assert_eq!(lost.uris(), [uri(1)]);
639 assert_eq!(next_publish(&mut reader).await.uri(), &uri(2));
640 assert_eq!(next_publish(&mut reader).await.uri(), &uri(1));
641 assert!(reader.recv().await.is_none());
642 }
643
644 fn small(max_pending: usize, max_bytes: usize) -> MailboxLimits {
645 MailboxLimits {
646 max_pending,
647 max_bytes,
648 }
649 }
650
651 async fn next_publish(reader: &mut PublishReader) -> BoundedPublish {
652 match reader.recv().await {
653 Some(PublishDelivery::Publish(publish)) => publish,
654 other => panic!("expected a publish, got {other:?}"),
655 }
656 }
657
658 #[tokio::test]
659 async fn burst_below_the_cap_is_delivered_whole_in_order() {
660 let (writer, mut reader) = mailbox(small(10, usize::MAX));
661 for n in 0..10 {
662 writer.publish(publish_of(n, "e"));
663 }
664 drop(writer);
665
666 for n in 0..10 {
667 assert_eq!(next_publish(&mut reader).await.uri(), &uri(n));
668 }
669 assert!(reader.recv().await.is_none());
670 }
671
672 #[tokio::test]
673 async fn later_publish_replaces_the_pending_one_of_the_same_file() {
674 let (writer, mut reader) = mailbox(small(10, usize::MAX));
675 writer.publish(publish_of(1, "old"));
676 writer.publish(publish_of(1, "new"));
677 drop(writer);
678
679 let delivered = next_publish(&mut reader).await;
680 assert_eq!(delivered.diagnostics().as_slice().len(), 1);
681 assert_eq!(
682 delivered
683 .diagnostics()
684 .as_slice()
685 .first()
686 .map(|d| d.message.clone()),
687 Some("new".to_owned().into())
688 );
689 assert!(reader.recv().await.is_none());
690 }
691
692 #[tokio::test]
693 async fn publish_over_the_file_cap_is_reported_lost_not_dropped_quietly() {
694 let (writer, mut reader) = mailbox(small(2, usize::MAX));
695 for n in 0..4 {
696 writer.publish(publish_of(n, "e"));
697 }
698
699 let Some(PublishDelivery::Lost(lost)) = reader.recv().await else {
700 panic!("the loss must be handed over first");
701 };
702 assert_eq!(lost.uris(), [uri(2), uri(3)]);
703 assert_eq!(next_publish(&mut reader).await.uri(), &uri(0));
704 assert_eq!(next_publish(&mut reader).await.uri(), &uri(1));
705 }
706
707 #[tokio::test]
710 async fn oversized_replacement_withdraws_the_older_publish_and_marks_the_file_lost() {
711 let one = bytes_of(1, "x");
712 let (writer, mut reader) = mailbox(small(10, one + one / 2));
713 writer.publish(publish_of(1, "x"));
714 writer.publish(publish_of(1, &"y".repeat(one * 2)));
715 drop(writer);
716
717 let Some(PublishDelivery::Lost(lost)) = reader.recv().await else {
718 panic!("expected the lost file");
719 };
720 assert_eq!(lost.uris(), [uri(1)]);
721 assert!(reader.recv().await.is_none());
722 }
723
724 #[tokio::test]
726 async fn arrival_queue_stays_bounded_under_withdrawals() {
727 let one = bytes_of(0, "x");
728 let (writer, _reader) = mailbox(small(4, one * 2));
729 for n in 0..1000 {
730 writer.publish(publish_of(n, "x"));
731 writer.publish(publish_of(n, &"y".repeat(one * 4)));
732 }
733
734 assert!(lock_std(&writer.shared.state).order.len() <= 16);
735 }
736
737 #[tokio::test]
738 async fn replacement_after_order_compaction_is_still_delivered() {
739 let one = bytes_of(0, "x");
740 let (writer, mut reader) = mailbox(small(4, one * 2));
741 for n in 0..16 {
742 writer.publish(publish_of(n, "x"));
743 writer.publish(publish_of(n, &"y".repeat(one * 4)));
744 }
745 writer.publish(publish_of(100, "x"));
746 writer.publish(publish_of(100, "z"));
747 drop(writer);
748
749 let mut delivered = None;
750 while let Some(delivery) = reader.recv().await {
751 if let PublishDelivery::Publish(params) = delivery {
752 delivered = Some(params);
753 }
754 }
755 let delivered = delivered.expect("the replaced publish must reach the reader");
756 assert_eq!(delivered.uri(), &uri(100));
757 assert_eq!(
758 delivered
759 .diagnostics()
760 .as_slice()
761 .first()
762 .map(|d| d.message.clone()),
763 Some("z".to_owned().into())
764 );
765 }
766
767 #[tokio::test]
768 async fn delivered_publish_carries_bounded_diagnostics() {
769 let (writer, mut reader) = mailbox(small(10, usize::MAX));
770 let oversized = 512 * 1024;
771 writer.publish(publish_of(1, &"m".repeat(oversized)));
772 drop(writer);
773
774 let delivered = next_publish(&mut reader).await;
775
776 let message = delivered
777 .diagnostics()
778 .as_slice()
779 .first()
780 .map(|d| d.message.clone());
781 assert_matches!(
782 message,
783 Some(lsp_types::Message::String(text)) if text.len() < oversized
784 );
785 }
786
787 #[test]
788 fn drop_counter_warns_once_per_period_and_counts_the_silenced() {
789 let mut counter = DropCounter::new();
790 let start = Instant::now();
791 let at = |secs| start + Duration::from_secs(secs);
792
793 assert_eq!(counter.record(at(0)), Some(0));
794 assert_eq!(counter.record(at(1)), None);
795 assert_eq!(counter.record(at(2)), None);
796 assert_eq!(counter.record(at(61)), Some(2));
797 }
798
799 #[test]
800 fn each_lane_has_its_own_drop_log() {
801 for lane in [Lane::Notification, Lane::Lifecycle] {
802 assert_eq!(DropLog::of(lane).lane, lane);
803 }
804 }
805
806 #[tokio::test]
807 async fn byte_cap_bounds_what_is_buffered() {
808 let one = bytes_of(0, "e");
809 let (writer, mut reader) = mailbox(small(100, one * 3));
810 for n in 0..10 {
811 writer.publish(publish_of(n, "e"));
812 }
813 drop(writer);
814
815 let mut delivered = 0;
816 let mut lost = 0;
817 while let Some(item) = reader.recv().await {
818 match item {
819 PublishDelivery::Publish(_) => delivered += 1,
820 PublishDelivery::Lost(files) => lost += files.uris().len(),
821 PublishDelivery::LostUnnamed => panic!("no unnamed losses expected"),
822 }
823 }
824 assert_eq!((delivered, lost), (3, 7));
825 }
826
827 #[tokio::test]
828 async fn lost_list_is_capped_and_then_only_the_fact_is_kept() {
829 let (writer, mut reader) = mailbox(small(0, usize::MAX));
830 for n in 0..=MAX_LOST_FILES {
831 writer.publish(publish_of(n, "e"));
832 }
833 drop(writer);
834
835 let Some(PublishDelivery::Lost(lost)) = reader.recv().await else {
836 panic!("expected losses");
837 };
838 assert_eq!(lost.uris().len(), MAX_LOST_FILES);
839 assert_matches!(reader.recv().await, Some(PublishDelivery::LostUnnamed));
840 assert!(reader.recv().await.is_none());
841 }
842
843 #[tokio::test]
844 async fn dropping_the_writer_ends_the_reader_and_dropping_the_reader_stops_buffering() {
845 let (writer, reader) = mailbox(small(10, usize::MAX));
846 writer.publish(publish_of(1, "e"));
847 drop(reader);
848 writer.publish(publish_of(2, "e"));
849 assert_eq!(lock_std(&writer.shared.state).bytes, 0);
850 assert!(lock_std(&writer.shared.state).pending.is_empty());
851
852 let (writer, mut reader) = mailbox(small(10, usize::MAX));
853 drop(writer);
854 assert!(reader.recv().await.is_none());
855 }
856
857 #[tokio::test]
858 async fn a_waiting_reader_wakes_for_a_new_publish() {
859 let (writer, mut reader) = mailbox(small(10, usize::MAX));
860 let waiting = tokio::spawn(async move { next_publish(&mut reader).await.uri });
861 tokio::task::yield_now().await;
862
863 writer.publish(publish_of(7, "e"));
864
865 assert_eq!(waiting.await.unwrap(), uri(7));
866 }
867}