Bundle the seqno and the lease together

Author: bkirwi

src/persist-client/src/fetch.rs

@@ -428,10 +428,19 @@ where
 ///
 /// Generally the state and lease are bundled together, as in [LeasedBatchPart]... but sometimes
 /// it's necessary to handle them separately, so this struct is exposed as well. Handle with care.
-#[derive(Clone, Debug, Default)]
-pub(crate) struct Lease(Arc<()>);
+#[derive(Clone, Debug)]
+pub(crate) struct Lease(Arc<SeqNo>);
 
 impl Lease {
+    pub fn new(seqno: SeqNo) -> Self {
+        Self(Arc::new(seqno))
+    }
+
+    #[cfg(test)]
+    pub fn seqno(&self) -> SeqNo {
+        *self.0
+    }
+
     /// Returns the number of live copies of this lease, including this one.
     pub fn count(&self) -> usize {
         Arc::strong_count(&self.0)
@@ -470,10 +479,6 @@ pub struct LeasedBatchPart<T> {
     pub(crate) filter: FetchBatchFilter<T>,
     pub(crate) desc: Description<T>,
     pub(crate) part: BatchPart<T>,
-    /// The `SeqNo` from which this part originated; we track this value as
-    /// to ensure the `SeqNo` isn't garbage collected while a
-    /// read still depends on it.
-    pub(crate) leased_seqno: SeqNo,
     /// The lease that prevents this part from being GCed. Code should ensure that this lease
     /// lives as long as the part is needed.
     pub(crate) lease: Option<Lease>,

src/persist-client/src/operators/shard_source.rs

@@ -696,6 +696,7 @@ mod tests {
     use super::*;
     use std::sync::Arc;
 
+    use mz_persist::location::SeqNo;
     use timely::dataflow::Scope;
     use timely::dataflow::operators::Leave;
     use timely::dataflow::operators::Probe;
@@ -706,7 +707,7 @@ mod tests {
 
     #[mz_ore::test]
     fn test_lease_manager() {
-        let lease = Lease::default();
+        let lease = Lease::new(SeqNo::minimum());
         let mut manager = LeaseManager::new();
         for t in 0u64..10 {
             manager.push_at(t, lease.clone());

src/persist-client/src/read.rs

@@ -726,7 +726,6 @@ where
             filter,
             desc,
             part,
-            leased_seqno: self.machine.seqno(),
             lease: Some(self.lease_seqno()),
             filter_pushdown_audit: false,
         }
@@ -752,7 +751,10 @@ where
     /// collected until its lease has been returned.
     fn lease_seqno(&mut self) -> Lease {
         let seqno = self.machine.seqno();
-        let lease = self.leased_seqnos.entry(seqno).or_default();
+        let lease = self
+            .leased_seqnos
+            .entry(seqno)
+            .or_insert_with(|| Lease::new(seqno));
         lease.clone()
     }
 
@@ -1448,7 +1450,7 @@ mod tests {
 
         // Repeat the same process as above, more or less, while fetching + returning parts
         for (mut i, part) in parts.into_iter().enumerate() {
-            let part_seqno = part.leased_seqno;
+            let part_seqno = part.lease.as_ref().unwrap().seqno();
             let last_seqno = this_seqno;
             this_seqno = part_seqno;
             assert!(this_seqno >= last_seqno);