Improve memory usage and speed.

src/common/light_weight_zone.rs

@@ -1,11 +1,12 @@
-/// A quick PoC to see if using a BTree compared to default in-memory zone
-/// store uses less memory, and it does, even with its dumb way of storing
-/// values in the tree. It's not a fair comparison either as the default
-/// in-memory store also supports proper answers to queries, versioning and
-/// IXFR diff generation.
+/// A simple zone store that trades the advanced functionality of the
+/// in-memory zone store for reduced memory usage, specifically: proper
+/// answers to queries, versioning and IXFR diff generation. Memory usage is
+/// reduced by having a less complex data structure, and also for signed zones
+/// not storing a copy of the unsigned records but instead referencing the
+/// actual unsigned zone store.
 use std::{
     any::Any,
-    collections::{btree_map::Entry, BTreeMap, HashSet},
+    collections::{hash_map::Entry, HashMap, HashSet},
     future::{ready, Future},
     ops::Deref,
     pin::Pin,
@@ -22,9 +23,10 @@ use domain::{
         name::Label,
         Name, NameBuilder, Rtype,
     },
+    rdata::ZoneRecordData,
     zonetree::{
-        error::OutOfZone, Answer, InMemoryZoneDiff, ReadableZone, SharedRrset, StoredName, WalkOp,
-        WritableZone, WritableZoneNode, ZoneStore,
+        error::OutOfZone, Answer, InMemoryZoneDiff, ReadableZone, Rrset, SharedRrset, StoredName,
+        WalkOp, WritableZone, WritableZoneNode, Zone, ZoneStore,
     },
 };
 use log::trace;
@@ -55,8 +57,22 @@ impl PartialEq for HashedByRtypeSharedRrset {
 #[derive(Clone, Debug)]
 struct SimpleZoneInner {
     root: StoredName,
-    tree: Arc<std::sync::RwLock<BTreeMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
+
+    /// Optional reference to a zone that contains the unsigned records of a
+    /// signed zone.
+    ///
+    /// Avoids copying the unsigned records into this zone thus reduces memory
+    /// usage.
+    unsigned_zone: Option<Zone>,
+
+    /// The records of the zone, either unsigned or signed.
+    tree: Arc<std::sync::RwLock<HashMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
+
+    /// The count of records skipped when skip_signed is true.
     skipped: Arc<AtomicUsize>,
+
+    /// Whether or not to prevent DNSSEC records being added to this zone via
+    /// its WritableZoneInterface.
     skip_signed: bool,
 }
 
@@ -66,10 +82,11 @@ pub struct LightWeightZone {
 }
 
 impl LightWeightZone {
-    pub fn new(root: StoredName, skip_signed: bool) -> Self {
+    pub fn new(root: StoredName, unsigned_zone: Option<Zone>, skip_signed: bool) -> Self {
         Self {
             inner: SimpleZoneInner {
                 root,
+                unsigned_zone,
                 tree: Default::default(),
                 skipped: Default::default(),
                 skip_signed,
@@ -128,12 +145,61 @@ impl ReadableZone for SimpleZoneInner {
 
     fn walk(&self, op: WalkOp) {
         trace!("WALK");
+
+        // Emit the records that we have, either unsigned or signed.
         for (name, rrsets) in self.tree.read().unwrap().iter() {
             for rrset in rrsets {
-                // TODO: Set false to proper value for "at zone cut or not"
-                (op)(name.clone(), rrset, false)
+                // If a DNSSEC RRSIG is encountered, ignore the RRSET TTL and
+                // instead emit it with the TTL of the original record the
+                // RRSIG covered, as for DNSSEC RRSIGs should not be grouped
+                // into RRSETs with a common TTL.
+                if rrset.rtype() == Rtype::RRSIG {
+                    for data in rrset.data() {
+                        let ZoneRecordData::Rrsig(rrsig) = data else {
+                            unreachable!();
+                        };
+                        let mut rrset = Rrset::new(Rtype::RRSIG, rrsig.original_ttl());
+                        rrset.push_data(data.clone());
+                        (op)(name.clone(), &rrset.into_shared(), false)
+                    }
+                } else {
+                    // TODO: Set false to proper value for "at zone cut or not"
+                    (op)(name.clone(), rrset, false)
+                }
             }
         }
+
+        /// A type for storing the passed op callback.
+        ///
+        /// A workaround for not being able to clone or store the op in an
+        /// Arc.
+        struct OpContainer {
+            op: WalkOp,
+        }
+
+        // Do we have an unsigned zone reference? Yes, walk over its records
+        // too.
+        //
+        // This allows us to store the signed records and refer to an
+        // existing copy of the unsigned records instead of duplicating them
+        // into our own storage.
+        if let Some(unsigned_zone) = &self.unsigned_zone {
+            let c = Arc::new(OpContainer { op });
+
+            // Wrap the real op inside a new one that allows us to omit the
+            // SOA RR when walking the unsigned records. We do this because
+            // the signed records contain the bumped SOA RR that should be
+            // used, not the original one in the unsigned records.
+            let op = Box::new(move |owner, rrset: &SharedRrset, _at_zone_cut| {
+                // Skip the SOA, use the new one that was part of the signed data.
+                if matches!(rrset.rtype(), Rtype::SOA) {
+                    return;
+                }
+
+                (c.op)(owner, rrset, _at_zone_cut)
+            });
+            unsigned_zone.read().walk(op);
+        }
         trace!("WALK FINISHED");
     }
 
@@ -173,15 +239,15 @@ impl WritableZone for SimpleZoneInner {
 }
 
 struct SimpleZoneNode {
-    pub tree: Arc<std::sync::RwLock<BTreeMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
+    pub tree: Arc<std::sync::RwLock<HashMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
     pub name: StoredName,
     pub skipped: Arc<AtomicUsize>,
     pub skip_signed: bool,
 }
 
 impl SimpleZoneNode {
     fn new(
-        tree: Arc<std::sync::RwLock<BTreeMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
+        tree: Arc<std::sync::RwLock<HashMap<StoredName, HashSet<HashedByRtypeSharedRrset>>>>,
         name: StoredName,
         skipped: Arc<AtomicUsize>,
         skip_signed: bool,