Merge pull request #142935 from wojtek-t/optimize_watch_event_allocations

Optimize watch event allocations

staging/src/k8s.io/apiserver/pkg/storage/cacher/cache_watcher.go

@@ -19,11 +19,14 @@ package cacher
 import (
 	"context"
 	"fmt"
+	"maps"
+	"reflect"
 	"sync"
 	"time"
 
 	"k8s.io/utils/clock"
 
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
@@ -354,6 +357,22 @@ func (c *cacheWatcher) isDoneChannelClosedLocked() bool {
 	return false
 }
 
+func getShallowCopyObject(object runtime.Object) runtime.Object {
+	if _, ok := object.(*cachingObject); ok {
+		return object
+	}
+	if u, ok := object.(*unstructured.Unstructured); ok {
+		return &unstructured.Unstructured{Object: maps.Clone(u.Object)}
+	}
+	v := reflect.ValueOf(object)
+	if v.Kind() == reflect.Pointer && v.Elem().Kind() == reflect.Struct {
+		cp := reflect.New(v.Type().Elem())
+		cp.Elem().Set(v.Elem())
+		return cp.Interface().(runtime.Object)
+	}
+	return object.DeepCopyObject()
+}
+
 func getMutableObject(object runtime.Object) runtime.Object {
 	if _, ok := object.(*cachingObject); ok {
 		// It is safe to return without deep-copy, because the underlying
@@ -394,9 +413,9 @@ func (c *cacheWatcher) convertToWatchEvent(event *watchCacheEvent) *watch.Event
 
 	switch {
 	case curObjPasses && !oldObjPasses:
-		return &watch.Event{Type: watch.Added, Object: getMutableObject(event.Object)}
+		return &watch.Event{Type: watch.Added, Object: getShallowCopyObject(event.Object)}
 	case curObjPasses && oldObjPasses:
-		return &watch.Event{Type: watch.Modified, Object: getMutableObject(event.Object)}
+		return &watch.Event{Type: watch.Modified, Object: getShallowCopyObject(event.Object)}
 	case !curObjPasses && oldObjPasses:
 		// return a delete event with the previous object content, but with the event's resource version
 		oldObj := getMutableObject(event.PrevObject)

staging/src/k8s.io/apiserver/pkg/storage/cacher/cache_watcher_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/fields"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -730,3 +731,57 @@ func gatherWithoutBuckets(gatherer compbasemetrics.Gatherer) testutil.GathererFu
 
 // filterAny allways passes the watch events as recursive on "" prefix with no predicate will always pass the filtering. GroupResource is just for logging.
 var filterAny watchEventFilter = newWatchEventFilter("", storage.ListOptions{Recursive: true, Predicate: storage.Everything}, schema.GroupResource{Resource: "pods"})
+
+func TestGetShallowCopyObject(t *testing.T) {
+	t.Run("typed struct", func(t *testing.T) {
+		orig := &v1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   "pod-1",
+				Labels: map[string]string{"app": "test"},
+			},
+		}
+		copied := getShallowCopyObject(orig).(*v1.Pod)
+		if copied == orig {
+			t.Fatalf("expected new struct pointer, got same pointer")
+		}
+		copied.GetObjectKind().SetGroupVersionKind(v1.SchemeGroupVersion.WithKind("Pod"))
+		if orig.APIVersion != "" || orig.Kind != "" {
+			t.Errorf("expected original TypeMeta to remain empty, got %v", orig.TypeMeta)
+		}
+		if reflect.ValueOf(copied.Labels).Pointer() != reflect.ValueOf(orig.Labels).Pointer() {
+			t.Errorf("expected shallow copy to share nested map")
+		}
+	})
+
+	t.Run("unstructured", func(t *testing.T) {
+		nested := map[string]interface{}{"key": "val"}
+		orig := &unstructured.Unstructured{
+			Object: map[string]interface{}{
+				"apiVersion": "v1",
+				"kind":       "Pod",
+				"spec":       nested,
+			},
+		}
+		copied := getShallowCopyObject(orig).(*unstructured.Unstructured)
+		if copied == orig {
+			t.Fatalf("expected new Unstructured pointer, got same pointer")
+		}
+		copied.SetGroupVersionKind(schema.GroupVersionKind{Group: "example.com", Version: "v2", Kind: "Custom"})
+		if orig.GetAPIVersion() != "v1" || orig.GetKind() != "Pod" {
+			t.Errorf("expected original GVK to remain v1/Pod, got %s/%s", orig.GetAPIVersion(), orig.GetKind())
+		}
+		if reflect.ValueOf(copied.Object["spec"]).Pointer() != reflect.ValueOf(orig.Object["spec"]).Pointer() {
+			t.Errorf("expected shallow copy to share nested spec map")
+		}
+	})
+
+	t.Run("cachingObject", func(t *testing.T) {
+		co, err := newCachingObject(&v1.Pod{})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got := getShallowCopyObject(co); got != co {
+			t.Errorf("expected same cachingObject pointer, got %#v", got)
+		}
+	})
+}

test/integration/apiserver/watchcache_test.go

@@ -17,22 +17,33 @@ limitations under the License.
 package apiserver
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
+	"net/http"
+	stdruntime "runtime"
+	"strings"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/go-logr/logr"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/serializer/protobuf"
 	"k8s.io/apimachinery/pkg/util/wait"
+	storagetesting "k8s.io/apiserver/pkg/storage/testing"
 	clientset "k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
 	"k8s.io/ktesting"
 	"k8s.io/kubernetes/cmd/kube-apiserver/app/options"
 	"k8s.io/kubernetes/pkg/controlplane"
 	"k8s.io/kubernetes/pkg/controlplane/reconcilers"
 	"k8s.io/kubernetes/test/integration/framework"
+	netutils "k8s.io/utils/net"
 )
 
 // setup create kube-apiserver backed up by two separate etcds,
@@ -228,3 +239,186 @@ func BenchmarkListFromWatchCache(b *testing.B) {
 		b.Logf("Number of secrets: %d", len(secrets.Items))
 	}
 }
+
+func BenchmarkPodListVsWatchList(b *testing.B) {
+	klog.SetLogger(logr.Discard())
+	tCtx := context.Background()
+	_, clientConfig, tearDownFn := framework.StartTestServer(tCtx, b, framework.TestServerSetup{
+		ModifyServerRunOptions: func(opts *options.ServerRunOptions) {
+			opts.GenericServerRunOptions.AdvertiseAddress = netutils.ParseIPSloppy("127.0.0.1")
+			opts.Admission.GenericAdmission.DisablePlugins = []string{"ServiceAccount"}
+		},
+		ModifyServerConfig: func(config *controlplane.Config) {
+			config.Extra.EndpointReconcilerType = reconcilers.NoneEndpointReconcilerType
+		},
+	})
+	defer tearDownFn()
+
+	cfg := rest.CopyConfig(clientConfig)
+	cfg.ContentConfig.AcceptContentTypes = runtime.ContentTypeProtobuf
+	cfg.ContentConfig.ContentType = runtime.ContentTypeProtobuf
+	cfg.QPS = -1
+	c, err := clientset.NewForConfig(cfg)
+	if err != nil {
+		b.Fatal(err)
+	}
+	httpClient, err := rest.HTTPClientFor(cfg)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	namespaces := 100
+	podsPerNamespace := 100
+	totalPods := namespaces * podsPerNamespace
+	data := storagetesting.PrepareBenchmarkData(namespaces, podsPerNamespace, 50)
+
+	for _, nsName := range data.NamespaceNames {
+		ns := &v1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{Name: nsName},
+		}
+		if _, err := c.CoreV1().Namespaces().Create(tCtx, ns, metav1.CreateOptions{}); err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	var wg sync.WaitGroup
+	errCh := make(chan error, totalPods)
+	sem := make(chan struct{}, 50)
+	for _, pod := range data.Pods {
+		wg.Add(1)
+		sem <- struct{}{}
+		go func(p *v1.Pod) {
+			defer wg.Done()
+			defer func() { <-sem }()
+			status := p.Status
+			created, err := c.CoreV1().Pods(p.Namespace).Create(tCtx, p, metav1.CreateOptions{FieldManager: "kube-controller-manager"})
+			if err != nil {
+				errCh <- err
+				return
+			}
+			created.Status = status
+			if _, err := c.CoreV1().Pods(p.Namespace).UpdateStatus(tCtx, created, metav1.UpdateOptions{FieldManager: "kubelet"}); err != nil {
+				errCh <- err
+				return
+			}
+		}(pod)
+	}
+	wg.Wait()
+	close(errCh)
+	for err := range errCh {
+		b.Fatal(err)
+	}
+
+	// Wait until all pods are present in the watch cache.
+	if err := wait.PollUntilContextTimeout(tCtx, 100*time.Millisecond, 30*time.Second, true, func(ctx context.Context) (bool, error) {
+		list, err := c.CoreV1().Pods("").List(ctx, metav1.ListOptions{ResourceVersion: "0"})
+		if err != nil {
+			return false, err
+		}
+		return len(list.Items) == totalPods, nil
+	}); err != nil {
+		b.Fatalf("failed waiting for watch cache to sync %d pods: %v", totalPods, err)
+	}
+
+	// Measure single pod protobuf wire size by fetching one pod raw.
+	firstPod := data.Pods[0]
+	singleRaw, err := c.CoreV1().RESTClient().Get().Namespace(firstPod.Namespace).Resource("pods").Name(firstPod.Name).DoRaw(tCtx)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.Logf("Single pod protobuf wire size: %d bytes (%d pods total = %.2f MB)", len(singleRaw), totalPods, float64(len(singleRaw)*totalPods)/(1024*1024))
+
+	listURL := strings.TrimRight(cfg.Host, "/") + "/api/v1/pods?resourceVersion=0"
+	watchListURL := strings.TrimRight(cfg.Host, "/") + "/api/v1/pods?watch=true&sendInitialEvents=true&allowWatchBookmarks=true&resourceVersion=0&resourceVersionMatch=NotOlderThan"
+
+	runList := func(copyBuf []byte) {
+		req, err := http.NewRequestWithContext(tCtx, http.MethodGet, listURL, nil)
+		if err != nil {
+			b.Fatal(err)
+		}
+		req.Header.Set("Accept", runtime.ContentTypeProtobuf)
+		req.Header.Set("Accept-Encoding", "identity")
+		resp, err := httpClient.Do(req)
+		if err != nil {
+			b.Fatal(err)
+		}
+		defer func() {
+			if err := resp.Body.Close(); err != nil {
+				b.Fatal(err)
+			}
+		}()
+		if resp.StatusCode != http.StatusOK {
+			b.Fatalf("unexpected status code: %d", resp.StatusCode)
+		}
+		n, err := io.CopyBuffer(io.Discard, resp.Body, copyBuf)
+		if err != nil {
+			b.Fatal(err)
+		}
+		if n < int64(totalPods*len(singleRaw)/2) {
+			b.Fatalf("unexpectedly small list response: %d bytes", n)
+		}
+	}
+
+	bookmarkMarker := []byte(metav1.InitialEventsAnnotationKey)
+	runWatchList := func(frameBuf []byte) {
+		reqCtx, cancel := context.WithCancel(tCtx)
+		defer cancel()
+		req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, watchListURL, nil)
+		if err != nil {
+			b.Fatal(err)
+		}
+		req.Header.Set("Accept", runtime.ContentTypeProtobuf)
+		req.Header.Set("Accept-Encoding", "identity")
+		resp, err := httpClient.Do(req)
+		if err != nil {
+			b.Fatal(err)
+		}
+		defer func() {
+			if err := resp.Body.Close(); err != nil {
+				b.Fatal(err)
+			}
+		}()
+		if resp.StatusCode != http.StatusOK {
+			b.Fatalf("unexpected status code: %d", resp.StatusCode)
+		}
+		frameReader := protobuf.LengthDelimitedFramer.NewFrameReader(resp.Body)
+		received := 0
+		for {
+			n, err := frameReader.Read(frameBuf)
+			if err != nil {
+				b.Fatalf("error reading watch frame after %d events: %v", received, err)
+			}
+			if bytes.Contains(frameBuf[:n], bookmarkMarker) {
+				break
+			}
+			received++
+		}
+		if received != totalPods {
+			b.Fatalf("expected %d initial events before bookmark, got %d", totalPods, received)
+		}
+	}
+
+	b.Run("List", func(b *testing.B) {
+		b.ReportAllocs()
+		copyBuf := make([]byte, 64*1024)
+		runList(copyBuf)
+		stdruntime.GC()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			runList(copyBuf)
+		}
+		b.ReportMetric(float64(len(singleRaw)), "pod-proto-B")
+	})
+
+	b.Run("WatchList", func(b *testing.B) {
+		b.ReportAllocs()
+		frameBuf := make([]byte, 64*1024)
+		runWatchList(frameBuf)
+		stdruntime.GC()
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			runWatchList(frameBuf)
+		}
+		b.ReportMetric(float64(len(singleRaw)), "pod-proto-B")
+	})
+}