kubectl-persistent-logger/logs/watcher.go

package logs

import (
	"context"
	"errors"
	"io"
	"log"
	"sync"
	"time"

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
)

// KubernetesClient provides both typed and untyped interfaces to the
// Kubernetes API.
type KubernetesClient struct {
	Typed   kubernetes.Interface
	Untyped dynamic.Interface
}

// concurrentWriter implements an io.Writer that can be safely written to from
// multiple goroutines.
type concurrentWriter struct {
	w  io.Writer
	mu sync.Mutex
}

// Write implements io.Writer.
func (cw *concurrentWriter) Write(p []byte) (int, error) {
	cw.mu.Lock()
	defer cw.mu.Unlock()

	return cw.w.Write(p)
}

type PodWatcherInterface interface {
	WatchPods(ctx context.Context) error
	Close()
}

// PodWatcherFunc builds a PodWatcher.
type PodWatcherFunc func(KubernetesClient, string, labels.Selector, io.Writer) PodWatcherInterface

// WatcherParams defines the input parameters of a Watcher.
type WatcherParams struct {
	Name        string
	Type        string
	Namespace   string
	Container   string
	StrictExist bool
}

// Watcher watches a deployment and tails the logs for its currently active
// pods.
type Watcher struct {
	params         WatcherParams
	client         KubernetesClient
	resourceUID    types.UID
	podSelector    labels.Selector
	podWatcher     PodWatcherInterface
	podWatcherFunc PodWatcherFunc
	errChan        chan error
	dst            *concurrentWriter
}

// NewWatcher creates a new Watcher.
func NewWatcher(params WatcherParams, client KubernetesClient, podWatcherFunc PodWatcherFunc, dst io.Writer) *Watcher {
	return &Watcher{
		params:         params,
		client:         client,
		podWatcherFunc: podWatcherFunc,
		errChan:        make(chan error),
		dst:            &concurrentWriter{w: dst},
	}
}

// Watch watches a deployment.
func (w *Watcher) Watch(ctx context.Context) error {
	ns := w.params.Namespace
	if ns == "" {
		ns = corev1.NamespaceDefault
	}

	// Supported resource types are deployments, statefulsets and replicasets
	// (all apps/v1).
	resourceID := schema.GroupVersionResource{
		Resource: w.params.Type,
		Group:    "apps",
		Version:  "v1",
	}

	if err := w.checkResourceExists(ctx, ns, resourceID); err != nil {
		return err
	}

	opts := metav1.ListOptions{Watch: true, FieldSelector: "metadata.name=" + w.params.Name}
	watcher, err := w.client.Untyped.Resource(resourceID).Namespace(ns).Watch(ctx, opts)
	if err != nil {
		return err
	}

	defer watcher.Stop()

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	resultChan := watcher.ResultChan()
	for {
		select {
		case evt, ok := <-resultChan:
			if !ok {
				resultChan = nil
				continue
			}
			switch evt.Type {
			case watch.Added, watch.Modified:
				resource := evt.Object.(*unstructured.Unstructured)
				uid := resource.GetUID()
				// TODO: handle matchExpressions
				selectorAsMap, ok, err := unstructured.NestedStringMap(resource.Object, "spec", "selector", "matchLabels")
				if !ok || err != nil {
					// matchLabels don't exist or cannot be parsed.
					// Should this be fatal?
					log.Printf("warning: unable to parse matchLabels: ok = %t, err = %v", ok, err)
					continue
				}
				selector := labels.SelectorFromSet(selectorAsMap)
				w.addDeployment(ctx, uid, selector)

			case watch.Deleted:
				w.removeDeployment()
			}
		case err := <-w.errChan:
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (w *Watcher) checkResourceExists(ctx context.Context, namespace string, resourceID schema.GroupVersionResource) error {
	_, err := w.client.Untyped.Resource(resourceID).Namespace(namespace).Get(ctx, w.params.Name, metav1.GetOptions{})
	var statusErr *apierrors.StatusError
	if !w.params.StrictExist && errors.As(err, &statusErr) && statusErr.Status().Reason == metav1.StatusReasonNotFound {
		log.Printf(`%s "%s" does not exist, waiting`, resourceID.Resource, w.params.Name)
		return nil
	}
	return err
}

func (w *Watcher) addDeployment(ctx context.Context, resourceUID types.UID, podSelector labels.Selector) {
	if w.resourceUID == resourceUID {
		return
	}

	w.removeDeployment()

	log.Println("[DeploymentWatcher] add podWatcher")

	w.resourceUID = resourceUID
	w.podSelector = podSelector
	w.podWatcher = w.podWatcherFunc(
		w.client,
		w.params.Container,
		w.podSelector,
		w.dst,
	)

	go func() {
		if err := w.podWatcher.WatchPods(ctx); err != nil {
			w.errChan <- err
		}
	}()
}

func (w *Watcher) removeDeployment() {
	if w.podWatcher != nil {
		log.Println("[DeploymentWatcher] remove podWatcher")
		w.podWatcher.Close()
		w.podWatcher = nil
	}
	w.resourceUID = ""
	w.podSelector = nil
}