featurebase/buffer/filebuffer.go

package buffer

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"sync"
)

// NewFileBuffer returns a file buffer which will use an in-memory buffer, until `max` bytes have been written, at which point it will write the contents of memory to a file, and continue writing future data to the file.
// The file will be written to `temp` directory. The buffer fulfills the io.Reader and io.Writer interface
func NewFileBuffer(max int, temp string) *FileBuffer {
	return &FileBuffer{max: max, tempDir: temp}
}

type FileBuffer struct {
	max     int
	buf     bytes.Buffer
	file    *os.File
	tempDir string
	reading bool
	files   []*os.File
	mu      sync.Mutex
}

func (fb *FileBuffer) Write(p []byte) (n int, err error) {
	if fb.reading {
		panic("cannot write after read")
	}
	if fb.file != nil {
		return fb.file.Write(p)
	}
	n, err = fb.buf.Write(p)
	if err != nil {
		return
	}
	if fb.buf.Len() > fb.max {
		fb.file, err = ioutil.TempFile(fb.tempDir, "filebuffer-")
		if err != nil {
			return
		}
		_, err = io.Copy(fb.file, &fb.buf)
		fb.buf.Reset()
	}
	return
}

func (fb *FileBuffer) Len() (int64, error) {
	if fb.file == nil {
		return int64(fb.buf.Len()), nil
	}
	fi, err := fb.file.Stat()
	if err != nil {
		return 0, err
	}

	return fi.Size(), nil
}

func (fb *FileBuffer) Read(p []byte) (n int, err error) {
	if fb.file != nil {
		if !fb.reading {
			fb.reading = true
			_, err = fb.file.Seek(0, 0)
			if err != nil {
				return
			}
		}
		return fb.file.Read(p)
	}
	fb.reading = true
	return fb.buf.Read(p)
}

func (fb *FileBuffer) Close() error {
	if fb.file != nil {
		name := fb.file.Name()
		if err := fb.file.Close(); err != nil {
			return err
		}
		for _, f := range fb.files {
			f.Close()
		}
		fb.files = fb.files[:0]
		fb.file = nil
		return os.Remove(name)
	}
	return nil
}

func (fb *FileBuffer) Reset() error {
	fb.mu.Lock()
	defer fb.mu.Unlock()
	fb.reading = false
	fb.buf.Reset()
	return fb.Close()
}

func (fb *FileBuffer) NewReader() (io.Reader, error) {
	fb.mu.Lock()
	defer fb.mu.Unlock()
	fb.reading = true
	if fb.file == nil {
		return bytes.NewReader(fb.buf.Bytes()), nil
	}
	f, err := os.OpenFile(fb.file.Name(), os.O_RDONLY, 0)
	fb.files = append(fb.files, f)
	return f, err
}