378 lines
9.2 KiB
Go
378 lines
9.2 KiB
Go
package common
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// BodyStorage 请求体存储接口
|
|
type BodyStorage interface {
|
|
io.ReadSeeker
|
|
io.Closer
|
|
// Bytes 获取全部内容
|
|
Bytes() ([]byte, error)
|
|
// Size 获取数据大小
|
|
Size() int64
|
|
// IsDisk 是否是磁盘存储
|
|
IsDisk() bool
|
|
// NewReader returns an independent reader positioned at the start of the
|
|
// stored payload. Each call returns a reader with its own cursor, so
|
|
// callers (e.g. http.Request.GetBody) can replay the body concurrently
|
|
// with, or after, other readers without sharing seek state. Closing the
|
|
// returned reader releases only that reader, never the storage itself;
|
|
// after the storage has been closed, NewReader returns ErrStorageClosed.
|
|
NewReader() (io.ReadCloser, error)
|
|
}
|
|
|
|
// ReplayableBody is an outbound request body that can report its byte size and
|
|
// create independent readers for transport-level retries.
|
|
type ReplayableBody interface {
|
|
io.Reader
|
|
Size() int64
|
|
NewReader() (io.ReadCloser, error)
|
|
}
|
|
|
|
// ErrStorageClosed 存储已关闭错误
|
|
var ErrStorageClosed = fmt.Errorf("body storage is closed")
|
|
|
|
// memoryStorage 内存存储实现
|
|
type memoryStorage struct {
|
|
data []byte
|
|
reader *bytes.Reader
|
|
size int64
|
|
closed int32
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func newMemoryStorage(data []byte) *memoryStorage {
|
|
size := int64(len(data))
|
|
IncrementMemoryBuffers(size)
|
|
return &memoryStorage{
|
|
data: data,
|
|
reader: bytes.NewReader(data),
|
|
size: size,
|
|
}
|
|
}
|
|
|
|
func (m *memoryStorage) Read(p []byte) (n int, err error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if atomic.LoadInt32(&m.closed) == 1 {
|
|
return 0, ErrStorageClosed
|
|
}
|
|
return m.reader.Read(p)
|
|
}
|
|
|
|
func (m *memoryStorage) Seek(offset int64, whence int) (int64, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if atomic.LoadInt32(&m.closed) == 1 {
|
|
return 0, ErrStorageClosed
|
|
}
|
|
return m.reader.Seek(offset, whence)
|
|
}
|
|
|
|
func (m *memoryStorage) Close() error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if atomic.CompareAndSwapInt32(&m.closed, 0, 1) {
|
|
DecrementMemoryBuffers(m.size)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *memoryStorage) Bytes() ([]byte, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if atomic.LoadInt32(&m.closed) == 1 {
|
|
return nil, ErrStorageClosed
|
|
}
|
|
return m.data, nil
|
|
}
|
|
|
|
func (m *memoryStorage) NewReader() (io.ReadCloser, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if atomic.LoadInt32(&m.closed) == 1 {
|
|
return nil, ErrStorageClosed
|
|
}
|
|
// A fresh bytes.Reader over the shared immutable backing array: an
|
|
// independent cursor at zero copy cost. NopCloser keeps Close a no-op, so
|
|
// the storage lifecycle stays owned by whoever holds the storage itself.
|
|
return io.NopCloser(bytes.NewReader(m.data)), nil
|
|
}
|
|
|
|
func (m *memoryStorage) Size() int64 {
|
|
return m.size
|
|
}
|
|
|
|
func (m *memoryStorage) IsDisk() bool {
|
|
return false
|
|
}
|
|
|
|
// diskStorage 磁盘存储实现
|
|
type diskStorage struct {
|
|
file *os.File
|
|
filePath string
|
|
size int64
|
|
closed int32
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func newDiskStorage(data []byte, cachePath string) (*diskStorage, error) {
|
|
// 使用统一的缓存目录管理
|
|
filePath, file, err := CreateDiskCacheFile(DiskCacheTypeBody)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 写入数据
|
|
n, err := file.Write(data)
|
|
if err != nil {
|
|
file.Close()
|
|
os.Remove(filePath)
|
|
return nil, fmt.Errorf("failed to write to temp file: %w", err)
|
|
}
|
|
|
|
// 重置文件指针
|
|
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
|
file.Close()
|
|
os.Remove(filePath)
|
|
return nil, fmt.Errorf("failed to seek temp file: %w", err)
|
|
}
|
|
|
|
size := int64(n)
|
|
IncrementDiskFiles(size)
|
|
|
|
return &diskStorage{
|
|
file: file,
|
|
filePath: filePath,
|
|
size: size,
|
|
}, nil
|
|
}
|
|
|
|
func newDiskStorageFromReader(reader io.Reader, maxBytes int64, cachePath string) (*diskStorage, error) {
|
|
// 使用统一的缓存目录管理
|
|
filePath, file, err := CreateDiskCacheFile(DiskCacheTypeBody)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 从 reader 读取并写入文件
|
|
written, err := io.Copy(file, io.LimitReader(reader, maxBytes+1))
|
|
if err != nil {
|
|
file.Close()
|
|
os.Remove(filePath)
|
|
return nil, fmt.Errorf("failed to write to temp file: %w", err)
|
|
}
|
|
|
|
if written > maxBytes {
|
|
file.Close()
|
|
os.Remove(filePath)
|
|
return nil, ErrRequestBodyTooLarge
|
|
}
|
|
|
|
// 重置文件指针
|
|
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
|
file.Close()
|
|
os.Remove(filePath)
|
|
return nil, fmt.Errorf("failed to seek temp file: %w", err)
|
|
}
|
|
|
|
IncrementDiskFiles(written)
|
|
|
|
return &diskStorage{
|
|
file: file,
|
|
filePath: filePath,
|
|
size: written,
|
|
}, nil
|
|
}
|
|
|
|
func (d *diskStorage) Read(p []byte) (n int, err error) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if atomic.LoadInt32(&d.closed) == 1 {
|
|
return 0, ErrStorageClosed
|
|
}
|
|
return d.file.Read(p)
|
|
}
|
|
|
|
func (d *diskStorage) Seek(offset int64, whence int) (int64, error) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if atomic.LoadInt32(&d.closed) == 1 {
|
|
return 0, ErrStorageClosed
|
|
}
|
|
return d.file.Seek(offset, whence)
|
|
}
|
|
|
|
func (d *diskStorage) Close() error {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if atomic.CompareAndSwapInt32(&d.closed, 0, 1) {
|
|
d.file.Close()
|
|
os.Remove(d.filePath)
|
|
DecrementDiskFiles(d.size)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (d *diskStorage) Bytes() ([]byte, error) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
|
|
if atomic.LoadInt32(&d.closed) == 1 {
|
|
return nil, ErrStorageClosed
|
|
}
|
|
|
|
// 保存当前位置
|
|
currentPos, err := d.file.Seek(0, io.SeekCurrent)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 移动到开头
|
|
if _, err := d.file.Seek(0, io.SeekStart); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 读取全部内容
|
|
data := make([]byte, d.size)
|
|
_, err = io.ReadFull(d.file, data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 恢复位置
|
|
if _, err := d.file.Seek(currentPos, io.SeekStart); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return data, nil
|
|
}
|
|
|
|
func (d *diskStorage) NewReader() (io.ReadCloser, error) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if atomic.LoadInt32(&d.closed) == 1 {
|
|
return nil, ErrStorageClosed
|
|
}
|
|
// A separate file descriptor over the same cache file: an independent
|
|
// cursor at zero copy cost. Closing the returned reader closes only that
|
|
// descriptor; the storage keeps owning the primary descriptor and the
|
|
// file's lifetime. Readers opened before Close stay usable even after the
|
|
// file is unlinked, as the descriptor keeps the inode alive.
|
|
file, err := os.Open(d.filePath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to open body cache file for replay: %w", err)
|
|
}
|
|
return file, nil
|
|
}
|
|
|
|
func (d *diskStorage) Size() int64 {
|
|
return d.size
|
|
}
|
|
|
|
func (d *diskStorage) IsDisk() bool {
|
|
return true
|
|
}
|
|
|
|
// CreateBodyStorage 根据数据大小创建合适的存储
|
|
func CreateBodyStorage(data []byte) (BodyStorage, error) {
|
|
size := int64(len(data))
|
|
threshold := GetDiskCacheThresholdBytes()
|
|
|
|
// 检查是否应该使用磁盘缓存
|
|
if IsDiskCacheEnabled() &&
|
|
size >= threshold &&
|
|
IsDiskCacheAvailable(size) {
|
|
storage, err := newDiskStorage(data, GetDiskCachePath())
|
|
if err != nil {
|
|
// 如果磁盘存储失败，回退到内存存储
|
|
SysError(fmt.Sprintf("failed to create disk storage, falling back to memory: %v", err))
|
|
return newMemoryStorage(data), nil
|
|
}
|
|
return storage, nil
|
|
}
|
|
|
|
return newMemoryStorage(data), nil
|
|
}
|
|
|
|
// CreateBodyStorageFromReader 从 Reader 创建存储（用于大请求的流式处理）
|
|
func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes int64) (BodyStorage, error) {
|
|
threshold := GetDiskCacheThresholdBytes()
|
|
|
|
// 如果启用了磁盘缓存且内容长度超过阈值，直接使用磁盘存储
|
|
if IsDiskCacheEnabled() &&
|
|
contentLength > 0 &&
|
|
contentLength >= threshold &&
|
|
IsDiskCacheAvailable(contentLength) {
|
|
storage, err := newDiskStorageFromReader(reader, maxBytes, GetDiskCachePath())
|
|
if err != nil {
|
|
if IsRequestBodyTooLargeError(err) {
|
|
return nil, err
|
|
}
|
|
// 磁盘存储失败，reader 已被消费，无法安全回退
|
|
// 直接返回错误而非尝试回退（因为 reader 数据已丢失）
|
|
return nil, fmt.Errorf("disk storage creation failed: %w", err)
|
|
}
|
|
IncrementDiskCacheHits()
|
|
return storage, nil
|
|
}
|
|
|
|
// 使用内存读取
|
|
data, err := io.ReadAll(io.LimitReader(reader, maxBytes+1))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if int64(len(data)) > maxBytes {
|
|
return nil, ErrRequestBodyTooLarge
|
|
}
|
|
|
|
storage, err := CreateBodyStorage(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// 如果最终使用内存存储，记录内存缓存命中
|
|
if !storage.IsDisk() {
|
|
IncrementMemoryCacheHits()
|
|
} else {
|
|
IncrementDiskCacheHits()
|
|
}
|
|
return storage, nil
|
|
}
|
|
|
|
type replayableBodyReader struct {
|
|
storage BodyStorage
|
|
}
|
|
|
|
func (r replayableBodyReader) Read(p []byte) (int, error) {
|
|
return r.storage.Read(p)
|
|
}
|
|
|
|
func (r replayableBodyReader) Size() int64 {
|
|
return r.storage.Size()
|
|
}
|
|
|
|
func (r replayableBodyReader) NewReader() (io.ReadCloser, error) {
|
|
return r.storage.NewReader()
|
|
}
|
|
|
|
// NewReplayableBodyReader exposes the replay capabilities of storage without
|
|
// exposing io.Closer. This keeps ownership of the storage lifecycle with the
|
|
// caller instead of allowing net/http to close it as the request body.
|
|
func NewReplayableBodyReader(storage BodyStorage) ReplayableBody {
|
|
return replayableBodyReader{storage: storage}
|
|
}
|
|
|
|
// CleanupOldCacheFiles 清理旧的缓存文件（用于启动时清理残留）
|
|
func CleanupOldCacheFiles() {
|
|
// 使用统一的缓存管理
|
|
CleanupOldDiskCacheFiles(5 * time.Minute)
|
|
}
|