waf-platform/EdgeNode/internal/caches/reader_file_mmap_plus.go

// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package caches

import (
	"encoding/binary"
	"errors"
	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
	memutils "github.com/TeaOSLab/EdgeNode/internal/utils/mem"
	"github.com/TeaOSLab/EdgeNode/internal/utils/mmap"
	rangeutils "github.com/TeaOSLab/EdgeNode/internal/utils/ranges"
	"github.com/iwind/TeaGo/types"
	"io"
	"os"
	"sync"
)

// 设置最大能映射的文件尺寸
var maxMMAPFileSize int64 = 8 << 20

func init() {
	var estimatedSize = int64((memutils.SystemMemoryGB() * 4) << 20)
	if estimatedSize > maxMMAPFileSize {
		maxMMAPFileSize = estimatedSize
	}
}

// MMAPFileReader 通过MMAP读取文件
type MMAPFileReader struct {
	BaseReader

	rawReader  *mmap.SharedMMAP
	path       string
	modifiedAt int64

	meta   []byte
	header []byte

	expiresAt    int64
	status       int
	headerOffset int64
	headerSize   int
	bodySize     int64
	bodyOffset   int64

	currentOffset int64

	isClosed bool

	once sync.Once
}

func NewMMAPFileReaderFromPath(filename string) (*MMAPFileReader, error) {
	reader, err := mmap.OpenSharedMMAP(filename)
	if err != nil {
		return nil, err
	}
	return NewMMAPFileReader(reader)
}

func NewMMAPFileReader(mmapReader *mmap.SharedMMAP) (*MMAPFileReader, error) {
	return &MMAPFileReader{
		path:       mmapReader.Name(),
		rawReader:  mmapReader,
		modifiedAt: mmapReader.Stat().ModTime().Unix(),
	}, nil
}

func (this *MMAPFileReader) Init() error {
	return this.InitAutoDiscard(true)
}

func (this *MMAPFileReader) InitAutoDiscard(autoDiscard bool) error {
	var isOk = false

	if autoDiscard {
		defer func() {
			if !isOk {
				_ = this.discard()
			}
		}()
	}

	var buf = this.meta
	if len(buf) == 0 {
		buf = make([]byte, SizeMeta)
		ok, err := this.readNext(buf)
		if err != nil {
			return err
		}
		if !ok {
			return ErrNotFound
		}
		this.meta = buf
	}

	this.expiresAt = int64(binary.BigEndian.Uint32(buf[:SizeExpiresAt]))

	var status = types.Int(string(buf[OffsetStatus : OffsetStatus+SizeStatus]))
	if status < 100 || status > 999 {
		return errors.New("invalid status")
	}
	this.status = status

	// URL
	var urlLength = binary.BigEndian.Uint32(buf[OffsetURLLength : OffsetURLLength+SizeURLLength])

	// header
	var headerSize = int(binary.BigEndian.Uint32(buf[OffsetHeaderLength : OffsetHeaderLength+SizeHeaderLength]))
	if headerSize == 0 {
		return nil
	}
	this.headerSize = headerSize
	this.headerOffset = int64(SizeMeta) + int64(urlLength)

	// body
	this.bodyOffset = this.headerOffset + int64(headerSize)
	var bodySize = int(binary.BigEndian.Uint64(buf[OffsetBodyLength : OffsetBodyLength+SizeBodyLength]))
	if bodySize == 0 {
		isOk = true
		return nil
	}
	this.bodySize = int64(bodySize)

	isOk = true

	return nil
}

func (this *MMAPFileReader) TypeName() string {
	return "disk"
}

func (this *MMAPFileReader) ExpiresAt() int64 {
	return this.expiresAt
}

func (this *MMAPFileReader) Status() int {
	return this.status
}

func (this *MMAPFileReader) LastModified() int64 {
	if this.modifiedAt > 0 {
		return this.modifiedAt
	}
	stat, err := os.Stat(this.path)
	if err != nil {
		return 0
	}
	return stat.ModTime().Unix()
}

func (this *MMAPFileReader) HeaderSize() int64 {
	return int64(this.headerSize)
}

func (this *MMAPFileReader) BodySize() int64 {
	return this.bodySize
}

func (this *MMAPFileReader) ReadHeader(buf []byte, callback ReaderFunc) error {
	// 使用缓存
	if len(this.header) > 0 && len(buf) >= len(this.header) {
		copy(buf, this.header)
		_, err := callback(len(this.header))
		if err != nil {
			return err
		}

		// 移动到Body位置
		this.moveTo(this.bodyOffset)
		return nil
	}

	var isOk = false

	defer func() {
		if !isOk {
			_ = this.discard()
		}
	}()

	this.moveTo(this.headerOffset)

	var headerSize = this.headerSize

	if len(buf) > headerSize {
		n, err := this.rawReader.ReadAt(buf[:headerSize], this.headerOffset)
		if err != nil {
			if err != io.EOF {
				return err
			}
		}

		_, err = callback(n)
		if err != nil {
			isOk = true
			return err
		}
	} else {
		for {
			n, err := this.read(buf)
			if n > 0 {
				if n < headerSize {
					goNext, e := callback(n)
					if e != nil {
						isOk = true
						return e
					}
					if !goNext {
						break
					}
					headerSize -= n
				} else {
					_, e := callback(headerSize)
					if e != nil {
						isOk = true
						return e
					}
					break
				}
			}
			if err != nil {
				if err != io.EOF {
					return err
				}
				break
			}
		}
	}

	isOk = true

	// 移动到Body位置
	this.moveTo(this.bodyOffset)

	return nil
}

func (this *MMAPFileReader) ReadBody(buf []byte, callback ReaderFunc) error {
	if this.bodySize == 0 {
		return nil
	}

	var isOk = false

	defer func() {
		if !isOk {
			_ = this.discard()
		}
	}()

	var offset = this.bodyOffset

	// 开始读Body部分
	this.moveTo(offset)

	for {
		n, err := this.read(buf)
		if n > 0 {
			goNext, e := callback(n)
			if e != nil {
				isOk = true
				return e
			}
			if !goNext {
				break
			}
		}
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}
	}

	isOk = true

	return nil
}

func (this *MMAPFileReader) Read(buf []byte) (n int, err error) {
	if this.bodySize == 0 {
		n = 0
		err = io.EOF
		return
	}

	n, err = this.read(buf)
	if err != nil && err != io.EOF {
		_ = this.discard()
	}

	return
}

func (this *MMAPFileReader) ReadBodyRange(buf []byte, start int64, end int64, callback ReaderFunc) error {
	var isOk = false

	defer func() {
		if !isOk {
			_ = this.discard()
		}
	}()

	var offset = start
	if start < 0 {
		offset = this.bodyOffset + this.bodySize + end
		end = this.bodyOffset + this.bodySize - 1
	} else if end < 0 {
		offset = this.bodyOffset + start
		end = this.bodyOffset + this.bodySize - 1
	} else {
		offset = this.bodyOffset + start
		end = this.bodyOffset + end
	}
	if offset < 0 || end < 0 || offset > end {
		isOk = true
		return ErrInvalidRange
	}

	this.moveTo(offset)

	for {
		n, err := this.read(buf)
		if n > 0 {
			var n2 = int(end-offset) + 1
			if n2 <= n {
				_, e := callback(n2)
				if e != nil {
					isOk = true
					return e
				}
				break
			} else {
				goNext, e := callback(n)
				if e != nil {
					isOk = true
					return e
				}
				if !goNext {
					break
				}
			}

			offset += int64(n)
			if offset > end {
				break
			}
		}
		if err != nil {
			if err != io.EOF {
				return err
			}
			break
		}
	}

	isOk = true

	return nil
}

// ContainsRange 是否包含某些区间内容
func (this *MMAPFileReader) ContainsRange(r rangeutils.Range) (r2 rangeutils.Range, ok bool) {
	return r, true
}

func (this *MMAPFileReader) CopyBodyTo(writer io.Writer) (int, error) {
	return this.rawReader.Write(writer, int(this.bodyOffset))
}

func (this *MMAPFileReader) Close() error {
	if this.isClosed {
		return nil
	}
	this.isClosed = true

	var err error
	this.once.Do(func() {
		err = this.rawReader.Close()
	})
	return err
}

func (this *MMAPFileReader) moveTo(offset int64) {
	this.currentOffset = offset
}

func (this *MMAPFileReader) readNext(buf []byte) (ok bool, err error) {
	n, err := this.rawReader.ReadAt(buf, this.currentOffset)
	if n > 0 {
		this.currentOffset += int64(n)
	}
	if err != nil {
		return false, err
	}
	ok = n == len(buf)
	return
}

func (this *MMAPFileReader) read(p []byte) (n int, err error) {
	n, err = this.rawReader.ReadAt(p, this.currentOffset)
	if n > 0 {
		this.currentOffset += int64(n)
	}
	return
}

func (this *MMAPFileReader) discard() error {
	this.once.Do(func() {
		_ = this.rawReader.Close()
	})

	this.isClosed = true

	// remove file
	return fsutils.Remove(this.path)
}