jldb/lib/wal/test-util_test.go

package wal

import (
	"bytes"
	crand "crypto/rand"
	"encoding/base32"
	"encoding/binary"
	"hash/crc32"
	"io"
	"math/rand"
	"os"
	"reflect"
	"testing"
	"time"
)

// ----------------------------------------------------------------------------

func randString() string {
	size := 8 + rand.Intn(92)
	buf := make([]byte, size)
	if _, err := crand.Read(buf); err != nil {
		panic(err)
	}
	return base32.StdEncoding.EncodeToString(buf)
}

// ----------------------------------------------------------------------------

type RawRecord struct {
	Record
	Data    []byte
	DataCRC uint32
}

func (rr *RawRecord) ReadFrom(t *testing.T, f *os.File, offset int64) {
	t.Helper()

	buf := make([]byte, recordHeaderSize)
	if _, err := f.ReadAt(buf, offset); err != nil {
		t.Fatal(err)
	}

	if err := rr.Record.readHeaderFrom(readerAtToReader(f, offset)); err != nil {
		t.Fatal(err)
	}

	rr.Data = make([]byte, rr.DataSize+4) // For data and CRC32.
	if _, err := f.ReadAt(rr.Data, offset+recordHeaderSize); err != nil {
		t.Fatal(err)
	}

	storedCRC := binary.LittleEndian.Uint32(rr.Data[rr.DataSize:])
	computedCRC := crc32.ChecksumIEEE(rr.Data[:rr.DataSize])

	if storedCRC != computedCRC {
		t.Fatal(storedCRC, computedCRC)
	}

	rr.Data = rr.Data[:rr.DataSize]
}

// ----------------------------------------------------------------------------

func appendRandomRecords(t *testing.T, w waLog, count int64) []RawRecord {
	t.Helper()

	recs := make([]RawRecord, count)

	for i := range recs {
		rec := RawRecord{
			Data: []byte(randString()),
		}
		rec.DataSize = int64(len(rec.Data))

		seqNum, _, err := w.Append(int64(len(rec.Data)), bytes.NewBuffer(rec.Data))
		if err != nil {
			t.Fatal(err)
		}

		rec.SeqNum = seqNum

		recs[i] = rec
	}

	// Check that sequence numbers are sequential.
	seqNum := recs[0].SeqNum
	for _, rec := range recs {
		if rec.SeqNum != seqNum {
			t.Fatal(seqNum, rec)
		}
		seqNum++
	}

	return recs
}

func checkIteratorMatches(t *testing.T, it Iterator, recs []RawRecord) {
	for i, expected := range recs {
		if !it.Next(time.Millisecond) {
			t.Fatal(i, "no next")
		}

		rec := it.Record()

		if rec.SeqNum != expected.SeqNum {
			t.Fatal(i, rec.SeqNum, expected.SeqNum)
		}

		if rec.DataSize != expected.DataSize {
			t.Fatal(i, rec.DataSize, expected.DataSize)
		}

		if rec.TimestampMS == 0 {
			t.Fatal(rec.TimestampMS)
		}

		data := make([]byte, rec.DataSize)
		if _, err := io.ReadFull(rec.Reader, data); err != nil {
			t.Fatal(err)
		}

		if !bytes.Equal(data, expected.Data) {
			t.Fatalf("%d %s != %s", i, data, expected.Data)
		}
	}

	if it.Error() != nil {
		t.Fatal(it.Error())
	}

	// Check that iterator is empty.
	if it.Next(0) {
		t.Fatal("extra", it.Record())
	}
}

func randData() []byte {
	data := make([]byte, 1+rand.Intn(128))
	crand.Read(data)
	return data
}

func writeRandomWithEOF(w waLog, dt time.Duration) error {
	tStart := time.Now()
	for time.Since(tStart) < dt {
		data := randData()
		_, _, err := w.Append(int64(len(data)), bytes.NewBuffer(data))
		if err != nil {
			return err
		}
		time.Sleep(time.Millisecond)
	}

	_, _, err := w.Append(3, bytes.NewBuffer([]byte("EOF")))
	return err
}

func waitForEOF(t *testing.T, w *WAL) {
	t.Helper()

	h := w.seg.Header()
	it, err := w.Iterator(h.FirstSeqNum)
	if err != nil {
		t.Fatal(err)
	}
	defer it.Close()

	for it.Next(time.Hour) {
		rec := it.Record()
		buf := make([]byte, rec.DataSize)
		if _, err := io.ReadFull(rec.Reader, buf); err != nil {
			t.Fatal(err)
		}
		if bytes.Equal(buf, []byte("EOF")) {
			return
		}
	}

	t.Fatal("waitForEOF", it.Error())
}

func checkWALsEqual(t *testing.T, w1, w2 *WAL) {
	t.Helper()

	info1 := w1.Info()
	info2 := w2.Info()

	if !reflect.DeepEqual(info1, info2) {
		t.Fatal(info1, info2)
	}

	it1, err := w1.Iterator(info1.FirstSeqNum)
	if err != nil {
		t.Fatal(err)
	}
	defer it1.Close()

	it2, err := w2.Iterator(info2.FirstSeqNum)
	if err != nil {
		t.Fatal(err)
	}
	defer it2.Close()

	for {
		ok1 := it1.Next(time.Second)
		ok2 := it2.Next(time.Second)
		if ok1 != ok2 {
			t.Fatal(ok1, ok2)
		}

		if !ok1 {
			return
		}

		rec1 := it1.Record()
		rec2 := it2.Record()

		data1, err := io.ReadAll(rec1.Reader)
		if err != nil {
			t.Fatal(err)
		}
		data2, err := io.ReadAll(rec2.Reader)
		if err != nil {
			t.Fatal(err)
		}

		if !bytes.Equal(data1, data2) {
			t.Fatal(data1, data2)
		}
	}
}