Untitled

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sort"
	"sync"
	"syscall"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	cache "k8s.io/client-go/tools/cache"

	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

type rec struct {
	created  time.Time
	sched    *time.Time
	ready    *time.Time
	reported bool
}

type ms struct{ Pending, Startup, Total int64 }

type agg struct {
	mu    sync.Mutex
	list  []ms
	sumP  int64
	sumS  int64
	sumT  int64
	minP  int64
	minS  int64
	minT  int64
	maxP  int64
	maxS  int64
	maxT  int64
	count int64
}

func newAgg() *agg {
	const big = int64(^uint64(0) >> 1) // max int64
	return &agg{minP: big, minS: big, minT: big}
}

func (a *agg) add(m ms) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.list = append(a.list, m)
	a.count++
	a.sumP += m.Pending
	a.sumS += m.Startup
	a.sumT += m.Total
	if m.Pending < a.minP {
		a.minP = m.Pending
	}
	if m.Startup < a.minS {
		a.minS = m.Startup
	}
	if m.Total < a.minT {
		a.minT = m.Total
	}
	if m.Pending > a.maxP {
		a.maxP = m.Pending
	}
	if m.Startup > a.maxS {
		a.maxS = m.Startup
	}
	if m.Total > a.maxT {
		a.maxT = m.Total
	}
}

func (a *agg) avg(x int64) int64 {
	if a.count == 0 {
		return 0
	}
	return x / a.count
}

func (a *agg) done(expect int) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	return expect > 0 && int(a.count) >= expect
}

func mustConfig() *rest.Config {
	// Try in-cluster, then kubeconfig
	if cfg, err := rest.InClusterConfig(); err == nil {
		return cfg
	}
	kc := os.Getenv("KUBECONFIG")
	if kc == "" {
		home, _ := os.UserHomeDir()
		kc = home + "/.kube/config"
	}
	cfg, err := clientcmd.BuildConfigFromFlags("", kc)
	if err != nil {
		panic(err)
	}
	return cfg
}

func main() {
	var (
		ns       = flag.String("namespace", "", "Namespace to watch (empty = all)")
		sel      = flag.String("selector", "app=kb-latency", "Label selector for pods to measure")
		expect   = flag.Int("expect", 0, "Stop after N pods measured (0 = never)")
		timeout  = flag.Duration("timeout", 0, "Hard stop after this duration (0 = no timeout)")
		showEach = flag.Bool("per-pod", true, "Print per-pod measurements")
	)
	flag.Parse()

	cfg := mustConfig()
	cs, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	if *timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, *timeout)
	}
	defer cancel()

	// Tweaked informer (namespace + label selector)
	opts := []informers.SharedInformerOption{}
	if *ns != "" {
		opts = append(opts, informers.WithNamespace(*ns))
	}
	selector, err := labels.Parse(*sel)
	if err != nil {
		panic(fmt.Errorf("bad selector: %w", err))
	}
	opts = append(opts, informers.WithTweakListOptions(func(lo *metav1.ListOptions) {
		lo.LabelSelector = selector.String()
	}))

	f := informers.NewSharedInformerFactoryWithOptions(cs, 0, opts...)
	podInf := f.Core().V1().Pods().Informer()

	var (
		lock   sync.Mutex
		byKey  = map[string]*rec{} // ns/name -> rec
		stats  = newAgg()
		sigCh  = make(chan os.Signal, 1)
	)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	handle := func(pod *corev1.Pod) {
		key := pod.Namespace + "/" + pod.Name

		lock.Lock()
		r, ok := byKey[key]
		if !ok {
			r = &rec{created: pod.CreationTimestamp.Time}
			byKey[key] = r
		}
		// Update created if zero (shouldn't happen) or earlier
		if pod.CreationTimestamp.Time.Before(r.created) {
			r.created = pod.CreationTimestamp.Time
		}

		// PodScheduled condition
		for _, c := range pod.Status.Conditions {
			if c.Type == corev1.PodScheduled && c.Status == corev1.ConditionTrue {
				t := c.LastTransitionTime.Time
				r.sched = &t
			}
			if c.Type == corev1.PodReady && c.Status == corev1.ConditionTrue {
				t := c.LastTransitionTime.Time
				r.ready = &t
			}
		}

		// If we have all timestamps and not reported → compute
		if !r.reported && r.sched != nil && r.ready != nil {
			pending := r.sched.Sub(r.created).Milliseconds()
			startup := r.ready.Sub(*r.sched).Milliseconds()
			total := r.ready.Sub(r.created).Milliseconds()
			stats.add(ms{Pending: pending, Startup: startup, Total: total})
			r.reported = true
			if *showEach {
				fmt.Printf("%s pending_ms=%d startup_ms=%d total_ms=%d (node=%s)\n",
					key, pending, startup, total, pod.Spec.NodeName)
			}
		}
		lock.Unlock()
	}

	podInf.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			if pod, ok := obj.(*corev1.Pod); ok {
				handle(pod)
			}
		},
		UpdateFunc: func(_, newObj interface{}) {
			if pod, ok := newObj.(*corev1.Pod); ok {
				handle(pod)
			}
		},
	})

	// Start informers
	f.Start(ctx.Done())
	f.WaitForCacheSync(ctx.Done())

	// Wait until: expect reached OR signal OR timeout
	tick := time.NewTicker(500 * time.Millisecond)
	defer tick.Stop()

loop:
	for {
		select {
		case <-ctx.Done():
			break loop
		case <-sigCh:
			break loop
		case <-tick.C:
			if stats.done(*expect) && *expect > 0 {
				break loop
			}
		}
	}

	// Final report
	stats.mu.Lock()
	defer stats.mu.Unlock()

	if stats.count == 0 {
		fmt.Println("no measurements collected (check selector/namespace or if pods reached Ready)")
		return
	}

	// Optional percentiles (nice to have)
	getP := func(values []int64, p float64) int64 {
		if len(values) == 0 {
			return 0
		}
		idx := int(float64(len(values)-1) * p)
		return values[idx]
	}
	var pendings, startups, totals []int64
	for _, m := range stats.list {
		pendings = append(pendings, m.Pending)
		startups = append(startups, m.Startup)
		totals = append(totals, m.Total)
	}
	sort.Slice(pendings, func(i, j int) bool { return pendings[i] < pendings[j] })
	sort.Slice(startups, func(i, j int) bool { return startups[i] < startups[j] })
	sort.Slice(totals, func(i, j int) bool { return totals[i] < totals[j] })

	fmt.Printf("\n--- SUMMARY (pods=%d) ---\n", stats.count)
	fmt.Printf("Pending ms: min=%d p50=%d p95=%d max=%d avg=%d\n",
		stats.minP, getP(pendings, 0.50), getP(pendings, 0.95), stats.maxP, stats.avg(stats.sumP))
	fmt.Printf("Startup ms: min=%d p50=%d p95=%d max=%d avg=%d\n",
		stats.minS, getP(startups, 0.50), getP(startups, 0.95), stats.maxS, stats.avg(stats.sumS))
	fmt.Printf("Total   ms: min=%d p50=%d p95=%d max=%d avg=%d\n",
		stats.minT, getP(totals, 0.50), getP(totals, 0.95), stats.maxT, stats.avg(stats.sumT))
}