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

     1  package v2
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"hash/crc32"
     7  	"io"
     8  	"unsafe"
     9  )
    10  
    11  const (
    12  	blockSize  = 64
    13  	headerSize = int(unsafe.Sizeof(header{}))
    14  
    15  	uint16SZ = 2
    16  	uint32SZ = 4
    17  	uint64SZ = 8
    18  )
    19  
    20  type header struct {
    21  	kind  uint16
    22  	crc32 uint32
    23  	size  uint64
    24  }
    25  
    26  func writeHeader(buf []byte, hdr *header) (int, error) {
    27  	if len(buf) < headerSize {
    28  		return 0, io.ErrShortWrite
    29  	}
    30  	var n int
    31  	binary.LittleEndian.PutUint16(buf[n:n+uint16SZ], hdr.kind)
    32  	n += uint16SZ
    33  	binary.LittleEndian.PutUint32(buf[n:n+uint32SZ], hdr.crc32)
    34  	n += uint32SZ
    35  	binary.LittleEndian.PutUint64(buf[n:n+uint64SZ], hdr.size)
    36  	n += uint64SZ
    37  	return n, nil
    38  }
    39  
    40  type Writer struct {
    41  	buf []byte
    42  	n   int
    43  	w   io.Writer
    44  }
    45  
    46  func NewWriter(w io.Writer, size int) *Writer {
    47  	if size%blockSize != 0 {
    48  		size = (size + blockSize - 1) &^ (blockSize - 1)
    49  	}
    50  	return &Writer{
    51  		buf: make([]byte, size),
    52  		n:   0,
    53  		w:   w,
    54  	}
    55  }
    56  
    57  var ErrDataOverflow = errors.New("data overflow--too big")
    58  
    59  func (w *Writer) undoWrite(n int) {
    60  	if n > 0 && n < w.n {
    61  		copy(w.buf[0:w.n-n], w.buf[n:w.n])
    62  	}
    63  	w.n -= n
    64  }
    65  
    66  func (w *Writer) Write(p []byte) (int, error) {
    67  	// error checking
    68  	if len(p)+headerSize > len(w.buf[w.n:]) {
    69  		return 0, ErrDataOverflow
    70  	}
    71  	var nn int
    72  
    73  	// write header
    74  	n, err := writeHeader(w.buf[w.n:], &header{
    75  		kind:  0xff,
    76  		crc32: crc32.ChecksumIEEE(p),
    77  		size:  uint64(len(p)),
    78  	})
    79  	if err != nil {
    80  		w.undoWrite(n)
    81  		return n, err
    82  	}
    83  
    84  	// update offsets
    85  	w.n += n
    86  	nn += n
    87  
    88  	// write data to buffer
    89  	n = copy(w.buf[w.n:], p)
    90  	w.n += n
    91  	nn += n
    92  
    93  	// take up some slack
    94  	if w.n%blockSize != 0 {
    95  		w.n = (w.n + blockSize - 1) &^ (blockSize - 1)
    96  	}
    97  
    98  	// persist
    99  	x, err := w.w.Write(w.buf[:w.n])
   100  	if err != nil {
   101  		w.undoWrite(x)
   102  		return x, err
   103  	}
   104  
   105  	// reset buffer
   106  	w.n = 0
   107  	return x, nil
   108  }