influxdb/tsdb/tsi1/measurement_block_test.go

package tsi1_test

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"reflect"
	"testing"

	"github.com/influxdata/platform/tsdb"
	"github.com/influxdata/platform/tsdb/tsi1"
)

func TestReadMeasurementBlockTrailer(t *testing.T) {
	// Build a trailer
	var (
		data                       = make([]byte, tsi1.MeasurementTrailerSize)
		blockversion               = uint16(1)
		blockOffset, blockSize     = uint64(1), uint64(2500)
		hashIdxOffset, hashIdxSize = uint64(2501), uint64(1000)
	)

	binary.BigEndian.PutUint64(data[0:], blockOffset)
	binary.BigEndian.PutUint64(data[8:], blockSize)
	binary.BigEndian.PutUint64(data[16:], hashIdxOffset)
	binary.BigEndian.PutUint64(data[24:], hashIdxSize)
	binary.BigEndian.PutUint64(data[32:], 0)
	binary.BigEndian.PutUint64(data[40:], 0)
	binary.BigEndian.PutUint64(data[48:], 0)
	binary.BigEndian.PutUint64(data[56:], 0)
	binary.BigEndian.PutUint16(data[64:], blockversion)

	trailer, err := tsi1.ReadMeasurementBlockTrailer(data)
	if err != nil {
		t.Logf("trailer is: %#v\n", trailer)
		t.Fatal(err)
	}

	ok := true &&
		trailer.Version == int(blockversion) &&
		trailer.Data.Offset == int64(blockOffset) &&
		trailer.Data.Size == int64(blockSize) &&
		trailer.HashIndex.Offset == int64(hashIdxOffset) &&
		trailer.HashIndex.Size == int64(hashIdxSize)

	if !ok {
		t.Fatalf("got %v\nwhich does not match expected", trailer)
	}
}

func TestMeasurementBlockTrailer_WriteTo(t *testing.T) {
	var trailer = tsi1.MeasurementBlockTrailer{
		Version: 1,
		Data: struct {
			Offset int64
			Size   int64
		}{Offset: 1, Size: 2},
		HashIndex: struct {
			Offset int64
			Size   int64
		}{Offset: 3, Size: 4},
	}

	var buf bytes.Buffer
	n, err := trailer.WriteTo(&buf)
	if got, exp := n, int64(tsi1.MeasurementTrailerSize); got != exp {
		t.Fatalf("got %v, exp %v", got, exp)
	}

	if got := err; got != nil {
		t.Fatalf("got %v, exp %v", got, nil)
	}

	// Verify trailer written correctly.
	exp := "" +
		"0000000000000001" + // data offset
		"0000000000000002" + // data size
		"0000000000000003" + // hash index offset
		"0000000000000004" + // hash index size
		"0000000000000000" + // legacy sketch offset
		"0000000000000000" + // legacy sketch size
		"0000000000000000" + // legacy tsketch offset
		"0000000000000000" + // legacy tsketch size
		"0001" // version

	if got, exp := fmt.Sprintf("%x", buf.String()), exp; got != exp {
		t.Fatalf("got %v, exp %v", got, exp)
	}
}

// Ensure measurement blocks can be written and opened.
func TestMeasurementBlockWriter(t *testing.T) {
	ms := Measurements{
		NewMeasurement([]byte("foo"), false, 100, 10, toSeriesIDs([]uint64{1, 3, 4})),
		NewMeasurement([]byte("bar"), false, 200, 20, toSeriesIDs([]uint64{2})),
		NewMeasurement([]byte("baz"), false, 300, 30, toSeriesIDs([]uint64{5, 6})),
	}

	// Write the measurements to writer.
	mw := tsi1.NewMeasurementBlockWriter()
	for _, m := range ms {
		mw.Add(m.Name, m.Deleted, m.Offset, m.Size, m.ids)
	}

	// Encode into buffer.
	var buf bytes.Buffer
	if n, err := mw.WriteTo(&buf); err != nil {
		t.Fatal(err)
	} else if n == 0 {
		t.Fatal("expected bytes written")
	}

	// Unmarshal into a block.
	var blk tsi1.MeasurementBlock
	if err := blk.UnmarshalBinary(buf.Bytes()); err != nil {
		t.Fatal(err)
	}

	// Verify data in block.
	if e, ok := blk.Elem([]byte("foo")); !ok {
		t.Fatal("expected element")
	} else if e.TagBlockOffset() != 100 || e.TagBlockSize() != 10 {
		t.Fatalf("unexpected offset/size: %v/%v", e.TagBlockOffset(), e.TagBlockSize())
	} else if !reflect.DeepEqual(e.SeriesIDs(), toSeriesIDs([]uint64{1, 3, 4})) {
		t.Fatalf("unexpected series data: %#v", e.SeriesIDs())
	}

	if e, ok := blk.Elem([]byte("bar")); !ok {
		t.Fatal("expected element")
	} else if e.TagBlockOffset() != 200 || e.TagBlockSize() != 20 {
		t.Fatalf("unexpected offset/size: %v/%v", e.TagBlockOffset(), e.TagBlockSize())
	} else if !reflect.DeepEqual(e.SeriesIDs(), toSeriesIDs([]uint64{2})) {
		t.Fatalf("unexpected series data: %#v", e.SeriesIDs())
	}

	if e, ok := blk.Elem([]byte("baz")); !ok {
		t.Fatal("expected element")
	} else if e.TagBlockOffset() != 300 || e.TagBlockSize() != 30 {
		t.Fatalf("unexpected offset/size: %v/%v", e.TagBlockOffset(), e.TagBlockSize())
	} else if !reflect.DeepEqual(e.SeriesIDs(), toSeriesIDs([]uint64{5, 6})) {
		t.Fatalf("unexpected series data: %#v", e.SeriesIDs())
	}

	// Verify non-existent measurement doesn't exist.
	if _, ok := blk.Elem([]byte("BAD_MEASUREMENT")); ok {
		t.Fatal("expected no element")
	}
}

type Measurements []Measurement

type Measurement struct {
	Name    []byte
	Deleted bool
	Offset  int64
	Size    int64
	ids     []tsdb.SeriesID
}

func NewMeasurement(name []byte, deleted bool, offset, size int64, ids []tsdb.SeriesID) Measurement {
	return Measurement{
		Name:    name,
		Deleted: deleted,
		Offset:  offset,
		Size:    size,
		ids:     ids,
	}
}

func toSeriesIDs(ids []uint64) []tsdb.SeriesID {
	sids := make([]tsdb.SeriesID, 0, len(ids))
	for _, id := range ids {
		sids = append(sids, tsdb.NewSeriesID(id))
	}
	return sids
}