From 487b9c520d4d9bb77531e78efc24af62d36f8af2 Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Sun, 28 May 2023 20:37:05 +0000
Subject: [PATCH] feat: add limit amount of kernels to fetch

---
 .github/workflows/debian-cache.yml |  2 +-
 distro.go                          | 14 ++++++++++++--
 distro/debian/kernel.go            | 21 +++++++++++++++++----
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/debian-cache.yml b/.github/workflows/debian-cache.yml
index 0582ec9..e639443 100644
--- a/.github/workflows/debian-cache.yml
+++ b/.github/workflows/debian-cache.yml
@@ -25,7 +25,7 @@ jobs:
       run: go build
 
     - name: Cache
-      run: ./out-of-tree --log-level=trace distro debian cache --refetch=0
+      run: ./out-of-tree --log-level=trace distro debian cache --refetch=0 --limit=64
 
     - name: Install s3cmd
       run: sudo apt install s3cmd
diff --git a/distro.go b/distro.go
index 12bf47c..09ee123 100644
--- a/distro.go
+++ b/distro.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -30,6 +31,7 @@ type DebianCmd struct {
 	Cache DebianCacheCmd `cmd:"" help:"populate cache"`
 	Fetch DebianFetchCmd `cmd:"" help:"download deb packages"`
 
+	Limit int    `help:"limit amount of kernels to fetch"`
 	Regex string `help:"match deb pkg names by regex" default:".*"`
 }
 
@@ -47,7 +49,11 @@ func (cmd *DebianCacheCmd) Run(dcmd *DebianCmd) (err error) {
 
 	log.Info().Msg("Fetching kernels...")
 
-	kernels, err := debian.GetKernels()
+	if dcmd.Limit == 0 {
+		dcmd.Limit = math.MaxInt32
+	}
+
+	kernels, err := debian.GetKernelsWithLimit(dcmd.Limit)
 	if err != nil {
 		log.Error().Err(err).Msg("")
 		return
@@ -152,7 +158,11 @@ func (cmd *DebianFetchCmd) Run(dcmd *DebianCmd) (err error) {
 	log.Info().Msg("will not download packages that exist on the mirror")
 	log.Info().Msg("use --ignore-mirror if you really need it")
 
-	kernels, err := debian.GetKernels()
+	if dcmd.Limit == 0 {
+		dcmd.Limit = math.MaxInt32
+	}
+
+	kernels, err := debian.GetKernelsWithLimit(dcmd.Limit)
 	if err != nil {
 		log.Error().Err(err).Msg("")
 		return
diff --git a/distro/debian/kernel.go b/distro/debian/kernel.go
index 4f82c9e..2d10d11 100644
--- a/distro/debian/kernel.go
+++ b/distro/debian/kernel.go
@@ -2,6 +2,7 @@ package debian
 
 import (
 	"errors"
+	"math"
 	"sort"
 	"strings"
 	"time"
@@ -260,7 +261,7 @@ func findKbuild(versions []string, kpkgver string) (
 }
 
 func getKernelsByVersion(slog zerolog.Logger, c *Cache, toolsVersions []string,
-	version string) (kernels []DebianKernel) {
+	version string) (kernels []DebianKernel, fromcache bool) {
 
 	var dk DebianKernel
 	dks, err := c.Get(version)
@@ -269,6 +270,7 @@ func getKernelsByVersion(slog zerolog.Logger, c *Cache, toolsVersions []string,
 		if !dk.Internal.Invalid {
 			slog.Trace().Msgf("found in cache")
 			kernels = append(kernels, dk)
+			fromcache = true
 			return
 		}
 	}
@@ -336,7 +338,9 @@ var (
 	RefetchDays int = 14
 )
 
-func GetKernels() (kernels []DebianKernel, err error) {
+// GetKernelsWithLimit is workaround for testing and building the
+// first cache, which is heavily rate limited by snapshot.debian.org
+func GetKernelsWithLimit(limit int) (kernels []DebianKernel, err error) {
 	if CachePath == "" {
 		CachePath = config.File("debian.cache")
 		log.Debug().Msgf("Use default kernels cache path: %s", CachePath)
@@ -381,10 +385,19 @@ func GetKernels() (kernels []DebianKernel, err error) {
 	for i, version := range versions {
 		slog := log.With().Str("version", version).Logger()
 		slog.Trace().Msgf("%03d/%03d", i, len(versions))
-		vkernels := getKernelsByVersion(slog, c, toolsVersions, version)
+		vkernels, fromcache := getKernelsByVersion(slog, c, toolsVersions, version)
 		kernels = append(kernels, vkernels...)
-
+		if !fromcache {
+			limit--
+		}
+		if limit <= 0 {
+			return
+		}
 	}
 
 	return
 }
+
+func GetKernels() (kernels []DebianKernel, err error) {
+	return GetKernelsWithLimit(math.MaxInt32)
+}