use informer pattern in ConfigMapSource

Co-authored-by: Devon Warshaw <[email protected]>
Signed-off-by: Matthew Conger-Eldeen <[email protected]>

Author: m25n

src/vizier/services/query_broker/script_runner/BUILD.bazel

@@ -48,9 +48,13 @@ go_library(
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@io_k8s_api//core/v1:core",
         "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
-        "@io_k8s_apimachinery//pkg/watch",
-        "@io_k8s_client_go//kubernetes/typed/core/v1:core",
+        "@io_k8s_apimachinery//pkg/labels",
+        "@io_k8s_client_go//informers",
+        "@io_k8s_client_go//informers/core/v1:core",
+        "@io_k8s_client_go//kubernetes",
+        "@io_k8s_client_go//listers/core/v1:core",
         "@io_k8s_client_go//rest",
+        "@io_k8s_client_go//tools/cache",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
@@ -84,11 +88,7 @@ pl_go_test(
         "@com_github_stretchr_testify//require",
         "@io_k8s_api//core/v1:core",
         "@io_k8s_apimachinery//pkg/apis/meta/v1:meta",
-        "@io_k8s_apimachinery//pkg/labels",
-        "@io_k8s_apimachinery//pkg/runtime",
-        "@io_k8s_apimachinery//pkg/watch",
         "@io_k8s_client_go//kubernetes/fake",
-        "@io_k8s_client_go//testing",
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",

src/vizier/services/query_broker/script_runner/config_map_source.go

@@ -20,111 +20,151 @@ package scriptrunner
 
 import (
 	"context"
+	"reflect"
+	"sync/atomic"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v2"
-	v1 "k8s.io/api/core/v1"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/watch"
-	clientv1 "k8s.io/client-go/kubernetes/typed/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/client-go/informers"
+	informercorev1 "k8s.io/client-go/informers/core/v1"
+	"k8s.io/client-go/kubernetes"
+	listercorev1 "k8s.io/client-go/listers/core/v1"
+	"k8s.io/client-go/tools/cache"
 
 	"px.dev/pixie/src/shared/cvmsgspb"
 	"px.dev/pixie/src/utils"
 )
 
 // ConfigMapSource pulls cron scripts from config maps.
 type ConfigMapSource struct {
-	stop   func()
-	client clientv1.ConfigMapInterface
+	stop     func()
+	informer informercorev1.ConfigMapInformer
 }
 
 // NewConfigMapSource constructs a [Source] that extracts cron scripts from config maps with the label "purpose=cron-script".
 // Each config map must contain
 //   - a script.pxl with the pixel script
 //   - a configs.yaml which will be stored in the Configs field of [cvmsgspb.CronScript]
 //   - a cron.yaml that contains a "frequency_s" key
-func NewConfigMapSource(client clientv1.ConfigMapInterface) *ConfigMapSource {
-	return &ConfigMapSource{client: client}
+func NewConfigMapSource(client kubernetes.Interface, namespace string) *ConfigMapSource {
+	return &ConfigMapSource{
+		informer: informers.NewSharedInformerFactoryWithOptions(
+			client,
+			12*time.Hour,
+			informers.WithNamespace(namespace),
+			informers.WithTweakListOptions(func(options *metav1.ListOptions) {
+				options.LabelSelector = "purpose=cron-script"
+			}),
+		).Core().V1().ConfigMaps(),
+	}
 }
 
 // Start watches for updates to matching configmaps and sends resulting updates on updatesCh.
-func (source *ConfigMapSource) Start(baseCtx context.Context, updatesCh chan<- *cvmsgspb.CronScriptUpdate) (map[string]*cvmsgspb.CronScript, error) {
-	options := metav1.ListOptions{LabelSelector: "purpose=cron-script"}
-	watcher, err := source.client.Watch(baseCtx, options)
+func (source *ConfigMapSource) Start(ctx context.Context, updatesCh chan<- *cvmsgspb.CronScriptUpdate) (map[string]*cvmsgspb.CronScript, error) {
+	stopCh := make(chan struct{})
+	isInitialized := &atomic.Bool{}
+	_, err := source.informer.Informer().AddEventHandler(configMapEventHandlers(isInitialized, updatesCh))
+	if err != nil {
+		return nil, err
+	}
+	go source.informer.Informer().Run(stopCh)
+	cache.WaitForCacheSync(ctx.Done(), source.informer.Informer().HasSynced)
+	initialScripts, err := getInitialScripts(source.informer.Lister())
 	if err != nil {
+		close(stopCh)
 		return nil, err
 	}
-	go configMapUpdater(watcher, updatesCh)
-	configmaps, err := source.client.List(baseCtx, options)
+	isInitialized.Store(true)
+	source.stop = func() { close(stopCh) }
+	return initialScripts, nil
+}
+
+func getInitialScripts(lister listercorev1.ConfigMapLister) (map[string]*cvmsgspb.CronScript, error) {
+	configMaps, err := lister.List(labels.Everything())
 	if err != nil {
-		watcher.Stop()
 		return nil, err
 	}
 	scripts := map[string]*cvmsgspb.CronScript{}
-	for _, configmap := range configmaps.Items {
-		id, cronScript, err := configmapToCronScript(&configmap)
+	for _, configMap := range configMaps {
+		id, cronScript, err := configmapToCronScript(configMap)
 		if err != nil {
 			logCronScriptParseError(err)
 			continue
 		}
 		scripts[id] = cronScript
 	}
-	source.stop = watcher.Stop
 	return scripts, nil
 }
 
-// Stop stops further updates from being sent.
-func (source *ConfigMapSource) Stop() {
-	source.stop()
-}
-
-func configMapUpdater(watcher watch.Interface, updatesCh chan<- *cvmsgspb.CronScriptUpdate) {
-	for event := range watcher.ResultChan() {
-		switch event.Type {
-		case watch.Modified, watch.Added:
-			configmap := event.Object.(*v1.ConfigMap)
-			id, script, err := configmapToCronScript(configmap)
-			if err != nil {
-				logCronScriptParseError(err)
-				continue
-			}
-			cronScriptUpdate := &cvmsgspb.CronScriptUpdate{
-				Msg: &cvmsgspb.CronScriptUpdate_UpsertReq{
-					UpsertReq: &cvmsgspb.RegisterOrUpdateCronScriptRequest{
-						Script: script,
-					},
-				},
-				RequestID: id,
-				Timestamp: time.Now().Unix(),
+func configMapEventHandlers(isInitialized *atomic.Bool, updatesCh chan<- *cvmsgspb.CronScriptUpdate) cache.ResourceEventHandlerFuncs {
+	return cache.ResourceEventHandlerFuncs{
+		AddFunc: func(obj interface{}) {
+			configmap := obj.(*corev1.ConfigMap)
+			if !isInitialized.Load() {
+				return
 			}
-			updatesCh <- cronScriptUpdate
-		case watch.Deleted:
-			configmap := event.Object.(*v1.ConfigMap)
-			id, script, err := configmapToCronScript(configmap)
-			if err != nil {
-				logCronScriptParseError(err)
-				continue
+			updatesCh <- makeUpdate(configmap)
+		},
+		UpdateFunc: func(oldObj, newObj interface{}) {
+			if reflect.DeepEqual(oldObj, newObj) {
+				return
 			}
-			cronScriptUpdate := &cvmsgspb.CronScriptUpdate{
-				Msg: &cvmsgspb.CronScriptUpdate_DeleteReq{
-					DeleteReq: &cvmsgspb.DeleteCronScriptRequest{
-						ScriptID: script.ID,
-					},
-				},
-				RequestID: id,
-				Timestamp: time.Now().Unix(),
-			}
-			updatesCh <- cronScriptUpdate
-		}
+			configmap := newObj.(*corev1.ConfigMap)
+			updatesCh <- makeUpdate(configmap)
+		},
+		DeleteFunc: func(obj interface{}) {
+			configmap := obj.(*corev1.ConfigMap)
+			updatesCh <- makeDelete(configmap)
+		},
 	}
 }
 
+func makeDelete(configmap *corev1.ConfigMap) *cvmsgspb.CronScriptUpdate {
+	_, script, err := configmapToCronScript(configmap)
+	if err != nil {
+		logCronScriptParseError(err)
+		return nil
+	}
+	return &cvmsgspb.CronScriptUpdate{
+		Msg: &cvmsgspb.CronScriptUpdate_DeleteReq{
+			DeleteReq: &cvmsgspb.DeleteCronScriptRequest{
+				ScriptID: script.ID,
+			},
+		},
+		Timestamp: time.Now().Unix(),
+	}
+}
+
+func makeUpdate(configmap *corev1.ConfigMap) *cvmsgspb.CronScriptUpdate {
+	_, script, err := configmapToCronScript(configmap)
+	if err != nil {
+		logCronScriptParseError(err)
+		return nil
+	}
+	return &cvmsgspb.CronScriptUpdate{
+		Msg: &cvmsgspb.CronScriptUpdate_UpsertReq{
+			UpsertReq: &cvmsgspb.RegisterOrUpdateCronScriptRequest{
+				Script: script,
+			},
+		},
+		Timestamp: time.Now().Unix(),
+	}
+}
+
+// Stop stops further updates from being sent.
+func (source *ConfigMapSource) Stop() {
+	source.stop()
+}
+
 func logCronScriptParseError(err error) {
 	log.WithError(err).Error("Failed to parse cron.yaml from configmap cron script")
 }
 
-func configmapToCronScript(configmap *v1.ConfigMap) (string, *cvmsgspb.CronScript, error) {
+func configmapToCronScript(configmap *corev1.ConfigMap) (string, *cvmsgspb.CronScript, error) {
 	id := string(configmap.UID)
 	cronScript := &cvmsgspb.CronScript{
 		ID:      utils.ProtoFromUUIDStrOrNil(id),