github.com/blend/go-sdk@v1.20220411.3/bufferutil/buffer.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package bufferutil
     9  
    10  import (
    11  	"bytes"
    12  	"sync"
    13  	"time"
    14  )
    15  
    16  // NewBuffer creates a new line writer from a given set of bytes.
    17  func NewBuffer(contents []byte) *Buffer {
    18  	lw := new(Buffer)
    19  	_, _ = lw.Write(contents)
    20  	return lw
    21  }
    22  
    23  // Buffer is a writer that accepts binary but splits out onto new lines.
    24  type Buffer struct {
    25  	sync.RWMutex
    26  	// Size is the size of all bytes written to the buffer.
    27  	Size int64
    28  	// Lines are the string lines broken up by newlines with associated timestamps
    29  	Chunks []BufferChunk `json:"chunks"`
    30  	// Handler is an optional listener for new line events.
    31  	Handler BufferChunkHandler `json:"-"`
    32  }
    33  
    34  // Write writes the contents to the output buffer.
    35  // An important gotcha here is the `contents` parameter is by reference, as a result
    36  // you can get into some bad loop capture states where this buffer will
    37  // be assigned to multiple chunks but be effectively the same value.
    38  // As a result, when you write to the output buffer, we fully copy this
    39  // contents parameter for storage in the buffer.
    40  func (b *Buffer) Write(contents []byte) (written int, err error) {
    41  	chunkData := make([]byte, len(contents))
    42  	copy(chunkData, contents)
    43  
    44  	chunk := BufferChunk{Timestamp: time.Now().UTC(), Data: chunkData}
    45  	written = len(chunkData)
    46  
    47  	// lock the buffer only to add the new chunk
    48  	b.Lock()
    49  	b.Size += int64(written)
    50  	b.Chunks = append(b.Chunks, chunk)
    51  	b.Unlock()
    52  
    53  	// called outside critical section
    54  	if b.Handler != nil {
    55  		// call the handler with the chunk.
    56  		b.Handler(chunk)
    57  	}
    58  	return
    59  }
    60  
    61  // Bytes returns the bytes written to the writer.
    62  func (b *Buffer) Bytes() []byte {
    63  	b.RLock()
    64  	defer b.RUnlock()
    65  	buffer := new(bytes.Buffer)
    66  	for _, chunk := range b.Chunks {
    67  		buffer.Write(chunk.Data)
    68  	}
    69  	return buffer.Bytes()
    70  }
    71  
    72  // String returns the current combined output as a string.
    73  func (b *Buffer) String() string {
    74  	return string(b.Bytes())
    75  }