From ab07b0aed57d46bd52e0544b210b7be5ce5db66e Mon Sep 17 00:00:00 2001
From: Michael Crenshaw <350466+crenshaw-dev@users.noreply.github.com>
Date: Wed, 18 Dec 2024 17:15:41 -0500
Subject: [PATCH] chore(controller): simplify sharding code (#21244)

Signed-off-by: Michael Crenshaw <350466+crenshaw-dev@users.noreply.github.com>
---
 controller/sharding/cache.go         | 20 ----------------
 controller/sharding/sharding_test.go | 36 +++++++++++++++-------------
 2 files changed, 19 insertions(+), 37 deletions(-)

diff --git a/controller/sharding/cache.go b/controller/sharding/cache.go
index 4a750e3545524..454b37e4cf5fe 100644
--- a/controller/sharding/cache.go
+++ b/controller/sharding/cache.go
@@ -19,7 +19,6 @@ type ClusterShardingCache interface {
 	UpdateApp(a *v1alpha1.Application)
 	IsManagedCluster(c *v1alpha1.Cluster) bool
 	GetDistribution() map[string]int
-	GetAppDistribution() map[string]int
 }
 
 type ClusterSharding struct {
@@ -244,22 +243,3 @@ func (sharding *ClusterSharding) UpdateApp(a *v1alpha1.Application) {
 		log.Debugf("Skipping sharding distribution update. No relevant changes")
 	}
 }
-
-// GetAppDistribution should be not be called from a DestributionFunction because
-// it could cause a deadlock when updateDistribution is called.
-func (sharding *ClusterSharding) GetAppDistribution() map[string]int {
-	sharding.lock.RLock()
-	clusters := sharding.Clusters
-	apps := sharding.Apps
-	sharding.lock.RUnlock()
-
-	appDistribution := make(map[string]int, len(clusters))
-
-	for _, a := range apps {
-		if _, ok := appDistribution[a.Spec.Destination.Server]; !ok {
-			appDistribution[a.Spec.Destination.Server] = 0
-		}
-		appDistribution[a.Spec.Destination.Server]++
-	}
-	return appDistribution
-}
diff --git a/controller/sharding/sharding_test.go b/controller/sharding/sharding_test.go
index b76741be92e16..c332e869c2789 100644
--- a/controller/sharding/sharding_test.go
+++ b/controller/sharding/sharding_test.go
@@ -988,35 +988,37 @@ func TestGetClusterSharding(t *testing.T) {
 }
 
 func TestAppAwareCache(t *testing.T) {
-	_, db, cluster1, cluster2, cluster3, cluster4, cluster5 := createTestClusters()
+	_, _, cluster1, cluster2, cluster3, cluster4, cluster5 := createTestClusters()
 	_, app1, app2, app3, app4, app5 := createTestApps()
 
-	clusterSharding := NewClusterSharding(db, 0, 1, "legacy")
+	clusterList := getClusterPointers([]v1alpha1.Cluster{cluster1, cluster2, cluster3, cluster4, cluster5})
+	appList := getAppPointers([]v1alpha1.Application{app1, app2, app3, app4, app5})
 
-	clusterList := &v1alpha1.ClusterList{Items: []v1alpha1.Cluster{cluster1, cluster2, cluster3, cluster4, cluster5}}
-	appList := &v1alpha1.ApplicationList{Items: []v1alpha1.Application{app1, app2, app3, app4, app5}}
-	clusterSharding.Init(clusterList, appList)
+	getClusters := func() []*v1alpha1.Cluster { return clusterList }
+	getApps := func() []*v1alpha1.Application { return appList }
 
-	appDistribution := clusterSharding.GetAppDistribution()
+	appDistribution := getAppDistribution(getClusters, getApps)
 
-	assert.Equal(t, 2, appDistribution["cluster1"])
-	assert.Equal(t, 2, appDistribution["cluster2"])
-	assert.Equal(t, 1, appDistribution["cluster3"])
+	assert.Equal(t, int64(2), appDistribution["cluster1"])
+	assert.Equal(t, int64(2), appDistribution["cluster2"])
+	assert.Equal(t, int64(1), appDistribution["cluster3"])
 
 	app6 := createApp("app6", "cluster4")
-	clusterSharding.AddApp(&app6)
+	appList = append(appList, &app6)
 
 	app1Update := createApp("app1", "cluster2")
-	clusterSharding.UpdateApp(&app1Update)
+	// replace app 1
+	appList[0] = &app1Update
 
-	clusterSharding.DeleteApp(&app3)
+	// Remove app 3
+	appList = append(appList[:2], appList[3:]...)
 
-	appDistribution = clusterSharding.GetAppDistribution()
+	appDistribution = getAppDistribution(getClusters, getApps)
 
-	assert.Equal(t, 1, appDistribution["cluster1"])
-	assert.Equal(t, 2, appDistribution["cluster2"])
-	assert.Equal(t, 1, appDistribution["cluster3"])
-	assert.Equal(t, 1, appDistribution["cluster4"])
+	assert.Equal(t, int64(1), appDistribution["cluster1"])
+	assert.Equal(t, int64(2), appDistribution["cluster2"])
+	assert.Equal(t, int64(1), appDistribution["cluster3"])
+	assert.Equal(t, int64(1), appDistribution["cluster4"])
 }
 
 func createTestApps() (appAccessor, v1alpha1.Application, v1alpha1.Application, v1alpha1.Application, v1alpha1.Application, v1alpha1.Application) {