github.com/aavshr/aws-sdk-go@v1.41.3/service/s3/s3manager/writer_read_from.go (about)

     1  package s3manager
     2  
     3  import (
     4  	"bufio"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/aavshr/aws-sdk-go/internal/sdkio"
     9  )
    10  
    11  // WriterReadFrom defines an interface implementing io.Writer and io.ReaderFrom
    12  type WriterReadFrom interface {
    13  	io.Writer
    14  	io.ReaderFrom
    15  }
    16  
    17  // WriterReadFromProvider provides an implementation of io.ReadFrom for the given io.Writer
    18  type WriterReadFromProvider interface {
    19  	GetReadFrom(writer io.Writer) (w WriterReadFrom, cleanup func())
    20  }
    21  
    22  type bufferedWriter interface {
    23  	WriterReadFrom
    24  	Flush() error
    25  	Reset(io.Writer)
    26  }
    27  
    28  type bufferedReadFrom struct {
    29  	bufferedWriter
    30  }
    31  
    32  func (b *bufferedReadFrom) ReadFrom(r io.Reader) (int64, error) {
    33  	n, err := b.bufferedWriter.ReadFrom(r)
    34  	if flushErr := b.Flush(); flushErr != nil && err == nil {
    35  		err = flushErr
    36  	}
    37  	return n, err
    38  }
    39  
    40  // PooledBufferedReadFromProvider is a WriterReadFromProvider that uses a sync.Pool
    41  // to manage allocation and reuse of *bufio.Writer structures.
    42  type PooledBufferedReadFromProvider struct {
    43  	pool sync.Pool
    44  }
    45  
    46  // NewPooledBufferedWriterReadFromProvider returns a new PooledBufferedReadFromProvider
    47  // Size is used to control the size of the underlying *bufio.Writer created for
    48  // calls to GetReadFrom.
    49  func NewPooledBufferedWriterReadFromProvider(size int) *PooledBufferedReadFromProvider {
    50  	if size < int(32*sdkio.KibiByte) {
    51  		size = int(64 * sdkio.KibiByte)
    52  	}
    53  
    54  	return &PooledBufferedReadFromProvider{
    55  		pool: sync.Pool{
    56  			New: func() interface{} {
    57  				return &bufferedReadFrom{bufferedWriter: bufio.NewWriterSize(nil, size)}
    58  			},
    59  		},
    60  	}
    61  }
    62  
    63  // GetReadFrom takes an io.Writer and wraps it with a type which satisfies the WriterReadFrom
    64  // interface/ Additionally a cleanup function is provided which must be called after usage of the WriterReadFrom
    65  // has been completed in order to allow the reuse of the *bufio.Writer
    66  func (p *PooledBufferedReadFromProvider) GetReadFrom(writer io.Writer) (r WriterReadFrom, cleanup func()) {
    67  	buffer := p.pool.Get().(*bufferedReadFrom)
    68  	buffer.Reset(writer)
    69  	r = buffer
    70  	cleanup = func() {
    71  		buffer.Reset(nil) // Reset to nil writer to release reference
    72  		p.pool.Put(buffer)
    73  	}
    74  	return r, cleanup
    75  }