chore: Fix flaky semantic tokens caching test (#23831)

The stderr stream from the LSP is consumed by a separate thread, so it
may not have processed the part we care about yet. Instead, wait until
you see the measure for the request you care about.

2 files changed, 33 insertions, 13 deletions

diff --git a/tests/integration/lsp_tests.rs b/tests/integration/lsp_tests.rs
index 47fefeafe..2d5bb2e24 100644
--- a/tests/integration/lsp_tests.rs
+++ b/tests/integration/lsp_tests.rs
@@ -12740,7 +12740,7 @@ fn lsp_semantic_token_caching() {
 
   assert_eq!(
     client
-      .perf()
+      .perf_wait_for_measure("lsp.semantic_tokens_range")
       .measure_count("tsc.request.getEncodedSemanticClassifications"),
     1,
   );
@@ -12755,7 +12755,7 @@ fn lsp_semantic_token_caching() {
 
   assert_eq!(
     client
-      .perf()
+      .perf_wait_for_measure("lsp.semantic_tokens_full")
       .measure_count("tsc.request.getEncodedSemanticClassifications"),
     2,
   );
@@ -12775,7 +12775,7 @@ fn lsp_semantic_token_caching() {
   // make sure we actually used the cache
   assert_eq!(
     client
-      .perf()
+      .perf_wait_for_measure("lsp.semantic_tokens_range")
       .measure_count("tsc.request.getEncodedSemanticClassifications"),
     2,
   );
diff --git a/tests/util/server/src/lsp.rs b/tests/util/server/src/lsp.rs
index ed9dd302f..7c48bae23 100644
--- a/tests/util/server/src/lsp.rs
+++ b/tests/util/server/src/lsp.rs
@@ -689,15 +689,33 @@ impl Perf {
       rx,
     }
   }
-  fn drain(&mut self) {
-    while let Ok(record) = self.rx.try_recv() {
-      if let PerfRecord::Measure(measure) = &record {
-        *self
-          .measures_counts
-          .entry(measure.name.clone())
-          .or_default() += 1;
+  fn drain_until(&mut self, f: impl Fn(&PerfRecord) -> bool) {
+    let timeout_time =
+      Instant::now().checked_add(Duration::from_secs(5)).unwrap();
+    let mut found = false;
+    loop {
+      while let Ok(record) = self.rx.try_recv() {
+        if let PerfRecord::Measure(measure) = &record {
+          *self
+            .measures_counts
+            .entry(measure.name.clone())
+            .or_default() += 1;
+        }
+        if f(&record) {
+          found = true;
+        }
+        self.records.push(record);
+      }
+
+      if found {
+        break;
+      }
+
+      std::thread::sleep(Duration::from_millis(20));
+
+      if Instant::now() > timeout_time {
+        panic!("timed out waiting for perf record");
       }
-      self.records.push(record);
     }
   }
   pub fn measures(&self) -> impl IntoIterator<Item = &PerfMeasure> {
@@ -757,12 +775,14 @@ impl LspClient {
     self.reader.pending_len()
   }
 
-  pub fn perf(&mut self) -> &Perf {
+  /// Collects performance records until a measure with the given name is
+  /// emitted.
+  pub fn perf_wait_for_measure(&mut self, name: &str) -> &Perf {
     let perf = self
       .perf
       .as_mut()
       .expect("must setup with client_builder.collect_perf()");
-    perf.drain();
+    perf.drain_until(|record| matches!(record, PerfRecord::Measure(measure) if measure.name == name));
     perf
   }