qb/storage/data_preparer.go

package storage

import (
	"bytes"
	"io"

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

type DataPreparer struct {
	pagePreparer *PagePreparer
	//dataPagePayloadBound int
	w            *bytes.Buffer
	indexPages   []PageToWrite
	dataPages    []PageToWrite
	writeResults []any
}

type DataPreparerOptions struct {
	MinWALBufferSize int
	PagePreparer     *PagePreparer
}

func NewDataPreparer(opt DataPreparerOptions) *DataPreparer {
	buf := make([]byte, opt.MinWALBufferSize)
	s := &DataPreparer{
		pagePreparer: opt.PagePreparer,
		w:            bytes.NewBuffer(buf),
	}
	return s
}

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

func (s *DataPreparer) 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 MetricAddRecord:
			x.Pack(w)
		// case DeletedMetric:
		// 	if len(x.FreePageNumbers) > 0 {
		// 		s.freeList.AddPageNumbers(x.FreePageNumbers)
		// 	}
		case AppendedMeasures:
			sealResult := s.pagePreparer.SealPages(SealPagesIn{
				LastPageNo:      x.LastPageNo,
				IndexLevelTails: x.IndexLevelTails,
				DataPages:       x.DataPages,
			})

			// сторінки для запису в index та data файли
			for _, p := range sealResult.DataPages {
				s.dataPages = append(s.dataPages, PageToWrite{
					PageNo:  p.PageNo,
					Content: p.Content,
				})
			}
			for _, level := range sealResult.IndexLevels {
				for _, p := range level.IndexPages {
					s.indexPages = append(s.indexPages, PageToWrite{
						PageNo:  p.PageNo,
						Content: p.Content,
					})
				}
			}

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

			for _, level := range sealResult.IndexLevels {
				if len(level.IndexPages) == 0 && level.SkipRecords == level.TailRecordsCount {
					break
				}
				var completedIndexPage *HeadIndexPage
				if len(level.IndexPages) > 0 {
					first := level.IndexPages[0]
					skipSize := level.SkipRecords * indexRecordSize
					completedIndexPage = &HeadIndexPage{
						PageNo:   first.PageNo,
						Reused:   first.Reused,
						Checksum: first.Checksum,
						Records:  first.Content[skipSize:],
					}
				}
				changedIndexPages = append(changedIndexPages, ChangedIndexLevel{
					HeadIndexPage: completedIndexPage,
					IndexPages:    level.IndexPages[1:],
					TailRecords:   level.TailRecords[:level.TailRecordsCount*indexRecordSize],
				})

				indexLevelTails = append(indexLevelTails, IndexLevelTail{
					Buffer:       level.TailRecords,
					RecordsCount: level.TailRecordsCount,
				})
			}

			first := sealResult.DataPages[0]

			rec := MeasuresAppendWithGrowRecord{
				MetricID: x.MetricID,
				HeadDataPage: HeadDataPage{
					PageNo:           first.PageNo,
					Reused:           first.Reused,
					PrevPageNo:       first.PrevPageNo,
					Checksum:         first.Checksum,
					TimestampsOffset: x.TimestampsOffset,
					Timestamps:       x.Timestamps,
					ValuesOffset:     x.ValuesOffset,
					Values:           x.Values,
				},
				DataPages:          sealResult.DataPages[1:],
				ChangedIndexLevels: changedIndexPages,
				TailTimestamps:     x.Timestamps,
				TailValues:         x.Values,
			}

			rec.WriteTo(s.w)

			// Дані для Worker
			s.writeResults = append(s.writeResults, AppendMeasuresSummary{
				MetricID:     x.MetricID,
				LastPageNo:   sealResult.DataPages[len(sealResult.DataPages)-1].PageNo,
				Index:        indexLevelTails,
				ResultCode:   x.ResultCode,
				WrittenCount: x.WrittenCount,
				ResultCh:     x.ResultCh,
			})

			//case DeletedMeasures:
			//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)
	// checksum
	bin.PutUint32(packet[9:], hasher.Sum32())
	//
	return PreparedData{
		Packet:       packet[start:],
		WriteToIndex: s.indexPages,
		WriteToData:  s.dataPages,
		WriteResults: s.writeResults,
	}
}

func (s *DataPreparer) Reset() {
	s.w.Reset()
	s.indexPages = nil
	s.dataPages = nil
	s.writeResults = nil
}

// MeasuresToWrite
type AppendedMeasures struct {
	MetricID         uint32
	LastPageNo       uint32
	TimestampsOffset int // offset on 1st page
	Timestamps       []byte
	ValuesOffset     int // offset on 1st page
	Values           []byte
	IndexLevelTails  []IndexLevelTail
	DataPages        []DataPayload
	TailTimestamps   []byte // data level tail
	TailValues       []byte // data level tail
	ResultCode       byte
	WrittenCount     int
	ResultCh         chan struct{}
}

// WRITE RESULTS

// результат
type AppendMeasuresSummary struct {
	MetricID     uint32
	LastPageNo   uint32
	Index        []IndexLevelTail
	ResultCode   byte
	WrittenCount int
	ResultCh     chan struct{}
}