Repository: dolthub/dolt
File: go/store/nbs/byte_sink.go (Go)

Code before:

// FlushToFile writes all the data that was written to the ByteSink to a file at the given path
func (sink *FixedBufferByteSink) FlushToFile(path string) (err error) {
	return flushSinkToFile(sink, path)
}

func (sink *FixedBufferByteSink) Reader() (io.ReadCloser, error) {
	return io.NopCloser(bytes.NewReader(sink.buff)), nil
}

// BlockBufferByteSink allocates blocks of data with a given block size to store the bytes written to the sink. New
// blocks are allocated as needed in order to handle all the data of the Write calls.
type BlockBufferByteSink struct {
	blocks    [][]byte
	blockSize int
	pos       uint64

	ctx context.Context
	q   MemoryQuotaProvider
}

// NewBlockBufferByteSink creates a BlockBufferByteSink with the provided block size.
func NewBlockBufferByteSink(ctx context.Context, blockSize int, q MemoryQuotaProvider) (*BlockBufferByteSink, error) {
	block, err := q.AcquireQuotaByteSlice(ctx, blockSize)
	if err != nil {
		return nil, err
	}
	block = block[:0]
	return &BlockBufferByteSink{
		ctx:       ctx,
		q:         q,
		pos:       0,
		blockSize: blockSize,
		blocks:    [][]byte{block},
	}, nil
}

// Write writes a byte array to the sink.
func (sink *BlockBufferByteSink) Write(src []byte) (int, error) {
	srcLen := len(src)
	currBlockIdx := len(sink.blocks) - 1
	currBlock := sink.blocks[currBlockIdx]
	remaining := cap(currBlock) - len(currBlock)

	if remaining >= srcLen {
		currBlock = append(currBlock, src...)
		sink.blocks[currBlockIdx] = currBlock
	} else {
		if remaining > 0 {
			currBlock = append(currBlock, src[:remaining]...)
			sink.blocks[currBlockIdx] = currBlock

Review comment (suggestion):
Maybe worth a comment here stating that you need to call Close() on this returned value. Not sure how to protect against leaking otherwise, and it's not immediately obvious that it's necessary. 

Code after:

// FlushToFile writes all the data that was written to the ByteSink to a file at the given path
func (sink *FixedBufferByteSink) FlushToFile(path string) (err error) {
	return flushSinkToFile(sink, path)
}

func (sink *FixedBufferByteSink) Reader() (io.ReadCloser, error) {
	return io.NopCloser(bytes.NewReader(sink.buff)), nil
}

// BlockBufferByteSink allocates blocks of data with a given block size to store the bytes written to the sink. New
// blocks are allocated as needed in order to handle all the data of the Write calls.
type BlockBufferByteSink struct {
	blocks    [][]byte
	blockSize int
	pos       uint64

	ctx context.Context
	q   MemoryQuotaProvider
}

// NewBlockBufferByteSink creates a BlockBufferByteSink with the provided block size.
//
// A BlockBufferByteSink acquires memory from |q|. In order to release the memory,
// it must be |Close|d.
func NewBlockBufferByteSink(ctx context.Context, blockSize int, q MemoryQuotaProvider) (*BlockBufferByteSink, error) {
	block, err := q.AcquireQuotaByteSlice(ctx, blockSize)
	if err != nil {
		return nil, err
	}
	block = block[:0]
	return &BlockBufferByteSink{
		ctx:       ctx,
		q:         q,
		pos:       0,
		blockSize: blockSize,
		blocks:    [][]byte{block},
	}, nil
}

// Write writes a byte array to the sink.
func (sink *BlockBufferByteSink) Write(src []byte) (int, error) {
	srcLen := len(src)
	currBlockIdx := len(sink.blocks) - 1
	currBlock := sink.blocks[currBlockIdx]
	remaining := cap(currBlock) - len(currBlock)

	if remaining >= srcLen {
		currBlock = append(currBlock, src...)
		sink.blocks[currBlockIdx] = currBlock
	} else {