From 6359820d3dc213eeceb0cbbcd97b14db590ff167 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Wed, 16 Oct 2024 16:35:13 -0700
Subject: [PATCH] startup: partially restore waiting for spool startup

We had a user report that launching kumod
post-7fbec8758720dad465d711458f6bf7923b778773 produced a panic about the
spool not being configured.

My supposition is that they were bringing the system up under load
(incoming traffic) and that a new message was being injected before we
had started the spool.

So what we do here is restore the spool_started check, but have it
be satisfied once the startup has been triggered, rather than
on completion of enumeration.

The result is that there is still a small window where we might turn
away a message during startup, but it should be very very small.
---
 .../src/http_server/check_liveness_v1.rs      |  3 +++
 crates/kumod/src/smtp_server.rs               | 20 +++++++++++++++++++
 crates/kumod/src/spool.rs                     | 10 +++++++++-
 docs/changelog/main.md                        |  6 ++++--
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/crates/kumod/src/http_server/check_liveness_v1.rs b/crates/kumod/src/http_server/check_liveness_v1.rs
index a665624d..8c43bee7 100644
--- a/crates/kumod/src/http_server/check_liveness_v1.rs
+++ b/crates/kumod/src/http_server/check_liveness_v1.rs
@@ -1,3 +1,4 @@
+use crate::spool::SpoolManager;
 use axum::response::Response;
 use kumo_server_lifecycle::Activity;
 
@@ -18,6 +19,8 @@ pub async fn check_liveness_v1() -> Response {
         Some(_activity) => {
             if kumo_server_memory::get_headroom() == 0 {
                 (503, "load shedding")
+            } else if !SpoolManager::get().spool_started() {
+                (503, "waiting for spool startup")
             } else if kumo_server_common::disk_space::is_over_limit() {
                 (503, "storage is too full")
             } else {
diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs
index 711260d5..c4f84625 100644
--- a/crates/kumod/src/smtp_server.rs
+++ b/crates/kumod/src/smtp_server.rs
@@ -4,6 +4,7 @@ use crate::http_server::admin_trace_smtp_server_v1::{
 use crate::logging::disposition::{log_disposition, LogDisposition, RecordType};
 use crate::logging::rejection::{log_rejection, LogRejection};
 use crate::queue::QueueManager;
+use crate::spool::SpoolManager;
 use anyhow::{anyhow, Context};
 use chrono::Utc;
 use cidr_map::CidrSet;
@@ -1041,6 +1042,25 @@ impl SmtpServer {
             return Ok(());
         }
 
+        if !SpoolManager::get().spool_started() {
+            // We don't bump the connection_denied_counter here, because
+            // startup is a normal condition and doesn't require an operator
+            // to respond.
+
+            // Can't accept any messages until the spool has started
+            // because we won't know where to put them.
+            self.write_response(
+                421,
+                format!(
+                    "4.3.2 {} waiting for spool startup. Try again soon!",
+                    self.params.hostname
+                ),
+                None,
+            )
+            .await?;
+            return Ok(());
+        }
+
         self.write_response(
             220,
             format!("{} {}", self.params.hostname, self.params.banner),
diff --git a/crates/kumod/src/spool.rs b/crates/kumod/src/spool.rs
index 58b6dcb1..5cadd6a8 100644
--- a/crates/kumod/src/spool.rs
+++ b/crates/kumod/src/spool.rs
@@ -15,7 +15,7 @@ use spool::rocks::{RocksSpool, RocksSpoolParams};
 use spool::{get_data_spool, get_meta_spool, Spool as SpoolTrait, SpoolEntry, SpoolId};
 use std::collections::HashMap;
 use std::path::PathBuf;
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, LazyLock, Mutex as StdMutex};
 use std::time::{Duration, Instant};
 use tokio::sync::Mutex;
@@ -129,12 +129,14 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> {
 
 pub struct SpoolManager {
     named: Mutex<HashMap<String, SpoolHandle>>,
+    started: AtomicBool,
 }
 
 impl SpoolManager {
     pub fn new() -> Self {
         Self {
             named: Mutex::new(HashMap::new()),
+            started: AtomicBool::new(false),
         }
     }
 
@@ -238,6 +240,10 @@ impl SpoolManager {
         (get_meta_spool(), get_data_spool())
     }
 
+    pub fn spool_started(&self) -> bool {
+        self.started.load(Ordering::SeqCst)
+    }
+
     pub async fn remove_from_spool(id: SpoolId) -> anyhow::Result<()> {
         let (data_spool, meta_spool) = Self::get_data_meta();
         let res_data = data_spool.remove(id).await;
@@ -406,6 +412,8 @@ impl SpoolManager {
     }
 
     pub async fn start_spool(&self, start_time: DateTime<Utc>) -> anyhow::Result<()> {
+        self.started.store(true, Ordering::SeqCst);
+
         let (tx, rx) = flume::bounded(1024);
         {
             let mut named = self.named.lock().await;
diff --git a/docs/changelog/main.md b/docs/changelog/main.md
index 87d46b76..0bacf2be 100644
--- a/docs/changelog/main.md
+++ b/docs/changelog/main.md
@@ -89,8 +89,10 @@
 * The bounce classifier will now automatically reload when the configuration
   epoch is updated. #298
 
-* Spool enumeration no longer prevents the reception of new messages and is
-  no longer considered as part of the liveness check.
+* Spool enumeration completion no longer prevents the reception of new messages.
+  The spool does need to have started before messages will be received, which
+  means that there is still a very small window during startup where the liveness
+  check and SMTP sessions can turn away incoming connections.
 
 * The dkim helper now supports passing the `expiration` value through to
   the underlying signer.