diff options
Diffstat (limited to 'weed/util/buffered_writer/buffered_writer.go')
| -rw-r--r-- | weed/util/buffered_writer/buffered_writer.go | 52 |
1 files changed, 52 insertions, 0 deletions
diff --git a/weed/util/buffered_writer/buffered_writer.go b/weed/util/buffered_writer/buffered_writer.go new file mode 100644 index 000000000..73d9f4995 --- /dev/null +++ b/weed/util/buffered_writer/buffered_writer.go @@ -0,0 +1,52 @@ +package buffered_writer + +import ( + "bytes" + "io" +) + +var _ = io.WriteCloser(&BufferedWriteCloser{}) + +type BufferedWriteCloser struct { + buffer bytes.Buffer + bufferLimit int + position int64 + nextFlushOffset int64 + FlushFunc func([]byte, int64) error + CloseFunc func() error +} + +func NewBufferedWriteCloser(bufferLimit int) *BufferedWriteCloser { + return &BufferedWriteCloser{ + bufferLimit: bufferLimit, + } +} + +func (b *BufferedWriteCloser) Write(p []byte) (n int, err error) { + + if b.buffer.Len()+len(p) >= b.bufferLimit { + if err := b.FlushFunc(b.buffer.Bytes(), b.nextFlushOffset); err != nil { + return 0, err + } + b.nextFlushOffset += int64(b.buffer.Len()) + b.buffer.Reset() + } + + return b.buffer.Write(p) + +} + +func (b *BufferedWriteCloser) Close() error { + if b.buffer.Len() > 0 { + if err := b.FlushFunc(b.buffer.Bytes(), b.nextFlushOffset); err != nil { + return err + } + } + if b.CloseFunc != nil { + if err := b.CloseFunc(); err != nil { + return err + } + } + + return nil +} |
