go-otel/internal/topology/watcher.go

package topology

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	networkingv1 "k8s.io/api/networking/v1"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"

	kube "git.tornberg.me/go-otel/internal/kubernetes"
)

// Registry serves the latest Snapshot. A nil *Registry resolves nothing, so
// callers need no cluster to run.
type Registry struct {
	snap   atomic.Pointer[Snapshot]
	dirty  atomic.Bool
	built  atomic.Int64 // unix seconds of the last rebuild
	client kubernetes.Interface // for on-demand pod lookups; nil outside a cluster
}

// Resolve returns the backends serving host+path.
func (r *Registry) Resolve(host, path string) []Target {
	if r == nil {
		return nil
	}
	s := r.snap.Load()
	if s == nil {
		return nil
	}
	return s.Resolve(host, path)
}

// Snapshot returns the current routing table.
func (r *Registry) Snapshot() Snapshot {
	if r == nil {
		return Snapshot{}
	}
	if s := r.snap.Load(); s != nil {
		return *s
	}
	return Snapshot{}
}

// NewFromEnv watches ingresses, services and deployments, following the same
// switches as the rollout watcher (OTEL_KUBERNETES_ENABLED / _NAMESPACE). It
// returns nil, nil when there is no cluster to watch.
func NewFromEnv(ctx context.Context) (*Registry, error) {
	if strings.EqualFold(os.Getenv("OTEL_KUBERNETES_ENABLED"), "false") {
		return nil, nil
	}
	cfg, err := kube.RESTConfig()
	if err != nil {
		if os.Getenv("OTEL_KUBERNETES_ENABLED") == "true" {
			return nil, err
		}
		return nil, nil
	}
	client, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return nil, fmt.Errorf("create kubernetes client: %w", err)
	}
	var opts []informers.SharedInformerOption
	if ns := os.Getenv("OTEL_KUBERNETES_NAMESPACE"); ns != "" {
		opts = append(opts, informers.WithNamespace(ns))
	}
	factory := informers.NewSharedInformerFactoryWithOptions(client, 30*time.Minute, opts...)
	ing := factory.Networking().V1().Ingresses().Informer()
	svc := factory.Core().V1().Services().Informer()
	dep := factory.Apps().V1().Deployments().Informer()

	r := &Registry{client: client}
	mark := cache.ResourceEventHandlerFuncs{
		AddFunc:    func(any) { r.dirty.Store(true) },
		UpdateFunc: func(_, _ any) { r.dirty.Store(true) },
		DeleteFunc: func(any) { r.dirty.Store(true) },
	}
	for _, i := range []cache.SharedIndexInformer{ing, svc, dep} {
		_, _ = i.AddEventHandler(mark)
	}
	rebuild := func() {
		var is []networkingv1.Ingress
		for _, o := range ing.GetStore().List() {
			if v, ok := o.(*networkingv1.Ingress); ok {
				is = append(is, *v)
			}
		}
		var ss []corev1.Service
		for _, o := range svc.GetStore().List() {
			if v, ok := o.(*corev1.Service); ok {
				ss = append(ss, *v)
			}
		}
		var ds []appsv1.Deployment
		for _, o := range dep.GetStore().List() {
			if v, ok := o.(*appsv1.Deployment); ok {
				ds = append(ds, *v)
			}
		}
		s := Build(is, ss, ds)
		r.snap.Store(&s)
		r.built.Store(time.Now().Unix())
	}
	go func() {
		factory.Start(ctx.Done())
		if !cache.WaitForCacheSync(ctx.Done(), ing.HasSynced, svc.HasSynced, dep.HasSynced) {
			log.Printf("topology: informers did not sync")
			return
		}
		rebuild()
		log.Printf("topology: %d ingress routes", len(r.Snapshot().Targets))
		// Coalesce bursts of changes (a rollout touches many objects).
		t := time.NewTicker(5 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				if r.dirty.Swap(false) {
					rebuild()
				}
			}
		}
	}()
	return r, nil
}

// HandleTopology serves GET /api/topology. With ?host=&path= it answers
// "what sits behind this request"; without them it lists every route.
func (r *Registry) HandleTopology(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	q := req.URL.Query()
	if host := q.Get("host"); host != "" {
		_ = json.NewEncoder(w).Encode(map[string]any{"targets": r.Resolve(host, q.Get("path"))})
		return
	}
	_ = json.NewEncoder(w).Encode(map[string]any{"targets": r.Snapshot().Targets, "updated": r.built.Load()})
}