Add caching to kuberesolver.

Only create one "watcher" per namespace/pod pair rather creating one per
gRPC connection.

Noticed from the logs that there are a lot of redundant resolutions.

I also removed the resolution from ResolveNow. The documentation
suggests it's just a hint and we already guarantee that we will send the
information as soon as we have it.

Depends on #11355

Author: vadimberezniker

server/util/kuberesolver/kuberesolver.go

@@ -66,137 +66,134 @@ func parsePodTarget(endpoint string) (podTarget, error) {
 	}, nil
 }
 
-type kubeResolverBuilder struct {
-	mu     sync.Mutex
-	client kubernetes.Interface
+type podResolution struct {
+	pod *corev1.Pod
+	err error
 }
 
-func (b *kubeResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
-	endpoint := target.Endpoint()
-	pt, err := parsePodTarget(endpoint)
-	if err != nil {
-		return nil, err
-	}
-
-	b.mu.Lock()
-	defer b.mu.Unlock()
-	if b.client == nil {
-		config, err := rest.InClusterConfig()
-		if err != nil {
-			return nil, fmt.Errorf("could not create k8s client: %w", err)
-		}
-		b.client, err = kubernetes.NewForConfig(config)
-		if err != nil {
-			return nil, fmt.Errorf("could not create k8s client: %w", err)
-		}
-	}
+// podWatcher owns a single Get+Watch loop for a specific (namespace, podName)
+// pair. Multiple kubeResolvers can subscribe to the same podWatcher.
+type podWatcher struct {
+	client    kubernetes.Interface
+	namespace string
+	podName   string
 
-	ctx, cancel := context.WithCancel(context.Background())
-	r := &kubeResolver{
-		cc:        cc,
-		client:    b.client,
-		podTarget: pt,
-		ctx:       ctx,
-		cancel:    cancel,
-		resolveCh: make(chan struct{}, 1),
-	}
+	// Only accessed from the watch goroutine.
+	resourceVersion string
 
-	go r.watch()
-	return r, nil
+	mu               sync.Mutex
+	subscribers      map[*kubeResolver]func(podResolution)
+	cancel           context.CancelFunc
+	latestResolution *podResolution
 }
 
-func (b *kubeResolverBuilder) Scheme() string {
-	return scheme
+func newPodWatcher(ctx context.Context, client kubernetes.Interface, namespace, podName string) *podWatcher {
+	ctx, cancel := context.WithCancel(ctx)
+	pw := &podWatcher{
+		client:      client,
+		namespace:   namespace,
+		podName:     podName,
+		subscribers: make(map[*kubeResolver]func(podResolution)),
+		cancel:      cancel,
+	}
+	go pw.watch(ctx)
+	return pw
 }
 
-type kubeResolver struct {
-	cc        resolver.ClientConn
-	client    kubernetes.Interface
-	podTarget podTarget
-	ctx       context.Context
-	cancel    context.CancelFunc
-	resolveCh chan struct{}
-
-	// resourceVersion is the last known resource version from a Get or Watch
-	// event. It is used to resume watches without missing events.
-	// Only accessed from the watch goroutine, so no mutex is needed.
-	resourceVersion string
+// subscribe adds a subscriber with callback cb. cb will be called immediately
+// with the latestResolution known state (pod or error) and then whenever the pod state
+// changes. cb must not block.
+func (pw *podWatcher) subscribe(r *kubeResolver, cb func(podResolution)) {
+	pw.mu.Lock()
+	pw.subscribers[r] = cb
+	if pw.latestResolution != nil {
+		cb(*pw.latestResolution)
+	}
+	pw.mu.Unlock()
 }
 
-func (r *kubeResolver) resolve() {
-	log.Infof("Resolving pod %s/%s", r.podTarget.namespace, r.podTarget.podName)
-	pod, err := r.client.CoreV1().Pods(r.podTarget.namespace).Get(r.ctx, r.podTarget.podName, metav1.GetOptions{})
-	if err != nil {
-		log.Warningf("Failed to get pod %s/%s from k8s: %s", r.podTarget.namespace, r.podTarget.podName, err)
-		r.cc.ReportError(fmt.Errorf("failed to get pod: %w", err))
-		return
-	}
-	r.updateStateFromPod(pod)
-	r.resourceVersion = pod.ResourceVersion
+// unsubscribe removes a subscriber. It returns true if no subscribers remain.
+func (pw *podWatcher) unsubscribe(r *kubeResolver) bool {
+	pw.mu.Lock()
+	defer pw.mu.Unlock()
+	delete(pw.subscribers, r)
+	return len(pw.subscribers) == 0
 }
 
-func (r *kubeResolver) updateStateFromPod(pod *corev1.Pod) {
-	ip := pod.Status.PodIP
-	if ip == "" {
-		log.Infof("Pod %s/%s doesn't have an IP yet", r.podTarget.namespace, r.podTarget.podName)
-		// Setting an empty address list will return an error, but it should
-		// still cause the balancer to close any existing connections.
-		_ = r.cc.UpdateState(resolver.State{Addresses: nil})
-		return
+func (pw *podWatcher) notifySubscribers(r podResolution) {
+	pw.mu.Lock()
+	pw.latestResolution = &r
+	callbacks := make([]func(podResolution), 0, len(pw.subscribers))
+	for _, cb := range pw.subscribers {
+		callbacks = append(callbacks, cb)
 	}
+	pw.mu.Unlock()
 
-	addr := ip
-	if r.podTarget.port != "" {
-		addr = ip + ":" + r.podTarget.port
+	for _, cb := range callbacks {
+		cb(r)
 	}
+}
 
-	log.Infof("Pod %s/%s resolved to IP %s", r.podTarget.namespace, r.podTarget.podName, ip)
-
-	err := r.cc.UpdateState(resolver.State{
-		Addresses: []resolver.Address{{Addr: addr}},
-	})
+// resolve does a one-shot Get and notifies subscribers. Returns true on success.
+func (pw *podWatcher) resolve(ctx context.Context) bool {
+	log.Infof("Resolving pod %s/%s", pw.namespace, pw.podName)
+	pod, err := pw.client.CoreV1().Pods(pw.namespace).Get(ctx, pw.podName, metav1.GetOptions{})
 	if err != nil {
-		log.Warningf("failed to update state: %s", err)
+		log.Warningf("Failed to get pod %s/%s from k8s: %s", pw.namespace, pw.podName, err)
+		pw.notifySubscribers(podResolution{err: fmt.Errorf("failed to get pod: %w", err)})
+		return false
 	}
+	pw.notifySubscribers(podResolution{pod: pod})
+	pw.resourceVersion = pod.ResourceVersion
+	return true
 }
 
-func (r *kubeResolver) watch() {
-	// Do an initial resolve before subscribing to async updates.
-	r.resolve()
-
+func (pw *podWatcher) watch(ctx context.Context) {
 	backoff := time.Second
 	maxBackoff := 30 * time.Second
 
+	// Resolve the pod before starting the watch loop. Retry until it
+	// succeeds so that the watch has a valid resourceVersion to start from.
+	for !pw.resolve(ctx) {
+		select {
+		case <-time.After(backoff):
+			backoff = min(backoff*2, maxBackoff)
+		case <-ctx.Done():
+			return
+		}
+	}
+	backoff = time.Second
+
 	for {
 		select {
-		case <-r.ctx.Done():
+		case <-ctx.Done():
 			return
 		default:
 		}
 
-		log.Infof("Watching pod %s/%s for updates.", r.podTarget.namespace, r.podTarget.podName)
-		watcher, err := r.client.CoreV1().Pods(r.podTarget.namespace).Watch(r.ctx, metav1.ListOptions{
-			FieldSelector:   "metadata.name=" + r.podTarget.podName,
-			ResourceVersion: r.resourceVersion,
+		log.Infof("Watching pod %s/%s for updates.", pw.namespace, pw.podName)
+		watcher, err := pw.client.CoreV1().Pods(pw.namespace).Watch(ctx, metav1.ListOptions{
+			FieldSelector:   "metadata.name=" + pw.podName,
+			ResourceVersion: pw.resourceVersion,
 		})
 		if err != nil {
-			log.Warningf("Failed to watch pod %s/%s: %s", r.podTarget.namespace, r.podTarget.podName, err)
+			log.Warningf("Failed to watch pod %s/%s: %s", pw.namespace, pw.podName, err)
 			select {
 			case <-time.After(backoff):
 				backoff = min(backoff*2, maxBackoff)
-			case <-r.ctx.Done():
+			case <-ctx.Done():
 				return
 			}
 			continue
 		}
 
-		if err := r.processWatchEvents(watcher); err != nil {
+		if err := pw.processWatchEvents(ctx, watcher); err != nil {
+			log.Warningf("Watch for pod %s/%s ended abruptly: %s", pw.namespace, pw.podName, err)
 			watcher.Stop()
-			log.Warningf("Watch for pod %s/%s ended abruptly: %s", r.podTarget.namespace, r.podTarget.podName, err)
 			select {
 			case <-time.After(backoff):
 				backoff = min(backoff*2, maxBackoff)
-			case <-r.ctx.Done():
+			case <-ctx.Done():
 				return
 			}
 			continue
@@ -209,18 +206,18 @@ func (r *kubeResolver) watch() {
 
 // processWatchEvents processes events received from the watch stream.
 // A non-nil error is returned when the watch ends abruptly.
-func (r *kubeResolver) processWatchEvents(watcher watch.Interface) error {
+func (pw *podWatcher) processWatchEvents(ctx context.Context, watcher watch.Interface) error {
 	for {
 		select {
-		case <-r.ctx.Done():
+		case <-ctx.Done():
 			return nil
 		case event, ok := <-watcher.ResultChan():
 			if !ok {
 				return fmt.Errorf("watch channel closed")
 			}
 			if obj, ok := event.Object.(metav1.ObjectMetaAccessor); ok {
 				if rv := obj.GetObjectMeta().GetResourceVersion(); rv != "" {
-					r.resourceVersion = rv
+					pw.resourceVersion = rv
 				}
 			}
 			switch event.Type {
@@ -229,32 +226,139 @@ func (r *kubeResolver) processWatchEvents(watcher watch.Interface) error {
 				if !ok {
 					continue
 				}
-				r.updateStateFromPod(pod)
+				pw.notifySubscribers(podResolution{pod: pod})
 			case watch.Deleted:
-				log.Warningf("Pod %s/%s was deleted", r.podTarget.namespace, r.podTarget.podName)
-				r.cc.ReportError(fmt.Errorf("pod %s/%s was deleted", r.podTarget.namespace, r.podTarget.podName))
+				log.Warningf("Pod %s/%s was deleted", pw.namespace, pw.podName)
+				pw.notifySubscribers(podResolution{err: fmt.Errorf("pod %s/%s was deleted", pw.namespace, pw.podName)})
 			case watch.Bookmark:
 				// No-op: bookmark events are informational.
 			case watch.Error:
 				err := errors.FromObject(event.Object)
-				log.Warningf("Watch error for pod %s/%s: %s", r.podTarget.namespace, r.podTarget.podName, err)
+				log.Warningf("Watch error for pod %s/%s: %s", pw.namespace, pw.podName, err)
 				return fmt.Errorf("watch error: %w", err)
 			}
-		case <-r.resolveCh:
-			r.resolve()
 		}
 	}
 }
 
-func (r *kubeResolver) ResolveNow(_ resolver.ResolveNowOptions) {
-	select {
-	case r.resolveCh <- struct{}{}:
-	default:
+type kubeResolverBuilder struct {
+	mu       sync.Mutex
+	client   kubernetes.Interface
+	watchers map[string]*podWatcher
+}
+
+func watcherKey(namespace, podName string) string {
+	return namespace + "/" + podName
+}
+
+func (b *kubeResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
+	endpoint := target.Endpoint()
+	pt, err := parsePodTarget(endpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.client == nil {
+		config, err := rest.InClusterConfig()
+		if err != nil {
+			return nil, fmt.Errorf("could not create k8s client: %w", err)
+		}
+		b.client, err = kubernetes.NewForConfig(config)
+		if err != nil {
+			return nil, fmt.Errorf("could not create k8s client: %w", err)
+		}
 	}
+
+	// gRPC creates one resolver per connection.
+	// To avoid making redundant calls to the k8s APIs we create a single
+	// watcher for each namespace/pod pair.
+	key := watcherKey(pt.namespace, pt.podName)
+	pw := b.watchers[key]
+	if pw == nil {
+		if b.watchers == nil {
+			b.watchers = make(map[string]*podWatcher)
+		}
+		pw = newPodWatcher(context.Background(), b.client, pt.namespace, pt.podName)
+		b.watchers[key] = pw
+	}
+
+	r := &kubeResolver{
+		cc:        cc,
+		podTarget: pt,
+		watcher:   pw,
+		builder:   b,
+	}
+
+	pw.subscribe(r, func(result podResolution) {
+		if result.err != nil {
+			r.cc.ReportError(result.err)
+			return
+		}
+		r.updateStateFromPod(result.pod)
+	})
+
+	return r, nil
+}
+
+func (b *kubeResolverBuilder) Scheme() string {
+	return scheme
+}
+
+// removeWatcher unsubscribes a resolver from its podWatcher and cleans up
+// the watcher if no subscribers remain.
+func (b *kubeResolverBuilder) removeWatcher(r *kubeResolver) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	pw := r.watcher
+	if pw.unsubscribe(r) {
+		pw.cancel()
+		delete(b.watchers, watcherKey(pw.namespace, pw.podName))
+	}
+}
+
+type kubeResolver struct {
+	cc        resolver.ClientConn
+	podTarget podTarget
+	watcher   *podWatcher
+	builder   *kubeResolverBuilder
+}
+
+func (r *kubeResolver) updateStateFromPod(pod *corev1.Pod) {
+	ip := pod.Status.PodIP
+	// This will happen when a pod restarts.
+	if ip == "" {
+		log.Infof("Pod %s/%s doesn't have an IP yet", r.podTarget.namespace, r.podTarget.podName)
+		// If we previously reported an IP, reporting an error should cause
+		// the balancer to close existing connections to the old IP.
+		r.cc.ReportError(fmt.Errorf("pod %s/%s doesn't have an IP yet", r.podTarget.namespace, r.podTarget.podName))
+		return
+	}
+
+	addr := ip
+	if r.podTarget.port != "" {
+		addr = ip + ":" + r.podTarget.port
+	}
+
+	log.Infof("Pod %s/%s resolved to IP %s", r.podTarget.namespace, r.podTarget.podName, ip)
+
+	err := r.cc.UpdateState(resolver.State{
+		Addresses: []resolver.Address{{Addr: addr}},
+	})
+	if err != nil {
+		log.Warningf("failed to update state: %s", err)
+	}
+}
+
+func (r *kubeResolver) ResolveNow(_ resolver.ResolveNowOptions) {
+	// Per the documentation, ResolveNow is a hint.
+	// We already resolve and watch the target for changes as soon as the
+	// resolver is created so there's no benefit in doing anything here.
 }
 
 func (r *kubeResolver) Close() {
-	r.cancel()
+	r.builder.removeWatcher(r)
 }
 
 func init() {