qb/storage/snapshot.go

package storage

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"

	bin "gordenko.dev/dima/bin/little"
	"gordenko.dev/dima/qb/util"
)

/*
Формат:
metricsQty - varsize
[metric]*
де metric - це:

index free list frozen pages - varsize
indexFreeList size - varsize
indexFreeList - Nb

data free list frozen pages - varsize
dataFreeList size - varsize
dataFreeList - Nb

CRC32 - 4b
*/

const ReadBufferSize = 8 * 1024 * 1024 // 8mb

type WritableMetric interface {
	WriteTo(io.Writer) error
}

type WriteSnapshotIn struct {
	SnapshotNumber        int
	Dir                   string
	WriteBufferSize       int
	Metrics               map[uint32]WritableMetric
	FrozenIndexPagesCount int
	IndexPageNumbers      []uint32
	FrozenDataPagesCount  int
	DataPageNumbers       []uint32
}

func WriteSnapshot(in WriteSnapshotIn) (err error) {
	var (
		fileName = filepath.Join(in.Dir, fmt.Sprintf("%d.snapshot", in.SnapshotNumber))
		hasher   = util.NewHasher()
	)

	file, err := os.OpenFile(fileName, os.O_CREATE|os.O_WRONLY, 0666) //0770)
	if err != nil {
		return
	}

	dst := io.MultiWriter(bufio.NewWriterSize(file, in.WriteBufferSize), hasher)

	_, err = bin.WriteVarSize(dst, len(in.Metrics))
	if err != nil {
		return
	}

	for metricID, metric := range in.Metrics {
		err = bin.WriteUint32(dst, metricID)
		if err != nil {
			return
		}
		err = metric.WriteTo(dst)
		if err != nil {
			return
		}
	}
	// free index pages
	err = writeFreePages(dst, in.FrozenIndexPagesCount, in.IndexPageNumbers)
	if err != nil {
		return
	}
	// free data pages
	err = writeFreePages(dst, in.FrozenDataPagesCount, in.DataPageNumbers)
	if err != nil {
		return
	}
	// CRC32
	bin.WriteUint32(file, hasher.Sum32())

	err = file.Close()
	if err != nil {
		return
	}

	// prevLogNumber := logNumber - 1
	// prevChanges := filepath.Join(s.dir, fmt.Sprintf("%d.changes", prevLogNumber))
	// prevSnapshot := filepath.Join(s.dir, fmt.Sprintf("%d.snapshot", prevLogNumber))

	// isExist, err := isFileExist(prevChanges)
	// if err != nil {
	// 	return
	// }

	// if isExist {
	// 	err = os.Remove(prevChanges)
	// 	if err != nil {
	// 		qb.Abort(qb.DeletePrevChangesFileFailed, err)
	// 	}
	// }

	// isExist, err = isFileExist(prevSnapshot)
	// if err != nil {
	// 	return
	// }

	// if isExist {
	// 	err = os.Remove(prevSnapshot)
	// 	if err != nil {
	// 		qb.Abort(qb.DeletePrevSnapshotFileFailed, err)
	// 	}
	// }
	return
}

type ReadSnapshotOut struct {
	Metrics               map[uint32]*ReplayMetric
	FrozenIndexPagesCount int
	IndexPageNumbers      []uint32
	FrozenDataPagesCount  int
	DataPageNumbers       []uint32
}

func ReadSnapshot(fileName string) (out ReadSnapshotOut, err error) {
	var (
		metrics               = make(map[uint32]*ReplayMetric)
		frozenIndexPagesCount int
		indexPageNumbers      []uint32
		frozenDataPagesCount  int
		dataPageNumbers       []uint32
		hasher                = util.NewHasher()
	)

	file, err := os.Open(fileName)
	if err != nil {
		return
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return
	}
	fileSize := stat.Size()

	if fileSize == 0 {
		return ReadSnapshotOut{
			Metrics: metrics,
		}, nil
	}

	if fileSize < 4 {
		err = fmt.Errorf("%s is corrupted", fileName)
		return
	}
	payloadSize := fileSize - 4

	// 2. Створюємо великий буфер для NVMe (4 MB)
	bufferedReader := bufio.NewReaderSize(file, readBufferSize)

	// 3. Обмежуємо читання лише розміром Payload
	limitReader := io.LimitReader(bufferedReader, payloadSize)

	// 4. Створюємо хеш та обгортку TeeReader, яка рахує CRC32 на льоту
	src := io.TeeReader(limitReader, hasher)

	// читаю payload
	metricsQty, err := bin.ReadVarSize(src)
	if err != nil {
		return
	}
	for range metricsQty {
		var (
			metricID uint32
			buf      = make([]byte, DataPageSize)
			// для простоти створюю буфер максимального розміру
			metric = &ReplayMetric{
				buf:     buf,
				databuf: buf[:DataPagePayloadSize],
			}
		)
		metricID, err = bin.ReadUint32(src)
		if err != nil {
			return
		}
		err = metric.ReadFrom(src)
		if err != nil {
			return
		}
		metrics[metricID] = metric
	}
	// index pages
	frozenIndexPagesCount, indexPageNumbers, err = readFreePages(src)
	if err != nil {
		return
	}
	// data pages
	frozenDataPagesCount, dataPageNumbers, err = readFreePages(src)
	if err != nil {
		return
	}
	// verify CRC32
	expectedCRC, err := bin.ReadUint32(bufferedReader)
	if err != nil {
		return
	}

	calculatedCRC := hasher.Sum32()

	if expectedCRC != calculatedCRC {
		err = fmt.Errorf("%s is corrupted. Calculated CRC %d not equal expected CRC %d",
			fileName, calculatedCRC, expectedCRC)
		return
	}

	return ReadSnapshotOut{
		Metrics:               metrics,
		FrozenIndexPagesCount: frozenIndexPagesCount,
		IndexPageNumbers:      indexPageNumbers,
		FrozenDataPagesCount:  frozenDataPagesCount,
		DataPageNumbers:       dataPageNumbers,
	}, nil
}

// HELPERS

func writeFreePages(dst io.Writer, frozenPagesCount int, pageNumbers []uint32) (err error) {
	_, err = bin.WriteVarSize(dst, frozenPagesCount)
	if err != nil {
		return
	}
	_, err = bin.WriteVarSize(dst, len(pageNumbers))
	if err != nil {
		return
	}
	for _, pageNo := range pageNumbers {
		err = bin.WriteUint32(dst, pageNo)
		if err != nil {
			return
		}
	}
	return
}

func readFreePages(src io.Reader) (_ int, _ []uint32, err error) {
	var (
		frozenPagesCount int
		pageNumbers      []uint32
	)
	frozenPagesCount, err = bin.ReadVarSize(src)
	if err != nil {
		return
	}
	pageNumbersCount, err := bin.ReadVarSize(src)
	if err != nil {
		return
	}
	for range pageNumbersCount {
		var pageNo uint32
		pageNo, err = bin.ReadUint32(src)
		if err != nil {
			return
		}
		pageNumbers = append(pageNumbers, pageNo)
	}
	return frozenPagesCount, pageNumbers, nil
}