tools/syz-aflow: support setting a custom cache size

This is useful to save time when iterating on syz-aflow changes.

Note that this also sets a cache size of 10GB by default

1 files changed, 30 insertions, 3 deletions

diff --git a/tools/syz-aflow/aflow.go b/tools/syz-aflow/aflow.go
index ea127b164..ee770e530 100644
--- a/tools/syz-aflow/aflow.go
+++ b/tools/syz-aflow/aflow.go
@@ -32,6 +32,7 @@ func main() {
 		flagInput       = flag.String("input", "", "input json file with workflow arguments")
 		flagWorkdir     = flag.String("workdir", "", "directory for kernel checkout, kernel builds, etc")
 		flagModel       = flag.String("model", aflow.DefaultModel, "use this LLM model")
+		flagCacheSize   = flag.String("cache-size", "10GB", "max cache size (e.g. 100MB, 5GB, 1TB)")
 		flagDownloadBug = flag.String("download-bug", "", "extid of a bug to download from the dashboard"+
 			" and save into -input file")
 	)
@@ -51,12 +52,38 @@ func main() {
 		}
 		return
 	}
-	if err := run(context.Background(), *flagModel, *flagFlow, *flagInput, *flagWorkdir); err != nil {
+	cacheSize, err := parseSize(*flagCacheSize)
+	if err != nil {
+		tool.Fail(err)
+	}
+	if err := run(context.Background(), *flagModel, *flagFlow, *flagInput, *flagWorkdir, cacheSize); err != nil {
 		tool.Fail(err)
 	}
 }
 
-func run(ctx context.Context, model, flowName, inputFile, workdir string) error {
+func parseSize(s string) (uint64, error) {
+	var size uint64
+	var suffix string
+	if _, err := fmt.Sscanf(s, "%d%s", &size, &suffix); err != nil {
+		return 0, fmt.Errorf("failed to parse cache size %q: %w", s, err)
+	}
+	switch suffix {
+	case "KB":
+		size <<= 10
+	case "MB":
+		size <<= 20
+	case "GB":
+		size <<= 30
+	case "TB":
+		size <<= 40
+	case "":
+	default:
+		return 0, fmt.Errorf("unknown size suffix %q", suffix)
+	}
+	return size, nil
+}
+
+func run(ctx context.Context, model, flowName, inputFile, workdir string, cacheSize uint64) error {
 	flow := aflow.Flows[flowName]
 	if flow == nil {
 		return fmt.Errorf("workflow %q is not found", flowName)
@@ -69,7 +96,7 @@ func run(ctx context.Context, model, flowName, inputFile, workdir string) error
 	if err := json.Unmarshal(inputData, &inputs); err != nil {
 		return err
 	}
-	cache, err := aflow.NewCache(filepath.Join(workdir, "cache"), 0)
+	cache, err := aflow.NewCache(filepath.Join(workdir, "cache"), cacheSize)
 	if err != nil {
 		return err
 	}