diff --git a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf.go b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf.go index 8358d77c39e..c63e6dc63f6 100644 --- a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf.go +++ b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf.go @@ -30,6 +30,7 @@ import ( "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/runtime/serializer/recognizer" "k8s.io/apimachinery/pkg/util/framer" + "k8s.io/klog/v2" ) var ( @@ -86,6 +87,7 @@ type Serializer struct { } var _ runtime.Serializer = &Serializer{} +var _ runtime.EncoderWithAllocator = &Serializer{} var _ recognizer.RecognizingDecoder = &Serializer{} const serializerIdentifier runtime.Identifier = "protobuf" @@ -161,22 +163,36 @@ func (s *Serializer) Decode(originalData []byte, gvk *schema.GroupVersionKind, i return unmarshalToObject(s.typer, s.creater, &actual, into, unk.Raw) } -// Encode serializes the provided object to the given writer. -func (s *Serializer) Encode(obj runtime.Object, w io.Writer) error { - if co, ok := obj.(runtime.CacheableObject); ok { - return co.CacheEncode(s.Identifier(), s.doEncode, w) - } - return s.doEncode(obj, w) +// EncodeWithAllocator writes an object to the provided writer. +// In addition, it allows for providing a memory allocator for efficient memory usage during object serialization. +func (s *Serializer) EncodeWithAllocator(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + return s.encode(obj, w, memAlloc) } -func (s *Serializer) doEncode(obj runtime.Object, w io.Writer) error { +// Encode serializes the provided object to the given writer. +func (s *Serializer) Encode(obj runtime.Object, w io.Writer) error { + return s.encode(obj, w, &runtime.SimpleAllocator{}) +} + +func (s *Serializer) encode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + if co, ok := obj.(runtime.CacheableObject); ok { + return co.CacheEncode(s.Identifier(), func(obj runtime.Object, w io.Writer) error { return s.doEncode(obj, w, memAlloc) }, w) + } + return s.doEncode(obj, w, memAlloc) +} + +func (s *Serializer) doEncode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + if memAlloc == nil { + klog.Error("a mandatory memory allocator wasn't provided, this might have a negative impact on performance, check invocations of EncodeWithAllocator method, falling back on runtime.SimpleAllocator") + memAlloc = &runtime.SimpleAllocator{} + } prefixSize := uint64(len(s.prefix)) var unk runtime.Unknown switch t := obj.(type) { case *runtime.Unknown: estimatedSize := prefixSize + uint64(t.Size()) - data := make([]byte, estimatedSize) + data := memAlloc.Allocate(estimatedSize) i, err := t.MarshalTo(data[prefixSize:]) if err != nil { return err @@ -196,11 +212,11 @@ func (s *Serializer) doEncode(obj runtime.Object, w io.Writer) error { switch t := obj.(type) { case bufferedMarshaller: - // this path performs a single allocation during write but requires the caller to implement - // the more efficient Size and MarshalToSizedBuffer methods + // this path performs a single allocation during write only when the Allocator wasn't provided + // it also requires the caller to implement the more efficient Size and MarshalToSizedBuffer methods encodedSize := uint64(t.Size()) estimatedSize := prefixSize + estimateUnknownSize(&unk, encodedSize) - data := make([]byte, estimatedSize) + data := memAlloc.Allocate(estimatedSize) i, err := unk.NestedMarshalTo(data[prefixSize:], t, encodedSize) if err != nil { @@ -221,7 +237,7 @@ func (s *Serializer) doEncode(obj runtime.Object, w io.Writer) error { unk.Raw = data estimatedSize := prefixSize + uint64(unk.Size()) - data = make([]byte, estimatedSize) + data = memAlloc.Allocate(estimatedSize) i, err := unk.MarshalTo(data[prefixSize:]) if err != nil { @@ -395,19 +411,33 @@ func unmarshalToObject(typer runtime.ObjectTyper, creater runtime.ObjectCreater, // Encode serializes the provided object to the given writer. Overrides is ignored. func (s *RawSerializer) Encode(obj runtime.Object, w io.Writer) error { - if co, ok := obj.(runtime.CacheableObject); ok { - return co.CacheEncode(s.Identifier(), s.doEncode, w) - } - return s.doEncode(obj, w) + return s.encode(obj, w, &runtime.SimpleAllocator{}) } -func (s *RawSerializer) doEncode(obj runtime.Object, w io.Writer) error { +// EncodeWithAllocator writes an object to the provided writer. +// In addition, it allows for providing a memory allocator for efficient memory usage during object serialization. +func (s *RawSerializer) EncodeWithAllocator(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + return s.encode(obj, w, memAlloc) +} + +func (s *RawSerializer) encode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + if co, ok := obj.(runtime.CacheableObject); ok { + return co.CacheEncode(s.Identifier(), func(obj runtime.Object, w io.Writer) error { return s.doEncode(obj, w, memAlloc) }, w) + } + return s.doEncode(obj, w, memAlloc) +} + +func (s *RawSerializer) doEncode(obj runtime.Object, w io.Writer, memAlloc runtime.MemoryAllocator) error { + if memAlloc == nil { + klog.Error("a mandatory memory allocator wasn't provided, this might have a negative impact on performance, check invocations of EncodeWithAllocator method, falling back on runtime.SimpleAllocator") + memAlloc = &runtime.SimpleAllocator{} + } switch t := obj.(type) { case bufferedReverseMarshaller: - // this path performs a single allocation during write but requires the caller to implement - // the more efficient Size and MarshalToSizedBuffer methods + // this path performs a single allocation during write only when the Allocator wasn't provided + // it also requires the caller to implement the more efficient Size and MarshalToSizedBuffer methods encodedSize := uint64(t.Size()) - data := make([]byte, encodedSize) + data := memAlloc.Allocate(encodedSize) n, err := t.MarshalToSizedBuffer(data) if err != nil { @@ -417,10 +447,10 @@ func (s *RawSerializer) doEncode(obj runtime.Object, w io.Writer) error { return err case bufferedMarshaller: - // this path performs a single allocation during write but requires the caller to implement - // the more efficient Size and MarshalTo methods + // this path performs a single allocation during write only when the Allocator wasn't provided + // it also requires the caller to implement the more efficient Size and MarshalTo methods encodedSize := uint64(t.Size()) - data := make([]byte, encodedSize) + data := memAlloc.Allocate(encodedSize) n, err := t.MarshalTo(data) if err != nil { diff --git a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf_test.go b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf_test.go index 001b0f64786..27f4496d18a 100644 --- a/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf_test.go +++ b/staging/src/k8s.io/apimachinery/pkg/runtime/serializer/protobuf/protobuf_test.go @@ -17,8 +17,12 @@ limitations under the License. package protobuf import ( + "bytes" + "reflect" "testing" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + testapigroupv1 "k8s.io/apimachinery/pkg/apis/testapigroup/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" runtimetesting "k8s.io/apimachinery/pkg/runtime/testing" @@ -66,3 +70,111 @@ func (t *mockTyper) ObjectKinds(obj runtime.Object) ([]schema.GroupVersionKind, func (t *mockTyper) Recognizes(_ schema.GroupVersionKind) bool { return false } + +func TestSerializerEncodeWithAllocator(t *testing.T) { + testCases := []struct { + name string + obj runtime.Object + }{ + { + name: "encode a bufferedMarshaller obj", + obj: &testapigroupv1.Carp{ + TypeMeta: metav1.TypeMeta{APIVersion: "group/version", Kind: "Carp"}, + ObjectMeta: metav1.ObjectMeta{ + Name: "name", + Namespace: "namespace", + }, + Spec: testapigroupv1.CarpSpec{ + Subdomain: "carp.k8s.io", + }, + }, + }, + + { + name: "encode a runtime.Unknown obj", + obj: &runtime.Unknown{TypeMeta: runtime.TypeMeta{APIVersion: "group/version", Kind: "Unknown"}, Raw: []byte("hello world")}, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + target := NewSerializer(nil, nil) + + writer := &bytes.Buffer{} + if err := target.Encode(tc.obj, writer); err != nil { + t.Fatal(err) + } + + writer2 := &bytes.Buffer{} + alloc := &testAllocator{} + if err := target.EncodeWithAllocator(tc.obj, writer2, alloc); err != nil { + t.Fatal(err) + } + if alloc.allocateCount != 1 { + t.Fatalf("expected the Allocate method to be called exactly 1 but it was executed: %v times ", alloc.allocateCount) + } + + // to ensure compatibility of the new method with the old one, serialized data must be equal + // also we are not testing decoding since "roundtripping" is tested elsewhere for all known types + if !reflect.DeepEqual(writer.Bytes(), writer2.Bytes()) { + t.Fatal("data mismatch, data serialized with the Encode method is different than serialized with the EncodeWithAllocator method") + } + }) + } +} + +func TestRawSerializerEncodeWithAllocator(t *testing.T) { + testCases := []struct { + name string + obj runtime.Object + }{ + { + name: "encode a bufferedReverseMarshaller obj", + obj: &testapigroupv1.Carp{ + TypeMeta: metav1.TypeMeta{APIVersion: "group/version", Kind: "Carp"}, + ObjectMeta: metav1.ObjectMeta{ + Name: "name", + Namespace: "namespace", + }, + Spec: testapigroupv1.CarpSpec{ + Subdomain: "carp.k8s.io", + }, + }, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + writer := &bytes.Buffer{} + target := NewRawSerializer(nil, nil) + + if err := target.Encode(tc.obj, writer); err != nil { + t.Fatal(err) + } + + writer2 := &bytes.Buffer{} + alloc := &testAllocator{} + if err := target.EncodeWithAllocator(tc.obj, writer2, alloc); err != nil { + t.Fatal(err) + } + if alloc.allocateCount != 1 { + t.Fatalf("expected the Allocate method to be called exactly 1 but it was executed: %v times ", alloc.allocateCount) + } + + // to ensure compatibility of the new method with the old one, serialized data must be equal + // also we are not testing decoding since "roundtripping" is tested elsewhere for all known types + if !reflect.DeepEqual(writer.Bytes(), writer2.Bytes()) { + t.Fatal("data mismatch, data serialized with the Encode method is different than serialized with the EncodeWithAllocator method") + } + }) + } +} + +type testAllocator struct { + buf []byte + allocateCount int +} + +func (ta *testAllocator) Allocate(n uint64) []byte { + ta.buf = make([]byte, n, n) + ta.allocateCount++ + return ta.buf +}