fastmirror/mirror.go

package fastmirror

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

	progressbar "github.com/schollz/progressbar/v3"
)

func findFastestMirror(suites []string) (string, time.Duration, []time.Duration, error) {
	log.Println("Fetching mirrors list")
	resp, err := http.Get("http://mirrors.ubuntu.com/mirrors.txt")
	if err != nil {
		return "", 0, nil, fmt.Errorf("failed to fetch mirrors list: %v", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", 0, nil, fmt.Errorf("failed to read mirrors list: %v", err)
	}

	mirrors := strings.Split(string(body), "\n")
	log.Printf("Found %d mirrors", len(mirrors))

	var fastestMirrorURL string
	var fastestTime time.Duration
	var latencies []time.Duration
	downCount := 0
	noIndexCount := 0

	if err != nil {
		return "", 0, nil, err
	}

	httpClient := http.Client{
		Timeout: 1 * time.Second,
	}

	log.Println("Testing mirrors for latency")
	bar := progressbar.Default(int64(len(mirrors)))

	for _, mirror := range mirrors {
		if err := bar.Add(1); err != nil {
			log.Printf("Error updating progress bar: %v", err)
		}
		if strings.HasPrefix(mirror, "https://") {
			mirror = strings.TrimSuffix(mirror, "/")
			startTime := time.Now()
			isValid := isValidMirror(httpClient, mirror, suites)
			elapsedTime := time.Since(startTime)

			if isValid {
				latencies = append(latencies, elapsedTime)
				if fastestMirrorURL == "" || elapsedTime < fastestTime {
					fastestMirrorURL = mirror
					fastestTime = elapsedTime
				}
			} else {
				noIndexCount++
			}
		} else {
			downCount++
		}
	}

	if fastestMirrorURL == "" {
		return "", 0, nil, fmt.Errorf("no suitable HTTPS mirror found")
	}

	displayLatencyStatistics(latencies, fastestMirrorURL, fastestTime)
	log.Printf("Number of mirrors that were down: %d", downCount)
	log.Printf("Number of mirrors without required package index: %d", noIndexCount)

	return fastestMirrorURL, fastestTime, latencies, nil
}

func isValidMirror(httpClient http.Client, mirrorURL string, suites []string) bool {
	for _, suite := range suites {
		uri := fmt.Sprintf("%s/dists/%s/Release", mirrorURL, suite)
		resp, err := httpClient.Get(uri)
		if err != nil || resp.StatusCode != http.StatusOK {
			return false
		}
		resp.Body.Close()
	}
	return true
}

func displayLatencyStatistics(latencies []time.Duration, fastestMirrorURL string, fastestTime time.Duration) {
	if len(latencies) == 0 {
		log.Println("No latencies recorded")
		return
	}

	var totalLatency time.Duration
	var minLatency time.Duration = latencies[0]
	var maxLatency time.Duration = latencies[0]

	for _, latency := range latencies {
		totalLatency += latency
		if latency < minLatency {
			minLatency = latency
		}
		if latency > maxLatency {
			maxLatency = latency
		}
	}

	averageLatency := totalLatency / time.Duration(len(latencies))

	log.Printf("Minimum latency: %s", minLatency)
	log.Printf("Maximum latency: %s", maxLatency)
	log.Printf("Average latency: %s", averageLatency)
	log.Printf("Chosen fastest mirror: %s with latency: %s", fastestMirrorURL, fastestTime)
}