qb/storage/writer.go

package storage

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"gordenko.dev/dima/qb"
	"gordenko.dev/dima/qb/freelist"
	"gordenko.dev/dima/qb/inbox"
)

const (
	lsnSize          = 4
	packetPrefixSize = 12 // 4 lsn + 4 packet length + 4 crc32

	lengthIdx   = 0
	checksumIdx = 4
	lsnIdx      = 8

	filePerm = 0770

	dumpSnapshotAfterNBytes = 1024 * 1024 * 1024 // 1 GB

	writeBufferSize = 4 * 1024 * 1024
)

const (
	CodeMetricAdd              byte = 1
	CodeMetricDelete           byte = 2
	CodeMeasuresAppendWithGrow byte = 5
	CodeMeasuresAppend         byte = 6
	CodeMeasuresDelete         byte = 7
)

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

type MetricsState struct {
	//Metrics []Metric
	WaitCh chan struct{}
}

type Changes struct {
	Commits []any
	// for snapshot only
	SnapshotNumberCh      chan int // log number
	FrozenIndexPagesCount int
	IndexPageNumbers      []uint32
	FrozenDataPagesCount  int
	DataPageNumbers       []uint32
}

type Writer struct {
	mutex           sync.Mutex
	inbox           *inbox.Inbox
	workerInbox     *inbox.Inbox
	dataPagesCount  uint32
	indexPagesCount uint32
	dataFreeList    *freelist.FreeList
	indexFreeList   *freelist.FreeList
	//atree               *atree.Atree
	dir           string
	databaseName  string
	w             *bytes.Buffer // wal буфер для упаковки даних перед записом на диск
	wal           *os.File
	dataFile      *os.File
	indexFile     *os.File
	input         []any
	writePreparer *WritePreparer
	written       int64
	isExited      bool
	exitCh        chan struct{}
	waitGroup     *sync.WaitGroup
}

type WriterOptions struct {
	Inbox         *inbox.Inbox
	WorkerInbox   *inbox.Inbox
	Dir           string
	DatabaseName  string
	WAL           string
	DataFreeList  *freelist.FreeList
	IndexFreeList *freelist.FreeList
	DataFile      *os.File
	IndexFile     *os.File
	ExitCh        chan struct{}
	WaitGroup     *sync.WaitGroup
}

func NewWriter(opt WriterOptions) (*Writer, error) {
	// if (opt.IndexPageSize % 2) != 0 {
	// 	return nil, errors.New("IndexPageSize must be multiple of 2")
	// }
	// if (opt.DataPageSize % 2) != 0 {
	// 	return nil, errors.New("DataPageSize must be multiple of 2")
	// }
	if opt.Inbox == nil {
		return nil, errors.New("Inbox option is required")
	}
	if opt.WorkerInbox == nil {
		return nil, errors.New("WorkerInbox option is required")
	}
	if opt.Dir == "" {
		return nil, errors.New("Dir option is required")
	}
	if opt.DatabaseName == "" {
		return nil, errors.New("DatabaseName option is required")
	}
	// if opt.AppendToWorkerQueue == nil {
	// 	return nil, errors.New("AppendToWorkerQueue option is required")
	// }
	if opt.DataFreeList == nil {
		return nil, errors.New("DataFreeList option is required")
	}
	if opt.IndexFreeList == nil {
		return nil, errors.New("IndexFreeList option is required")
	}
	if opt.DataFile == nil {
		return nil, errors.New("DataFile option is required")
	}
	if opt.IndexFile == nil {
		return nil, errors.New("IndexFile option is required")
	}
	if opt.ExitCh == nil {
		return nil, errors.New("ExitCh option is required")
	}
	if opt.WaitGroup == nil {
		return nil, errors.New("WaitGroup option is required")
	}

	indexPageNoProvider, err := NewPageNoProvider(PageNoProviderOptions{
		FilePath: qb.GetIndexFilePath(opt.Dir, opt.DatabaseName),
		PageSize: IndexPageSize,
		FreeList: opt.IndexFreeList,
	})

	dataPageNoProvider, err := NewPageNoProvider(PageNoProviderOptions{
		FilePath: qb.GetDataFilePath(opt.Dir, opt.DatabaseName),
		PageSize: DataPageSize,
		FreeList: opt.DataFreeList,
	})

	pageIngester, err := NewPageIngester(PageIngesterOptions{
		GetIndexPageNumber: indexPageNoProvider.GetPageNumber,
		GetDataPageNumber:  dataPageNoProvider.GetPageNumber,
	})
	if err != nil {
		return nil, err
	}

	writePreparer := NewWritePreparer(WritePreparerOptions{
		//MinWALBufferSize: 128,
		PageIngester: pageIngester,
	})

	s := &Writer{
		inbox:         opt.Inbox,
		workerInbox:   opt.WorkerInbox,
		dir:           opt.Dir,
		databaseName:  opt.DatabaseName,
		dataFile:      opt.DataFile,
		indexFile:     opt.IndexFile,
		dataFreeList:  opt.DataFreeList,
		indexFreeList: opt.IndexFreeList,
		writePreparer: writePreparer,
		exitCh:        opt.ExitCh,
		waitGroup:     opt.WaitGroup,
	}

	if opt.WAL != "" {
		s.wal, err = os.OpenFile(filepath.Join(opt.Dir, opt.WAL), os.O_APPEND|os.O_WRONLY, filePerm)
		if err != nil {
			return nil, err
		}
	} else {
		s.wal, err = os.OpenFile(
			qb.GetWALFilePath(opt.Dir, opt.DatabaseName, 0), // .wal_0
			os.O_CREATE|os.O_APPEND|os.O_WRONLY,
			filePerm,
		)
		if err != nil {
			return nil, err
		}
	}

	return s, nil
}

func (s *Writer) Run() {
	for {
		select {
		case <-s.inbox.Ready():
			if err := s.packAndWrite(); err != nil {
				qb.Abort(qb.FailedWriteToTxLog, err)
			}

		case <-s.exitCh:
			s.exit()
			return
		}
	}
}

func (s *Writer) packAndWrite() (err error) {
	//fmt.Println("packAndWrite")
	//s.mutex.Lock()
	//isExited := s.isExited

	// var exitWaitGroup *sync.WaitGroup
	// if s.isExited {
	// 	exitWaitGroup = s.waitGroup
	// }
	//s.mutex.Unlock()

	input := s.inbox.Drain()

	//fmt.Printf("Drain: %#v\n", input)

	prepared := s.writePreparer.Prepare(input)

	//fmt.Printf("prepared: %#v\n", prepared)
	fmt.Println("prepared")

	// 3. Пишу на диск WAL (append)
	n, err := s.wal.Write(prepared.Packet)
	if err != nil {
		return
	}
	//fmt.Printf("written to wal: % x\n", prepared.Packet)
	if n != len(prepared.Packet) {
		return fmt.Errorf("written %d != packet size %d", n, len(prepared.Packet))
	}
	if err = s.wal.Sync(); err != nil {
		return
	}
	fmt.Println("synced to wal")

	// 4. Пишу в atree сторінки
	if len(prepared.WriteToData) > 0 {
		err = WriteDataPages(s.dataFile, prepared.WriteToData)
		if err != nil {
			return
		}
		if err = s.dataFile.Sync(); err != nil {
			return
		}
		fmt.Println("written to data")
	}
	if len(prepared.WriteToIndex) > 0 {
		err = WriteIndexPages(s.indexFile, prepared.WriteToIndex)
		if err != nil {
			return
		}
		if err = s.indexFile.Sync(); err != nil {
			return
		}
		fmt.Println("written to index")
	}
	// 6. відправляю input - worker-у

	var (
		forceSnapshot bool
	)

	if s.written > dumpSnapshotAfterNBytes {
		forceSnapshot = true
	}

	// if isExited && s.written > 0 {
	// 	forceSnapshot = true
	// }

	if forceSnapshot {
		if err := s.wal.Close(); err != nil {
			return fmt.Errorf("close wal file: %s", err)
		}
		s.wal = nil
		s.written = 0

		snapshotNumberCh := make(chan int, 1)

		s.workerInbox.Push(Changes{
			Commits:               prepared.Commits,
			SnapshotNumberCh:      snapshotNumberCh,
			FrozenIndexPagesCount: s.indexFreeList.Pages(),
			IndexPageNumbers:      s.indexFreeList.Cached(),
			FrozenDataPagesCount:  s.dataFreeList.Pages(),
			DataPageNumbers:       s.dataFreeList.Cached(),
		})

		snapshotNumber := <-snapshotNumberCh

		// копіюю сторінки із delta файла в base файл і потім роблю Truncate
		err = s.indexFreeList.Merge()
		if err != nil {
			return
		}
		err = s.dataFreeList.Merge()
		if err != nil {
			return
		}

		var err error
		s.wal, err = os.OpenFile(
			qb.GetWALFilePath(s.dir, s.databaseName, snapshotNumber),
			os.O_CREATE|os.O_WRONLY,
			filePerm,
		)
		if err != nil {
			return fmt.Errorf("create new changes file: %s", err)
		}
	} else {
		s.workerInbox.Push(Changes{
			Commits: prepared.Commits,
		})
		fmt.Println("pushed commits")
	}
	return nil
}

// func (s *Writer) reset() {
// 	s.buf.Reset()
// 	s.buf.Write([]byte{
// 		0, 0, 0, 0, // packet length
// 		0, 0, 0, 0, // crc32
// 		0, 0, 0, 0, // lsn

// 	})
// 	s.pagesToWrite = nil
// 	s.workerReqs = nil
// 	s.waitCh = make(chan struct{})
// }

// func (s *Writer) flush(w *bytes.Buffer) error {

// 	return nil
// }

func (s *Writer) exit() {
	s.mutex.Lock()
	s.isExited = true
	s.mutex.Unlock()

	if err := s.packAndWrite(); err != nil {
		qb.Abort(qb.FailedWriteToTxLog, err)
	}
}

// type IndexPayloadToWrite struct {
// 	PageNo    uint32
// 	Payload   []byte
// 	ZeroLevel bool
// }

type PageToWrite struct {
	PageNo  uint32
	Content []byte
}

func WriteDataPages(file *os.File, pages []PageToWrite) (err error) {
	for _, p := range pages {
		if len(p.Content) != DataPageSize {
			return fmt.Errorf("wrong data page size: %d", len(p.Content))
		}
		var (
			off = int(p.PageNo-1) * DataPageSize
			n   int
		)
		n, err = file.WriteAt(p.Content, int64(off))
		if err != nil {
			return fmt.Errorf("file.WriteAt: %s; off=%d", err, off)
		}
		if n != DataPageSize {
			return fmt.Errorf("write %d instead of %d", n, DataPageSize)
		}
	}
	return
}

func WriteIndexPages(file *os.File, pages []PageToWrite) (err error) {
	for _, p := range pages {
		if len(p.Content) != IndexPageSize {
			return fmt.Errorf("wrong index page size: %d", len(p.Content))
		}
		var (
			off = int(p.PageNo-1) * IndexPageSize
			n   int
		)
		n, err = file.WriteAt(p.Content, int64(off))
		if err != nil {
			return
		}
		if n != IndexPageSize {
			return fmt.Errorf("write %d instead of %d", n, IndexPageSize)
		}
	}
	return
}

// API

// FIX - add
// type DeletedMeasuresSince struct {
// 	MetricID       uint32
// 	LastPageNo     uint32
// 	IsRootChanged  bool
// 	RootPageNo     uint32
// 	FreeDataPages  []uint32
// 	FreeIndexPages []uint32
// 	TimestampsBuf  []byte
// 	ValuesBuf      []byte
// }