qb/storage/write_preparer.go

package storage

import (
	"bytes"
	"fmt"
	"io"

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

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
}

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

	hasher := util.NewHasher()

	w := 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(w)
			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(w)
			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(w)
			// Дані для 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]
						skipSize := level.SkipRecords * indexRecordSize
						indexPageTail = &IndexPageTail{
							PageNo: first.PageNo,
							Reused: first.Reused,
							// в 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:],
				ChangedIndexLevels: changedIndexPages,
				TailTimestamps:     x.Timestamps,
				TailValues:         x.Values,
			}

			rec.Pack(w)
			// Дані для 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(w)
			s.commits = append(s.commits, MeasuresDeleteCommited{
				MetricID: x.MetricID,
				ResultCh: x.ResultCh,
			})
			//case DeletedMeasuresSince:
		}
	}
	// 2. Додаю розмір пакету і чексуму

	//s.written += int64(len(packet)) + 12

	packet := s.w.Bytes()
	// write size
	payloadSize := len(packet) - 13
	start := 9 - bin.CountVarSize(payloadSize)
	bin.PutVarSize(packet[start:], payloadSize)
	// CRC32
	bin.PutUint32(packet[9:], hasher.Sum32())
	//
	return PreparedData{
		Packet:       packet[start:],
		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
}