mirror of
https://github.com/hanzoai/forge.git
synced 2026-10-30 15:00:39 +00:00
zeekayandHanzo Dev
aa87442fdd
Source files keep their SPDX-License-Identifier and Hanzo copyright line; the Gitea, Gogs and LFS Test Server copyrights they used to repeat now live once in NOTICE, alongside LICENSE. The header lint accepts any copyright line above the SPDX line, so it no longer demands an upstream one. Co-authored-by: Hanzo Dev <[email protected]>
134 lines
2.7 KiB
Go
134 lines
2.7 KiB
Go
// Copyright 2026 Hanzo AI, Inc. All rights reserved.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
package filebuffer
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"io"
|
|
"os"
|
|
)
|
|
|
|
var ErrWriteAfterRead = errors.New("write is unsupported after a read operation") // occurs if Write is called after a read operation
|
|
|
|
type readAtSeeker interface {
|
|
io.ReadSeeker
|
|
io.ReaderAt
|
|
}
|
|
|
|
// FileBackedBuffer uses a memory buffer with a fixed size.
|
|
// If more data is written a temporary file is used instead.
|
|
// It implements io.ReadWriteCloser, io.ReadSeekCloser and io.ReaderAt
|
|
type FileBackedBuffer struct {
|
|
maxMemorySize int64
|
|
size int64
|
|
buffer bytes.Buffer
|
|
tempDir string
|
|
file *os.File
|
|
reader readAtSeeker
|
|
}
|
|
|
|
// New creates a file backed buffer with a specific maximum memory size
|
|
func New(maxMemorySize int, tempDir string) *FileBackedBuffer {
|
|
return &FileBackedBuffer{
|
|
maxMemorySize: int64(maxMemorySize),
|
|
tempDir: tempDir,
|
|
}
|
|
}
|
|
|
|
// Write implements io.Writer
|
|
func (b *FileBackedBuffer) Write(p []byte) (int, error) {
|
|
if b.reader != nil {
|
|
return 0, ErrWriteAfterRead
|
|
}
|
|
|
|
var n int
|
|
var err error
|
|
|
|
if b.file != nil {
|
|
n, err = b.file.Write(p)
|
|
} else {
|
|
if b.size+int64(len(p)) > b.maxMemorySize {
|
|
b.file, err = os.CreateTemp(b.tempDir, "forge-buffer-")
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
_, err = io.Copy(b.file, &b.buffer)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return b.Write(p)
|
|
}
|
|
|
|
n, err = b.buffer.Write(p)
|
|
}
|
|
|
|
if err != nil {
|
|
return n, err
|
|
}
|
|
b.size += int64(n)
|
|
return n, nil
|
|
}
|
|
|
|
// Size returns the byte size of the buffered data
|
|
func (b *FileBackedBuffer) Size() int64 {
|
|
return b.size
|
|
}
|
|
|
|
func (b *FileBackedBuffer) switchToReader() error {
|
|
if b.reader != nil {
|
|
return nil
|
|
}
|
|
|
|
if b.file != nil {
|
|
if _, err := b.file.Seek(0, io.SeekStart); err != nil {
|
|
return err
|
|
}
|
|
b.reader = b.file
|
|
} else {
|
|
b.reader = bytes.NewReader(b.buffer.Bytes())
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Read implements io.Reader
|
|
func (b *FileBackedBuffer) Read(p []byte) (int, error) {
|
|
if err := b.switchToReader(); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return b.reader.Read(p)
|
|
}
|
|
|
|
// ReadAt implements io.ReaderAt
|
|
func (b *FileBackedBuffer) ReadAt(p []byte, off int64) (int, error) {
|
|
if err := b.switchToReader(); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return b.reader.ReadAt(p, off)
|
|
}
|
|
|
|
// Seek implements io.Seeker
|
|
func (b *FileBackedBuffer) Seek(offset int64, whence int) (int64, error) {
|
|
if err := b.switchToReader(); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return b.reader.Seek(offset, whence)
|
|
}
|
|
|
|
// Close implements io.Closer
|
|
func (b *FileBackedBuffer) Close() error {
|
|
if b.file != nil {
|
|
err := b.file.Close()
|
|
_ = os.Remove(b.file.Name())
|
|
b.file = nil
|
|
return err
|
|
}
|
|
return nil
|
|
}
|