From 81cf09675166a082949c00d8c853400f26fbf0c9 Mon Sep 17 00:00:00 2001
From: Lukasz Szaszkiewicz <lukasz.szaszkiewicz@gmail.com>
Date: Thu, 17 Feb 2022 16:05:25 +0100
Subject: [PATCH] codec: exposes EncodeWithAllocator method

The new method is implemented by the protobuf serializer and helps to reduce memory footprint during object serialization.
---
 .../serializer/versioning/versioning.go       | 43 ++++++++++++++-----
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/versioning/versioning.go b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/versioning/versioning.go
index ea7c580bd6b..e637f8b8e00 100644
--- a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/versioning/versioning.go
+++ b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/versioning/versioning.go
@@ -89,6 +89,8 @@ type codec struct {
 	originalSchemeName string
 }
 
+var _ runtime.EncoderWithAllocator = &codec{}
+
 var identifiersMap sync.Map
 
 type codecIdentifier struct {
@@ -182,19 +184,40 @@ func (c *codec) Decode(data []byte, defaultGVK *schema.GroupVersionKind, into ru
 	return out, gvk, strictDecodingErr
 }
 
+// EncodeWithAllocator ensures the provided object is output in the appropriate group and version, invoking
+// conversion if necessary. Unversioned objects (according to the ObjectTyper) are output as is.
+// In addition, it allows for providing a memory allocator for efficient memory usage during object serialization.
+func (c *codec) EncodeWithAllocator(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error {
+	return c.encode(obj, w, memAlloc)
+}
+
 // Encode ensures the provided object is output in the appropriate group and version, invoking
 // conversion if necessary. Unversioned objects (according to the ObjectTyper) are output as is.
 func (c *codec) Encode(obj runtime.Object, w io.Writer) error {
-	if co, ok := obj.(runtime.CacheableObject); ok {
-		return co.CacheEncode(c.Identifier(), c.doEncode, w)
-	}
-	return c.doEncode(obj, w)
+	return c.encode(obj, w, nil)
 }
 
-func (c *codec) doEncode(obj runtime.Object, w io.Writer) error {
+func (c *codec) encode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error {
+	if co, ok := obj.(runtime.CacheableObject); ok {
+		return co.CacheEncode(c.Identifier(), func(obj runtime.Object, w io.Writer) error { return c.doEncode(obj, w, memAlloc) }, w)
+	}
+	return c.doEncode(obj, w, memAlloc)
+}
+
+func (c *codec) doEncode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error {
+	encodeFn := c.encoder.Encode
+	if memAlloc != nil {
+		if encoder, supportsAllocator := c.encoder.(runtime.EncoderWithAllocator); supportsAllocator {
+			encodeFn = func(obj runtime.Object, w io.Writer) error {
+				return encoder.EncodeWithAllocator(obj, w, memAlloc)
+			}
+		} else {
+			klog.V(4).Infof("a memory allocator was provided but the encoder %T doesn't implement the runtime.EncoderWithAllocator, using regular encoder.Encode method")
+		}
+	}
 	switch obj := obj.(type) {
 	case *runtime.Unknown:
-		return c.encoder.Encode(obj, w)
+		return encodeFn(obj, w)
 	case runtime.Unstructured:
 		// An unstructured list can contain objects of multiple group version kinds. don't short-circuit just
 		// because the top-level type matches our desired destination type. actually send the object to the converter
@@ -203,14 +226,14 @@ func (c *codec) doEncode(obj runtime.Object, w io.Writer) error {
 			// avoid conversion roundtrip if GVK is the right one already or is empty (yes, this is a hack, but the old behaviour we rely on in kubectl)
 			objGVK := obj.GetObjectKind().GroupVersionKind()
 			if len(objGVK.Version) == 0 {
-				return c.encoder.Encode(obj, w)
+				return encodeFn(obj, w)
 			}
 			targetGVK, ok := c.encodeVersion.KindForGroupVersionKinds([]schema.GroupVersionKind{objGVK})
 			if !ok {
 				return runtime.NewNotRegisteredGVKErrForTarget(c.originalSchemeName, objGVK, c.encodeVersion)
 			}
 			if targetGVK == objGVK {
-				return c.encoder.Encode(obj, w)
+				return encodeFn(obj, w)
 			}
 		}
 	}
@@ -232,7 +255,7 @@ func (c *codec) doEncode(obj runtime.Object, w io.Writer) error {
 			}
 		}
 		objectKind.SetGroupVersionKind(gvks[0])
-		return c.encoder.Encode(obj, w)
+		return encodeFn(obj, w)
 	}
 
 	// Perform a conversion if necessary
@@ -248,7 +271,7 @@ func (c *codec) doEncode(obj runtime.Object, w io.Writer) error {
 	}
 
 	// Conversion is responsible for setting the proper group, version, and kind onto the outgoing object
-	return c.encoder.Encode(out, w)
+	return encodeFn(out, w)
 }
 
 // Identifier implements runtime.Encoder interface.