feat: shared informer priority queue improvements

scrapers/incremental.go

@@ -15,7 +15,6 @@ import (
 	"github.com/flanksource/config-db/scrapers/kubernetes"
 	"github.com/flanksource/config-db/utils/kube"
 	"github.com/flanksource/duty/job"
-	"github.com/samber/lo"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
 
@@ -48,8 +47,9 @@ func ConsumeKubernetesWatchJobFunc(sc api.ScrapeContext, config v1.Kubernetes, q
 			config.BaseScraper = config.BaseScraper.ApplyPlugins(plugins...)
 
 			var (
-				objs       []*unstructured.Unstructured
-				queuedTime = map[string]time.Time{}
+				objs           []*unstructured.Unstructured
+				deletedObjects []string
+				queuedTime     = map[string]time.Time{}
 
 				seenObjects       = map[string]struct{}{}
 				objectsFromEvents = map[string]v1.InvolvedObject{}
@@ -73,6 +73,11 @@ func ConsumeKubernetesWatchJobFunc(sc api.ScrapeContext, config v1.Kubernetes, q
 				}
 				obj := queueItem.Obj
 
+				if queueItem.Operation == kubernetes.QueueItemOperationDelete {
+					deletedObjects = append(deletedObjects, string(obj.GetUID()))
+					continue
+				}
+
 				if obj.GetKind() == "Event" {
 					involvedObjectRaw, ok, _ := unstructured.NestedMap(obj.Object, "involvedObject")
 					if !ok {
@@ -124,7 +129,7 @@ func ConsumeKubernetesWatchJobFunc(sc api.ScrapeContext, config v1.Kubernetes, q
 			// a way that no two objects in a batch have the same id.
 
 			objs = dedup(objs)
-			if err := consumeResources(ctx, *sc.ScrapeConfig(), *config, objs); err != nil {
+			if err := consumeResources(ctx, *sc.ScrapeConfig(), *config, objs, deletedObjects); err != nil {
 				ctx.History.AddErrorf("failed to consume resources: %v", err)
 				return err
 			}
@@ -140,7 +145,7 @@ func ConsumeKubernetesWatchJobFunc(sc api.ScrapeContext, config v1.Kubernetes, q
 	}
 }
 
-func consumeResources(ctx job.JobRuntime, scrapeConfig v1.ScrapeConfig, config v1.Kubernetes, objs []*unstructured.Unstructured) error {
+func consumeResources(ctx job.JobRuntime, scrapeConfig v1.ScrapeConfig, config v1.Kubernetes, objs []*unstructured.Unstructured, deletedResourcesIDs []string) error {
 	cc := api.NewScrapeContext(ctx.Context).WithScrapeConfig(&scrapeConfig).WithJobHistory(ctx.History).AsIncrementalScrape()
 	cc.Context = cc.Context.WithoutName().WithName(fmt.Sprintf("watch[%s/%s]", cc.GetNamespace(), cc.GetName()))
 	results, err := processObjects(cc, config, objs)
@@ -162,15 +167,7 @@ func consumeResources(ctx job.JobRuntime, scrapeConfig v1.ScrapeConfig, config v
 		}
 	}
 
-	_deleteCh, ok := kubernetes.DeleteResourceBuffer.Load(config.Hash())
-	if !ok {
-		return fmt.Errorf("no resource watcher channel found for config (scrapeconfig: %s)", config.Hash())
-	}
-	deleteChan := _deleteCh.(chan string)
-
-	if len(deleteChan) > 0 {
-		deletedResourcesIDs, _, _, _ := lo.Buffer(deleteChan, len(deleteChan))
-
+	if len(deletedResourcesIDs) > 0 {
 		total, err := db.SoftDeleteConfigItems(ctx.Context, deletedResourcesIDs...)
 		if err != nil {
 			return fmt.Errorf("failed to delete %d resources: %w", len(deletedResourcesIDs), err)

scrapers/kubernetes/informers.go

@@ -29,10 +29,6 @@ var (
 	// WatchQueue stores a sync buffer per kubernetes config
 	WatchQueue = sync.Map{}
 
-	// DeleteResourceBuffer stores a buffer per kubernetes config
-	// that contains the ids of resources that have been deleted.
-	DeleteResourceBuffer = sync.Map{}
-
 	informerLagBuckets = []float64{1_000, 5_000, 30_000, 120_000, 300_000, 600_000, 900_000, 1_800_000}
 )
 
@@ -43,9 +39,6 @@ func WatchResources(ctx api.ScrapeContext, config v1.Kubernetes) (*pq.Queue, err
 		priorityQueue = loaded.(*pq.Queue)
 	}
 
-	deleteBuffer := make(chan string, BufferSize)
-	DeleteResourceBuffer.Store(config.Hash(), deleteBuffer)
-
 	if config.Kubeconfig != nil {
 		var err error
 		c, err := ctx.WithKubeconfig(*config.Kubeconfig)
@@ -56,7 +49,7 @@ func WatchResources(ctx api.ScrapeContext, config v1.Kubernetes) (*pq.Queue, err
 	}
 
 	for _, watchResource := range lo.Uniq(config.Watch) {
-		if err := globalSharedInformerManager.Register(ctx, watchResource, priorityQueue, deleteBuffer); err != nil {
+		if err := globalSharedInformerManager.Register(ctx, watchResource, priorityQueue); err != nil {
 			return nil, fmt.Errorf("failed to register informer: %w", err)
 		}
 	}
@@ -92,7 +85,7 @@ type SharedInformerManager struct {
 
 type DeleteObjHandler func(ctx context.Context, id string) error
 
-func (t *SharedInformerManager) Register(ctx api.ScrapeContext, watchResource v1.KubernetesResourceToWatch, queue *pq.Queue, deleteBuffer chan<- string) error {
+func (t *SharedInformerManager) Register(ctx api.ScrapeContext, watchResource v1.KubernetesResourceToWatch, queue *pq.Queue) error {
 	start := time.Now()
 
 	apiVersion, kind := watchResource.ApiVersion, watchResource.Kind
@@ -123,7 +116,7 @@ func (t *SharedInformerManager) Register(ctx api.ScrapeContext, watchResource v1
 				return
 			}
 
-			queue.Enqueue(NewQueueItem(u))
+			queue.Enqueue(NewQueueItem(u, QueueItemOperationAdd))
 
 			if ctx.Properties().On(false, "scraper.log.items") {
 				ctx.Logger.V(4).Infof("added: %s %s %s", u.GetUID(), u.GetKind(), u.GetName())
@@ -166,7 +159,7 @@ func (t *SharedInformerManager) Register(ctx api.ScrapeContext, watchResource v1
 				).Record(time.Duration(time.Since(*lastUpdatedTime).Milliseconds()))
 			}
 
-			queue.Enqueue(NewQueueItem(u))
+			queue.Enqueue(NewQueueItem(u, QueueItemOperationUpdate))
 		},
 		DeleteFunc: func(obj any) {
 			u, err := getUnstructuredFromInformedObj(watchResource, obj)
@@ -197,7 +190,7 @@ func (t *SharedInformerManager) Register(ctx api.ScrapeContext, watchResource v1
 				).Record(time.Duration(time.Since(u.GetDeletionTimestamp().Time).Milliseconds()))
 			}
 
-			deleteBuffer <- string(u.GetUID())
+			queue.Enqueue(NewQueueItem(u, QueueItemOperationDelete))
 		},
 	})
 	if err != nil {
@@ -329,28 +322,79 @@ func kubeConfigIdentifier(ctx api.ScrapeContext) string {
 	return rs.Host
 }
 
+type QueueItemOperation int
+
+const (
+	QueueItemOperationAdd QueueItemOperation = iota + 1
+	QueueItemOperationUpdate
+	QueueItemOperationDelete
+)
+
+func (t *QueueItemOperation) Priority() int {
+	// smaller value represents higher priority
+	priority := map[QueueItemOperation]int{
+		QueueItemOperationAdd:    1,
+		QueueItemOperationUpdate: 2,
+		QueueItemOperationDelete: 3,
+	}
+
+	return priority[*t]
+}
+
 type QueueItem struct {
 	Timestamp time.Time // Queued time
 	Obj       *unstructured.Unstructured
+	Operation QueueItemOperation
 }
 
-func NewQueueItem(obj *unstructured.Unstructured) *QueueItem {
+func NewQueueItem(obj *unstructured.Unstructured, operation QueueItemOperation) *QueueItem {
 	return &QueueItem{
 		Timestamp: time.Now(),
 		Obj:       obj,
+		Operation: operation,
 	}
 }
 
 func pqComparator(a, b any) int {
-	var aTimestamp, bTimestamp time.Time
 	qa := a.(*QueueItem)
 	qb := b.(*QueueItem)
 
-	if qa.Obj.GetCreationTimestamp().Time.Before(qb.Obj.GetCreationTimestamp().Time) {
+	if opResult := pqCompareOperation(qa.Operation, qb.Operation); opResult != 0 {
+		return opResult
+	}
+
+	if opResult := pqCompareKind(qa.Obj.GetKind(), qb.Obj.GetKind()); opResult != 0 {
+		return opResult
+	}
+
+	lastUpdatedTimeA := *health.GetLastUpdatedTime(qa.Obj)
+	lastUpdatedTimeB := *health.GetLastUpdatedTime(qb.Obj)
+
+	if lastUpdatedTimeA.Before(lastUpdatedTimeB) {
 		return -1
-	} else if aTimestamp.Equal(bTimestamp) {
+	} else if lastUpdatedTimeA.Equal(lastUpdatedTimeB) {
 		return 0
 	} else {
 		return 1
 	}
 }
+
+func pqCompareOperation(a, b QueueItemOperation) int {
+	return a.Priority() - b.Priority()
+}
+
+func pqCompareKind(a, b string) int {
+	// smaller means earlier in the queue
+	priority := map[string]int{
+		"Namespace":  1,
+		"Deployment": 2,
+		"ReplicaSet": 3,
+		"Pod":        4,
+		"Event":      5,
+	}
+
+	pa := lo.CoalesceOrEmpty(priority[a], 3)
+	pb := lo.CoalesceOrEmpty(priority[b], 3)
+
+	return pa - pb
+}

scrapers/kubernetes/informers_test.go

@@ -0,0 +1,171 @@
+package kubernetes
+
+import (
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/emirpasic/gods/queues/priorityqueue"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+)
+
+func TestPqComparator(t *testing.T) {
+	now := time.Now()
+
+	tests := []struct {
+		name     string
+		a        QueueItem
+		b        QueueItem
+		expected []string
+	}{
+		{
+			name: "add should have higher priority than update",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "a", now),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationUpdate,
+				Obj:       getUnstructured("Pod", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "update should have higher priority than delete",
+			a: QueueItem{
+				Operation: QueueItemOperationUpdate,
+				Obj:       getUnstructured("Pod", "a", now),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationDelete,
+				Obj:       getUnstructured("Pod", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "same operation should compare by kind - Namespace vs Pod",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Namespace", "a", now),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "same operation and kind should compare by timestamp - earlier first",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "a", now.Add(-1*time.Hour)),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "namespace comes first even before a pod created earlier",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "a", now.Add(-1*time.Hour)),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Namespace", "b", now),
+			},
+			expected: []string{"b", "a"},
+		},
+		{
+			name: "operation priority should override kind priority",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "a", now),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationDelete,
+				Obj:       getUnstructured("Namespace", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "unknown kind should use default priority",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Canary", "a", now),
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructured("Pod", "b", now),
+			},
+			expected: []string{"a", "b"},
+		},
+		{
+			name: "events with managed fields",
+			a: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructuredEvent("Event", "a", now.Add(-2*time.Hour), now.Add(time.Hour)), // created ealier but re-created later
+			},
+			b: QueueItem{
+				Operation: QueueItemOperationAdd,
+				Obj:       getUnstructuredEvent("Event", "b", now.Add(-time.Hour), now.Add(time.Minute)),
+			},
+			expected: []string{"b", "a"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			q := priorityqueue.NewWith(pqComparator)
+			q.Enqueue(&tt.a)
+			q.Enqueue(&tt.b)
+
+			var result []string
+			for {
+				v, ok := q.Dequeue()
+				if !ok {
+					break
+				}
+
+				item := v.(*QueueItem)
+				result = append(result, item.Obj.GetName())
+			}
+
+			if !reflect.DeepEqual(result, tt.expected) {
+				t.Errorf("Test %s failed\nExpected: %v\nGot: %v", tt.name, tt.expected, result)
+			}
+		})
+	}
+}
+
+func getUnstructuredEvent(kind, name string, creationTimestamp, recreationTimestamp time.Time) *unstructured.Unstructured {
+	return &unstructured.Unstructured{
+		Object: map[string]any{
+			"kind": kind,
+			"metadata": map[string]any{
+				"name":              name,
+				"creationTimestamp": creationTimestamp.Format(time.RFC3339),
+				"managedFields": []any{
+					map[string]any{
+						"operation": "Update",
+						"time":      recreationTimestamp.Format(time.RFC3339),
+					},
+				},
+			},
+		},
+	}
+}
+
+func getUnstructured(kind, name string, creationTimestamp time.Time) *unstructured.Unstructured {
+	return &unstructured.Unstructured{
+		Object: map[string]any{
+			"kind": kind,
+			"metadata": map[string]any{
+				"name":              name,
+				"creationTimestamp": creationTimestamp.Format(time.RFC3339),
+			},
+		},
+	}
+}