From ec6fd2befac30ebb7dfb55b94d8c9874489a3902 Mon Sep 17 00:00:00 2001
From: Gaurav Ghildiyal <gauravkg@google.com>
Date: Mon, 4 Mar 2024 14:07:45 -0800
Subject: [PATCH] Add options construct to EndpointSlice NewReconciler for the
 new trafficDistributionEnabled field

---
 .../endpointslice/endpointslice_controller.go |  2 +-
 .../src/k8s.io/endpointslice/reconciler.go    | 35 +++++++++++++------
 .../k8s.io/endpointslice/reconciler_test.go   |  1 -
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/pkg/controller/endpointslice/endpointslice_controller.go b/pkg/controller/endpointslice/endpointslice_controller.go
index 687e08b9fc2..2e137fafdda 100644
--- a/pkg/controller/endpointslice/endpointslice_controller.go
+++ b/pkg/controller/endpointslice/endpointslice_controller.go
@@ -173,9 +173,9 @@ func NewController(ctx context.Context, podInformer coreinformers.PodInformer,
 		c.maxEndpointsPerSlice,
 		c.endpointSliceTracker,
 		c.topologyCache,
-		utilfeature.DefaultFeatureGate.Enabled(features.ServiceTrafficDistribution),
 		c.eventRecorder,
 		controllerName,
+		endpointslicerec.WithTrafficDistributionEnabled(utilfeature.DefaultFeatureGate.Enabled(features.ServiceTrafficDistribution)),
 	)
 
 	return c
diff --git a/staging/src/k8s.io/endpointslice/reconciler.go b/staging/src/k8s.io/endpointslice/reconciler.go
index 1ec466c503f..9242a76c0bc 100644
--- a/staging/src/k8s.io/endpointslice/reconciler.go
+++ b/staging/src/k8s.io/endpointslice/reconciler.go
@@ -59,6 +59,16 @@ type Reconciler struct {
 	controllerName string
 }
 
+type ReconcilerOption func(*Reconciler)
+
+// WithTrafficDistributionEnabled controls whether the Reconciler considers the
+// `trafficDistribution` field while reconciling EndpointSlices.
+func WithTrafficDistributionEnabled(enabled bool) ReconcilerOption {
+	return func(r *Reconciler) {
+		r.trafficDistributionEnabled = enabled
+	}
+}
+
 // endpointMeta includes the attributes we group slices on, this type helps with
 // that logic in Reconciler
 type endpointMeta struct {
@@ -327,18 +337,21 @@ func (r *Reconciler) reconcileByAddressType(logger klog.Logger, service *corev1.
 
 }
 
-func NewReconciler(client clientset.Interface, nodeLister corelisters.NodeLister, maxEndpointsPerSlice int32, endpointSliceTracker *endpointsliceutil.EndpointSliceTracker, topologyCache *topologycache.TopologyCache, trafficDistributionEnabled bool, eventRecorder record.EventRecorder, controllerName string) *Reconciler {
-	return &Reconciler{
-		client:                     client,
-		nodeLister:                 nodeLister,
-		maxEndpointsPerSlice:       maxEndpointsPerSlice,
-		endpointSliceTracker:       endpointSliceTracker,
-		metricsCache:               metrics.NewCache(maxEndpointsPerSlice),
-		topologyCache:              topologyCache,
-		trafficDistributionEnabled: trafficDistributionEnabled,
-		eventRecorder:              eventRecorder,
-		controllerName:             controllerName,
+func NewReconciler(client clientset.Interface, nodeLister corelisters.NodeLister, maxEndpointsPerSlice int32, endpointSliceTracker *endpointsliceutil.EndpointSliceTracker, topologyCache *topologycache.TopologyCache, eventRecorder record.EventRecorder, controllerName string, options ...ReconcilerOption) *Reconciler {
+	r := &Reconciler{
+		client:               client,
+		nodeLister:           nodeLister,
+		maxEndpointsPerSlice: maxEndpointsPerSlice,
+		endpointSliceTracker: endpointSliceTracker,
+		metricsCache:         metrics.NewCache(maxEndpointsPerSlice),
+		topologyCache:        topologyCache,
+		eventRecorder:        eventRecorder,
+		controllerName:       controllerName,
 	}
+	for _, option := range options {
+		option(r)
+	}
+	return r
 }
 
 // placeholderSliceCompare is a conversion func for comparing two placeholder endpoint slices.
diff --git a/staging/src/k8s.io/endpointslice/reconciler_test.go b/staging/src/k8s.io/endpointslice/reconciler_test.go
index 007b3a9293e..0f931835b13 100644
--- a/staging/src/k8s.io/endpointslice/reconciler_test.go
+++ b/staging/src/k8s.io/endpointslice/reconciler_test.go
@@ -2199,7 +2199,6 @@ func newReconciler(client *fake.Clientset, nodes []*corev1.Node, maxEndpointsPer
 		maxEndpointsPerSlice,
 		endpointsliceutil.NewEndpointSliceTracker(),
 		nil,
-		false,
 		eventRecorder,
 		controllerName,
 	)