qb/storage/write_preparer.go

package storage

import (
	"bytes"
	"io"

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

type Packable interface {
	Pack(io.Writer)
}

type WritePreparer struct {
	pagePreparer *PageIngester
	//dataPagePayloadBound int
	w          *bytes.Buffer
	indexPages []PageToWrite
	dataPages  []PageToWrite
	commits    []any
}

type WritePreparerOptions struct {
	MinWALBufferSize int
	PageIngester     *PageIngester
}

func NewWritePreparer(opt WritePreparerOptions) *WritePreparer {
	//buf := make([]byte, opt.MinWALBufferSize)
	s := &WritePreparer{
		pagePreparer: opt.PageIngester,
		w:            bytes.NewBuffer(nil),
	}
	return s
}

type PreparedData struct {
	Packet       []byte // пакет для запису в WAL
	WriteToIndex []PageToWrite
	WriteToData  []PageToWrite
	Commits      []any
}

// const (
// 	prefixSize = 9
// 	crc32Size  = 4
// )

func (s *WritePreparer) Prepare(input []any) PreparedData {
	s.w.Reset()
	s.w.Write([]byte{
		0, 0, 0, 0, 0, 0, 0, 0, 0, // size (max 9 byte)
	})

	hasher := util.NewHasher()

	multi := io.MultiWriter(s.w, hasher)

	// 1. Пакую всі дані в  WAL буфер запису
	for _, untyped := range input {
		switch x := untyped.(type) {
		case MetricAdd:
			rec := MetricAddRecord{
				MetricID:   x.MetricID,
				MetricType: x.MetricType,
				FracDigits: x.FracDigits,
			}
			rec.Pack(multi)
			s.commits = append(s.commits, MetricAddCommited{
				MetricID: x.MetricID,
				ResultCh: x.ResultCh,
			})

		case MetricDelete:
			// if len(x.FreeIndexPages) > 0 {
			// 	s.indexFreeList.AddPageNumbers(x.FreeIndexPages)
			// }
			// if len(x.FreeIndexPages) > 0 {
			// 	s.dataFreeList.AddPageNumbers(x.FreeDataPages)
			// }
			rec := MetricDeleteRecord{
				MetricID:       x.MetricID,
				FreeIndexPages: x.FreeIndexPages,
				FreeDataPages:  x.FreeDataPages,
			}
			rec.Pack(multi)
			s.commits = append(s.commits, MetricDeleteCommited{
				MetricID: x.MetricID,
				ResultCh: x.ResultCh,
			})

		case MeasuresAppend:
			rec := MeasuresAppendRecord{
				MetricID:               x.MetricID,
				TimestampsRewindOffset: x.TimestampsRewindOffset,
				Timestamps:             x.Timestamps,
				ValuesRewindOffset:     x.ValuesRewindOffset,
				Values:                 x.Values,
			}
			rec.Pack(multi)
			// Дані для Worker
			s.commits = append(s.commits, MeasuresAppendCommited{
				MetricID:     x.MetricID,
				ResultCode:   x.ResultCode,
				WrittenCount: x.WrittenCount,
				ResultCh:     x.ResultCh,
			})

		case MeasuresAppendWithGrow:
			sealedDataPages, sealedIndexLevels := s.pagePreparer.Ingest(IngestIn{
				LastPageNo:      x.LastPageNo,
				IndexLevelTails: x.IndexLevelTails,
				DataPages:       x.DataPages,
			})

			// сторінки для запису в index та data файли
			for _, p := range sealedDataPages {
				s.dataPages = append(s.dataPages, PageToWrite{
					PageNo:  p.PageNo,
					Content: p.Content,
				})
			}
			for _, level := range sealedIndexLevels {
				for _, p := range level.IndexPages {
					//fmt.Printf("records: %v\n", getIndexRecords(p.Content, 2))
					s.indexPages = append(s.indexPages, PageToWrite{
						PageNo:  p.PageNo,
						Content: p.Content,
					})
				}
			}

			var (
				// дані для запису в WAL
				changedIndexPages []ChangedIndexLevel
				// дані для Воркера
				indexLevelTails []IndexLevelTail
			)

			for _, level := range sealedIndexLevels {
				if len(level.IndexPages) > 0 || level.SkipRecords != level.TailRecordsCount {
					// на індексному внесені зміни
					var (
						indexPageTail *IndexPageTail
						indexPages    []SealedIndexPage
					)
					if len(level.IndexPages) > 0 {
						first := level.IndexPages[0]
						firstCRC32, _ := bin.GetUint32(first.Content[indexCRC32Idx:])
						skipSize := level.SkipRecords * IndexRecordSize
						indexPageTail = &IndexPageTail{
							PageNo: first.PageNo,
							Reused: first.Reused,
							CRC32:  firstCRC32,
							// в WAL файл попадає лише payload
							Records: first.Content[skipSize : maxRecordsOnIndexPage*IndexRecordSize],
						}
						indexPages = level.IndexPages[1:]
					} else {
						indexPages = level.IndexPages
					}
					changedIndexPages = append(changedIndexPages, ChangedIndexLevel{
						IndexPageTail: indexPageTail,
						IndexPages:    indexPages,
						// в WAL файл попадає лише payload
						TailRecords: level.TailRecords[:level.TailRecordsCount*IndexRecordSize],
					})
				}
				// Worker-у відправляю повний index
				indexLevelTails = append(indexLevelTails, IndexLevelTail{
					// Worker отримує повний буфер
					Buffer:       level.TailRecords,
					RecordsCount: level.TailRecordsCount,
				})
			}

			first := sealedDataPages[0]
			firstCRC32, _ := bin.GetUint32(first.Content[dataCRC32Idx:])

			rec := MeasuresAppendWithGrowRecord{
				MetricID: x.MetricID,
				DataPageTail: DataPageTail{
					PageNo:                 first.PageNo,
					Reused:                 first.Reused,
					PrevPageNo:             x.LastPageNo,
					CRC32:                  firstCRC32,
					TimestampsRewindOffset: x.TimestampsRewindOffset,
					Timestamps:             x.Timestamps,
					ValuesRewindOffset:     x.ValuesRewindOffset,
					Values:                 x.Values,
				},
				DataPages:          sealedDataPages[1:],
				TailTimestamps:     x.TailTimestamps,
				TailValues:         x.TailValues,
				ChangedIndexLevels: changedIndexPages,
			}

			//fmt.Printf("%#v\ns", rec.ChangedIndexLevels[0].IndexPageTail)

			rec.Pack(multi)
			// Дані для Worker
			s.commits = append(s.commits, MeasuresAppendWithGrowCommited{
				MetricID:     x.MetricID,
				LastPageNo:   sealedDataPages[len(sealedDataPages)-1].PageNo,
				Index:        indexLevelTails,
				ResultCode:   x.ResultCode,
				WrittenCount: x.WrittenCount,
				ResultCh:     x.ResultCh,
			})

		case MeasuresDelete:
			// if len(x.FreeIndexPages) > 0 {
			// 	s.indexFreeList.AddPageNumbers(x.FreeIndexPages)
			// }
			// if len(x.FreeIndexPages) > 0 {
			// 	s.dataFreeList.AddPageNumbers(x.FreeDataPages)
			// }
			rec := MeasuresDeleteRecord{
				MetricID:       x.MetricID,
				FreeIndexPages: x.FreeIndexPages,
				FreeDataPages:  x.FreeDataPages,
			}
			rec.Pack(multi)
			s.commits = append(s.commits, MeasuresDeleteCommited{
				MetricID: x.MetricID,
				ResultCh: x.ResultCh,
			})
			//case DeletedMeasuresSince:
		}
	}
	// 2. Додаю розмір пакету і чексуму
	packet := finalizePacket(s.w, hasher.Sum32())
	//s.written += int64(len(packet)) + 12
	//
	return PreparedData{
		Packet:       packet,
		WriteToIndex: s.indexPages,
		WriteToData:  s.dataPages,
		Commits:      s.commits,
	}
}

func (s *WritePreparer) Reset() {
	s.w.Reset()
	s.indexPages = nil
	s.dataPages = nil
	s.commits = nil
}

func getIndexRecords(buf []byte, count int) (list []IndexRecord) {
	//fmt.Printf("ipage: % x\n", buf)
	i := 0
	for range count {
		var rec IndexRecord
		rec.Timestamp, _ = bin.GetUint32(buf[i:])
		rec.PageNo, _ = bin.GetUint32(buf[i+4:])
		list = append(list, rec)
		i += IndexRecordSize
	}
	return
}

func finalizePacket(w *bytes.Buffer, checksum uint32) []byte {
	bin.WriteUint32(w, checksum)
	packet := w.Bytes()
	bodySize := len(packet) - 13 // 9 bytes prefix + 4 bytes crc32
	// cut empty bytes at start
	start := 9 - bin.CountVarSize(bodySize)
	bin.PutVarSize(packet[start:], bodySize)

	//fmt.Printf("checksum: %d\n", checksum)
	//fmt.Printf("bodySize: %d\n", bodySize)
	//fmt.Printf("start: %d\n", start)
	//fmt.Printf("body: % x\n", packet[9:len(packet)-4])

	return packet[start:]
}