EdgeNode/internal/caches/writer_partial_file.go

// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.

package caches

import (
	"encoding/binary"
	"github.com/iwind/TeaGo/types"
	"io"
	"os"
	"strings"
	"sync"
)

type PartialFileWriter struct {
	rawWriter  *os.File
	key        string
	headerSize int64
	bodySize   int64
	expiredAt  int64
	endFunc    func()
	once       sync.Once

	isNew      bool
	isPartial  bool
	bodyOffset int64
}

func NewPartialFileWriter(rawWriter *os.File, key string, expiredAt int64, isNew bool, isPartial bool, bodyOffset int64, endFunc func()) *PartialFileWriter {
	return &PartialFileWriter{
		key:        key,
		rawWriter:  rawWriter,
		expiredAt:  expiredAt,
		endFunc:    endFunc,
		isNew:      isNew,
		isPartial:  isPartial,
		bodyOffset: bodyOffset,
	}
}

// WriteHeader 写入数据
func (this *PartialFileWriter) WriteHeader(data []byte) (n int, err error) {
	if !this.isNew {
		return
	}
	n, err = this.rawWriter.Write(data)
	this.headerSize += int64(n)
	if err != nil {
		_ = this.Discard()
	}
	return
}

// WriteHeaderLength 写入Header长度数据
func (this *PartialFileWriter) WriteHeaderLength(headerLength int) error {
	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
	}
	_, err = this.rawWriter.Write(bytes4)
	if err != nil {
		_ = this.Discard()
		return err
	}
	return nil
}

// Write 写入数据
func (this *PartialFileWriter) Write(data []byte) (n int, err error) {
	n, err = this.rawWriter.Write(data)
	this.bodySize += int64(n)
	if err != nil {
		_ = this.Discard()
	}
	return
}

// WriteAt 在指定位置写入数据
func (this *PartialFileWriter) WriteAt(data []byte, offset int64) error {
	if this.bodyOffset == 0 {
		this.bodyOffset = SizeMeta + int64(len(this.key)) + this.headerSize
	}
	_, err := this.rawWriter.WriteAt(data, this.bodyOffset+offset)
	return err
}

// WriteBodyLength 写入Body长度数据
func (this *PartialFileWriter) WriteBodyLength(bodyLength int64) error {
	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
	}
	_, err = this.rawWriter.Write(bytes8)
	if err != nil {
		_ = this.Discard()
		return err
	}
	return nil
}

// Close 关闭
func (this *PartialFileWriter) Close() error {
	defer this.once.Do(func() {
		this.endFunc()
	})

	var path = this.rawWriter.Name()

	if this.isNew {
		err := this.WriteHeaderLength(types.Int(this.headerSize))
		if err != nil {
			_ = this.rawWriter.Close()
			_ = os.Remove(path)
			return err
		}
		err = this.WriteBodyLength(this.bodySize)
		if err != nil {
			_ = this.rawWriter.Close()
			_ = os.Remove(path)
			return err
		}
	}

	err := this.rawWriter.Close()
	if err != nil {
		_ = os.Remove(path)
	} else if !this.isPartial {
		err = os.Rename(path, strings.Replace(path, ".tmp", "", 1))
		if err != nil {
			_ = os.Remove(path)
		}
	}

	return err
}

// Discard 丢弃
func (this *PartialFileWriter) Discard() error {
	defer this.once.Do(func() {
		this.endFunc()
	})

	_ = this.rawWriter.Close()

	err := os.Remove(this.rawWriter.Name())
	return err
}

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

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

func (this *PartialFileWriter) ExpiredAt() int64 {
	return this.expiredAt
}

func (this *PartialFileWriter) Key() string {
	return this.key
}

// ItemType 获取内容类型
func (this *PartialFileWriter) ItemType() ItemType {
	return ItemTypeFile
}