flock/pkg/agent/podinfo.go

package agent

import (
	"context"
	"fmt"
	"log/slog"
	"time"

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

// PodCache exposes a Get(ns, name) lookup against a node-scoped Pod
// informer. ADD/DEL handlers consult it to read annotations + labels for
// IPAM and (later) NetworkPolicy.
type PodCache struct {
	lister cache.GenericLister
	logger *slog.Logger
	store  cache.Store
}

// StartPodInformer launches a Pod informer filtered to spec.nodeName ==
// node. Returns a PodCache once the cache is synced. Blocks on initial
// list/watch sync.
func StartPodInformer(ctx context.Context, cfg *rest.Config, node string, logger *slog.Logger) (*PodCache, error) {
	cs, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		return nil, fmt.Errorf("kubernetes client: %w", err)
	}
	tweak := func(opts *metav1.ListOptions) {
		opts.FieldSelector = fields.OneTermEqualSelector("spec.nodeName", node).String()
	}
	factory := informers.NewSharedInformerFactoryWithOptions(cs, 10*time.Minute,
		informers.WithTweakListOptions(tweak))
	inf := factory.Core().V1().Pods().Informer()

	logger.Info("Pod informer starting", "node", node, "field_selector", "spec.nodeName="+node)
	factory.Start(ctx.Done())
	if !cache.WaitForCacheSync(ctx.Done(), inf.HasSynced) {
		return nil, fmt.Errorf("pod informer cache failed to sync")
	}
	logger.Info("Pod informer synced", "node", node, "items", len(inf.GetStore().ListKeys()))
	return &PodCache{store: inf.GetStore(), logger: logger}, nil
}

// Get looks up a Pod by namespace and name. Returns (nil, false) if absent.
func (c *PodCache) Get(namespace, name string) (*corev1.Pod, bool) {
	key := namespace + "/" + name
	obj, ok, err := c.store.GetByKey(key)
	if err != nil || !ok || obj == nil {
		return nil, false
	}
	pod, ok := obj.(*corev1.Pod)
	if !ok {
		return nil, false
	}
	return pod, true
}

// WaitForPod polls the cache for up to `timeout` for a pod to appear.
// kubelet may invoke CNI ADD slightly before the informer has observed the
// PodSpec, so this helper smooths the race.
func (c *PodCache) WaitForPod(ctx context.Context, namespace, name string, timeout time.Duration) (*corev1.Pod, error) {
	deadline := time.Now().Add(timeout)
	for {
		if pod, ok := c.Get(namespace, name); ok {
			return pod, nil
		}
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("pod %s/%s not found in informer cache after %s", namespace, name, timeout)
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(50 * time.Millisecond):
		}
	}
}