go-cart-actor/pkg/discovery/discovery.go

package discovery

import (
	"context"
	"log"
	"os"
	"slices"
	"strings"
	"sync"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	toolsWatch "k8s.io/client-go/tools/watch"
)

// InClusterNamespace returns the namespace the current pod runs in: POD_NAMESPACE
// (downward API) if set, else the service-account namespace file mounted into
// every pod. Empty means "not in a cluster" — callers can treat that as
// all-namespaces or single-node. Pass the result to NewK8sDiscoveryInNamespace
// to keep pod discovery (and thus the required RBAC) scoped to one namespace.
func InClusterNamespace() string {
	if ns := os.Getenv("POD_NAMESPACE"); ns != "" {
		return ns
	}
	if b, err := os.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/namespace"); err == nil {
		return strings.TrimSpace(string(b))
	}
	return ""
}

type K8sDiscovery struct {
	ctx         context.Context
	client      *kubernetes.Clientset
	listOptions metav1.ListOptions
	// namespace scopes the pod list/watch. Empty means all namespaces (requires
	// a cluster-scoped RBAC role); a specific namespace only needs a namespaced
	// Role, which is the least-privilege option.
	namespace string
}

func (k *K8sDiscovery) Discover() ([]string, error) {
	return k.DiscoverInNamespace(k.namespace)
}
func (k *K8sDiscovery) DiscoverInNamespace(namespace string) ([]string, error) {
	pods, err := k.client.CoreV1().Pods(namespace).List(k.ctx, k.listOptions)
	if err != nil {
		return nil, err
	}
	hosts := make([]string, 0, len(pods.Items))
	for _, pod := range pods.Items {
		if hasReadyCondition(&pod) {
			hosts = append(hosts, pod.Status.PodIP)
		}
	}
	return hosts, nil
}

func hasReadyCondition(pod *v1.Pod) bool {
	return slices.ContainsFunc(pod.Status.Conditions, func(condition v1.PodCondition) bool {
		return condition.Type == v1.PodReady && condition.Status == v1.ConditionTrue
	})
}

func (k *K8sDiscovery) Watch() (<-chan HostChange, error) {
	ipsThatAreReady := make(map[string]bool)
	m := sync.Mutex{}
	watcherFn := func(options metav1.ListOptions) (watch.Interface, error) {
		return k.client.CoreV1().Pods(k.namespace).Watch(k.ctx, k.listOptions)
	}
	watcher, err := toolsWatch.NewRetryWatcherWithContext(k.ctx, "1", &cache.ListWatch{WatchFunc: watcherFn})
	if err != nil {
		return nil, err
	}
	ch := make(chan HostChange)
	go func() {
		for event := range watcher.ResultChan() {

			pod := event.Object.(*v1.Pod)
			isReady := hasReadyCondition(pod)
			m.Lock()
			oldState := ipsThatAreReady[pod.Status.PodIP]
			ipsThatAreReady[pod.Status.PodIP] = isReady
			m.Unlock()
			if oldState != isReady {
				ch <- HostChange{
					Host:    pod.Status.PodIP,
					IsReady: isReady,
				}
			}
			ch <- HostChange{
				Host:    pod.Status.PodIP,
				IsReady: isReady,
			}
		}
	}()
	return ch, nil
}

// NewK8sDiscovery watches pods across ALL namespaces (label-filtered via
// listOptions). It requires a cluster-scoped RBAC role for pod list/watch.
func NewK8sDiscovery(client *kubernetes.Clientset, listOptions metav1.ListOptions) *K8sDiscovery {
	return &K8sDiscovery{
		ctx:         context.Background(),
		client:      client,
		listOptions: listOptions,
	}
}

// NewK8sDiscoveryInNamespace watches pods only within namespace, so it needs
// just a namespaced Role (pods: get/list/watch) rather than a ClusterRole.
func NewK8sDiscoveryInNamespace(client *kubernetes.Clientset, namespace string, listOptions metav1.ListOptions) *K8sDiscovery {
	return &K8sDiscovery{
		ctx:         context.Background(),
		client:      client,
		listOptions: listOptions,
		namespace:   namespace,
	}
}

// StartK8sDiscovery initializes kubernetes discovery for target.
// If podIp is empty, it returns early.
func StartK8sDiscovery(podIp string, labelSelector string, timeoutSeconds int64, target DiscoveryTarget) {
	if podIp == "" {
		log.Print("No discovery service available")
		return
	}

	config, kerr := rest.InClusterConfig()
	if kerr != nil {
		log.Fatalf("Error creating kubernetes client: %v\n", kerr)
	}
	client, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatalf("Error creating client: %v\n", err)
	}

	hw := NewK8sDiscoveryInNamespace(client, InClusterNamespace(), metav1.ListOptions{
		LabelSelector:  labelSelector,
		TimeoutSeconds: &timeoutSeconds,
	})

	go func() {
		ch, err := hw.Watch()
		if err != nil {
			log.Printf("Discovery error: %v", err)
			return
		}
		for evt := range ch {
			if evt.Host == "" {
				continue
			}
			switch evt.IsReady {
			case false:
				if target.IsKnown(evt.Host) {
					log.Printf("Host %s is not ready, removing", evt.Host)
					target.RemoveHost(evt.Host)
				}
			default:
				if !target.IsKnown(evt.Host) {
					log.Printf("Discovered host %s", evt.Host)
					target.AddRemoteHost(evt.Host)
				}
			}
		}
	}()
}