qb/atree/atree.go

package atree

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

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

const (
	PageTypeData  = 1
	PageTypeIndex = 2

	filePerm = 0770

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

	// index page
	indexCRC32Idx      = IndexPageSize - 4
	isLastLevelIdx     = IndexPageSize - 5
	indexRecordsQtyIdx = IndexPageSize - 7

	indexRecordSize = 8

	// data page
	dataCRC32Idx      = DataPageSize - 4
	prevPageIdx       = DataPageSize - 8
	valuesSizeIdx     = DataPageSize - 10
	timestampsSizeIdx = DataPageSize - 12

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

	PageSize      = 8192
	DataPageSize  = 8192
	IndexPageSize = 1024
	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 {
	file              *os.File
	mutex             sync.Mutex
	freelist          FreeList
	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)

		// fix
		if buf[isLastLevelIdx] == 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)

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

// APPEND DATA PAGE

//MetricID         uint32
//Timestamp        uint32
//Value            float64
//RootPageNo       uint32
//PrevPageNo       uint32

type AppendDataPageReq struct {
	Since            uint32
	TimestampsChunks [][]byte
	TimestampsSize   uint16
	ValuesChunks     [][]byte
	ValuesSize       uint16
}

// 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 будуть відбуватись в одному потоці
// NotLinkedDataPage - це data сторінка із payload, але без встановленого prevPageNo та без розрахованого CRC32
type NotLinkedDataPage struct {
	Since uint32
	Data  []byte
}

func (s NotLinkedDataPage) SetPrevPageNo(prevPageNo uint32) {
	bin.PutUint32(s.Data[prevPageIdx:], prevPageNo)
	bin.PutUint32(s.Data[crc32Idx:], calcChecksum(s.Data[:crc32Idx]))
}

type AppendDataPagesReq struct {
	LastPageNo uint32
	Legs       []PathLeg
	DataPages  []NotLinkedDataPage
}

type Report struct {
	NewRootPageNo uint32
	LastPageNo    uint32
	Pages         []PageToWrite
}

// Можливо atree не треба блокувати при читанні. Я можу створити копії index сторінок,
// а потім під 1 мутексом замінити в pages.
func (s *Atree) AppendDataPages(req AppendDataPagesReq) Report {
	var (
		//pagesToRelease []uint32
		newRootPageNo uint32
		lastPageNo    = req.LastPageNo
		pages         []PageToWrite
		legs          = req.Legs
		changed       = make(map[uint32]PathLeg)
	)

	for _, p := range req.DataPages {
		newDataPageNo, isReused := s.allocPageNumber() // alloc only number
		//	pagesToRelease = append(pagesToRelease, newDataPage.PageNo)

		// set prevPageNo
		p.SetPrevPageNo(lastPageNo)

		pages = append(pages, PageToWrite{
			PageNo:   newDataPageNo,
			Data:     p.Data,
			IsReused: isReused,
		})

		// FIX - після додавання index page потрібно модифікувати path,
		// або будувати щось типу дерева знизу вверх

		if len(legs) > 0 {
			newPageNo := newDataPageNo
			lastIdx := len(legs) - 1

			for legIdx := lastIdx; legIdx >= 0; legIdx-- {
				leg := req.Legs[legIdx]
				ok := appendPair(leg.Data, p.Since, newPageNo)
				if ok {
					// index FIX
					// потрібно запам'ятати змінені сторінки, але їх можуть змінювати
					// кілька ітерацій, тому додавати в Pages не можна.
					// на індексній сторінці достатньо місця. Запис вставлено.
					changed[leg.PageNo] = leg
					break
				}
				// на індексній сторінці НЕ достатньо місця. Створюю нову.
				newIndexPage := s.allocPage()
				//pagesToRelease = append(pagesToRelease, newIndexPage.PageNo)
				appendPair(newIndexPage.Data, p.Since, newPageNo)
				// ставлю мітку що всі pageNo на сторінці - це data pageNo
				// fix - єдина оптимізація від існування isDataPageNumbersIdx - getAllPages не завантажує data pages
				// if legIdx == lastIdx {
				// 	newIndexPage.Data[isDataPageNumbersIdx] = 1
				// }
				pages = append(pages, PageToWrite{
					PageNo:   newIndexPage.PageNo,
					Data:     newIndexPage.Data,
					IsReused: newIndexPage.IsReused,
				})
				// замінюю крок в path на новий
				legs[legIdx] = PathLeg{
					PageNo: newIndexPage.PageNo,
					Data:   newIndexPage.Data,
				}
				//
				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, p.Since, newIndexPage.PageNo)

					// Фиксирую новый root в REDO логе
					pages = append(pages, PageToWrite{
						PageNo:   newRoot.PageNo,
						Data:     newRoot.Data,
						IsReused: newRoot.IsReused,
					})
					newRootPageNo = newRoot.PageNo
					// додаю в початок списку кроків нову root сторінку
					legs = append([]PathLeg{
						{
							PageNo: newRoot.PageNo,
							Data:   newRoot.Data,
						},
					}, legs...)
					break
				}
			}
		} else {
			// індексну root сторінку створюю одразу для першої data сторінки,
			// root data сторінки, як в B+Tree не буває.
			newRoot := s.allocPage()
			//pagesToRelease = append(pagesToRelease, newRoot.PageNo)
			//newRoot.Data[isDataPageNumbersIdx] = 1
			appendPair(newRoot.Data, p.Since, newDataPageNo)

			pages = append(pages, PageToWrite{
				PageNo:   newRoot.PageNo,
				Data:     newRoot.Data,
				IsReused: newRoot.IsReused,
			})
			newRootPageNo = newRoot.PageNo
			// додаю в початок списку кроків нову root сторінку
			legs = append([]PathLeg{
				{
					PageNo: newRoot.PageNo,
					Data:   newRoot.Data,
				},
			}, legs...)
		}

		// На данний момен схема - наступна. Всі сторінки - data та index - зафіксовані в кеші.
		// Отже запис на диск пройде максимально швидко. Після цього ReferenceCount кожної
		// сторінки зменшиться на 1. Оскільки на метрику утримується XLock, сторінки мають
		// ReferenceCount = 1 (немає інших читачів).
		// for _, pageNo := range indexPagesToRelease {
		// 	s.releasePage(pageNo)
		// }
		lastPageNo = newDataPageNo
	}
	for _, leg := range changed {
		// fix recalc checksum
		pages = append(pages, PageToWrite{
			PageNo: leg.PageNo,
			Data:   leg.Data,
		})
	}
	return Report{
		NewRootPageNo: newRootPageNo,
		LastPageNo:    lastPageNo,
		Pages:         pages,
	}
}

// DELETE

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

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

	buf, err := s.fetchIndexPage(rootPageNo)
	if err != nil {
		return nil, fmt.Errorf("fetchIndexPage(%d): %s", rootPageNo, err)
	}
	pageNumbers = append(pageNumbers, 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 pageNumbers, 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.fetchPage(pageNo)
			if err != nil {
				return nil, fmt.Errorf("fetchPage(%d): %s", pageNo, err)
			}
			pageNumbers = append(pageNumbers, pageNo)

			if buf[pageTypeIdx] == PageTypeData {
				//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]
		}
	}
}