Refactor to use k8s client directly instead of an informer

Author: carreter

cmd/agent/app/server.go

@@ -153,7 +153,7 @@ func (a *Agent) runProxyConnection(o *options.GrpcProxyAgentOptions, drainCh, st
 			return nil, fmt.Errorf("invalid server count lease selector: %w", err)
 		}
 		serverLeaseCounter := agent.NewServerLeaseCounter(
-			sharedInformerFactory.Coordination().V1().Leases().Lister(),
+			sharedInformerFactory,
 			serverLeaseSelector,
 		)
 		if err != nil {

pkg/agent/clientset.go

@@ -184,15 +184,15 @@ func (cs *ClientSet) resetBackoff() *wait.Backoff {
 }
 
 // sync makes sure that #clients >= #proxy servers
-func (cs *ClientSet) sync() {
+func (cs *ClientSet) sync(ctx context.Context) {
 	defer cs.shutdown()
 	backoff := cs.resetBackoff()
 	var duration time.Duration
 	for {
-		if err := cs.connectOnce(); err != nil {
+		if err := cs.connectOnce(ctx); err != nil {
 			if dse, ok := err.(*DuplicateServerError); ok {
 				clientsCount := cs.ClientsCount()
-				serverCount := cs.ServerCount()
+				serverCount := cs.ServerCount(ctx)
 				klog.V(4).InfoS("duplicate server", "serverID", dse.ServerID, "serverCount", serverCount, "clientsCount", clientsCount)
 				if serverCount != 0 && clientsCount >= serverCount {
 					duration = backoff.Step()
@@ -217,10 +217,10 @@ func (cs *ClientSet) sync() {
 	}
 }
 
-func (cs *ClientSet) ServerCount() int {
+func (cs *ClientSet) ServerCount(ctx context.Context) int {
 	var serverCount int
 	if cs.serverCounter != nil {
-		serverCount = cs.serverCounter.Count()
+		serverCount = cs.serverCounter.Count(ctx)
 		if serverCount == 0 {
 			klog.Warningf("server lease counter could not find any leases")
 		}
@@ -232,8 +232,8 @@ func (cs *ClientSet) ServerCount() int {
 	return serverCount
 }
 
-func (cs *ClientSet) connectOnce() error {
-	serverCount := cs.ServerCount()
+func (cs *ClientSet) connectOnce(ctx context.Context) error {
+	serverCount := cs.ServerCount(ctx)
 
 	if !cs.syncForever && serverCount != 0 && cs.ClientsCount() >= serverCount {
 		return nil
@@ -263,7 +263,7 @@ func (cs *ClientSet) Serve() {
 		"agentIdentifiers", cs.agentIdentifiers,
 		"serverAddress", cs.address,
 	)
-	go runpprof.Do(context.Background(), labels, func(context.Context) { cs.sync() })
+	go runpprof.Do(context.Background(), labels, func(ctx context.Context) { cs.sync(ctx) })
 }
 
 func (cs *ClientSet) shutdown() {

pkg/agent/lease_counter.go

@@ -1,30 +1,33 @@
 package agent
 
 import (
+	"context"
 	"time"
 
 	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
 
 	coordinationv1api "k8s.io/api/coordination/v1"
-	coordinationv1listers "k8s.io/client-go/listers/coordination/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	coordinationv1 "k8s.io/client-go/kubernetes/typed/coordination/v1"
 )
 
 var timeNow = time.Now
 
 // A ServerLeaseCounter counts leases in the k8s apiserver to determine the
 // current proxy server count.
 type ServerLeaseCounter struct {
-	lister        coordinationv1listers.LeaseLister
+	leaseClient   coordinationv1.LeaseInterface
 	selector      labels.Selector
 	fallbackCount int
 }
 
 // NewServerLeaseCounter creates a server counter that counts valid leases that match the label
 // selector and provides the fallback count (initially 0) if this fails.
-func NewServerLeaseCounter(lister coordinationv1listers.LeaseLister, labelSelector labels.Selector) *ServerLeaseCounter {
+func NewServerLeaseCounter(k8sClient kubernetes.Interface, labelSelector labels.Selector) *ServerLeaseCounter {
 	return &ServerLeaseCounter{
-		lister:        lister,
+		leaseClient:   k8sClient.CoordinationV1().Leases(""),
 		selector:      labelSelector,
 		fallbackCount: 0,
 	}
@@ -36,20 +39,17 @@ func NewServerLeaseCounter(lister coordinationv1listers.LeaseLister, labelSelect
 // In the event that no valid leases are found or lease listing fails, the
 // fallback count is returned. This fallback count is updated upon successful
 // discovery of valid leases.
-func (lc *ServerLeaseCounter) Count() int {
+func (lc *ServerLeaseCounter) Count(ctx context.Context) int {
 	// Since the number of proxy servers is generally small (1-10), we opted against
-	// using a LIST and WATCH pattern and instead list all leases in the informer.
-	// The informer still uses LIST and WATCH under the hood, so this doesn't result
-	// in additional calls in the apiserver, and checking whether a lease is valid
-	// is cheap.
-	leases, err := lc.lister.List(lc.selector)
+	// using a LIST and WATCH pattern and instead list all leases on each call.
+	leases, err := lc.leaseClient.List(ctx, metav1.ListOptions{LabelSelector: lc.selector.String()})
 	if err != nil {
 		klog.Errorf("could not list leases to update server count, using fallback count (%v): %v", lc.fallbackCount, err)
 		return lc.fallbackCount
 	}
 
 	count := 0
-	for _, lease := range leases {
+	for _, lease := range leases.Items {
 		if isLeaseValid(lease) {
 			count++
 		} else {
@@ -67,7 +67,7 @@ func (lc *ServerLeaseCounter) Count() int {
 	return count
 }
 
-func isLeaseValid(lease *coordinationv1api.Lease) bool {
+func isLeaseValid(lease coordinationv1api.Lease) bool {
 	var lastRenewTime time.Time
 	if lease.Spec.RenewTime != nil {
 		lastRenewTime = lease.Spec.RenewTime.Time

pkg/agent/lease_counter_test.go

@@ -1,13 +1,19 @@
 package agent
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/fake"
+	k8stesting "k8s.io/client-go/testing"
+
 	coordinationv1 "k8s.io/api/coordination/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/labels"
 	coordinationv1listers "k8s.io/client-go/listers/coordination/v1"
 )
 
@@ -34,6 +40,7 @@ func newLeaseFromTemplate(template leaseTemplate) *coordinationv1.Lease {
 	lease := &coordinationv1.Lease{
 		TypeMeta: metav1.TypeMeta{},
 		ObjectMeta: metav1.ObjectMeta{
+			Name:   uuid.New().String(),
 			Labels: template.labels,
 		},
 		Spec: coordinationv1.LeaseSpec{},
@@ -103,7 +110,7 @@ func TestIsLeaseValid(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			lease := newLeaseFromTemplate(tc.template)
 
-			got := isLeaseValid(lease)
+			got := isLeaseValid(*lease)
 			if got != tc.want {
 				t.Errorf("incorrect lease validity (got: %v, want: %v)", got, tc.want)
 			}
@@ -234,20 +241,18 @@ func TestServerLeaseCounter(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			ct := &controlledTime{t: time.Unix(10000000, 0)}
 			timeNow = ct.Now
-			leases := make([]*coordinationv1.Lease, len(tc.templates))
+			leases := make([]runtime.Object, len(tc.templates))
 			for i, template := range tc.templates {
 				leases[i] = newLeaseFromTemplate(template)
 			}
-			lister := &fakeLeaseLister{
-				leases: leases,
-				err:    tc.leaseListerError,
-			}
 			ct.Advance(time.Millisecond)
 
+			k8sClient := fake.NewSimpleClientset(leases...)
 			selector, _ := labels.Parse(tc.labelSelector)
-			counter := NewServerLeaseCounter(lister, selector)
 
-			got := counter.Count()
+			counter := NewServerLeaseCounter(k8sClient, selector)
+
+			got := counter.Count(context.Background())
 			if tc.want != got {
 				t.Errorf("incorrect server count (got: %v, want: %v)", got, tc.want)
 			}
@@ -269,37 +274,41 @@ func TestServerLeaseCounter_FallbackCount(t *testing.T) {
 
 	ct := &controlledTime{t: time.Unix(1000, 0)}
 	timeNow = ct.Now
-	leases := []*coordinationv1.Lease{}
-	leases = append(leases, newLeaseFromTemplate(validLease), newLeaseFromTemplate(validLease), newLeaseFromTemplate(validLease), newLeaseFromTemplate(invalidLease))
+	leases := []runtime.Object{newLeaseFromTemplate(validLease), newLeaseFromTemplate(validLease), newLeaseFromTemplate(validLease), newLeaseFromTemplate(invalidLease)}
 	ct.Advance(time.Millisecond)
 
-	lister := &fakeLeaseLister{
-		leases: leases,
-		err:    fmt.Errorf("dummy lister error"),
-	}
+	k8sClient := fake.NewSimpleClientset(leases...)
+	callShouldFail := true
 
 	selector, _ := labels.Parse("label=value")
-	counter := NewServerLeaseCounter(lister, selector)
 
-	// First call should return fallback count of 0 because of lister error.
-	got := counter.Count()
+	counter := NewServerLeaseCounter(k8sClient, selector)
+
+	// First call should return fallback count of 0 because of leaseClient error.
+	k8sClient.PrependReactor("*", "*", func(action k8stesting.Action) (handled bool, ret runtime.Object, err error) {
+		if callShouldFail {
+			return true, nil, fmt.Errorf("dummy lease client error")
+		}
+		return false, nil, nil
+	})
+	ctx := context.Background()
+	got := counter.Count(ctx)
 	if got != 0 {
-		t.Errorf("lease counter did not return fallback count on lister error (got: %v, want: 0", got)
+		t.Errorf("lease counter did not return fallback count on leaseClient error (got: %v, want: 0)", got)
 	}
 
-	// Second call should return the actual count (3) upon lister success.
+	// Second call should return the actual count (3) upon leaseClient success.
+	callShouldFail = false
 	actualCount := 3
-	lister.err = nil
-	got = counter.Count()
+	got = counter.Count(ctx)
 	if got != actualCount {
-		t.Errorf("lease counter did not return actual count on lister success (got: %v, want: %v)", got, actualCount)
+		t.Errorf("lease counter did not return actual count on leaseClient success (got: %v, want: %v)", got, actualCount)
 	}
 
-	// Third call should return updated fallback count (3) upon lister failure.
-	lister.err = fmt.Errorf("dummy lister error")
-	lister.leases = append(lister.leases, newLeaseFromTemplate(validLease)) // Change actual count just in case.
-	got = counter.Count()
+	// Third call should return updated fallback count (3) upon leaseClient failure.
+	callShouldFail = true
+	got = counter.Count(ctx)
 	if got != actualCount {
-		t.Errorf("lease counter did not update fallback count after lister success, returned incorrect count on subsequent lister error (got: %v, want: %v)", got, actualCount)
+		t.Errorf("lease counter did not update fallback count after leaseClient success, returned incorrect count on subsequent leaseClient error (got: %v, want: %v)", got, actualCount)
 	}
 }