mirror of
https://github.com/pocketbase/pocketbase.git
synced 2026-03-27 18:53:23 -04:00
126 lines
3.0 KiB
Go
126 lines
3.0 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"io"
|
|
"os"
|
|
)
|
|
|
|
var _ io.ReadWriteCloser = (*bufferWithFile)(nil)
|
|
|
|
// newBufferWithFile initializes and returns a new bufferWithFile with the specified memoryLimit.
|
|
//
|
|
// If memoryLimit is negative or zero, defaults to [DefaultMaxMemory].
|
|
func newBufferWithFile(memoryLimit int64) *bufferWithFile {
|
|
if memoryLimit <= 0 {
|
|
memoryLimit = DefaultMaxMemory
|
|
}
|
|
|
|
return &bufferWithFile{
|
|
buf: new(bytes.Buffer),
|
|
memoryLimit: memoryLimit,
|
|
}
|
|
}
|
|
|
|
// bufferWithFile is similar to [bytes.Buffer] but after the limit it
|
|
// fallbacks to a temporary file to minimize excessive memory usage.
|
|
type bufferWithFile struct {
|
|
buf *bytes.Buffer
|
|
file *os.File
|
|
memoryLimit int64
|
|
fileReadOffset int64
|
|
}
|
|
|
|
// Read implements the standard [io.Reader] interface by reading
|
|
// up to len(p) bytes into p.
|
|
func (b *bufferWithFile) Read(p []byte) (n int, err error) {
|
|
if b.buf == nil {
|
|
return 0, errors.New("[bufferWithFile.Read] not initialized or already closed")
|
|
}
|
|
|
|
// eagerly get length because bytes.Buffer may resize and change it
|
|
maxToRead := len(p)
|
|
|
|
// read first from the memory buffer
|
|
if b.buf.Len() > 0 {
|
|
n, err = b.buf.Read(p)
|
|
if err != nil && err != io.EOF {
|
|
return n, err
|
|
}
|
|
}
|
|
|
|
// continue reading from the file to fill the remaining bytes
|
|
if n < maxToRead && b.file != nil {
|
|
fileN, fileErr := b.file.ReadAt(p[n:maxToRead], b.fileReadOffset)
|
|
b.fileReadOffset += int64(fileN)
|
|
n += fileN
|
|
err = fileErr
|
|
}
|
|
|
|
// return EOF if the buffers are empty and nothing has been read
|
|
// (to minimize potential breaking changes and for consistency with the bytes.Buffer rules)
|
|
if n == 0 && maxToRead > 0 && err == nil {
|
|
return 0, io.EOF
|
|
}
|
|
|
|
return n, err
|
|
}
|
|
|
|
// Write implements the standard [io.Writer] interface by writing the
|
|
// content of p into the buffer.
|
|
//
|
|
// If the current memory buffer doesn't have enough space to hold len(p),
|
|
// it write p into a temp disk file.
|
|
func (b *bufferWithFile) Write(p []byte) (int, error) {
|
|
if b.buf == nil {
|
|
return 0, errors.New("[bufferWithFile.Write] not initialized or already closed")
|
|
}
|
|
|
|
// already above the limit -> continue with the file
|
|
if b.file != nil {
|
|
return b.file.Write(p)
|
|
}
|
|
|
|
// above limit -> create and write to file
|
|
if int64(b.buf.Len()+len(p)) > b.memoryLimit {
|
|
if b.file == nil {
|
|
var err error
|
|
b.file, err = os.CreateTemp("", "pb_buffer_file_*")
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
return b.file.Write(p)
|
|
}
|
|
|
|
// write in memory
|
|
return b.buf.Write(p)
|
|
}
|
|
|
|
// Close implements the standard [io.Closer] interface.
|
|
//
|
|
// It unsets the memory buffer and will cleanup after the fallback
|
|
// temporary file (if exists).
|
|
//
|
|
// It is safe to call Close multiple times.
|
|
// Once Close is invoked the buffer no longer can be used and should be discarded.
|
|
func (b *bufferWithFile) Close() error {
|
|
if b.file != nil {
|
|
err := errors.Join(
|
|
b.file.Close(),
|
|
os.Remove(b.file.Name()),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
b.file = nil
|
|
}
|
|
|
|
b.buf = nil
|
|
|
|
return nil
|
|
}
|