feat: add engine key ring with lazy tenant keys, deadlines and reconnect

crates/control/src/error.rs

@@ -27,19 +27,15 @@ impl fmt::Display for Error {
             Self::Keys {
                 tenant,
                 generation,
-                error: KeyError::Unavailable,
-            } => write!(
-                f,
-                "keys unavailable for tenant {tenant} generation {generation}"
-            ),
-            Self::Keys {
-                tenant,
-                generation,
-                error: KeyError::Unknown,
-            } => write!(
-                f,
-                "keys do not know tenant {tenant} generation {generation}"
-            ),
+                error,
+            } => {
+                let what = match error {
+                    KeyError::Unavailable => "unavailable",
+                    KeyError::Unknown => "unknown",
+                    KeyError::Failed => "failed",
+                };
+                write!(f, "keys {what} for tenant {tenant} generation {generation}")
+            }
         }
     }
 }

crates/control/src/lib.rs

@@ -11,7 +11,7 @@ mod writer;
 pub use cache::{Cache, CacheStats};
 pub use error::Error;
 pub use leases::{Ask, Gcra, IDLE, LIFETIME, LeaseStats, Leases, Taken, Tokens, WAIT, limit_core};
-pub use names::{Call, Loaded as LoadedQueues, Naming, NamingStats, Queues};
+pub use names::{Call, Loaded as LoadedQueues, Naming, NamingStats, Queues, Unnamed};
 pub use queue::{
     KeyError, KeySource, LABEL_SIZE, Label, MAX_DOMAIN, Name, Queue, Ready, Tenancy, domain_ok,
     mix, naming_shard, queue_shard, tenant_keys,

crates/control/src/names.rs

@@ -22,7 +22,7 @@ pub enum Call {
     Name(Name),
     Named {
         name: Name,
-        id: Option<u32>,
+        id: Result<u32, Unnamed>,
     },
     Forget {
         tenant: u32,
@@ -41,16 +41,25 @@ pub struct Naming {
     pub cache: usize,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Unnamed {
+    Refused,
+    Failed,
+    Keys(KeyError),
+}
+
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub struct NamingStats {
     pub cache: CacheStats,
     pub allocated: u64,
     pub failed: u64,
+    pub keys: u64,
     pub refused: u64,
     pub skipped: u64,
 }
 
 type Lock = (u32, Box<str>);
+type Named = Result<u32, Unnamed>;
 
 struct Ready {
     name: Name,
@@ -204,7 +213,7 @@ struct Inner {
     changes: Changes,
     owned: RefCell<Owned>,
     cache: RefCell<Cache>,
-    waiting: RefCell<BTreeMap<Name, Vec<Sender<Option<u32>>>>>,
+    waiting: RefCell<BTreeMap<Name, Vec<Sender<Named>>>>,
     forgetting: RefCell<BTreeMap<u64, (u16, Sender<()>)>>,
     tickets: Cell<u64>,
     peers: OnceCell<Peers<Call>>,
@@ -266,8 +275,13 @@ impl Queues {
                         hidden: parts.sealed.to_vec(),
                     })
                 });
+                let fits = |head: &Head| {
+                    head.id == id
+                        && head.tenant == item.tenant
+                        && queue_shard(id, naming.shards) == shard
+                };
                 match stored {
-                    Some(stored) if stored.head.id == id => found
+                    Some(stored) if fits(&stored.head) => found
                         .entry((stored.head.tenant, stored.head.generation))
                         .or_default()
                         .push(stored),
@@ -324,7 +338,15 @@ impl Queues {
             let core = self.naming_core(name.tenant, &name.domain);
             self.send(core, Call::Name(name.clone()));
         }
-        answer.recv().await.flatten().ok_or(Error::Refused)
+        match answer.recv().await {
+            Some(Ok(id)) => Ok(id),
+            Some(Err(Unnamed::Keys(error))) => Err(Error::Keys {
+                tenant: name.tenant,
+                generation: name.generation,
+                error,
+            }),
+            _ => Err(Error::Refused),
+        }
     }
 
     pub async fn rebind(&self, name: &Name, pool: u32) -> Result<u32, Error> {
@@ -443,7 +465,7 @@ impl Queues {
         match call {
             Call::Name(name) => self.asked(name, from),
             Call::Named { name, id } => {
-                if let Some(id) = id {
+                if let Ok(id) = id {
                     self.inner.cache.borrow_mut().put(name.clone(), id);
                 }
                 let waiting = self.inner.waiting.borrow_mut().remove(&name);
@@ -495,7 +517,7 @@ impl Queues {
             let mut owned = self.inner.owned.borrow_mut();
             if let Some(&id) = owned.names.get(&name) {
                 drop(owned);
-                self.send(from, Call::Named { name, id: Some(id) });
+                self.send(from, Call::Named { name, id: Ok(id) });
                 return false;
             }
             let shard = naming_shard(&naming.key, name.tenant, &name.domain, naming.shards);
@@ -507,7 +529,13 @@ impl Queues {
         };
         let Some((shard, id)) = id else {
             self.count(|s| s.refused += 1);
-            self.send(from, Call::Named { name, id: None });
+            self.send(
+                from,
+                Call::Named {
+                    name,
+                    id: Err(Unnamed::Refused),
+                },
+            );
             return false;
         };
         self.inner.ready.borrow_mut().push_back(Ready {
@@ -574,13 +602,20 @@ impl Queues {
             name: name.clone(),
             draining: false,
         };
-        let id = if self.write(shard, &queue).await.is_ok() {
-            self.inner.owned.borrow_mut().insert(queue);
-            self.count(|s| s.allocated += 1);
-            Some(id)
-        } else {
-            self.count(|s| s.failed += 1);
-            None
+        let id = match self.write(shard, &queue).await {
+            Ok(()) => {
+                self.inner.owned.borrow_mut().insert(queue);
+                self.count(|s| s.allocated += 1);
+                Ok(id)
+            }
+            Err(Error::Keys { error, .. }) => {
+                self.count(|s| s.keys += 1);
+                Err(Unnamed::Keys(error))
+            }
+            Err(_) => {
+                self.count(|s| s.failed += 1);
+                Err(Unnamed::Failed)
+            }
         };
         self.send(from, Call::Named { name, id });
         drop(held);

crates/control/src/queue.rs

@@ -144,6 +144,7 @@ impl Label {
 pub enum KeyError {
     Unavailable,
     Unknown,
+    Failed,
 }
 
 pub type Ready<'a> = Pin<Box<dyn Future<Output = Result<TenantKeys, KeyError>> + 'a>>;

crates/control/tests/queues.rs

@@ -54,6 +54,7 @@ const DOMAINS: [&str; 12] = [
 struct Toy {
     wrong: BTreeSet<u32>,
     unknown: BTreeSet<u32>,
+    failed: BTreeSet<u32>,
     unavailable: Cell<u32>,
     asked: Cell<u32>,
 }
@@ -68,6 +69,9 @@ impl Toy {
         if self.unknown.contains(&tenant) {
             return Err(KeyError::Unknown);
         }
+        if self.failed.contains(&tenant) {
+            return Err(KeyError::Failed);
+        }
         let seed = if self.wrong.contains(&tenant) { 9 } else { 1 };
         let root = Tenant::new(&Secret::new([seed; 32]));
         Ok(TenantKeys::derive(tenant, generation, &root, 1))
@@ -502,7 +506,7 @@ fn records_that_do_not_open_or_decode_are_skipped_at_start() {
         let shard = naming_shard(&KEY, 1, "b.example", SHARDS);
         let ids: Vec<u32> = (900..)
             .filter(|&id| queue_shard(id, SHARDS) == shard)
-            .take(4)
+            .take(5)
             .collect();
         let real = |id| Queue {
             id,
@@ -515,6 +519,17 @@ fn records_that_do_not_open_or_decode_are_skipped_at_start() {
             (ids[1], 1, real(ids[1]).encode(), b"not sealed".to_vec()),
             (ids[2], 3, real(ids[2]).encode(), seal_for(ids[2])),
             (ids[3], 1, real(ids[2]).encode(), seal_for(ids[2])),
+            (
+                ids[4],
+                1,
+                Queue {
+                    id: ids[4],
+                    name: name(9, 1, 1, "b.example"),
+                    draining: false,
+                }
+                .encode(),
+                seal_for(ids[4]),
+            ),
         ] {
             let batch = Batch {
                 shard,
@@ -539,12 +554,13 @@ fn records_that_do_not_open_or_decode_are_skipped_at_start() {
     machine.block_on(0, node.stop());
     let refusing = Rc::new(Toy {
         wrong: BTreeSet::from([2]),
+        unknown: BTreeSet::from([9]),
         ..Toy::default()
     });
     let driver = Rc::new(SimDriver::with_disk(2, Config::default(), driver.crash()));
     let machine = self::machine(&driver, cores);
     let node = Node::start(&machine, cores, &refusing).expect("restart");
-    assert_eq!(node.skipped, 5);
+    assert_eq!(node.skipped, 6);
     assert_eq!(node.loaded, 1);
     assert!(node.queues[0].entry(kept).is_some());
     assert!(node.queues[0].entry(gone).is_none());
@@ -645,6 +661,23 @@ fn start_waits_for_keys_asks_once_per_tenant_and_fails_on_unknown_ones() {
         ),
         "{failed:?}"
     );
+    let broken = Rc::new(Toy {
+        failed: BTreeSet::from([1]),
+        ..Toy::default()
+    });
+    let driver = Rc::new(SimDriver::with_disk(6, Config::default(), disk.clone()));
+    let failed = Node::try_start(&self::machine(&driver, cores), cores, &broken).err();
+    assert!(
+        matches!(
+            failed,
+            Some(Some(Error::Keys {
+                tenant: 1,
+                generation: 1,
+                error: KeyError::Failed,
+            }))
+        ),
+        "{failed:?}"
+    );
     let slow = Rc::new(Toy {
         unavailable: Cell::new(3),
         ..Toy::default()
@@ -657,6 +690,28 @@ fn start_waits_for_keys_asks_once_per_tenant_and_fails_on_unknown_ones() {
     assert_eq!(node.skipped, 0);
     let groups = 2 * SHARDS;
     assert!(slow.asked.get() <= groups + 3, "{} asks", slow.asked.get());
+    slow.unavailable.set(1_000);
+    let queues = node.queues[0].clone();
+    let refused = machine.block_on(1, async move {
+        let named = queues.name(&name(1, 1, 1, "new.example")).await;
+        (named, queues.stats().keys)
+    });
+    assert!(
+        matches!(
+            refused,
+            (
+                Err(Error::Keys {
+                    tenant: 1,
+                    generation: 1,
+                    error: KeyError::Unavailable,
+                }),
+                _
+            )
+        ),
+        "{refused:?}"
+    );
+    let keys: u64 = node.queues.iter().map(|q| q.stats().keys).sum();
+    assert_eq!(keys, 1);
     machine.block_on(0, node.stop());
 }
 

crates/crypto/src/keys.rs

@@ -1,9 +1,9 @@
-use std::rc::Rc;
+use std::sync::Arc;
 
 use crate::{DataKey, Error, Secret, Tenant};
 
 #[derive(Clone)]
-pub struct TenantKeys(Rc<Held>);
+pub struct TenantKeys(Arc<Held>);
 
 struct Held {
     tenant: u32,
@@ -20,7 +20,7 @@ impl TenantKeys {
         version: u8,
         hash: Secret<16>,
     ) -> Self {
-        Self(Rc::new(Held {
+        Self(Arc::new(Held {
             tenant,
             generation,
             data: DataKey::new(data, version),