[dbnode] Cross-block series reader

src/dbnode/generated/mocks/generate.go

@@ -20,7 +20,7 @@
 
 // mockgen rules for generating mocks for exported interfaces (reflection mode)
 
-//go:generate sh -c "mockgen -package=fs $PACKAGE/src/dbnode/persist/fs DataFileSetWriter,DataFileSetReader,DataFileSetSeeker,IndexFileSetWriter,IndexFileSetReader,IndexSegmentFileSetWriter,IndexSegmentFileSet,IndexSegmentFile,SnapshotMetadataFileWriter,DataFileSetSeekerManager,ConcurrentDataFileSetSeeker,MergeWith | genclean -pkg $PACKAGE/src/dbnode/persist/fs -out $GOPATH/src/$PACKAGE/src/dbnode/persist/fs/fs_mock.go"
+//go:generate sh -c "mockgen -package=fs $PACKAGE/src/dbnode/persist/fs DataFileSetWriter,DataFileSetReader,DataFileSetSeeker,IndexFileSetWriter,IndexFileSetReader,IndexSegmentFileSetWriter,IndexSegmentFileSet,IndexSegmentFile,SnapshotMetadataFileWriter,DataFileSetSeekerManager,ConcurrentDataFileSetSeeker,MergeWith,CrossBlockReader | genclean -pkg $PACKAGE/src/dbnode/persist/fs -out $GOPATH/src/$PACKAGE/src/dbnode/persist/fs/fs_mock.go"
 //go:generate sh -c "mockgen -package=xio $PACKAGE/src/dbnode/x/xio SegmentReader,SegmentReaderPool | genclean -pkg $PACKAGE/src/dbnode/x/xio -out $GOPATH/src/$PACKAGE/src/dbnode/x/xio/io_mock.go"
 //go:generate sh -c "mockgen -package=digest -destination=$GOPATH/src/$PACKAGE/src/dbnode/digest/digest_mock.go $PACKAGE/src/dbnode/digest ReaderWithDigest"
 //go:generate sh -c "mockgen -package=series $PACKAGE/src/dbnode/storage/series DatabaseSeries,QueryableBlockRetriever | genclean -pkg $PACKAGE/src/dbnode/storage/series -out $GOPATH/src/$PACKAGE/src/dbnode/storage/series/series_mock.go"

src/dbnode/persist/fs/cross_block_reader.go

@@ -0,0 +1,168 @@
+package fs
+
+import (
+	"bytes"
+	"container/heap"
+	"errors"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/m3db/m3/src/x/checked"
+	xerrors "github.com/m3db/m3/src/x/errors"
+	"github.com/m3db/m3/src/x/ident"
+)
+
+var (
+	errReaderNotOrderedByIndex = errors.New("CrossBlockReader can only use DataFileSetReaders ordered by index")
+	_ CrossBlockReader = (*crossBlockReader)(nil)
+	_ heap.Interface   = (*minHeap)(nil)
+)
+
+type crossBlockReader struct {
+	dataFileSetReaders []DataFileSetReader
+	initialized        bool
+	minHeap            minHeap
+	err                error
+}
+
+type minHeapEntry struct {
+	dataFileSetReaderIndex int
+	id                     ident.ID
+	tags                   ident.TagIterator
+	data                   checked.Bytes
+	checksum               uint32
+}
+
+type minHeap []*minHeapEntry
+
+// NewCrossBlockReader constructs a new CrossBlockReader based on given DataFileSetReaders.
+// DataFileSetReaders must be configured to return the data in the order of index, and must be
+// provided in a slice sorted by block start time.
+// Callers are responsible for closing the DataFileSetReaders.
+func NewCrossBlockReader(dataFileSetReaders []DataFileSetReader) (*crossBlockReader, error) {
+	var previousStart time.Time
+	for _, dataFileSetReader := range dataFileSetReaders {
+		if !dataFileSetReader.IsOrderedByIndex() {
+			return nil, errReaderNotOrderedByIndex
+		}
+		currentStart := dataFileSetReader.Range().Start
+		if !currentStart.After(previousStart) {
+			return nil, fmt.Errorf("dataFileSetReaders are not ordered by time (%s followed by %s)", previousStart, currentStart)
+		}
+		previousStart = currentStart
+	}
+
+	return &crossBlockReader{dataFileSetReaders: dataFileSetReaders}, nil
+}
+
+func (r *crossBlockReader) Read() (id ident.ID, tags ident.TagIterator, data checked.Bytes, checksum uint32, err error) {
+	if !r.initialized {
+		r.initialized = true
+		err := r.init()
+		if err != nil {
+			return nil, nil, nil, 0, err
+		}
+	}
+
+	if len(r.minHeap) == 0 {
+		return nil, nil, nil, 0, io.EOF
+	}
+
+	entry := heap.Pop(&r.minHeap).(*minHeapEntry)
+	if r.dataFileSetReaders[entry.dataFileSetReaderIndex] != nil {
+		nextEntry, err := r.readFromDataFileSet(entry.dataFileSetReaderIndex)
+		if err == io.EOF {
+			// will no longer read from this one
+			r.dataFileSetReaders[entry.dataFileSetReaderIndex] = nil
+		} else if err != nil {
+			return nil, nil, nil, 0, err
+		} else {
+			heap.Push(&r.minHeap, nextEntry)
+		}
+	}
+
+	return entry.id, entry.tags, entry.data, entry.checksum, nil
+}
+
+func (r *crossBlockReader) init() error {
+	r.minHeap = make([]*minHeapEntry, 0, len(r.dataFileSetReaders))
+
+	for i := range r.dataFileSetReaders {
+		entry, err := r.readFromDataFileSet(i)
+		if err == io.EOF {
+			continue
+		}
+		if err != nil {
+			return err
+		}
+		r.minHeap = append(r.minHeap, entry)
+	}
+
+	heap.Init(&r.minHeap)
+
+	return nil
+}
+
+func (r *crossBlockReader) readFromDataFileSet(index int) (*minHeapEntry, error) {
+	id, tags, data, checksum, err := r.dataFileSetReaders[index].Read()
+
+	if err == io.EOF {
+		return nil, err
+	}
+
+	if err != nil {
+		closeErr := r.Close()
+		if closeErr != nil {
+			return nil, xerrors.NewMultiError().Add(err).Add(closeErr)
+		}
+		return nil, err
+	}
+
+	return &minHeapEntry{
+		dataFileSetReaderIndex: index,
+		id:                     id,
+		tags:                   tags,
+		data:                   data,
+		checksum:               checksum,
+	}, nil
+}
+
+func (r *crossBlockReader) Close() error {
+	for _, entry := range r.minHeap {
+		entry.id.Finalize()
+		entry.tags.Close()
+		entry.data.DecRef()
+		entry.data.Finalize()
+	}
+	r.minHeap = r.minHeap[:0]
+	return nil
+}
+
+func (h minHeap) Len() int {
+	return len(h)
+}
+
+func (h minHeap) Less(i, j int) bool {
+	idsCmp := bytes.Compare(h[i].id.Bytes(), h[j].id.Bytes())
+	if idsCmp == 0 {
+		return h[i].dataFileSetReaderIndex < h[j].dataFileSetReaderIndex
+	}
+	return idsCmp < 0
+}
+
+func (h minHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+}
+
+func (h *minHeap) Push(x interface{}) {
+	*h = append(*h, x.(*minHeapEntry))
+}
+
+func (h *minHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}

src/dbnode/persist/fs/cross_block_reader_test.go

@@ -0,0 +1,153 @@
+package fs
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/m3db/m3/src/x/ident"
+	xtest "github.com/m3db/m3/src/x/test"
+	xtime "github.com/m3db/m3/src/x/time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var expectedError = errors.New("expected error")
+
+func TestCrossBlockReaderRejectMisconfiguredInputs(t *testing.T) {
+	ctrl := xtest.NewController(t)
+	defer ctrl.Finish()
+
+	dfsReader := NewMockDataFileSetReader(ctrl)
+	dfsReader.EXPECT().IsOrderedByIndex().Return(false)
+
+	_, err := NewCrossBlockReader([]DataFileSetReader{dfsReader})
+
+	assert.Equal(t, errReaderNotOrderedByIndex, err)
+}
+
+func TestCrossBlockReaderRejectMisorderedInputs(t *testing.T) {
+	ctrl := xtest.NewController(t)
+	defer ctrl.Finish()
+
+	now := time.Now().Truncate(time.Hour)
+	dfsReader1 := NewMockDataFileSetReader(ctrl)
+	dfsReader1.EXPECT().IsOrderedByIndex().Return(true)
+	dfsReader1.EXPECT().Range().Return(xtime.Range{Start: now})
+
+	later := now.Add(time.Hour)
+	dfsReader2 := NewMockDataFileSetReader(ctrl)
+	dfsReader2.EXPECT().IsOrderedByIndex().Return(true)
+	dfsReader2.EXPECT().Range().Return(xtime.Range{Start: later})
+
+	_, err := NewCrossBlockReader([]DataFileSetReader{dfsReader2, dfsReader1})
+
+	expectedErr := fmt.Errorf("dataFileSetReaders are not ordered by time (%s followed by %s)", later, now)
+
+	assert.Equal(t, expectedErr, err)
+}
+
+func TestCrossBlockReader(t *testing.T) {
+	tests := []struct {
+		name           string
+		blockSeriesIds [][]string
+	}{
+		{"no readers", [][]string{}},
+		{"empty readers", [][]string{{}, {}, {}}},
+		{"one reader, one series", [][]string{{"id1"}}},
+		{"one reader, many series", [][]string{{"id1", "id2", "id3"}}},
+		{"many readers with same series", [][]string{{"id1"}, {"id1"}, {"id1"}}},
+		{"many readers with different series", [][]string{{"id1"}, {"id2"}, {"id3"}}},
+		{"many readers with unordered series", [][]string{{"id3"}, {"id1"}, {"id2"}}},
+		{"complex case", [][]string{{"id2", "id3", "id5"}, {"id1", "id2", "id4"}, {"id1", "id4"}}},
+		{"reader error", [][]string{{"id1", "id2"}, {"id1", "error"}}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			testCrossBlockReader(t, tt.blockSeriesIds)
+		})
+	}
+}
+
+func testCrossBlockReader(t *testing.T, blockSeriesIds [][]string) {
+	ctrl := xtest.NewController(t)
+	defer ctrl.Finish()
+
+	now := time.Now().Truncate(time.Hour)
+	var dfsReaders []DataFileSetReader
+	expectedCount := 0
+
+	for blockIndex, ids := range blockSeriesIds {
+		dfsReader := NewMockDataFileSetReader(ctrl)
+		dfsReader.EXPECT().IsOrderedByIndex().Return(true)
+		dfsReader.EXPECT().Range().Return(xtime.Range{Start: now.Add(time.Hour * time.Duration(blockIndex))})
+
+		blockHasError := false
+		for j, id := range ids {
+			tags := ident.NewTags(ident.StringTag("foo", string(j)))
+			data := checkedBytes([]byte{byte(j)})
+			checksum := uint32(blockIndex) // somewhat hacky - using checksum to propagate block index value for assertions
+			if id == "error" {
+				dfsReader.EXPECT().Read().Return(nil, nil, nil, uint32(0), expectedError)
+				blockHasError = true
+			} else {
+				dfsReader.EXPECT().Read().Return(ident.StringID(id), ident.NewTagsIterator(tags), data, checksum, nil)
+			}
+		}
+
+		if !blockHasError {
+			dfsReader.EXPECT().Read().Return(nil, nil, nil, uint32(0), io.EOF).MaxTimes(1)
+		}
+
+		dfsReaders = append(dfsReaders, dfsReader)
+		expectedCount += len(ids)
+	}
+
+	cbReader, err := NewCrossBlockReader(dfsReaders)
+	require.NoError(t, err)
+	defer cbReader.Close()
+
+	hadError := false
+	actualCount := 0
+	previousId := ""
+	var previousBlockIndex uint32
+	for {
+		id, tags, data, checksum, err := cbReader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err == expectedError {
+			hadError = true
+			break
+		}
+		require.NoError(t, err)
+
+		strId := id.String()
+		id.Finalize()
+		blockIndex := checksum // see the comment above
+		assert.True(t, strId >= previousId, "series must be read in non-decreasing id order")
+		if strId == previousId {
+			assert.True(t, blockIndex >= previousBlockIndex, "same id blocks must be read in temporal order")
+		}
+
+		assert.NotNil(t, tags)
+		tags.Close()
+
+		assert.NotNil(t, data)
+		data.DecRef()
+		data.Finalize()
+
+		previousId = strId
+		previousBlockIndex = blockIndex
+
+		actualCount++
+	}
+
+	if !hadError {
+		assert.Equal(t, expectedCount, actualCount, "count of series read")
+	}
+}