From 3a7af1e701353a9f632de5a452c63c9f9e1e0640 Mon Sep 17 00:00:00 2001
From: Gilbert Gilb's <gilbsgilbert@gmail.com>
Date: Thu, 13 Feb 2025 19:33:32 +0100
Subject: [PATCH] feat(backends/s3): add warmup support for check command

Follow-up from https://github.com/restic/restic/pull/5173
See also https://github.com/restic/restic/issues/3202
---
 changelog/unreleased/issue-3202 |  6 +++---
 cmd/restic/cmd_check.go         |  5 +----
 doc/faq.rst                     |  1 +
 internal/checker/checker.go     | 34 ++++++++++++++++++++++++++-------
 4 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/changelog/unreleased/issue-3202 b/changelog/unreleased/issue-3202
index e6fca05e1..2985d1062 100644
--- a/changelog/unreleased/issue-3202
+++ b/changelog/unreleased/issue-3202
@@ -1,10 +1,10 @@
 Enhancement: Add warmup support on S3 backend before repacks and restores
 
 Introduce S3 backend options for transitioning pack files from cold to hot
-storage on S3 and S3-compatible providers. Note: only works before repacks
-(prune/copy) and restore for now, and gated behind a new "s3-restore" feature
-flag.
+storage on S3 and S3-compatible providers. Note: only works with prune, copy,
+check and restore commands, and gated behind a new "s3-restore" feature flag.
 
 https://github.com/restic/restic/pull/5173
+https://github.com/restic/restic/pull/5248
 https://github.com/restic/restic/issues/3202
 https://github.com/restic/restic/issues/2504
diff --git a/cmd/restic/cmd_check.go b/cmd/restic/cmd_check.go
index f87303933..34842d510 100644
--- a/cmd/restic/cmd_check.go
+++ b/cmd/restic/cmd_check.go
@@ -375,11 +375,9 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 	}
 
 	doReadData := func(packs map[restic.ID]int64) {
-		p := printer.NewCounter("packs")
-		p.SetMax(uint64(len(packs)))
 		errChan := make(chan error)
 
-		go chkr.ReadPacks(ctx, packs, p, errChan)
+		go chkr.ReadPacks(ctx, packs, printer, errChan)
 
 		for err := range errChan {
 			errorsFound = true
@@ -389,7 +387,6 @@ func runCheck(ctx context.Context, opts CheckOptions, gopts GlobalOptions, args
 				salvagePacks.Insert(err.PackID)
 			}
 		}
-		p.Done()
 	}
 
 	switch {
diff --git a/doc/faq.rst b/doc/faq.rst
index 3b62f641d..a92371102 100644
--- a/doc/faq.rst
+++ b/doc/faq.rst
@@ -269,6 +269,7 @@ Archive** storage classes is available:
 - Currently, only the following commands are known to work:
 
   - `backup`
+  - `check`
   - `copy`
   - `prune`
   - `restore`
diff --git a/internal/checker/checker.go b/internal/checker/checker.go
index 12020891a..164c264e5 100644
--- a/internal/checker/checker.go
+++ b/internal/checker/checker.go
@@ -10,6 +10,7 @@ import (
 	"github.com/klauspost/compress/zstd"
 	"github.com/restic/restic/internal/debug"
 	"github.com/restic/restic/internal/errors"
+	"github.com/restic/restic/internal/feature"
 	"github.com/restic/restic/internal/repository"
 	"github.com/restic/restic/internal/repository/index"
 	"github.com/restic/restic/internal/repository/pack"
@@ -441,15 +442,39 @@ func (c *Checker) GetPacks() map[restic.ID]int64 {
 
 // ReadData loads all data from the repository and checks the integrity.
 func (c *Checker) ReadData(ctx context.Context, errChan chan<- error) {
-	c.ReadPacks(ctx, c.packs, nil, errChan)
+	c.ReadPacks(ctx, c.packs, &progress.NoopPrinter{}, errChan)
 }
 
 const maxStreamBufferSize = 4 * 1024 * 1024
 
 // ReadPacks loads data from specified packs and checks the integrity.
-func (c *Checker) ReadPacks(ctx context.Context, packs map[restic.ID]int64, p *progress.Counter, errChan chan<- error) {
+func (c *Checker) ReadPacks(ctx context.Context, packs map[restic.ID]int64, printer progress.Printer, errChan chan<- error) {
 	defer close(errChan)
 
+	p := printer.NewCounter("packs")
+	p.SetMax(uint64(len(packs)))
+	defer p.Done()
+
+	packSet := restic.NewIDSet()
+	for pack := range packs {
+		packSet.Insert(pack)
+	}
+
+	if feature.Flag.Enabled(feature.S3Restore) {
+		job, err := c.repo.StartWarmup(ctx, packSet)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		if job.HandleCount() != 0 {
+			printer.P("warming up %d packs from cold storage, this may take a while...", job.HandleCount())
+			if err := job.Wait(ctx); err != nil {
+				errChan <- err
+				return
+			}
+		}
+	}
+
 	g, ctx := errgroup.WithContext(ctx)
 	type checkTask struct {
 		id    restic.ID
@@ -497,11 +522,6 @@ func (c *Checker) ReadPacks(ctx context.Context, packs map[restic.ID]int64, p *p
 		})
 	}
 
-	packSet := restic.NewIDSet()
-	for pack := range packs {
-		packSet.Insert(pack)
-	}
-
 	// push packs to ch
 	for pbs := range c.repo.ListPacksFromIndex(ctx, packSet) {
 		size := packs[pbs.PackID]