github.com/scottcagno/storage@v1.8.0/pkg/bw/writer.go (about)

     1  package bw
     2  
     3  import (
     4  	"encoding/binary"
     5  	"hash/crc32"
     6  	"io"
     7  )
     8  
     9  const (
    10  	defaultBufSize = 4096
    11  	headerSize     = 16
    12  
    13  	pageSize = 64
    14  	pageMask = pageSize - headerSize
    15  
    16  	kindGeneric = 0x01
    17  )
    18  
    19  // Writer implements buffering for an io.Writer object.
    20  // If an error occurs writing to a Writer, no more data will be
    21  // accepted and all subsequent writes, and Flush, will return the error.
    22  // After all data has been written, the client should call the
    23  // Flush method to guarantee all data has been forwarded to
    24  // the underlying io.Writer.
    25  type Writer struct {
    26  	err   error
    27  	buf   []byte
    28  	n     int
    29  	wr    io.Writer
    30  	align int
    31  }
    32  
    33  // NewWriterSize returns a new Writer whose buffer has at least the specified
    34  // size. If the argument io.Writer is already a Writer with large enough
    35  // size, it returns the underlying Writer.
    36  func NewWriterSize(w io.Writer, size int, align int) *Writer {
    37  	if size <= 0 {
    38  		size = defaultBufSize
    39  	}
    40  	return &Writer{
    41  		buf:   make([]byte, size),
    42  		wr:    w,
    43  		align: align,
    44  	}
    45  }
    46  
    47  // NewWriter returns a new Writer whose buffer has the default size.
    48  func NewWriter(w io.Writer) *Writer {
    49  	return NewWriterSize(w, defaultBufSize, 0)
    50  }
    51  
    52  // Size returns the size of the underlying buffer in bytes.
    53  func (b *Writer) Size() int { return len(b.buf) }
    54  
    55  // Reset discards any unflushed buffered data, clears any error, and
    56  // resets b to write its output to w.
    57  func (b *Writer) Reset(w io.Writer) {
    58  	b.err = nil
    59  	b.n = 0
    60  	b.wr = w
    61  }
    62  
    63  // Flush writes any buffered data to the underlying io.Writer.
    64  func (b *Writer) Flush() error {
    65  	if b.err != nil {
    66  		return b.err
    67  	}
    68  	if b.n == 0 {
    69  		return nil
    70  	}
    71  	n, err := b.wr.Write(b.buf[0:b.n])
    72  	if n < b.n && err == nil {
    73  		err = io.ErrShortWrite
    74  	}
    75  	if err != nil {
    76  		if n > 0 && n < b.n {
    77  			copy(b.buf[0:b.n-n], b.buf[n:b.n])
    78  		}
    79  		b.n -= n
    80  		b.err = err
    81  		return err
    82  	}
    83  	b.n = 0
    84  	return nil
    85  }
    86  
    87  func (b *Writer) Available() int {
    88  	return len(b.buf) - b.n
    89  }
    90  
    91  func (b *Writer) Buffered() int {
    92  	return b.n
    93  }
    94  
    95  type header struct {
    96  	kind  uint16
    97  	crc32 uint32
    98  	size  uint64
    99  }
   100  
   101  func (b *Writer) WriteHeader(buf []byte, hdr *header) (int, error) {
   102  	if len(buf) < headerSize {
   103  		return 0, io.ErrShortWrite
   104  	}
   105  	var n int
   106  	binary.LittleEndian.PutUint16(buf[n:n+2], hdr.kind)
   107  	n += 2
   108  	binary.LittleEndian.PutUint32(buf[n:n+4], hdr.crc32)
   109  	n += 4
   110  	binary.LittleEndian.PutUint64(buf[n:n+8], hdr.size)
   111  	n += 8
   112  	return n, nil
   113  }
   114  
   115  func (b *Writer) Write(p []byte) (int, error) {
   116  	nn := 0
   117  	n, err := b.WriteHeader(b.buf[b.n:b.n+headerSize], &header{
   118  		kind:  0xff,
   119  		crc32: crc32.ChecksumIEEE(p),
   120  		size:  uint64(len(p)),
   121  	})
   122  	if err != nil {
   123  		return nn, err
   124  	}
   125  	b.n += n
   126  	nn += n
   127  	for len(p)+headerSize > b.Available() && b.err == nil {
   128  		n = copy(b.buf[b.n:], p)
   129  		b.n += n
   130  		nn += n
   131  		p = p[n:]
   132  		b.Flush()
   133  	}
   134  	if b.err != nil {
   135  		return nn, b.err
   136  	}
   137  	n = copy(b.buf[b.n:], p)
   138  	b.n += n
   139  	nn += n
   140  
   141  	if pageSize != 0 && b.n%pageSize != 0 {
   142  		b.n = (b.n + pageSize - 1) &^ (pageSize - 1)
   143  	}
   144  	return b.n, nil
   145  }
   146  
   147  func (b *Writer) WriteV2(p []byte) (int, error) {
   148  	nn := 0
   149  	for len(p) > b.Available() && b.err == nil {
   150  		n := copy(b.buf[b.n:], p)
   151  		b.n += n
   152  		nn += n
   153  		p = p[n:]
   154  		b.Flush()
   155  	}
   156  	if b.err != nil {
   157  		return nn, b.err
   158  	}
   159  	n := copy(b.buf[b.n:], p)
   160  	b.n += n
   161  	nn += n
   162  	return nn, nil
   163  }