github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/topic/topicreaderinternal/message_content_pool.go (about)

     1  package topicreaderinternal
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    10  )
    11  
    12  // Pool is interface for sync.Pool and may be extended by follow to original type
    13  //
    14  //go:generate mockgen -destination=pool_interface_mock_test.go -write_package_comment=false -package=topicreaderinternal github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreaderinternal Pool
    15  type Pool interface {
    16  	Get() interface{}
    17  	Put(x interface{})
    18  }
    19  
    20  // CallbackWithMessageContentFunc is callback function for work with message content
    21  // data bytes MUST NOT be used after f returned
    22  // if you need content longer - copy content to other slice
    23  type CallbackWithMessageContentFunc func(data []byte) error
    24  
    25  const (
    26  	minInitializeBufferSize = bytes.MinRead * 2
    27  	maxInitialBufferSize    = 1024 * 1024 * 50 // protection from bad UncompressedSize from stream
    28  )
    29  
    30  var globalReadMessagePool = &sync.Pool{}
    31  
    32  func callbackOnReaderContent(
    33  	p Pool,
    34  	reader io.Reader,
    35  	estinamedSize int,
    36  	consumer PublicMessageContentUnmarshaler,
    37  ) error {
    38  	var buf *bytes.Buffer
    39  	var ok bool
    40  	if buf, ok = p.Get().(*bytes.Buffer); !ok {
    41  		buf = &bytes.Buffer{}
    42  	}
    43  	defer func() {
    44  		buf.Reset()
    45  		p.Put(buf)
    46  	}()
    47  
    48  	// + bytes.MinRead need for prevent additional allocation for read io.EOF after read message content
    49  	targetSize := estinamedSize + bytes.MinRead
    50  	switch {
    51  	case targetSize < minInitializeBufferSize:
    52  		buf.Grow(minInitializeBufferSize)
    53  	case targetSize > maxInitialBufferSize:
    54  		buf.Grow(maxInitialBufferSize)
    55  	default:
    56  		buf.Grow(targetSize)
    57  	}
    58  
    59  	if _, err := buf.ReadFrom(reader); err != nil {
    60  		return xerrors.WithStackTrace(fmt.Errorf("ydb: error read from data reader: %w", err))
    61  	}
    62  
    63  	if err := consumer.UnmarshalYDBTopicMessage(buf.Bytes()); err != nil {
    64  		return xerrors.WithStackTrace(fmt.Errorf("ydb: error unmarshal data: %w", err))
    65  	}
    66  
    67  	return nil
    68  }