[datastore] Lookup expunged Oximeter returns Ok(None) (#9651)

This is a small change pulled out of upcoming blueprint pruning work.
Reconfigurator needs to be able to look up a potentially-expunged
Oximeter instance and know whether the cleanup work that marks it
expunged in the `oximeter` crdb table is complete. It's easy to know if
it's _not_ complete: we'll get back some Oximeter info. But if the
cleanup work has happened, prior to this PR, we'd get back an error, and
the details get squashed down by way of `public_error_from_diesel()`
into an `Error::InternalError { ... }`. We need to be able to tell the
difference between "lookup failed because there is no non-expunged
Oximeter by this ID" and "query failed for $DATABASE_REASON", so this PR
changes the lookup to return `Result<Option<OximeterInfo>, _>`, allowing
us to return `Ok(None)` for the "query succeeded but no non-expunged
Oximeter exists" case.

This required changing both callers, and in doing so I think the (rare)
case of trying to operate on an expunged Oximeter would fail prior to
this PR but will now succeed, because both are trying to notify Oximeter
of no-longer-present producers.

While I was there, I added a tiny bit of in-memory caching (thrown away
on each activation) to `ExpiredProducers::new()`. Previously, it would
issue a db lookup for the assigned Oximeter ID of every expired
producer. However:

* If we have any expired producers, we're likely to have several (e.g.,
from a sled going down - all its producers will expire at close to the
same time)
* We only run one Oximeter instance at a time, so every producer is
almost certain to be assigned to the same Oximeter

which meant we'd be issuing "num expired producers" queries for the same
information. We'll now only issue one query for each unique Oximeter ID
(almost always 1).

Author: jgallagher

nexus/db-queries/src/db/datastore/oximeter.rs

@@ -41,18 +41,19 @@ pub enum CollectorReassignment {
 impl DataStore {
     /// Lookup an oximeter instance by its ID.
     ///
-    /// Fails if the instance has been expunged.
+    /// Returns `Ok(None)` if the instance does not exist or has been expunged.
     pub async fn oximeter_lookup(
         &self,
         opctx: &OpContext,
         id: &Uuid,
-    ) -> Result<OximeterInfo, Error> {
+    ) -> Result<Option<OximeterInfo>, Error> {
         use nexus_db_schema::schema::oximeter::dsl;
         dsl::oximeter
             .filter(dsl::time_expunged.is_null())
             .find(*id)
             .first_async(&*self.pool_connection_authorized(opctx).await?)
             .await
+            .optional()
             .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
     }
 

nexus/metrics-producer-gc/src/lib.rs

@@ -17,6 +17,7 @@ use futures::stream::FuturesUnordered;
 use nexus_db_queries::context::OpContext;
 use nexus_db_queries::db::DataStore;
 use nexus_db_queries::db::identity::Asset;
+use nexus_db_queries::db::model::OximeterInfo;
 use nexus_db_queries::db::model::ProducerEndpoint;
 use omicron_common::api::external::Error as DbError;
 use oximeter_client::Client as OximeterClient;
@@ -42,7 +43,7 @@ pub enum Error {
     #[error("failed to list expired producers")]
     ListExpiredProducers(#[source] DbError),
     #[error("failed to get Oximeter info for {id}")]
-    GetOximterInfo {
+    GetOximeterInfo {
         id: Uuid,
         #[source]
         err: DbError,
@@ -62,12 +63,21 @@ pub async fn prune_expired_producers(
     // Build a FuturesUnordered to prune each expired producer.
     let mut all_prunes = expired_producers
         .producer_client_pairs()
-        .map(|(producer, client)| async {
-            let result = unregister_producer(
-                opctx, datastore, producer, client, &opctx.log,
-            )
-            .await;
-            (producer.id(), result)
+        .map(|(producer, maybe_client)| async move {
+            match maybe_client {
+                Some(client) => {
+                    let result = unregister_producer(
+                        opctx, datastore, producer, client, &opctx.log,
+                    )
+                    .await;
+                    (producer.id(), result)
+                }
+                // Treat "no client" as success: the only way to have no client
+                // from `expired_producers` is if this producer's Oximeter no
+                // longer exists (in which case we don't need to tell it about
+                // this expired producer!).
+                None => (producer.id(), Ok(())),
+            }
         })
         .collect::<FuturesUnordered<_>>();
 
@@ -139,18 +149,27 @@ impl ExpiredProducers {
             .map_err(Error::ListExpiredProducers)?;
 
         let mut clients = BTreeMap::new();
+        let mut oximeter_by_id = OximeterInfoById::default();
         for producer in &producers {
             let entry = match clients.entry(producer.oximeter_id) {
                 Entry::Vacant(entry) => entry,
                 Entry::Occupied(_) => continue,
             };
-            let info = datastore
-                .oximeter_lookup(opctx, &producer.oximeter_id)
-                .await
-                .map_err(|err| Error::GetOximterInfo {
-                    id: producer.oximeter_id,
-                    err,
-                })?;
+            let Some(info) = oximeter_by_id
+                .get(opctx, datastore, producer.oximeter_id)
+                .await?
+            else {
+                // If the Oximeter for this producer doesn't exist, that's fine
+                // - we don't need to tell it this producer is gone.
+                info!(
+                    opctx.log,
+                    "Oximeter instance not found (presumed expunged); \
+                     skipping notification of expired producer";
+                    "collector-id" => %producer.oximeter_id,
+                    "producer-id" => %producer.id(),
+                );
+                continue;
+            };
             let client_log =
                 opctx.log.new(o!("oximeter-collector" => info.id.to_string()));
             let address = SocketAddr::new(info.ip.ip(), *info.port);
@@ -164,16 +183,42 @@ impl ExpiredProducers {
 
     fn producer_client_pairs(
         &self,
-    ) -> impl Iterator<Item = (&ProducerEndpoint, &OximeterClient)> {
+    ) -> impl Iterator<Item = (&ProducerEndpoint, Option<&OximeterClient>)>
+    {
         self.producers.iter().map(|producer| {
-            // In `new()` we add a client for every producer.oximeter_id, so we
-            // can unwrap this lookup.
-            let client = self.clients.get(&producer.oximeter_id).unwrap();
-            (producer, client)
+            let maybe_client = self.clients.get(&producer.oximeter_id);
+            (producer, maybe_client)
         })
     }
 }
 
+#[derive(Debug, Default)]
+struct OximeterInfoById(BTreeMap<Uuid, Option<OximeterInfo>>);
+
+impl OximeterInfoById {
+    async fn get(
+        &mut self,
+        opctx: &OpContext,
+        datastore: &DataStore,
+        id: Uuid,
+    ) -> Result<Option<&OximeterInfo>, Error> {
+        // If we've already looked up this Oximeter's info, return it.
+        let vacant_entry = match self.0.entry(id) {
+            Entry::Vacant(entry) => entry,
+            Entry::Occupied(entry) => return Ok(entry.into_mut().as_ref()),
+        };
+
+        // Otherwise, try to look it up. If this succeeds, we'll insert the
+        // `info` into `vacant_entry` so we don't have to look it up again.
+        let info = datastore
+            .oximeter_lookup(opctx, &id)
+            .await
+            .map_err(|err| Error::GetOximeterInfo { id, err })?;
+
+        Ok(vacant_entry.insert(info).as_ref())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -186,6 +231,7 @@ mod tests {
     use nexus_db_model::OximeterInfo;
     use nexus_db_queries::db::pub_test_utils::TestDatabase;
     use nexus_types::internal_api::params;
+    use omicron_common::api::external::DataPageParams;
     use omicron_common::api::internal::nexus;
     use omicron_test_utils::dev;
     use std::time::Duration;
@@ -293,6 +339,102 @@ mod tests {
         logctx.cleanup_successful();
     }
 
+    #[tokio::test]
+    async fn test_prune_expired_producers_expunged_oximeter() {
+        // Setup
+        let logctx = dev::test_setup_log(
+            "test_prune_expired_producers_expunged_oximeter",
+        );
+        let db = TestDatabase::new_with_datastore(&logctx.log).await;
+        let (opctx, datastore) = (db.opctx(), db.datastore());
+
+        // Insert two Oximeter collectors
+        let collector_info1 = OximeterInfo::new(&params::OximeterInfo {
+            collector_id: Uuid::new_v4(),
+            address: "[::1]:0".parse().unwrap(),
+        });
+        let collector_info2 = OximeterInfo::new(&params::OximeterInfo {
+            collector_id: Uuid::new_v4(),
+            address: "[::1]:0".parse().unwrap(),
+        });
+        datastore
+            .oximeter_create(&opctx, &collector_info1)
+            .await
+            .expect("failed to insert collector");
+        datastore
+            .oximeter_create(&opctx, &collector_info2)
+            .await
+            .expect("failed to insert collector");
+
+        // Insert several producers; we hope to get at least one assigned to
+        // each collector, so pick a sufficiently large number that almost all
+        // test runs will do so.
+        let total_producers = 32;
+        let mut producers = Vec::new();
+        for _ in 0..total_producers {
+            let producer = nexus::ProducerEndpoint {
+                id: Uuid::new_v4(),
+                kind: nexus::ProducerKind::Service,
+                address: "[::1]:0".parse().unwrap(), // unused
+                interval: Duration::from_secs(0),    // unused
+            };
+            datastore
+                .producer_endpoint_upsert_and_assign(&opctx, &producer)
+                .await
+                .expect("failed to insert producer");
+            producers.push(producer);
+        }
+
+        // Confirm we created `total_producers`, with some split between the two
+        // collectors.
+        let num_assigned_to_collector_1 = datastore
+            .producers_list_by_oximeter_id(
+                opctx,
+                collector_info1.id,
+                &DataPageParams::max_page(),
+            )
+            .await
+            .expect("listed producers")
+            .len();
+        let num_assigned_to_collector_2 = datastore
+            .producers_list_by_oximeter_id(
+                opctx,
+                collector_info2.id,
+                &DataPageParams::max_page(),
+            )
+            .await
+            .expect("listed producers")
+            .len();
+        eprintln!("assigned to collector 1: {num_assigned_to_collector_1}");
+        eprintln!("assigned to collector 2: {num_assigned_to_collector_2}");
+        assert_eq!(
+            num_assigned_to_collector_1 + num_assigned_to_collector_2,
+            total_producers
+        );
+
+        // Expunge one of the collectors.
+        datastore
+            .oximeter_expunge(opctx, collector_info1.id)
+            .await
+            .expect("expunged collector");
+
+        // Attempt to prune; this should report success for all
+        // `total_producers` producers, even though some of them were assigned
+        // to a now-expunged Oximeter.
+        let pruned = prune_expired_producers(
+            &opctx,
+            &datastore,
+            Utc::now() + Duration::from_secs(1),
+        )
+        .await
+        .expect("failed to prune expired producers");
+        assert_eq!(pruned.successes.len(), total_producers);
+        assert!(pruned.failures.is_empty());
+
+        db.terminate().await;
+        logctx.cleanup_successful();
+    }
+
     #[tokio::test]
     async fn test_prune_expired_producers_notifies_collector() {
         // Setup

nexus/src/app/oximeter.rs

@@ -201,8 +201,23 @@ pub(crate) async fn unassign_producer(
             "producer_id" => %id,
             "collector_id" => %collector_id,
         );
-        let oximeter_info =
-            datastore.oximeter_lookup(opctx, &collector_id).await?;
+
+        let Some(oximeter_info) =
+            datastore.oximeter_lookup(opctx, &collector_id).await?
+        else {
+            // If we can't look up the Oximeter info, it must have been
+            // expunged; that's sufficient to consider this producer
+            // "unassigned".
+            info!(
+                log,
+                "attempted to unassign metric producer, but Oximeter instance \
+                 not found - assuming it was expunged";
+                "producer_id" => %id,
+                "collector_id" => %collector_id,
+            );
+            return Ok(());
+        };
+
         let address =
             SocketAddr::new(oximeter_info.ip.ip(), *oximeter_info.port);
         let client = build_oximeter_client(&log, &id, address);