From f50872f1bcf2592c31d3de2fb1cf0b18120adab3 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Fri, 6 Oct 2023 15:05:30 +0300
Subject: [PATCH] [#58] object: Allow to set marshal data

Now it is possible set marshaled data to reduce memory
allocations.

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
---
 object/marshal.go        | 48 ++++++++++++++++++++++++++++++++++++++++
 object/types.go          |  6 +++++
 util/proto/marshal.go    | 25 +++++++++++++++++++++
 util/signature/buffer.go |  7 ++++++
 util/signature/data.go   | 12 ++++++----
 5 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/object/marshal.go b/object/marshal.go
index fe0e232a..18172e67 100644
--- a/object/marshal.go
+++ b/object/marshal.go
@@ -322,6 +322,14 @@ func (o *Object) StableMarshal(buf []byte) []byte {
 		return []byte{}
 	}
 
+	if o.marshalData != nil {
+		if buf == nil {
+			return o.marshalData
+		}
+		copy(buf, o.marshalData)
+		return buf
+	}
+
 	if buf == nil {
 		buf = make([]byte, o.StableSize())
 	}
@@ -336,6 +344,20 @@ func (o *Object) StableMarshal(buf []byte) []byte {
 	return buf
 }
 
+// SetMarshalData sets marshal data to reduce memory allocations.
+//
+// It is unsafe to modify object data after setting marshal data.
+func (o *Object) SetMarshalData(data []byte) {
+	if o == nil {
+		return
+	}
+	if data == nil {
+		o.marshalData = o.StableMarshal(nil)
+	} else {
+		o.marshalData = data
+	}
+}
+
 func (o *Object) StableSize() (size int) {
 	if o == nil {
 		return 0
@@ -1071,6 +1093,15 @@ func (r *PutSingleRequestBody) StableMarshal(buf []byte) []byte {
 	if r == nil {
 		return []byte{}
 	}
+
+	if r.marshalData != nil {
+		if buf == nil {
+			return r.marshalData
+		}
+		copy(buf, r.marshalData)
+		return buf
+	}
+
 	if buf == nil {
 		buf = make([]byte, r.StableSize())
 	}
@@ -1082,6 +1113,23 @@ func (r *PutSingleRequestBody) StableMarshal(buf []byte) []byte {
 	return buf
 }
 
+// SetMarshalData sets marshal data to reduce memory allocations.
+//
+// It is unsafe to modify request data after setting marshal data.
+func (r *PutSingleRequestBody) SetMarshalData(data []byte) {
+	if r == nil {
+		return
+	}
+
+	if data == nil {
+		r.marshalData = r.StableMarshal(nil)
+	} else {
+		r.marshalData = data
+	}
+
+	proto.NestedStructureSetMarshalData(putSingleReqObjectField, r.marshalData, r.object)
+}
+
 func (r *PutSingleRequestBody) StableSize() int {
 	if r == nil {
 		return 0
diff --git a/object/types.go b/object/types.go
index c28bf357..895afb0b 100644
--- a/object/types.go
+++ b/object/types.go
@@ -75,6 +75,9 @@ type Object struct {
 	header *Header
 
 	payload []byte
+
+	// marshalData holds marshaled data, must not be marshaled by StableMarshal
+	marshalData []byte
 }
 
 type SplitInfo struct {
@@ -304,6 +307,9 @@ type GetRangeHashResponse struct {
 type PutSingleRequestBody struct {
 	object  *Object
 	copyNum []uint32
+
+	// marshalData holds marshaled data, must not be marshaled by StableMarshal
+	marshalData []byte
 }
 
 type PutSingleRequest struct {
diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index 606086aa..b16375ad 100644
--- a/util/proto/marshal.go
+++ b/util/proto/marshal.go
@@ -19,6 +19,11 @@ type (
 		StableMarshal([]byte) []byte
 		StableSize() int
 	}
+
+	setMarshalData interface {
+		SetMarshalData([]byte)
+		StableSize() int
+	}
 )
 
 func BytesMarshal(field int, buf, v []byte) int {
@@ -263,6 +268,26 @@ func NestedStructureMarshal[T stableMarshaller](field int64, buf []byte, v T) in
 	return offset + n
 }
 
+// NestedStructureSetMarshalData calculates offset for field in parentData
+// and calls SetMarshalData for nested structure.
+//
+// Returns marshalled data length of nested structure.
+func NestedStructureSetMarshalData(field int64, parentData []byte, v setMarshalData) int {
+	n := v.StableSize()
+	if n == 0 {
+		return 0
+	}
+
+	buf := make([]byte, binary.MaxVarintLen64)
+	prefix := protowire.EncodeTag(protowire.Number(field), protowire.BytesType)
+	offset := binary.PutUvarint(buf, prefix)
+	offset += binary.PutUvarint(buf, uint64(n))
+
+	v.SetMarshalData(parentData[offset : offset+n])
+
+	return offset + n
+}
+
 func NestedStructureSize[T stableMarshaller](field int64, v T) (size int) {
 	n := v.StableSize()
 	if n == 0 {
diff --git a/util/signature/buffer.go b/util/signature/buffer.go
index f9b003db..8102495d 100644
--- a/util/signature/buffer.go
+++ b/util/signature/buffer.go
@@ -14,6 +14,13 @@ var buffersPool = sync.Pool{
 	},
 }
 
+func tryGetNewBufferFromPool(size int) (*buffer, bool) {
+	if size > poolSliceMaxSize {
+		return &buffer{}, false
+	}
+	return newBufferFromPool(size), true
+}
+
 func newBufferFromPool(size int) *buffer {
 	result := buffersPool.Get().(*buffer)
 	if cap(result.data) < size {
diff --git a/util/signature/data.go b/util/signature/data.go
index 1db46dce..4170e6c7 100644
--- a/util/signature/data.go
+++ b/util/signature/data.go
@@ -35,8 +35,10 @@ func SignDataWithHandler(key *ecdsa.PrivateKey, src DataSource, handler KeySigna
 		opts[i](cfg)
 	}
 
-	buffer := newBufferFromPool(src.SignedDataSize())
-	defer returnBufferToPool(buffer)
+	buffer, ok := tryGetNewBufferFromPool(src.SignedDataSize())
+	if ok {
+		defer returnBufferToPool(buffer)
+	}
 
 	data, err := src.ReadSignedData(buffer.data)
 	if err != nil {
@@ -64,8 +66,10 @@ func VerifyDataWithSource(dataSrc DataSource, sigSrc KeySignatureSource, opts ..
 		opts[i](cfg)
 	}
 
-	buffer := newBufferFromPool(dataSrc.SignedDataSize())
-	defer returnBufferToPool(buffer)
+	buffer, ok := tryGetNewBufferFromPool(dataSrc.SignedDataSize())
+	if ok {
+		defer returnBufferToPool(buffer)
+	}
 
 	data, err := dataSrc.ReadSignedData(buffer.data)
 	if err != nil {