refactor!: remove index from trace

sk-core/src/constants.rs

@@ -11,7 +11,7 @@ use crate::k8s::{
 };
 
 // Trace Format Version
-pub const CURRENT_TRACE_FORMAT_VERSION: u16 = 2;
+pub const CURRENT_TRACE_FORMAT_VERSION: u16 = 3;
 
 // Well-known labels, annotations, and taints
 pub const KUBERNETES_IO_METADATA_NAME_KEY: &str = "kubernetes.io/metadata.name";

sk-core/src/k8s/mod.rs

@@ -14,7 +14,10 @@ use k8s_openapi::api::core::v1 as corev1;
 use k8s_openapi::apimachinery::pkg::apis::meta::v1 as metav1;
 use kube::api::TypeMeta;
 pub use lease::*;
-pub use owners::OwnersCache;
+pub use owners::{
+    OwnersCache,
+    PodOwner,
+};
 use serde::{
     Deserialize,
     Serialize,

sk-core/src/k8s/owners.rs

@@ -21,6 +21,8 @@ use crate::k8s::{
 };
 
 
+pub type PodOwner = (GVK, String);
+
 // TODO I really want a way to mock out the OwnersCache, because
 // any tests that depend on it implicitly now have to depend on tokio
 // and also the fake_apiserver, which is cumbersome to deal with; unfortunately,
@@ -31,7 +33,7 @@ use crate::k8s::{
 // this we'll have to implement the mock ourselves.
 pub struct OwnersCache {
     apiset: DynamicApiSet,
-    owners: HashMap<(GVK, String), Vec<metav1::OwnerReference>>,
+    owners: HashMap<PodOwner, Vec<metav1::OwnerReference>>,
 }
 
 impl OwnersCache {
@@ -41,7 +43,7 @@ impl OwnersCache {
 
     pub fn new_from_parts(
         apiset: DynamicApiSet,
-        owners: HashMap<(GVK, String), Vec<metav1::OwnerReference>>,
+        owners: HashMap<PodOwner, Vec<metav1::OwnerReference>>,
     ) -> OwnersCache {
         OwnersCache { apiset, owners }
     }

sk-core/src/lib.rs

@@ -35,7 +35,6 @@ pub mod prelude {
         KubeResourceExt,
         OpenApiResourceExt,
     };
-    pub use crate::trace::Trace;
     pub use crate::trace::config::{
         TracerConfig,
         TrackedObjectConfig,
@@ -49,4 +48,8 @@ pub mod prelude {
         PodLifecyclesMap,
         PodOwnersMap,
     };
+    pub use crate::trace::{
+        ResourceMetadata,
+        Trace,
+    };
 }

sk-core/src/trace/mod.rs

@@ -1,6 +1,5 @@
 pub mod config;
 pub mod event;
-pub mod index;
 pub mod pod_owners_map;
 
 use std::collections::HashMap;
@@ -23,11 +22,11 @@ use crate::external_storage::{
 use crate::k8s::{
     GVK,
     PodLifecycleData,
+    PodOwner,
 };
 use crate::time::duration_to_ts_from;
 use crate::trace::config::TracerConfig;
 use crate::trace::event::TraceEvent;
-use crate::trace::index::TraceIndex;
 use crate::trace::pod_owners_map::PodLifecyclesMap;
 
 #[derive(Debug, Error)]
@@ -40,12 +39,15 @@ pub enum TraceError {
 }
 
 #[derive(Clone, Deserialize, Serialize)]
+pub struct ResourceMetadata {}
+
+#[derive(Deserialize, Serialize)]
 pub struct Trace {
     pub version: u16,
     pub config: TracerConfig,
     pub events: Vec<TraceEvent>,
-    pub index: TraceIndex,
-    pub pod_lifecycles: HashMap<(GVK, String), PodLifecyclesMap>,
+    pub tracked_objects: HashMap<GVK, HashMap<String, ResourceMetadata>>,
+    pub pod_lifecycles: HashMap<PodOwner, PodLifecyclesMap>,
 }
 
 impl Default for Trace {
@@ -54,7 +56,7 @@ impl Default for Trace {
             version: CURRENT_TRACE_FORMAT_VERSION,
             config: TracerConfig::default(),
             events: vec![],
-            index: TraceIndex::default(),
+            tracked_objects: HashMap::default(),
             pod_lifecycles: HashMap::default(),
         }
     }
@@ -141,7 +143,9 @@ impl Trace {
     }
 
     pub fn has_obj(&self, gvk: &GVK, ns_name: &str) -> bool {
-        self.index.contains(gvk, ns_name)
+        self.tracked_objects
+            .get(gvk)
+            .is_some_and(|objects| objects.contains_key(ns_name))
     }
 
     pub fn get_object(&self, event_idx: usize, obj_idx: usize) -> Option<&DynamicObject> {

sk-core/src/trace/pod_owners_map.rs

@@ -1,15 +1,18 @@
-use std::collections::HashMap;
 use std::collections::hash_map::Entry;
+use std::collections::{
+    HashMap,
+    HashSet,
+};
 
 use tracing::*;
 
 use crate::errors::*;
 use crate::k8s::{
     GVK,
     PodLifecycleData,
+    PodOwner,
     format_gvk_name,
 };
-use crate::trace::index::TraceIndex;
 
 // The PodOwnersMap tracks lifecycle data for all pods that are owned by some object that we care
 // about (e.g., if we are tracking Deployments, the owners map will track the lifecycle data for
@@ -55,8 +58,8 @@ pub type PodLifecyclesMap = HashMap<u64, Vec<PodLifecycleData>>;
 
 #[derive(Clone, Default)]
 pub struct PodOwnersMap {
-    m: HashMap<(GVK, String), PodLifecyclesMap>,
-    index: HashMap<String, ((GVK, String), u64, usize)>,
+    m: HashMap<PodOwner, PodLifecyclesMap>,
+    index: HashMap<String, (PodOwner, u64, usize)>,
 }
 
 impl PodOwnersMap {
@@ -128,23 +131,25 @@ impl PodOwnersMap {
 
     // Given an index of "owning objects", get a list of all the pods between a given start and end
     // time that belong to one of those owning objects.
-    pub fn filter(&self, start_ts: i64, end_ts: i64, index: &TraceIndex) -> HashMap<(GVK, String), PodLifecyclesMap> {
+    pub fn filter(
+        &self,
+        start_ts: i64,
+        end_ts: i64,
+        owning_objects: &HashSet<PodOwner>,
+    ) -> HashMap<PodOwner, PodLifecyclesMap> {
         self.m
             .iter()
             // The filtering is a little complicated here; if the owning object isn't in the index,
             // we discard it.  Also, if none of the pods belonging to the owning object land
             // within the given time window, we want to discard it.  Otherwise, we want to filter
             // down the list of pods to the ones that fall between the given time window.
-            .filter_map(|((owner_gvk, owner_ns_name), lifecycles_map)| {
-                if !index.contains(owner_gvk, owner_ns_name) {
+            .filter_map(|(owner, lifecycles_map)| {
+                if !owning_objects.contains(owner) {
                     return None;
                 }
 
                 // Note the question mark here, doing a bunch of heavy lifting
-                Some((
-                    (owner_gvk.clone(), owner_ns_name.clone()),
-                    filter_lifecycles_map(start_ts, end_ts, lifecycles_map)?,
-                ))
+                Some((owner.clone(), filter_lifecycles_map(start_ts, end_ts, lifecycles_map)?))
             })
             .collect()
     }
@@ -200,13 +205,13 @@ impl PodOwnersMap {
     }
 
     pub fn new_from_parts(
-        m: HashMap<(GVK, String), PodLifecyclesMap>,
-        index: HashMap<String, ((GVK, String), u64, usize)>,
+        m: HashMap<PodOwner, PodLifecyclesMap>,
+        index: HashMap<String, (PodOwner, u64, usize)>,
     ) -> PodOwnersMap {
         PodOwnersMap { m, index }
     }
 
-    pub fn pod_owner_meta(&self, pod_ns_name: &str) -> Option<&((GVK, String), u64, usize)> {
+    pub fn pod_owner_meta(&self, pod_ns_name: &str) -> Option<&(PodOwner, u64, usize)> {
         self.index.get(pod_ns_name)
     }
 }

sk-core/src/trace/tests/pod_owners_map_test.rs

@@ -1,9 +1,11 @@
-use std::collections::HashMap;
+use std::collections::{
+    HashMap,
+    HashSet,
+};
 
 use super::*;
 use crate::constants::*;
 use crate::k8s::PodLifecycleData;
-use crate::trace::TraceIndex;
 use crate::trace::pod_owners_map::{
     PodLifecyclesMap,
     PodOwnersMap,
@@ -57,9 +59,9 @@ fn test_store_new_pod_lifecycle(mut owners_map: PodOwnersMap) {
 
 #[rstest]
 fn test_filter_owners_map() {
-    let mut index = TraceIndex::new();
-    index.insert(DEPLOYMENT_GVK.clone(), "test/deployment1".into(), 9876);
-    index.insert(DEPLOYMENT_GVK.clone(), "test/deployment2".into(), 5432);
+    let mut owning_objects = HashSet::new();
+    owning_objects.insert((DEPLOYMENT_GVK.clone(), "test/deployment1".into()));
+    owning_objects.insert((DEPLOYMENT_GVK.clone(), "test/deployment2".into()));
     let owners_map = PodOwnersMap::new_from_parts(
         HashMap::from([
             (
@@ -78,7 +80,7 @@ fn test_filter_owners_map() {
         HashMap::new(),
     );
 
-    let res = owners_map.filter(START_TS, END_TS, &index);
+    let res = owners_map.filter(START_TS, END_TS, &owning_objects);
     assert_eq!(
         res,
         HashMap::from([(

sk-core/src/trace/tests/trace_test.rs

@@ -26,8 +26,7 @@ fn test_lookup_pod_lifecycle_no_owner(test_trace: Trace) {
 }
 
 #[rstest]
-fn test_lookup_pod_lifecycle_no_hash(mut test_trace: Trace) {
-    test_trace.index.insert(DEPLOYMENT_GVK.clone(), TEST_DEPLOYMENT.into(), 1234);
+fn test_lookup_pod_lifecycle_no_hash(test_trace: Trace) {
     let res = test_trace.lookup_pod_lifecycle(&DEPLOYMENT_GVK, TEST_DEPLOYMENT, EMPTY_POD_SPEC_HASH, 0);
     assert_eq!(res, PodLifecycleData::Empty);
 }
@@ -37,7 +36,6 @@ fn test_lookup_pod_lifecycle(mut test_trace: Trace) {
     let owner_ns_name = format!("{TEST_NAMESPACE}/{TEST_DEPLOYMENT}");
     let pod_lifecycle = PodLifecycleData::Finished(1, 2);
 
-    test_trace.index.insert(DEPLOYMENT_GVK.clone(), owner_ns_name.clone(), 1234);
     test_trace.pod_lifecycles = HashMap::from([(
         (DEPLOYMENT_GVK.clone(), owner_ns_name.clone()),
         HashMap::from([(EMPTY_POD_SPEC_HASH, vec![pod_lifecycle.clone()])]),

sk-driver/src/tests/mutation_test.rs

@@ -296,7 +296,11 @@ mod itest {
                 (DEPLOYMENT_GVK.clone(), owner_ns_name.clone()),
                 PodLifecyclesMap::from([(pod_spec_hash, vec![PodLifecycleData::Finished(0, 42)])]),
             );
-            trace.index.insert(DEPLOYMENT_GVK.clone(), owner_ns_name.clone(), 1234);
+            trace
+                .tracked_objects
+                .entry(DEPLOYMENT_GVK.clone())
+                .or_insert(HashMap::new())
+                .insert(owner_ns_name.clone(), ResourceMetadata {});
         }
 
         let owners = vec![root_owner_ref, depl_owner_ref];

sk-skel/src/lib.rs

@@ -59,7 +59,7 @@ pub async fn apply_skel(trace: &Trace, skel_file: &str, update_channel: mpsc::Se
         version: trace.version,
         config: trace.config.clone(),
         events: new_events,
-        index: trace.index.clone(),
+        tracked_objects: trace.tracked_objects.clone(),
         pod_lifecycles: trace.pod_lifecycles.clone(),
     };
 

sk-core/src/trace/index.rs → sk-tracer/src/index.rs

@@ -1,38 +1,17 @@
 use std::collections::HashMap;
-use std::mem::take;
 
-use serde::{
-    Deserialize,
-    Serialize,
-};
+use sk_core::k8s::GVK;
 
-use crate::k8s::{
-    GVK,
-    format_gvk_name,
-};
-
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Default)]
 pub struct TraceIndex {
-    #[serde(flatten)]
     index: HashMap<GVK, HashMap<String, u64>>,
 }
 
 impl TraceIndex {
-    pub fn new() -> TraceIndex {
-        TraceIndex::default()
-    }
-
     pub fn contains(&self, gvk: &GVK, ns_name: &str) -> bool {
         self.index.get(gvk).is_some_and(|gvk_hash| gvk_hash.contains_key(ns_name))
     }
 
-    pub fn flattened_keys(&self) -> Vec<String> {
-        self.index
-            .iter()
-            .flat_map(|(gvk, gvk_hash)| gvk_hash.keys().map(move |ns_name| format_gvk_name(gvk, ns_name)))
-            .collect()
-    }
-
     pub fn get(&self, gvk: &GVK, ns_name: &str) -> Option<u64> {
         self.index.get(gvk)?.get(ns_name).cloned()
     }
@@ -41,21 +20,16 @@ impl TraceIndex {
         self.index.entry(gvk).or_default().insert(ns_name, hash);
     }
 
-    pub fn is_empty(&self) -> bool {
-        self.index.values().all(|gvk_hash| gvk_hash.is_empty())
-    }
-
-    pub fn len(&self) -> usize {
-        self.index.values().map(|gvk_hash| gvk_hash.len()).sum()
-    }
-
     pub fn remove(&mut self, gvk: GVK, ns_name: &str) {
         self.index.entry(gvk).and_modify(|gvk_hash| {
             gvk_hash.remove(ns_name);
         });
     }
+}
 
-    pub fn take_gvk_index(&mut self, gvk: &GVK) -> HashMap<String, u64> {
-        take(self.index.get_mut(gvk).unwrap_or(&mut HashMap::new()))
+#[cfg(test)]
+impl TraceIndex {
+    pub fn len(&self) -> usize {
+        self.index.values().map(|gvk_hash| gvk_hash.len()).sum()
     }
 }

sk-tracer/src/main.rs

@@ -1,6 +1,7 @@
 #![cfg_attr(coverage, feature(coverage_attribute))]
 
 mod errors;
+mod index;
 mod manager;
 mod store;
 mod watchers;