qb/atree/atree.go

package atree

import (
	"errors"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sync"

	diploma "gordenko.dev/dima/qb"
	"gordenko.dev/dima/qb/bin"
)

const (
	PageTypeData  = 1
	PageTypeIndex = 2

	filePerm = 0770

	// common
	crc32Idx = PageSize - 4
	pageType = PageSize - 5

	// index page
	indexRecordsQtyIdx   = PageSize - 8
	isDataPageNumbersIdx = PageSize - 6

	// data page
	timestampsSizeIdx = PageSize - 13
	valuesSizeIdx     = PageSize - 11
	prevPageIdx       = PageSize - 9

	timestampSize  = 4
	pairSize       = timestampSize + PageNoSize
	indexFooterIdx = indexRecordsQtyIdx
	dataFooterIdx  = timestampsSizeIdx

	PageSize   = 8192
	PageNoSize = 4
	//
	DataPagePayloadSize int = dataFooterIdx
)

const (
	FlagReused  byte = 1 // сторінка із FreeList
	FlagNewRoot byte = 2 // новая страница
)

var (
	castagnoliTable = crc32.MakeTable(crc32.Castagnoli)
)

func calcChecksum(page []byte) uint32 {
	return crc32.Checksum(page, castagnoliTable)
}

type PageToWrite struct {
	PageNo   uint32
	Data     []byte
	IsReused bool
}

type FreeList interface {
	// використовується в allocPage
	ReservePage() uint32
}

type _page struct {
	PageNo         uint32
	Buf            []byte
	ReferenceCount int
}

type Atree struct {
	freelist          FreeList
	file              *os.File
	mutex             sync.Mutex
	allocatedPagesQty uint32
	pages             map[uint32]*_page
	pageWaits         map[uint32][]chan readResult
	pagesToRead       []uint32
	readSignalCh      chan struct{}
	writeSignalCh     chan struct{}
	writeTasksQueue   []WriteTask
}

type Options struct {
	Dir          string
	DatabaseName string
	FreeList     FreeList
}

func New(opt Options) (*Atree, error) {
	if opt.Dir == "" {
		return nil, errors.New("Dir option is required")
	}
	// if opt.RedoDir == "" {
	// 	return nil, errors.New("RedoDir option is required")
	// }
	if opt.DatabaseName == "" {
		return nil, errors.New("DatabaseName option is required")
	}
	if opt.FreeList == nil {
		return nil, errors.New("FreeList option is required")
	}
	// открываю или создаю dbName.data и dbName.index файлы
	var (
		fileName          = filepath.Join(opt.Dir, opt.DatabaseName+".db")
		file              *os.File
		allocatedPagesQty uint32
	)

	// При создании data файла сразу создается индекс, поэтому корректное
	// состояние БД: либо оба файла есть, либо ни одного файла нет.
	isDataExist, err := isFileExist(fileName)
	if err != nil {
		return nil, fmt.Errorf("check data file is exist: %s", err)
	}
	if isDataExist {
		file, allocatedPagesQty, err = openFile(fileName, PageSize)
		if err != nil {
			return nil, fmt.Errorf("open data file: %s", err)
		}
	} else {
		// нет файла
		file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, filePerm)
		if err != nil {
			return nil, err
		}
	}

	tree := &Atree{
		//freelist:          opt.FreeList,
		file:              file,
		allocatedPagesQty: allocatedPagesQty,
		pages:             make(map[uint32]*_page),
		pageWaits:         make(map[uint32][]chan readResult),
		readSignalCh:      make(chan struct{}, 1),
		writeSignalCh:     make(chan struct{}, 1),
	}

	return tree, nil
}

func (s *Atree) Run() {
	go s.pageWriter()
	go s.pageReader()
}

// FIND

func (s *Atree) findDataPage(rootPageNo uint32, timestamp uint32) (uint32, []byte, error) {
	indexPageNo := rootPageNo
	for {
		buf, err := s.fetchIndexPage(indexPageNo)
		if err != nil {
			return 0, nil, fmt.Errorf("fetchIndexPage(%d): %s", indexPageNo, err)
		}

		foundPageNo := findPageNo(buf, timestamp)
		s.releasePage(indexPageNo)

		if buf[isDataPageNumbersIdx] == 1 {
			buf, err := s.fetchDataPage(foundPageNo)
			if err != nil {
				return 0, nil, fmt.Errorf("fetchDataPage(%d): %s", foundPageNo, err)
			}
			return foundPageNo, buf, nil
		}
		// вглубь
		indexPageNo = foundPageNo
	}
}

type pathLeg struct {
	PageNo uint32
	Data   []byte
}

type pathToDataPage struct {
	Legs       []pathLeg
	LastPageNo uint32
}

func (s *Atree) findPathToLastPage(rootPageNo uint32) (_ pathToDataPage, err error) {
	var (
		pageNo = rootPageNo
		legs   []pathLeg
	)

	for {
		var buf []byte
		buf, err = s.fetchIndexPage(pageNo)
		if err != nil {
			err = fmt.Errorf("FetchIndexPage(%d): %s", pageNo, err)
			return
		}

		legs = append(legs, pathLeg{
			PageNo: pageNo,
			Data:   buf,
			// childIdx не нужен
		})

		foundPageNo := getLastPageNo(buf)

		if buf[isDataPageNumbersIdx] == 1 {
			return pathToDataPage{
				Legs:       legs,
				LastPageNo: foundPageNo,
			}, nil
		}
		// вглубь
		pageNo = foundPageNo
	}
}

// APPEND DATA PAGE

type AppendDataPageReq struct {
	MetricID         uint32
	Timestamp        uint32
	Value            float64
	Since            uint32
	RootPageNo       uint32
	PrevPageNo       uint32
	TimestampsChunks [][]byte
	TimestampsSize   uint16
	ValuesChunks     [][]byte
	ValuesSize       uint16
}
type Report struct {
	//IsDataPageReused bool
	//DataPageNo       uint32
	IsRootChanged bool
	NewRootPageNo uint32
	//ReusedIndexPages []uint32
	Pages []PageToWrite
}

// type ChangedPage struct {
// 	PageNo   uint32
// 	Data     []byte
// 	IsReused bool
// }

// AppendDataPage - метод не записує дані в data-файл, а лише змінює дані
// в page cache та freeList і повертає звіт що змінено.
// Цей звіт txlog має записати в transaction log і лише потім можна змінювати data файл.
// Є ідея - записати у index файли заглушки 255,255,255,255 замість номерів сторінок і зберегти зміщення.
// А потім одним викликом отримати із FreeList список вільних сторінок.
// Тому що є проблема із відновленням FreeList після збою, якщо з нього будуть паралельно
// забирати та добавляти номери сторінок інші потоки.
// Це буде працювати, якщо додавання в txlog і маніпуляції із freeList будуть відбуватись в одному потоці
func (s *Atree) AppendDataPage(req AppendDataPageReq) (_ Report, err error) {
	var (
		pagesToRelease []uint32
		report         Report
	)

	newDataPage := s.allocPage()
	pagesToRelease = append(pagesToRelease, newDataPage.PageNo)

	chunksToDataPage(newDataPage.Data, chunksToDataPageReq{
		PrevPageNo:       req.PrevPageNo,
		TimestampsChunks: req.TimestampsChunks,
		TimestampsSize:   req.TimestampsSize,
		ValuesChunks:     req.ValuesChunks,
		ValuesSize:       req.ValuesSize,
	})

	report.Pages = append(report.Pages, PageToWrite{
		PageNo:   newDataPage.PageNo,
		Data:     newDataPage.Data,
		IsReused: newDataPage.IsReused,
	})

	// redoWriter, err := NewWriter(WriterOptions{
	// 	MetricID:         req.MetricID,
	// 	Timestamp:        req.Timestamp,
	// 	Value:            req.Value,
	// 	IsDataPageReused: newDataPage.IsReused,
	// 	DataPageNo:       newDataPage.PageNo,
	// 	Page:             newDataPage.Data,
	// })
	// if err != nil {
	// 	return
	// }

	if req.RootPageNo > 0 {
		var path pathToDataPage
		path, err = s.findPathToLastPage(req.RootPageNo)
		if err != nil {
			return
		}
		for _, leg := range path.Legs {
			pagesToRelease = append(pagesToRelease, leg.PageNo)
		}

		if path.LastPageNo != req.PrevPageNo {
			diploma.Abort(
				diploma.WrongPrevPageNo,
				fmt.Errorf("bug: last pageNo %d in tree != prev pageNo %d in _metric",
					path.LastPageNo, req.PrevPageNo),
			)
		}

		newPageNo := newDataPage.PageNo
		lastIdx := len(path.Legs) - 1

		for legIdx := lastIdx; legIdx >= 0; legIdx-- {
			leg := path.Legs[legIdx]

			ok := appendPair(leg.Data, req.Since, newPageNo)
			if ok {
				// index
				report.Pages = append(report.Pages, PageToWrite{
					PageNo: leg.PageNo,
					Data:   leg.Data,
				})
				// err = redoWriter.AppendIndexPage(leg.PageNo, leg.Data, 0)
				// if err != nil {
				// 	return
				// }
				break
			}

			newIndexPage := s.allocPage()
			pagesToRelease = append(pagesToRelease, newIndexPage.PageNo)
			appendPair(newIndexPage.Data, req.Since, newPageNo)
			// ставлю мітку що всі pageNo на сторінці - це data pageNo
			if legIdx == lastIdx {
				newIndexPage.Data[isDataPageNumbersIdx] = 1
			}

			// flags = 0
			// if newIndexPage.IsReused {
			// 	flags |= FlagReused
			// }
			report.Pages = append(report.Pages, PageToWrite{
				PageNo:   newIndexPage.PageNo,
				Data:     newIndexPage.Data,
				IsReused: newIndexPage.IsReused,
			})
			// err = redoWriter.AppendIndexPage(newIndexPage.PageNo, newIndexPage.Data, flags)
			// if err != nil {
			// 	return
			// }
			//
			newPageNo = newIndexPage.PageNo

			if legIdx == 0 {
				newRoot := s.allocPage()
				pagesToRelease = append(pagesToRelease, newRoot.PageNo)
				appendPair(newRoot.Data, getSince(leg.Data), leg.PageNo) // old rootPageNo
				appendPair(newRoot.Data, req.Since, newIndexPage.PageNo)

				// Фиксирую новый root в REDO логе

				report.Pages = append(report.Pages, PageToWrite{
					PageNo:   newRoot.PageNo,
					Data:     newRoot.Data,
					IsReused: newRoot.IsReused,
				})
				report.NewRootPageNo = newRoot.PageNo
				// flags = FlagNewRoot
				// if newRoot.IsReused {
				// 	flags |= FlagReused
				// }
				// err = redoWriter.AppendIndexPage(newRoot.PageNo, newRoot.Data, flags)
				// if err != nil {
				// 	return
				// }
				break
			}
		}
	} else {
		newRoot := s.allocPage()
		pagesToRelease = append(pagesToRelease, newRoot.PageNo)
		newRoot.Data[isDataPageNumbersIdx] = 1
		appendPair(newRoot.Data, req.Since, newDataPage.PageNo)

		report.Pages = append(report.Pages, PageToWrite{
			PageNo:   newRoot.PageNo,
			Data:     newRoot.Data,
			IsReused: newRoot.IsReused,
		})
		report.NewRootPageNo = newRoot.PageNo
		// flags = FlagNewRoot
		// if newRoot.IsReused {
		// 	flags |= FlagReused
		// }
		// err = redoWriter.AppendIndexPage(newRoot.PageNo, newRoot.Data, flags)
		// if err != nil {
		// 	return
		// }
	}

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

	// На данний момен схема - наступна. Всі сторінки - data та index - зафіксовані в кеші.
	// Отже запис на диск пройде максимально швидко. Після цього ReferenceCount кожної
	// сторінки зменшиться на 1. Оскільки на метрику утримується XLock, сторінки мають
	// ReferenceCount = 1 (немає інших читачів).
	// waitCh := make(chan struct{})

	// task := WriteTask{
	// 	WaitCh: waitCh,
	// 	Pages:  report.Pages,
	// }

	// s.appendWriteTaskToQueue(task)

	// <-waitCh

	// for _, pageNo := range dataPagesToRelease {
	// 	s.releasePage(pageNo)
	// }
	// for _, pageNo := range indexPagesToRelease {
	// 	s.releasePage(pageNo)
	// }
	return report, nil
}

// DELETE

type PageLists struct {
	DataPages  []uint32
	IndexPages []uint32
}

type Level struct {
	PageNo   uint32
	PageData []byte
	Idx      int
	ChildQty int
}

func (s *Atree) GetAllPages(rootPageNo uint32) (_ PageLists, err error) {
	var (
		dataPages  []uint32
		indexPages []uint32
		levels     []*Level
	)

	buf, err := s.fetchIndexPage(rootPageNo)
	if err != nil {
		err = fmt.Errorf("fetchIndexPage(%d): %s", rootPageNo, err)
		return
	}
	indexPages = append(indexPages, rootPageNo)

	if buf[isDataPageNumbersIdx] == 1 {
		pageNumbers := listPageNumbers(buf)
		dataPages = append(dataPages, pageNumbers...)

		s.releasePage(rootPageNo)

		return PageLists{
			DataPages:  dataPages,
			IndexPages: indexPages,
		}, nil
	}

	levels = append(levels, &Level{
		PageNo:   rootPageNo,
		PageData: buf,
		Idx:      0,
		ChildQty: bin.GetUint16AsInt(buf[indexRecordsQtyIdx:]),
	})

	for {
		if len(levels) == 0 {
			return PageLists{
				DataPages:  dataPages,
				IndexPages: indexPages,
			}, nil
		}

		lastIdx := len(levels) - 1
		level := levels[lastIdx]

		if level.Idx < level.ChildQty {
			pageNo := getPageNo(level.PageData, level.Idx)
			level.Idx++

			var buf []byte
			buf, err = s.fetchIndexPage(pageNo)
			if err != nil {
				err = fmt.Errorf("fetchIndexPage(%d): %s", pageNo, err)
				return
			}
			indexPages = append(indexPages, pageNo)

			if buf[isDataPageNumbersIdx] == 1 {
				pageNumbers := listPageNumbers(buf)
				dataPages = append(dataPages, pageNumbers...)

				s.releasePage(pageNo)
			} else {
				levels = append(levels, &Level{
					PageNo:   pageNo,
					PageData: buf,
					Idx:      0,
					ChildQty: bin.GetUint16AsInt(buf[indexRecordsQtyIdx:]),
				})
			}
		} else {
			s.releasePage(level.PageNo)
			levels = levels[:lastIdx]
		}
	}
}