From 157848d87ae57ab24c02aee910f361e35a99df6c Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Tue, 14 Jul 2026 10:16:00 +0800 Subject: [PATCH] main: bound startup sweep with context timeout Sweep() used to run synchronously at startup without any time bound. On a large uploads directory this can block the server for 30s or more, keeping /healthz returning 503 while the reaper works. Add a context timeout so startup continues if the sweep cannot finish within budget. The timeout is set to 5 seconds: long enough to clear typical leftover state, short enough that a stuck sweep will not meaningfully delay health checks or startup readiness. Change Store.Sweep to accept a context.Context and check ctx.Err() at the top of each iteration. If the context is cancelled or times out, Sweep returns the files deleted so far and ctx.Err(). Add TestStore_Sweep_HonorsContextCancel to ensure a cancelled context causes Sweep to exit early instead of running to completion. Deliberately unchanged: there is still no periodic reaper. The server only sweeps once at startup; this change just bounds that single sweep. Co-Authored-By: Claude --- internal/uploads/store.go | 6 ++++- internal/uploads/store_test.go | 43 +++++++++++++++++++++++++++++++--- main.go | 10 ++++++-- 3 files changed, 53 insertions(+), 6 deletions(-) diff --git a/internal/uploads/store.go b/internal/uploads/store.go index 13e1bea..766ae23 100644 --- a/internal/uploads/store.go +++ b/internal/uploads/store.go @@ -4,6 +4,7 @@ package uploads import ( + "context" "crypto/rand" "crypto/sha256" "encoding/hex" @@ -160,7 +161,7 @@ func (s *Store) Validate(absPath string) (string, error) { // Sweep deletes data files (and their sidecars) whose uploaded_at is older // than ttl, as well as orphan sidecars that have no matching data file. For // data files whose sidecar is missing, the file's mtime is used as a fallback. -func (s *Store) Sweep(ttl time.Duration) (deleted int, err error) { +func (s *Store) Sweep(ctx context.Context, ttl time.Duration) (deleted int, err error) { entries, err := os.ReadDir(s.Root) if err != nil { return 0, fmt.Errorf("uploads: read root: %w", err) @@ -197,6 +198,9 @@ func (s *Store) Sweep(ttl time.Duration) (deleted int, err error) { cutoff := time.Now().Add(-ttl) for fileID, it := range items { + if err := ctx.Err(); err != nil { + return deleted, err + } dataPath := filepath.Join(s.Root, fileID) metaPath := dataPath + ".meta.json" diff --git a/internal/uploads/store_test.go b/internal/uploads/store_test.go index 336bf4e..00ed5a4 100644 --- a/internal/uploads/store_test.go +++ b/internal/uploads/store_test.go @@ -1,6 +1,7 @@ package uploads import ( + "context" "crypto/sha256" "encoding/hex" "encoding/json" @@ -201,7 +202,7 @@ func TestStore_Sweep_DeletesExpiredKeepsRecent(t *testing.T) { create(freshID, timePtr(now), true) create(noMetaID, nil, true) - deleted, err := store.Sweep(1 * time.Hour) + deleted, err := store.Sweep(context.Background(), 1*time.Hour) if err != nil { t.Fatalf("sweep: %v", err) } @@ -241,7 +242,7 @@ func TestStore_Sweep_OrphanSidecarReaped(t *testing.T) { t.Fatalf("create orphan sidecar: %v", err) } - deleted, err := store.Sweep(1 * time.Hour) + deleted, err := store.Sweep(context.Background(), 1*time.Hour) if err != nil { t.Fatalf("sweep: %v", err) } @@ -277,6 +278,42 @@ func TestStore_Validate_RejectsSidecarPath(t *testing.T) { } } +func TestStore_Sweep_HonorsContextCancel(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + + id := "00000000000000000000000000000006" + if err := os.WriteFile(filepath.Join(store.Root, id), []byte("x"), 0o640); err != nil { + t.Fatalf("create data file: %v", err) + } + sc := map[string]any{ + "name": "x", + "size": 1, + "sha256": "abc", + "uploaded_at": time.Now().Add(-2 * time.Hour).Format(time.RFC3339Nano), + } + b, _ := json.Marshal(sc) + if err := os.WriteFile(filepath.Join(store.Root, id+".meta.json"), b, 0o600); err != nil { + t.Fatalf("create sidecar: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + deleted, err := store.Sweep(ctx, 1*time.Hour) + if err != ctx.Err() { + t.Fatalf("sweep err=%v, want %v", err, ctx.Err()) + } + if deleted != 0 { + t.Errorf("deleted=%d, want 0", deleted) + } + if _, existErr := os.Stat(filepath.Join(store.Root, id)); os.IsNotExist(existErr) { + t.Logf("cancelled sweep managed to delete the file before checking ctx; this is timing-dependent and acceptable") + } +} + func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) { store, err := New(t.TempDir()) if err != nil { @@ -300,7 +337,7 @@ func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) { t.Fatalf("set mtime: %v", err) } - deleted, err := store.Sweep(1 * time.Hour) + deleted, err := store.Sweep(context.Background(), 1*time.Hour) if err != nil { t.Fatalf("sweep: %v", err) } diff --git a/main.go b/main.go index c99b5e1..c0275d0 100644 --- a/main.go +++ b/main.go @@ -30,6 +30,8 @@ import ( "spark-mcp-go/internal/uploads" ) +const startupSweepTimeout = 5 * time.Second + func main() { if err := run(); err != nil { // run() handles its own logging; this is the last-chance report. @@ -71,8 +73,12 @@ func run() error { if err != nil { return err } - n, err := uploadStore.Sweep(cfg.UploadTTL) - if err != nil { + sweepCtx, sweepCancel := context.WithTimeout(context.Background(), startupSweepTimeout) + n, err := uploadStore.Sweep(sweepCtx, cfg.UploadTTL) + sweepCancel() + if errors.Is(err, context.DeadlineExceeded) { + logger.Warn("uploads.sweep.timeout", "deleted", n, "err", err) + } else if err != nil { logger.Error("uploads.sweep", "err", err) } else { logger.Info("uploads.sweep", "deleted", n)