qb/recovery/recovery.go

package recovery

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"gordenko.dev/dima/qb"
	"gordenko.dev/dima/qb/freelist"
	"gordenko.dev/dima/qb/storage"
	"gordenko.dev/dima/qb/worker"
)

// Result зберігає результати для обох типів файлів
type FileVersion struct {
	Name    string
	Version int
}

// FindMaxSnapshotFiles шукає файли з розширеннями .snapshot_XXX та .wal_XXX
// і повертає структуру з максимальними номерами.
func findFileLatestVersion(dir string, extPrefix string) (FileVersion, error) {
	var x FileVersion
	if !strings.HasPrefix(extPrefix, ".") {
		extPrefix = "." + extPrefix
	}
	entries, err := os.ReadDir(dir)
	if err != nil {
		return x, err
	}
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		ext := filepath.Ext(name)

		if strings.HasPrefix(ext, extPrefix) {
			version, err := strconv.Atoi(strings.TrimPrefix(ext, extPrefix))
			if err == nil {
				if version > x.Version || x.Name == "" {
					x.Name = name
					x.Version = version
				}
			}
		}
	}
	return x, nil
}

var (
	ErrLostSnapshot = errors.New("lost snapshot")
)

type RecoveryFiles struct {
	Snapshot       string
	WAL            string
	SnapshotNumber int
}

func resolveRecoveryFiles(dir string) (_ RecoveryFiles, err error) {
	wal, err := findFileLatestVersion(dir, "wal_")
	if err != nil {
		return
	}
	snapshot, err := findFileLatestVersion(dir, "snapshot_")
	if err != nil {
		return
	}
	if snapshot.Name != "" {
		if wal.Name != "" {
			if snapshot.Version == wal.Version {
				return RecoveryFiles{
					Snapshot:       snapshot.Name,
					WAL:            wal.Name,
					SnapshotNumber: snapshot.Version,
				}, nil
			} else if snapshot.Version > wal.Version {
				return RecoveryFiles{
					Snapshot:       snapshot.Name,
					SnapshotNumber: snapshot.Version,
				}, nil
			} else {
				// WAL version > snapshot version
				err = ErrLostSnapshot
				return
			}
		} else {
			return RecoveryFiles{
				Snapshot:       snapshot.Name,
				SnapshotNumber: snapshot.Version,
			}, nil
		}
	} else {
		if wal.Name != "" {
			if wal.Version == 0 {
				return RecoveryFiles{
					WAL: wal.Name,
				}, nil
			} else {
				err = ErrLostSnapshot
				return
			}
		} else {
			// no files - ok, 1st start of database
			return
		}
	}
}

type RecoveryReport struct {
	SnapshotNumber int
	WAL            string
	ReplayMetrics  map[uint32]*storage.ReplayMetric
	IndexFreeList  *freelist.FreeList
	DataFreeList   *freelist.FreeList
}

func Recovery(dir string, databaseName string) (_ RecoveryReport, err error) {
	recoveryFiles, err := resolveRecoveryFiles(dir)
	if err != nil {
		return
	}
	fmt.Printf("recovery files: %#v\n", recoveryFiles)
	var (
		frozenIndexPageCount int
		frozenDataPageCount  int
		freeIndexPages       []uint32
		freeDataPages        []uint32
		metrics              map[uint32]*storage.ReplayMetric
	)
	if recoveryFiles.Snapshot != "" {
		var snapshot worker.ReadSnapshotOut
		snapshot, err = worker.ReadSnapshot(filepath.Join(dir, recoveryFiles.Snapshot))
		if err != nil {
			return
		}
		fmt.Printf("snapshot: %#v\n", snapshot)
		frozenIndexPageCount = snapshot.FrozenIndexPagesCount
		frozenDataPageCount = snapshot.FrozenDataPagesCount
		freeIndexPages = snapshot.IndexPageNumbers
		freeDataPages = snapshot.DataPageNumbers
		metrics = snapshot.Metrics
	} else {
		metrics = make(map[uint32]*storage.ReplayMetric)
	}

	if recoveryFiles.WAL != "" {
		var (
			wal         *os.File
			walReader   *storage.WALReader
			walReplayer *storage.WALReplayer
		)
		wal, err = os.Open(filepath.Join(dir, recoveryFiles.WAL))
		if err != nil {
			return
		}
		defer wal.Close()

		walReader, err = storage.NewWALReader(wal)
		if err != nil {
			return
		}

		walReplayer, err = storage.NewWALReplayer(storage.WALReplayerOptions{
			WALReader:      walReader,
			Metrics:        metrics,
			FreeIndexPages: freeIndexPages,
			FreeDataPages:  freeDataPages,
		})
		if err != nil {
			return
		}

		fmt.Printf("before replay\n")
		err = walReplayer.Replay()
		if err != nil {
			return
		}
		fmt.Printf("after replay\n")
		// перезаписую сторінки із останнього комміта
		var dataFile *os.File
		dataFile, err = os.OpenFile(qb.GetDataFilePath(dir, databaseName), os.O_CREATE|os.O_WRONLY, 0666)
		if err != nil {
			return
		}
		err = storage.WriteDataPages(dataFile, walReplayer.DataPagesToRewrite())
		if err != nil {
			return
		}
		err = dataFile.Close()
		if err != nil {
			return
		}
		var indexFile *os.File
		indexFile, err = os.OpenFile(qb.GetIndexFilePath(dir, databaseName), os.O_CREATE|os.O_WRONLY, 0666)
		if err != nil {
			return
		}
		err = storage.WriteIndexPages(indexFile, walReplayer.IndexPagesToRewrite())
		if err != nil {
			return
		}
		err = indexFile.Close()
		if err != nil {
			return
		}

		freeIndexPages = walReplayer.FreeIndexPages()
		freeDataPages = walReplayer.FreeDataPages()
	}

	indexFreeList, err := freelist.New(freelist.Options{
		PageSize:        2048,
		BaseFilePath:    qb.GetIndexFreeListFilePath(dir, databaseName),
		DeltaFilePath:   qb.GetIndexFreeListDeltaFilePath(dir, databaseName),
		FrozenPageCount: frozenIndexPageCount,
	})
	indexFreeList.AddPageNumbers(freeIndexPages)

	dataFreeList, err := freelist.New(freelist.Options{
		PageSize:        2048,
		BaseFilePath:    qb.GetDataFreeListFilePath(dir, databaseName),
		DeltaFilePath:   qb.GetDataFreeListDeltaFilePath(dir, databaseName),
		FrozenPageCount: frozenDataPageCount,
	})
	dataFreeList.AddPageNumbers(freeDataPages)

	return RecoveryReport{
		SnapshotNumber: recoveryFiles.SnapshotNumber,
		WAL:            recoveryFiles.WAL,
		ReplayMetrics:  metrics,
		IndexFreeList:  indexFreeList,
		DataFreeList:   dataFreeList,
	}, nil
}