Increment write backpressure before deferred encryption (#1444)

(Staged on top of #1443 )

This PR tightens the feedback loop for backpressure, by incrementing our
counters _before_ encryption is complete. It uses the new RAII
`BackpressureGuard` type, so we don't need to worry about weird
accounting failures.

Before:
```
                ┌────────────┐ incr             
             ┌──┤backpressure│◄───┐             
             │  └────────────┘    │             
┌─────┐  ┌───▼──┐    ┌──────────┐ │    ┌───────┐
│Guest├─►│submit├───►│encryption├─┼───►│clients│
└───▲─┘  └──────┘    └──────────┘ │    └───────┘
    │                             │             
    └─────────────────────────────┘             
              ack                               
```

After:
```
                ┌────────────┐                  
             ┌──┤backpressure│                  
             │  └─▲──────────┘                  
             │    │incr                         
┌─────┐  ┌───▼──┐ │  ┌──────────┐      ┌───────┐
│Guest├─►│submit├─┴─►│encryption├─┬───►│clients│
└───▲─┘  └──────┘    └──────────┘ │    └───────┘
    │                             │             
    └─────────────────────────────┘             
              ack                               
```

Author: mkeeter

upstairs/src/backpressure.rs

@@ -59,6 +59,22 @@ impl Drop for BackpressureGuard {
     }
 }
 
+impl BackpressureGuard {
+    #[cfg(test)]
+    pub fn dummy() -> Self {
+        let counter = Arc::new(BackpressureCountersInner {
+            write_bytes: 0.into(),
+            io_bytes: 0.into(),
+            jobs: 1.into(),
+        });
+        Self {
+            counter,
+            write_bytes: 0,
+            io_bytes: 0,
+        }
+    }
+}
+
 impl BackpressureCounters {
     pub fn new() -> Self {
         Self(Arc::new(BackpressureCountersInner {
@@ -80,6 +96,20 @@ impl BackpressureCounters {
         self.0.jobs.load(Ordering::Relaxed)
     }
 
+    /// Stores write / IO bytes (and 1 job) for a pending write
+    #[must_use]
+    pub fn early_write_increment(&mut self, bytes: u64) -> BackpressureGuard {
+        self.0.write_bytes.fetch_add(bytes, Ordering::Relaxed);
+        self.0.io_bytes.fetch_add(bytes, Ordering::Relaxed);
+        self.0.jobs.fetch_add(1, Ordering::Relaxed);
+        BackpressureGuard {
+            counter: self.0.clone(),
+            write_bytes: bytes,
+            io_bytes: bytes,
+            // implicit 1 job
+        }
+    }
+
     /// Stores write / IO bytes (and 1 job) in the backpressure counters
     #[must_use]
     pub fn increment(&mut self, io: &IOop) -> BackpressureGuard {

upstairs/src/client.rs

@@ -379,7 +379,10 @@ impl DownstairsClient {
             // backpressure.  Remove the backpressure guard for this client,
             // which decrements backpressure counters on drop.
             job.backpressure_guard.take(&self.client_id);
-        } else if is_running && !was_running {
+        } else if is_running
+            && !was_running
+            && !job.backpressure_guard.contains(&self.client_id)
+        {
             // This should only happen if a job is replayed, but that still
             // counts!
             job.backpressure_guard.insert(
@@ -935,7 +938,8 @@ impl DownstairsClient {
                 IOState::New
             }
         };
-        if r == IOState::New {
+        if r == IOState::New && !io.backpressure_guard.contains(&self.client_id)
+        {
             io.backpressure_guard.insert(
                 self.client_id,
                 self.backpressure_counters.increment(&io.work),

upstairs/src/deferred.rs

@@ -3,8 +3,9 @@
 use std::sync::Arc;
 
 use crate::{
-    client::ConnectionId, upstairs::UpstairsConfig, BlockContext, BlockOp,
-    BlockRes, ClientId, ImpactedBlocks, Message, RawWrite, Validation,
+    backpressure::BackpressureGuard, client::ConnectionId,
+    upstairs::UpstairsConfig, BlockContext, BlockOp, BlockRes, ClientData,
+    ClientId, ImpactedBlocks, Message, RawWrite, Validation,
 };
 use bytes::BytesMut;
 use crucible_common::{integrity_hash, CrucibleError, RegionDefinition};
@@ -114,6 +115,7 @@ pub(crate) struct DeferredWrite {
     pub res: BlockRes,
     pub is_write_unwritten: bool,
     pub cfg: Arc<UpstairsConfig>,
+    pub guard: ClientData<BackpressureGuard>,
 }
 
 /// Result of a deferred `BlockOp`
@@ -135,6 +137,7 @@ pub(crate) struct EncryptedWrite {
     pub impacted_blocks: ImpactedBlocks,
     pub res: BlockRes,
     pub is_write_unwritten: bool,
+    pub guard: ClientData<BackpressureGuard>,
 }
 
 impl DeferredWrite {
@@ -183,6 +186,7 @@ impl DeferredWrite {
             impacted_blocks: self.impacted_blocks,
             res: self.res,
             is_write_unwritten: self.is_write_unwritten,
+            guard: self.guard,
         }
     }
 }

upstairs/src/downstairs.rs

@@ -7,6 +7,7 @@ use std::{
 };
 
 use crate::{
+    backpressure::BackpressureGuard,
     cdt,
     client::{ClientAction, ClientStopReason, DownstairsClient},
     guest::GuestWork,
@@ -1822,6 +1823,7 @@ impl Downstairs {
             GuestWorkId(10),
             request,
             is_write_unwritten,
+            ClientData::from_fn(|_| BackpressureGuard::dummy()),
         )
     }
 
@@ -1831,6 +1833,7 @@ impl Downstairs {
         gw_id: GuestWorkId,
         write: RawWrite,
         is_write_unwritten: bool,
+        bp_guard: ClientData<BackpressureGuard>,
     ) -> JobId {
         let ds_id = self.next_id();
         let dependencies = self.ds_active.deps_for_write(ds_id, blocks);
@@ -1872,7 +1875,7 @@ impl Downstairs {
             replay: false,
             data: None,
             read_validations: Vec::new(),
-            backpressure_guard: ClientMap::new(),
+            backpressure_guard: bp_guard.into(),
         };
         self.enqueue(io);
         ds_id
@@ -2453,6 +2456,7 @@ impl Downstairs {
         blocks: ImpactedBlocks,
         write: RawWrite,
         is_write_unwritten: bool,
+        backpressure_guard: ClientData<BackpressureGuard>,
     ) -> JobId {
         // If there is a live-repair in progress that intersects with this read,
         // then reserve job IDs for those jobs.
@@ -2463,6 +2467,7 @@ impl Downstairs {
             guest_id,
             write,
             is_write_unwritten,
+            backpressure_guard,
         )
     }
 
@@ -3674,6 +3679,7 @@ impl Downstairs {
                 data,
             },
             is_write_unwritten,
+            ClientData::from_fn(|_| BackpressureGuard::dummy()),
         )
     }
 
@@ -4422,6 +4428,19 @@ impl Downstairs {
             }
         });
     }
+
+    /// Assign the given number of write bytes to the backpressure counters
+    #[must_use]
+    pub(crate) fn early_write_backpressure(
+        &mut self,
+        bytes: u64,
+    ) -> ClientData<BackpressureGuard> {
+        ClientData::from_fn(|i| {
+            self.clients[i]
+                .backpressure_counters
+                .early_write_increment(bytes)
+        })
+    }
 }
 
 /// Configuration for per-client backpressure

upstairs/src/lib.rs

@@ -427,6 +427,15 @@ impl<T> ClientData<T> {
         std::mem::swap(&mut self[c], &mut v);
         v
     }
+
+    /// Builds a `ClientData` from a builder function
+    pub fn from_fn<F: FnMut(ClientId) -> T>(mut f: F) -> Self {
+        Self([
+            f(ClientId::new(0)),
+            f(ClientId::new(1)),
+            f(ClientId::new(2)),
+        ])
+    }
 }
 
 /// Map of data associated with clients, keyed by `ClientId`
@@ -466,6 +475,12 @@ impl<T> std::ops::Index<ClientId> for ClientMap<T> {
     }
 }
 
+impl<T> From<ClientData<T>> for ClientMap<T> {
+    fn from(c: ClientData<T>) -> Self {
+        Self(ClientData(c.0.map(Option::Some)))
+    }
+}
+
 /*
  * These counts describe the various states that a Downstairs IO can
  * be in.

upstairs/src/upstairs.rs

@@ -1439,13 +1439,16 @@ impl Upstairs {
         let impacted_blocks =
             extent_from_offset(&ddef, offset, ddef.bytes_to_blocks(data.len()));
 
+        let guard = self.downstairs.early_write_backpressure(data.len() as u64);
+
         Some(DeferredWrite {
             ddef,
             impacted_blocks,
             data,
             res,
             is_write_unwritten,
             cfg: self.cfg.clone(),
+            guard,
         })
     }
 
@@ -1473,6 +1476,7 @@ impl Upstairs {
                     write.impacted_blocks,
                     write.data,
                     write.is_write_unwritten,
+                    write.guard,
                 )
             },
             Some(GuestBlockRes::Other(write.res)),