Improve the BOLD Challenge Cache

staker/challenge-cache/cache.go

@@ -10,15 +10,15 @@ store them in a filesystem cache to avoid recomputing them and for hierarchical
 Each file contains a list of 32 byte hashes, concatenated together as bytes.
 Using this structure, we can namespace hashes by message number and by challenge level.
 
-Once a validator receives a full list of computed machine hashes for the first time from a validatio node,
+Once a validator receives a full list of computed machine hashes for the first time from a validation node,
 it will write the hashes to this filesystem hierarchy for fast access next time these hashes are needed.
 
 Example uses:
 - Obtain all the hashes for the execution of message num 70 to 71 for a given wavm module root.
 - Obtain all the hashes from step 100 to 101 at subchallenge level 1 for the execution of message num 70.
 
 	  wavm-module-root-0xab/
-		rollup-block-hash-0x12...-message-num-70/
+		message-num-70-rollup-block-hash-0x12.../
 			hashes.bin
 			subchallenge-level-1-big-step-100/
 				hashes.bin
@@ -31,18 +31,21 @@ package challengecache
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"regexp"
+	"strconv"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 )
 
 var (
-	ErrNotFoundInCache    = errors.New("no found in challenge cache")
+	ErrNotFoundInCache    = errors.New("not found in challenge cache")
 	ErrFileAlreadyExists  = errors.New("file already exists")
 	ErrNoHashes           = errors.New("no hashes being written")
 	hashesFileName        = "hashes.bin"
@@ -59,30 +62,46 @@ type HistoryCommitmentCacher interface {
 	Put(lookup *Key, hashes []common.Hash) error
 }
 
+// Key for cache lookups includes the wavm module root of a challenge, as well
+// as the heights for messages and big steps as needed.
+type Key struct {
+	RollupBlockHash common.Hash
+	WavmModuleRoot  common.Hash
+	MessageHeight   uint64
+	StepHeights     []uint64
+}
+
 // Cache for history commitments on disk.
 type Cache struct {
-	baseDir string
+	baseDir       string
+	tempWritesDir string
 }
 
 // New cache from a base directory path.
 func New(baseDir string) (*Cache, error) {
-	if _, err := os.Stat(baseDir); err != nil {
-		if err := os.MkdirAll(baseDir, os.ModePerm); err != nil {
-			return nil, fmt.Errorf("could not make base cache directory %s: %w", baseDir, err)
-		}
-	}
 	return &Cache{
-		baseDir: baseDir,
+		baseDir:       baseDir,
+		tempWritesDir: "",
 	}, nil
 }
 
-// Key for cache lookups includes the wavm module root of a challenge, as well
-// as the heights for messages and big steps as needed.
-type Key struct {
-	RollupBlockHash common.Hash
-	WavmModuleRoot  common.Hash
-	MessageHeight   uint64
-	StepHeights     []uint64
+// Init a cache by verifying its base directory exists.
+func (c *Cache) Init(_ context.Context) error {
+	if _, err := os.Stat(c.baseDir); err != nil {
+		if err := os.MkdirAll(c.baseDir, os.ModePerm); err != nil {
+			return fmt.Errorf("could not make initialize challenge cache directory %s: %w", c.baseDir, err)
+		}
+	}
+	// We create a temp directory to write our hashes to first when putting to the cache.
+	// Once writing succeeds, we rename in an atomic operation to the correct file name
+	// in the cache directory hierarchy in the `Put` function. All of these temporary writes
+	// will occur in a subdir of the base directory called temp.
+	tempWritesDir, err := os.MkdirTemp(c.baseDir, "temp")
+	if err != nil {
+		return err
+	}
+	c.tempWritesDir = tempWritesDir
+	return nil
 }
 
 // Get a list of hashes from the cache from index 0 up to a certain index. Hashes are saved as files in the directory
@@ -122,24 +141,14 @@ func (c *Cache) Put(lookup *Key, hashes []common.Hash) error {
 	if len(hashes) == 0 {
 		return ErrNoHashes
 	}
-	fName, err := determineFilePath(c.baseDir, lookup)
-	if err != nil {
-		return err
+	if c.tempWritesDir == "" {
+		return fmt.Errorf("cache not initialized by calling .Init(ctx)")
 	}
-	// We create a tmp file to write our hashes to first. If writing fails,
-	// we don't want to leave a half-written file in our cache directory.
-	// Once writing succeeds, we rename in an atomic operation to the correct file name
-	// in the cache directory hierarchy.
-	tmp, err := os.MkdirTemp(c.baseDir, "tmpdir")
+	fName, err := determineFilePath(c.baseDir, lookup)
 	if err != nil {
 		return err
 	}
-	tmpFName := filepath.Join(tmp, fName)
-	dir := filepath.Dir(tmpFName)
-	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
-		return fmt.Errorf("could not make tmp directory %s: %w", dir, err)
-	}
-	f, err := os.Create(tmpFName)
+	f, err := os.CreateTemp(c.tempWritesDir, fmt.Sprintf("%s-*", hashesFileName))
 	if err != nil {
 		return err
 	}
@@ -154,11 +163,57 @@ func (c *Cache) Put(lookup *Key, hashes []common.Hash) error {
 	if err := os.MkdirAll(filepath.Dir(fName), os.ModePerm); err != nil {
 		return fmt.Errorf("could not make file directory %s: %w", fName, err)
 	}
-	// If the file writing was successful, we rename the file from the tmp directory
+	// If the file writing was successful, we rename the file from the temp directory
 	// into our cache directory. This is an atomic operation.
 	// For more information on this atomic write pattern, see:
 	// https://stackoverflow.com/questions/2333872/how-to-make-file-creation-an-atomic-operation
-	return os.Rename(tmpFName /*old */, fName /* new */)
+	return os.Rename(f.Name() /*old */, fName /* new */)
+}
+
+// Prune all entries in the cache with a message number <= a specified value.
+func (c *Cache) Prune(ctx context.Context, messageNumber uint64) error {
+	// Define a regex pattern to extract the message number
+	numPruned := 0
+	messageNumPattern := fmt.Sprintf(`%s-(\d+)-`, messageNumberPrefix)
+	pattern := regexp.MustCompile(messageNumPattern)
+	pathsToDelete := make([]string, 0)
+	if err := filepath.WalkDir(c.baseDir, func(path string, info os.DirEntry, err error) error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		if err != nil {
+			return err
+		}
+		if info.IsDir() {
+			matches := pattern.FindStringSubmatch(info.Name())
+			if len(matches) > 1 {
+				dirNameMessageNum, err := strconv.Atoi(matches[1])
+				if err != nil {
+					return err
+				}
+				// Collect the directory path if the message number is <= the specified value.
+				if dirNameMessageNum <= int(messageNumber) {
+					pathsToDelete = append(pathsToDelete, path)
+				}
+			}
+		}
+		return nil
+	}); err != nil {
+		return err
+	}
+	// We delete separately from collecting the paths, as deleting while walking
+	// a dir can cause issues with the filepath.Walk function.
+	for _, path := range pathsToDelete {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		if err := os.RemoveAll(path); err != nil {
+			return fmt.Errorf("could not prune directory with path %s: %w", path, err)
+		}
+		numPruned += 1
+	}
+	log.Info("Pruned challenge cache", "numDirsPruned", numPruned, "messageNumber", messageNumPattern)
+	return nil
 }
 
 // Reads 32 bytes at a time from a reader up to a specified height. If none, then read all.
@@ -217,15 +272,15 @@ for the data requested within the cache directory hierarchy. The folder structur
 for a given filesystem challenge cache will look as follows:
 
 	  wavm-module-root-0xab/
-		rollup-block-hash-0x12...-message-num-70/
+		message-num-70-rollup-block-hash-0x12.../
 			hashes.bin
 			subchallenge-level-1-big-step-100/
 				hashes.bin
 */
 func determineFilePath(baseDir string, lookup *Key) (string, error) {
 	key := make([]string, 0)
 	key = append(key, fmt.Sprintf("%s-%s", wavmModuleRootPrefix, lookup.WavmModuleRoot.Hex()))
-	key = append(key, fmt.Sprintf("%s-%s-%s-%d", rollupBlockHashPrefix, lookup.RollupBlockHash.Hex(), messageNumberPrefix, lookup.MessageHeight))
+	key = append(key, fmt.Sprintf("%s-%d-%s-%s", messageNumberPrefix, lookup.MessageHeight, rollupBlockHashPrefix, lookup.RollupBlockHash.Hex()))
 	for challengeLevel, height := range lookup.StepHeights {
 		key = append(key, fmt.Sprintf(
 			"%s-%d-%s-%d",