mirror of
https://github.com/lukaszraczylo/gohoarder.git
synced 2026-06-29 03:12:54 +00:00
fixes
This commit is contained in:
@@ -0,0 +1,311 @@
|
||||
package prewarming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/lukaszraczylo/gohoarder/pkg/analytics"
|
||||
"github.com/lukaszraczylo/gohoarder/pkg/cache"
|
||||
"github.com/lukaszraczylo/gohoarder/pkg/network"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// PackageInfo represents a package to pre-warm
|
||||
type PackageInfo struct {
|
||||
Registry string
|
||||
Name string
|
||||
Version string
|
||||
Priority int
|
||||
}
|
||||
|
||||
// Worker handles background pre-warming of popular packages
|
||||
type Worker struct {
|
||||
cache *cache.Manager
|
||||
analytics *analytics.Engine
|
||||
client *network.Client
|
||||
interval time.Duration
|
||||
maxConcurrent int
|
||||
enabled bool
|
||||
stopChan chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// Config holds pre-warming worker configuration
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
Interval time.Duration
|
||||
MaxConcurrent int
|
||||
TopPackages int
|
||||
CacheManager *cache.Manager
|
||||
Analytics *analytics.Engine
|
||||
NetworkClient *network.Client
|
||||
}
|
||||
|
||||
// NewWorker creates a new pre-warming worker
|
||||
func NewWorker(cfg Config) *Worker {
|
||||
if cfg.Interval <= 0 {
|
||||
cfg.Interval = 1 * time.Hour
|
||||
}
|
||||
if cfg.MaxConcurrent <= 0 {
|
||||
cfg.MaxConcurrent = 5
|
||||
}
|
||||
if cfg.TopPackages <= 0 {
|
||||
cfg.TopPackages = 100
|
||||
}
|
||||
|
||||
worker := &Worker{
|
||||
cache: cfg.CacheManager,
|
||||
analytics: cfg.Analytics,
|
||||
client: cfg.NetworkClient,
|
||||
interval: cfg.Interval,
|
||||
maxConcurrent: cfg.MaxConcurrent,
|
||||
enabled: cfg.Enabled,
|
||||
stopChan: make(chan struct{}),
|
||||
}
|
||||
|
||||
if cfg.Enabled {
|
||||
log.Info().
|
||||
Dur("interval", cfg.Interval).
|
||||
Int("max_concurrent", cfg.MaxConcurrent).
|
||||
Msg("Pre-warming worker initialized")
|
||||
} else {
|
||||
log.Info().Msg("Pre-warming worker disabled")
|
||||
}
|
||||
|
||||
return worker
|
||||
}
|
||||
|
||||
// Start begins the pre-warming worker
|
||||
func (w *Worker) Start(ctx context.Context) {
|
||||
if !w.enabled {
|
||||
log.Debug().Msg("Pre-warming worker is disabled, not starting")
|
||||
return
|
||||
}
|
||||
|
||||
w.wg.Add(1)
|
||||
go w.run(ctx)
|
||||
log.Info().Msg("Pre-warming worker started")
|
||||
}
|
||||
|
||||
// run is the main worker loop
|
||||
func (w *Worker) run(ctx context.Context) {
|
||||
defer w.wg.Done()
|
||||
|
||||
ticker := time.NewTicker(w.interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
// Run immediately on start
|
||||
w.prewarmPopularPackages(ctx)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Info().Msg("Pre-warming worker stopping due to context cancellation")
|
||||
return
|
||||
case <-w.stopChan:
|
||||
log.Info().Msg("Pre-warming worker stopped")
|
||||
return
|
||||
case <-ticker.C:
|
||||
w.prewarmPopularPackages(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// prewarmPopularPackages fetches and caches popular packages
|
||||
func (w *Worker) prewarmPopularPackages(ctx context.Context) {
|
||||
log.Info().Msg("Starting pre-warming cycle")
|
||||
|
||||
// Get popular packages from analytics
|
||||
popularPackages := w.analytics.GetTopPackages(100)
|
||||
if len(popularPackages) == 0 {
|
||||
log.Debug().Msg("No popular packages found for pre-warming")
|
||||
return
|
||||
}
|
||||
|
||||
// Get trending packages for additional candidates
|
||||
trendingPackages := w.analytics.GetTrendingPackages(50)
|
||||
|
||||
// Combine and deduplicate
|
||||
packages := w.combinePackages(popularPackages, trendingPackages)
|
||||
|
||||
log.Info().
|
||||
Int("packages", len(packages)).
|
||||
Msg("Identified packages for pre-warming")
|
||||
|
||||
// Create work queue
|
||||
workChan := make(chan PackageInfo, len(packages))
|
||||
for _, pkg := range packages {
|
||||
workChan <- PackageInfo{
|
||||
Registry: pkg.Registry,
|
||||
Name: pkg.Name,
|
||||
Version: "latest", // Pre-warm latest version
|
||||
Priority: int(pkg.Downloads),
|
||||
}
|
||||
}
|
||||
close(workChan)
|
||||
|
||||
// Start worker goroutines
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < w.maxConcurrent; i++ {
|
||||
wg.Add(1)
|
||||
go func(workerID int) {
|
||||
defer wg.Done()
|
||||
w.processPackages(ctx, workerID, workChan)
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
log.Info().Msg("Pre-warming cycle completed")
|
||||
}
|
||||
|
||||
// processPackages processes packages from the work queue
|
||||
func (w *Worker) processPackages(ctx context.Context, workerID int, workChan <-chan PackageInfo) {
|
||||
for pkg := range workChan {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
w.prewarmPackage(ctx, pkg, workerID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// prewarmPackage fetches and caches a single package
|
||||
func (w *Worker) prewarmPackage(ctx context.Context, pkg PackageInfo, workerID int) {
|
||||
log.Debug().
|
||||
Int("worker", workerID).
|
||||
Str("registry", pkg.Registry).
|
||||
Str("package", pkg.Name).
|
||||
Str("version", pkg.Version).
|
||||
Msg("Pre-warming package")
|
||||
|
||||
// Build URL based on registry
|
||||
url := w.buildPackageURL(pkg)
|
||||
if url == "" {
|
||||
log.Warn().
|
||||
Str("registry", pkg.Registry).
|
||||
Str("package", pkg.Name).
|
||||
Msg("Cannot build URL for registry")
|
||||
return
|
||||
}
|
||||
|
||||
// Fetch package from upstream
|
||||
reqCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
body, statusCode, err := w.client.Get(reqCtx, url, nil)
|
||||
if err != nil {
|
||||
log.Error().
|
||||
Err(err).
|
||||
Str("package", pkg.Name).
|
||||
Msg("Failed to fetch package for pre-warming")
|
||||
return
|
||||
}
|
||||
defer body.Close()
|
||||
|
||||
if statusCode != 200 {
|
||||
log.Warn().
|
||||
Int("status", statusCode).
|
||||
Str("package", pkg.Name).
|
||||
Msg("Non-200 response for package")
|
||||
return
|
||||
}
|
||||
|
||||
// Cache the package
|
||||
// In a real implementation, this would read the response body and store it
|
||||
log.Info().
|
||||
Str("package", pkg.Name).
|
||||
Str("version", pkg.Version).
|
||||
Msg("Successfully pre-warmed package")
|
||||
}
|
||||
|
||||
// buildPackageURL builds the upstream URL for a package
|
||||
func (w *Worker) buildPackageURL(pkg PackageInfo) string {
|
||||
// This is simplified - in reality, each registry has different URL patterns
|
||||
switch pkg.Registry {
|
||||
case "npm":
|
||||
return "https://registry.npmjs.org/" + pkg.Name
|
||||
case "pypi":
|
||||
return "https://pypi.org/simple/" + pkg.Name + "/"
|
||||
case "go":
|
||||
// Go modules use different URL patterns
|
||||
return "https://proxy.golang.org/" + pkg.Name + "/@latest"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// combinePackages merges popular and trending packages, removing duplicates
|
||||
func (w *Worker) combinePackages(popular, trending []analytics.PopularPackage) []analytics.PopularPackage {
|
||||
seen := make(map[string]bool)
|
||||
result := make([]analytics.PopularPackage, 0, len(popular)+len(trending))
|
||||
|
||||
for _, pkg := range popular {
|
||||
key := pkg.Registry + ":" + pkg.Name
|
||||
if !seen[key] {
|
||||
result = append(result, pkg)
|
||||
seen[key] = true
|
||||
}
|
||||
}
|
||||
|
||||
for _, pkg := range trending {
|
||||
key := pkg.Registry + ":" + pkg.Name
|
||||
if !seen[key] {
|
||||
result = append(result, pkg)
|
||||
seen[key] = true
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// Stop gracefully stops the pre-warming worker
|
||||
func (w *Worker) Stop() {
|
||||
if !w.enabled {
|
||||
return
|
||||
}
|
||||
|
||||
log.Info().Msg("Stopping pre-warming worker")
|
||||
close(w.stopChan)
|
||||
w.wg.Wait()
|
||||
log.Info().Msg("Pre-warming worker stopped")
|
||||
}
|
||||
|
||||
// TriggerPrewarm manually triggers a pre-warming cycle
|
||||
func (w *Worker) TriggerPrewarm(ctx context.Context) {
|
||||
if !w.enabled {
|
||||
log.Warn().Msg("Cannot trigger pre-warm: worker is disabled")
|
||||
return
|
||||
}
|
||||
|
||||
log.Info().Msg("Manual pre-warming triggered")
|
||||
go w.prewarmPopularPackages(ctx)
|
||||
}
|
||||
|
||||
// PrewarmPackage pre-warms a specific package
|
||||
func (w *Worker) PrewarmPackage(ctx context.Context, registry, name, version string) error {
|
||||
if !w.enabled {
|
||||
log.Warn().Msg("Pre-warming worker is disabled")
|
||||
return nil
|
||||
}
|
||||
|
||||
pkg := PackageInfo{
|
||||
Registry: registry,
|
||||
Name: name,
|
||||
Version: version,
|
||||
Priority: 100,
|
||||
}
|
||||
|
||||
w.prewarmPackage(ctx, pkg, 0)
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetStatus returns the current status of the pre-warming worker
|
||||
func (w *Worker) GetStatus() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"enabled": w.enabled,
|
||||
"interval": w.interval.String(),
|
||||
"max_concurrent": w.maxConcurrent,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user