Initial commit (code only without large binaries)

EdgeNode/internal/caches/writer_file.go

@@ -0,0 +1,261 @@
+package caches
+
+import (
+	"bufio"
+	"encoding/binary"
+	"errors"
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
+	"github.com/iwind/TeaGo/types"
+	"io"
+	"strings"
+	"sync"
+)
+
+const (
+	DefaultWriteBufferSize = 64 << 10 // 默认写入缓冲区大小：64KB
+)
+
+type FileWriter struct {
+	storage   StorageInterface
+	rawWriter *fsutils.File
+	bufWriter *bufio.Writer // 优化：引入 bufio.Writer 减少系统调用
+	key       string
+
+	metaHeaderSize int
+	headerSize     int64
+
+	metaBodySize int64 // 写入前的内容长度
+	bodySize     int64
+
+	expiredAt int64
+	maxSize   int64
+	endFunc   func()
+	once      sync.Once
+
+	modifiedBytes int
+}
+
+func NewFileWriter(storage StorageInterface, rawWriter *fsutils.File, key string, expiredAt int64, metaHeaderSize int, metaBodySize int64, maxSize int64, endFunc func()) *FileWriter {
+	writer := &FileWriter{
+		storage:        storage,
+		key:            key,
+		rawWriter:      rawWriter,
+		expiredAt:      expiredAt,
+		maxSize:        maxSize,
+		endFunc:        endFunc,
+		metaHeaderSize: metaHeaderSize,
+		metaBodySize:   metaBodySize,
+	}
+	// 优化：创建 bufio.Writer，减少系统调用
+	writer.bufWriter = bufio.NewWriterSize(rawWriter, DefaultWriteBufferSize)
+	return writer
+}
+
+// WriteHeader 写入数据
+func (this *FileWriter) WriteHeader(data []byte) (n int, err error) {
+	// 优化：使用 bufio.Writer 减少系统调用
+	// 注意：Header 写入通常较小，但为了保持一致性也使用缓冲
+	n, err = this.bufWriter.Write(data)
+	this.headerSize += int64(n)
+	if err != nil {
+		_ = this.Discard()
+	}
+	return
+}
+
+// WriteHeaderLength 写入Header长度数据
+func (this *FileWriter) WriteHeaderLength(headerLength int) error {
+	if this.metaHeaderSize > 0 && this.metaHeaderSize == headerLength {
+		return nil
+	}
+	// 优化：先 Flush 缓冲区，确保之前的数据已写入
+	err := this.bufWriter.Flush()
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	var bytes4 = make([]byte, 4)
+	binary.BigEndian.PutUint32(bytes4, uint32(headerLength))
+	_, err = this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength, io.SeekStart)
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	// 注意：Seek 后需要重新创建 bufio.Writer，因为底层文件位置已改变
+	// 但这里只写入 4 字节，直接写入即可
+	_, err = this.rawWriter.Write(bytes4)
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	// Seek 后重新创建 bufio.Writer（因为文件位置已改变）
+	this.bufWriter.Reset(this.rawWriter)
+	return nil
+}
+
+// Write 写入数据
+func (this *FileWriter) Write(data []byte) (n int, err error) {
+	// split LARGE data
+	var l = len(data)
+	if l > (2 << 20) {
+		var offset = 0
+		const bufferSize = 64 << 10
+		for {
+			var end = offset + bufferSize
+			if end > l {
+				end = l
+			}
+			n1, err1 := this.write(data[offset:end])
+			n += n1
+			if err1 != nil {
+				return n, err1
+			}
+			if end >= l {
+				return n, nil
+			}
+			offset = end
+		}
+	}
+
+	// write NORMAL size data
+	return this.write(data)
+}
+
+// WriteAt 在指定位置写入数据
+func (this *FileWriter) WriteAt(offset int64, data []byte) error {
+	_ = data
+	_ = offset
+	return errors.New("not supported")
+}
+
+// WriteBodyLength 写入Body长度数据
+func (this *FileWriter) WriteBodyLength(bodyLength int64) error {
+	if this.metaBodySize >= 0 && bodyLength == this.metaBodySize {
+		return nil
+	}
+	// 优化：先 Flush 缓冲区，确保之前的数据已写入
+	err := this.bufWriter.Flush()
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	var bytes8 = make([]byte, 8)
+	binary.BigEndian.PutUint64(bytes8, uint64(bodyLength))
+	_, err = this.rawWriter.Seek(SizeExpiresAt+SizeStatus+SizeURLLength+SizeHeaderLength, io.SeekStart)
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	// 注意：Seek 后需要重新创建 bufio.Writer，因为底层文件位置已改变
+	// 但这里只写入 8 字节，直接写入即可
+	_, err = this.rawWriter.Write(bytes8)
+	if err != nil {
+		_ = this.Discard()
+		return err
+	}
+	// Seek 后重新创建 bufio.Writer（因为文件位置已改变）
+	this.bufWriter.Reset(this.rawWriter)
+	return nil
+}
+
+// Close 关闭
+func (this *FileWriter) Close() error {
+	defer this.once.Do(func() {
+		this.endFunc()
+	})
+
+	var path = this.rawWriter.Name()
+
+	// 优化：确保所有缓冲数据都已写入
+	err := this.bufWriter.Flush()
+	if err != nil {
+		_ = this.rawWriter.Close()
+		_ = fsutils.Remove(path)
+		return err
+	}
+
+	// check content length
+	if this.metaBodySize > 0 && this.bodySize != this.metaBodySize {
+		_ = this.rawWriter.Close()
+		_ = fsutils.Remove(path)
+		return ErrUnexpectedContentLength
+	}
+
+	err = this.WriteHeaderLength(types.Int(this.headerSize))
+	if err != nil {
+		_ = this.rawWriter.Close()
+		_ = fsutils.Remove(path)
+		return err
+	}
+	err = this.WriteBodyLength(this.bodySize)
+	if err != nil {
+		_ = this.rawWriter.Close()
+		_ = fsutils.Remove(path)
+		return err
+	}
+
+	err = this.rawWriter.Close()
+	if err != nil {
+		_ = fsutils.Remove(path)
+	} else if strings.HasSuffix(path, FileTmpSuffix) {
+		err = fsutils.Rename(path, strings.Replace(path, FileTmpSuffix, "", 1))
+		if err != nil {
+			_ = fsutils.Remove(path)
+		}
+	}
+
+	return err
+}
+
+// Discard 丢弃
+func (this *FileWriter) Discard() error {
+	defer this.once.Do(func() {
+		this.endFunc()
+	})
+
+	_ = this.rawWriter.Close()
+
+	err := fsutils.Remove(this.rawWriter.Name())
+	return err
+}
+
+func (this *FileWriter) HeaderSize() int64 {
+	return this.headerSize
+}
+
+func (this *FileWriter) BodySize() int64 {
+	return this.bodySize
+}
+
+func (this *FileWriter) ExpiredAt() int64 {
+	return this.expiredAt
+}
+
+func (this *FileWriter) Key() string {
+	return this.key
+}
+
+// ItemType 获取内容类型
+func (this *FileWriter) ItemType() ItemType {
+	return ItemTypeFile
+}
+
+func (this *FileWriter) write(data []byte) (n int, err error) {
+	// 优化：使用 bufio.Writer 减少系统调用
+	n, err = this.bufWriter.Write(data)
+	this.bodySize += int64(n)
+
+	if this.maxSize > 0 && this.bodySize > this.maxSize {
+		err = ErrEntityTooLarge
+
+		if this.storage != nil {
+			this.storage.IgnoreKey(this.key, this.maxSize)
+		}
+	}
+
+	if err != nil {
+		_ = this.Discard()
+	}
+
+	return
+}