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 }