EdgeNode/internal/caches/writer_memory.go

package caches

import (
	"errors"
	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
	"github.com/cespare/xxhash/v2"
	"github.com/iwind/TeaGo/types"
	"sync"
	"sync/atomic"
)

type MemoryWriter struct {
	storage *MemoryStorage

	key        string
	expiredAt  int64
	headerSize int64
	bodySize   int64
	status     int
	isDirty    bool

	expectedBodySize int64
	maxSize          int64

	hash    uint64
	item    *MemoryItem
	endFunc func(valueItem *MemoryItem)
	once    sync.Once
}

func NewMemoryWriter(memoryStorage *MemoryStorage, key string, expiredAt int64, status int, isDirty bool, expectedBodySize int64, maxSize int64, endFunc func(valueItem *MemoryItem)) *MemoryWriter {
	var valueItem = &MemoryItem{
		ExpiresAt:  expiredAt,
		ModifiedAt: fasttime.Now().Unix(),
		Status:     status,
	}

	if expectedBodySize > 0 {
		valueItem.BodyValue = make([]byte, 0, expectedBodySize)
	}

	var w = &MemoryWriter{
		storage:          memoryStorage,
		key:              key,
		expiredAt:        expiredAt,
		item:             valueItem,
		status:           status,
		isDirty:          isDirty,
		expectedBodySize: expectedBodySize,
		maxSize:          maxSize,
		endFunc:          endFunc,
	}

	w.hash = w.calculateHash(key)

	return w
}

// WriteHeader 写入数据
func (this *MemoryWriter) WriteHeader(data []byte) (n int, err error) {
	this.headerSize += int64(len(data))
	this.item.HeaderValue = append(this.item.HeaderValue, data...)
	return len(data), nil
}

// Write 写入数据
func (this *MemoryWriter) Write(data []byte) (n int, err error) {
	var l = len(data)
	if l == 0 {
		return
	}

	if this.item.IsPrepared {
		if this.item.WriteOffset+int64(l) > this.expectedBodySize {
			err = ErrWritingUnavailable
			return
		}
		copy(this.item.BodyValue[this.item.WriteOffset:], data)
		this.item.WriteOffset += int64(l)
	} else {
		this.item.BodyValue = append(this.item.BodyValue, data...)
	}

	this.bodySize += int64(l)

	// 检查尺寸
	if this.maxSize > 0 && this.bodySize > this.maxSize {
		err = ErrEntityTooLarge
		this.storage.IgnoreKey(this.key, this.maxSize)
		return l, err
	}

	return l, nil
}

// WriteAt 在指定位置写入数据
func (this *MemoryWriter) WriteAt(offset int64, b []byte) error {
	_ = b
	_ = offset
	return errors.New("not supported")
}

// HeaderSize 数据尺寸
func (this *MemoryWriter) HeaderSize() int64 {
	return this.headerSize
}

// BodySize 主体内容尺寸
func (this *MemoryWriter) BodySize() int64 {
	return this.bodySize
}

// Close 关闭
func (this *MemoryWriter) Close() error {
	// 需要在Locker之外
	defer this.once.Do(func() {
		this.endFunc(this.item)
	})

	if this.item == nil {
		return nil
	}

	// check content length
	if this.expectedBodySize > 0 && this.bodySize != this.expectedBodySize {
		this.storage.locker.Lock()
		delete(this.storage.valuesMap, this.hash)
		this.storage.locker.Unlock()
		return ErrUnexpectedContentLength
	}

	this.storage.locker.Lock()
	this.item.IsDone = true
	var err error
	if this.isDirty {
		if this.storage.parentStorage != nil {
			this.storage.valuesMap[this.hash] = this.item

			select {
			case this.storage.dirtyChan <- types.String(this.bodySize) + "@" + this.key:
				atomic.AddInt64(&this.storage.totalDirtySize, this.bodySize)
			default:
				// remove from values map
				delete(this.storage.valuesMap, this.hash)

				err = ErrWritingQueueFull
			}
		} else {
			this.storage.valuesMap[this.hash] = this.item
		}
	} else {
		this.storage.valuesMap[this.hash] = this.item
	}

	this.storage.locker.Unlock()

	return err
}

// Discard 丢弃
func (this *MemoryWriter) Discard() error {
	// 需要在Locker之外
	defer this.once.Do(func() {
		this.endFunc(this.item)
	})

	this.storage.locker.Lock()
	delete(this.storage.valuesMap, this.hash)
	this.storage.locker.Unlock()
	return nil
}

// Key 获取Key
func (this *MemoryWriter) Key() string {
	return this.key
}

// ExpiredAt 过期时间
func (this *MemoryWriter) ExpiredAt() int64 {
	return this.expiredAt
}

// ItemType 内容类型
func (this *MemoryWriter) ItemType() ItemType {
	return ItemTypeMemory
}

// 计算Key Hash
func (this *MemoryWriter) calculateHash(key string) uint64 {
	return xxhash.Sum64String(key)
}