Update kstat sampler futures when a target is updated (#8915)

- When a target is added to the `KstatSampler`, we also add a future
that resolves when we need to generate a sample from that target. Prior
to this, when a target was updated or removed, that future was never
touched! That means we could be polling for samples for removed targets
or that we might have multiple futures for the same target. This ensures
that we always update or remove the futures, along with the targets
themselves.
- Fixes #8889

Author: bnaecker

oximeter/instruments/src/kstat/link.rs

@@ -740,4 +740,76 @@ mod tests {
         // kept the final samples.
         assert!(count.value() > 4096);
     }
+
+    // Regression for https://github.com/oxidecomputer/omicron/issues/8889
+    #[tokio::test]
+    async fn updating_target_changes_existing_sampling_interval() {
+        let log = test_logger();
+        let sampler = KstatSampler::new(&log).unwrap();
+        let link = TestEtherstub::new();
+        info!(log, "created test etherstub"; "name" => &link.name);
+        let target = SledDataLinkTarget {
+            rack_id: RACK_ID,
+            sled_id: SLED_ID,
+            sled_serial: SLED_SERIAL.into(),
+            link_name: link.name.clone().into(),
+            kind: KIND.into(),
+            sled_model: SLED_MODEL.into(),
+            sled_revision: SLED_REVISION,
+            zone_name: ZONE_NAME.into(),
+        };
+        let dl = SledDataLink::new(target.clone(), true);
+        let collection_interval = Duration::from_millis(10);
+        let details = CollectionDetails::never(collection_interval);
+        let id = sampler.add_target(dl.clone(), details).await.unwrap();
+
+        // Update the target.
+        let new_duration = Duration::from_millis(15);
+        sampler
+            .update_target(dl, CollectionDetails::never(new_duration))
+            .await
+            .unwrap();
+
+        // Get the futures that the sampler knows about and ensure the value has
+        // been updated.
+        let futs = sampler.future_details().await;
+        assert_eq!(futs.len(), 1, "should have updated the only target");
+        assert_eq!(
+            futs[0],
+            (id, new_duration),
+            "failed to correctly update target"
+        );
+    }
+
+    #[tokio::test]
+    async fn no_futures_to_await_after_removing_target() {
+        let log = test_logger();
+        let sampler = KstatSampler::new(&log).unwrap();
+        let link = TestEtherstub::new();
+        info!(log, "created test etherstub"; "name" => &link.name);
+        let target = SledDataLinkTarget {
+            rack_id: RACK_ID,
+            sled_id: SLED_ID,
+            sled_serial: SLED_SERIAL.into(),
+            link_name: link.name.clone().into(),
+            kind: KIND.into(),
+            sled_model: SLED_MODEL.into(),
+            sled_revision: SLED_REVISION,
+            zone_name: ZONE_NAME.into(),
+        };
+        let dl = SledDataLink::new(target.clone(), true);
+        let collection_interval = Duration::from_millis(100);
+        let details = CollectionDetails::never(collection_interval);
+        let id = sampler.add_target(dl.clone(), details).await.unwrap();
+
+        // And remove right away.
+        sampler.remove_target(id).await.unwrap();
+
+        // And ensure there are zero actual futures
+        let futs = sampler.future_details().await;
+        assert!(
+            futs.is_empty(),
+            "should have zero futures to poll after removing target"
+        );
+    }
 }

oximeter/instruments/src/kstat/sampler.rs

@@ -187,6 +187,10 @@ enum Request {
     CreationTimes {
         reply_tx: oneshot::Sender<BTreeMap<KstatPath, DateTime<Utc>>>,
     },
+    /// Return the list of IDs and intervals in the set of futures the sampler
+    /// is tracking.
+    #[cfg(all(test, target_os = "illumos"))]
+    FutureDetails { reply_tx: oneshot::Sender<Vec<(TargetId, Duration)>> },
 }
 
 /// Data about a single kstat target.
@@ -267,6 +271,17 @@ impl core::future::Future for YieldIdAfter {
     }
 }
 
+// The operation we want to take on a future in our set, after handling an inbox
+// message.
+enum Operation {
+    // We want to add a new future.
+    Add(YieldIdAfter),
+    // Remove a future with the existing ID.
+    Remove(TargetId),
+    // We want to update an existing future.
+    Update((TargetId, Duration)),
+}
+
 /// An owned type used to keep track of the creation time for each kstat in
 /// which interest has been signaled.
 #[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
@@ -345,6 +360,9 @@ struct KstatSamplerWorker {
     /// at construction time. In that case, we'll try again the next time we
     /// need it.
     self_stats: Option<self_stats::SelfStats>,
+
+    /// The futures that resolve when it's time to sample the next target.
+    sample_timeouts: FuturesUnordered<YieldIdAfter>,
 }
 
 fn hostname() -> Option<String> {
@@ -358,7 +376,7 @@ fn hostname() -> Option<String> {
 
 /// Stores the number of samples taken, used for testing.
 #[cfg(all(test, target_os = "illumos"))]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Default)]
 pub(crate) struct SampleCounts {
     pub total: usize,
     pub overflow: usize,
@@ -393,6 +411,7 @@ impl KstatSamplerWorker {
             sample_limit,
             self_stat_queue,
             self_stats,
+            sample_timeouts: FuturesUnordered::new(),
         })
     }
 
@@ -405,7 +424,6 @@ impl KstatSamplerWorker {
         #[cfg(all(test, target_os = "illumos"))]
         sample_count_tx: mpsc::UnboundedSender<SampleCounts>,
     ) {
-        let mut sample_timeouts = FuturesUnordered::new();
         let mut creation_prune_interval =
             interval(CREATION_TIME_PRUNE_INTERVAL);
         creation_prune_interval.tick().await; // Completes immediately.
@@ -420,7 +438,7 @@ impl KstatSamplerWorker {
                         );
                     }
                 }
-                maybe_id = sample_timeouts.next(), if !sample_timeouts.is_empty() => {
+                maybe_id = self.sample_timeouts.next(), if !self.sample_timeouts.is_empty() => {
                     let Some((id, interval)) = maybe_id else {
                         unreachable!();
                     };
@@ -430,7 +448,7 @@ impl KstatSamplerWorker {
                         #[cfg(all(test, target_os = "illumos"))]
                         &sample_count_tx,
                     ) {
-                        sample_timeouts.push(next_timeout);
+                        self.sample_timeouts.push(next_timeout);
                     }
                 }
                 maybe_request = self.inbox.recv() => {
@@ -443,19 +461,53 @@ impl KstatSamplerWorker {
                         "received request on inbox";
                         "request" => ?request,
                     );
-                    if let Some(next_timeout) = self.handle_inbox_request(request) {
-                        sample_timeouts.push(next_timeout);
+                    if let Some(next_op) = self.handle_inbox_request(request) {
+                        self.update_sample_timeouts(next_op);
+                    }
+                }
+            }
+        }
+    }
+
+    fn update_sample_timeouts(&mut self, next_op: Operation) {
+        match next_op {
+            Operation::Add(fut) => self.sample_timeouts.push(fut),
+            Operation::Remove(id) => {
+                // Swap out all futures, and then filter out the one we're now
+                // removing.
+                let old = std::mem::take(&mut self.sample_timeouts);
+                self.sample_timeouts
+                    .extend(old.into_iter().filter(|fut| fut.id != id));
+            }
+            Operation::Update((new_id, new_interval)) => {
+                // Update just the one future, if it exists, or insert one.
+                //
+                // NOTE: we update the _interval_, not the sleep object itself,
+                // which means this won't take effect until the next tick.
+                match self
+                    .sample_timeouts
+                    .iter_mut()
+                    .find(|fut| fut.id == new_id)
+                {
+                    Some(old) => old.interval = new_interval,
+                    None => {
+                        warn!(
+                            &self.log,
+                            "attempting to update the samping future \
+                            for a target, but no active future found \
+                            in the set, it will be added directly";
+                            "id" => %&new_id,
+                        );
+                        self.sample_timeouts
+                            .push(YieldIdAfter::new(new_id, new_interval));
                     }
                 }
             }
         }
     }
 
     // Handle a message on the worker's inbox.
-    fn handle_inbox_request(
-        &mut self,
-        request: Request,
-    ) -> Option<YieldIdAfter> {
+    fn handle_inbox_request(&mut self, request: Request) -> Option<Operation> {
         match request {
             Request::AddTarget { target, details, reply_tx } => {
                 match self.add_target(target, details) {
@@ -475,7 +527,10 @@ impl KstatSamplerWorker {
                                 "error" => ?e,
                             ),
                         }
-                        Some(YieldIdAfter::new(id, details.interval))
+                        Some(Operation::Add(YieldIdAfter::new(
+                            id,
+                            details.interval,
+                        )))
                     }
                     Err(e) => {
                         error!(
@@ -513,7 +568,7 @@ impl KstatSamplerWorker {
                                 "error" => ?e,
                             ),
                         }
-                        Some(YieldIdAfter::new(id, details.interval))
+                        Some(Operation::Update((id, details.interval)))
                     }
                     Err(e) => {
                         error!(
@@ -534,7 +589,7 @@ impl KstatSamplerWorker {
                 }
             }
             Request::RemoveTarget { id, reply_tx } => {
-                self.targets.remove(&id);
+                let do_remove = self.targets.remove(&id).is_some();
                 if let Some(remaining_samples) =
                     self.samples.lock().unwrap().remove(&id)
                 {
@@ -555,7 +610,7 @@ impl KstatSamplerWorker {
                         "error" => ?e,
                     ),
                 }
-                None
+                if do_remove { Some(Operation::Remove(id)) } else { None }
             }
             Request::TargetStatus { id, reply_tx } => {
                 trace!(
@@ -594,6 +649,18 @@ impl KstatSamplerWorker {
                 debug!(self.log, "sent reply for creation times");
                 None
             }
+            #[cfg(all(test, target_os = "illumos"))]
+            Request::FutureDetails { reply_tx } => {
+                debug!(self.log, "request for future details");
+                let details = self
+                    .sample_timeouts
+                    .iter()
+                    .map(|fut| (fut.id, fut.interval))
+                    .collect();
+                reply_tx.send(details).unwrap();
+                debug!(self.log, "sent reply for future details");
+                None
+            }
         }
     }
 
@@ -1296,6 +1363,15 @@ impl KstatSampler {
         self.outbox.send(request).await.map_err(|_| Error::SendError).unwrap();
         reply_rx.await.map_err(|_| Error::RecvError).unwrap()
     }
+
+    /// Return the IDs and sampling intervals for all futures in the sampler.
+    #[cfg(all(test, target_os = "illumos"))]
+    pub(crate) async fn future_details(&self) -> Vec<(TargetId, Duration)> {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        let request = Request::FutureDetails { reply_tx };
+        self.outbox.send(request).await.map_err(|_| Error::SendError).unwrap();
+        reply_rx.await.map_err(|_| Error::RecvError).unwrap()
+    }
 }
 
 impl oximeter::Producer for KstatSampler {