[dbnode] Add wide filter (#2949)

.golangci.yml

@@ -223,6 +223,8 @@ linters:
     # New line required before return would require a large fraction of the
     # code base to need updating, it's not worth the perceived benefit.
     - nlreturn 
+    # Opinionated and sometimes wrong.
+    - paralleltest
   disable-all: false
   presets:
     # bodyclose, errcheck, gosec, govet, scopelint, staticcheck, typecheck

src/dbnode/persist/fs/retriever.go

@@ -40,6 +40,7 @@ import (
 	"time"
 
 	"github.com/m3db/m3/src/dbnode/namespace"
+	"github.com/m3db/m3/src/dbnode/persist/schema"
 	"github.com/m3db/m3/src/dbnode/sharding"
 	"github.com/m3db/m3/src/dbnode/storage/block"
 	"github.com/m3db/m3/src/dbnode/storage/limits"
@@ -314,7 +315,7 @@ func (r *blockRetriever) filterAndCompleteWideReqs(
 			retrieverResources.dataReqs = append(retrieverResources.dataReqs, req)
 
 		case streamWideEntryReq:
-			entry, err := seeker.SeekWideEntry(req.id, seekerResources)
+			entry, err := seeker.SeekWideEntry(req.id, req.wideFilter, seekerResources)
 			if err != nil {
 				if errors.Is(err, errSeekIDNotFound) {
 					// Missing, return empty result, successful lookup.
@@ -329,7 +330,7 @@ func (r *blockRetriever) filterAndCompleteWideReqs(
 
 			// Enqueue for fetch in batch in offset ascending order.
 			req.wideEntry = entry
-			entry.Shard = req.shard
+			req.wideEntry.Shard = req.shard
 			retrieverResources.appendWideEntryReq(req)
 
 		default:
@@ -663,10 +664,12 @@ func (r *blockRetriever) StreamWideEntry(
 	shard uint32,
 	id ident.ID,
 	startTime time.Time,
+	filter schema.WideEntryFilter,
 	nsCtx namespace.Context,
 ) (block.StreamedWideEntry, error) {
 	req := r.reqPool.Get()
 	req.streamReqType = streamWideEntryReq
+	req.wideFilter = filter
 
 	found, err := r.streamRequest(ctx, req, shard, id, startTime, nsCtx)
 	if err != nil {
@@ -788,6 +791,7 @@ type retrieveRequest struct {
 	streamReqType streamReqType
 	indexEntry    IndexEntry
 	wideEntry     xio.WideEntry
+	wideFilter    schema.WideEntryFilter
 	reader        xio.SegmentReader
 
 	err error
@@ -965,6 +969,7 @@ func (req *retrieveRequest) resetForReuse() {
 	req.streamReqType = streamInvalidReq
 	req.indexEntry = IndexEntry{}
 	req.wideEntry = xio.WideEntry{}
+	req.wideFilter = nil
 	req.reader = nil
 	req.err = nil
 	req.notFound = false

src/dbnode/persist/fs/seek.go

@@ -477,6 +477,7 @@ func (s *seeker) SeekIndexEntry(
 //        far we know it does not exist.
 func (s *seeker) SeekWideEntry(
 	id ident.ID,
+	filter schema.WideEntryFilter,
 	resources ReusableSeekerResources,
 ) (xio.WideEntry, error) {
 	offset, err := s.indexLookup.getNearestIndexFileOffset(id, resources)
@@ -508,6 +509,16 @@ func (s *seeker) SeekWideEntry(
 			return xio.WideEntry{}, instrument.InvariantErrorf(err.Error())
 		}
 
+		if filter != nil {
+			filtered, err := filter(entry)
+			if err != nil || filtered {
+				// NB: this entry is not being taken, can free memory.
+				resources.decodeIndexEntryBytesPool.Put(entry.ID)
+				resources.decodeIndexEntryBytesPool.Put(entry.EncodedTags)
+				return xio.WideEntry{}, err
+			}
+		}
+
 		if status != xmsgpack.MatchedLookupStatus {
 			// No longer being used so we can return to the pool.
 			resources.decodeIndexEntryBytesPool.Put(entry.ID)

src/dbnode/persist/fs/types.go

@@ -29,6 +29,7 @@ import (
 	"github.com/m3db/m3/src/dbnode/namespace"
 	"github.com/m3db/m3/src/dbnode/persist"
 	"github.com/m3db/m3/src/dbnode/persist/fs/msgpack"
+	"github.com/m3db/m3/src/dbnode/persist/schema"
 	"github.com/m3db/m3/src/dbnode/runtime"
 	"github.com/m3db/m3/src/dbnode/sharding"
 	"github.com/m3db/m3/src/dbnode/storage/block"
@@ -221,7 +222,11 @@ type DataFileSetSeeker interface {
 
 	// SeekWideEntry seeks in a manner similar to SeekIndexEntry, but
 	// instead yields a wide entry checksum of the series.
-	SeekWideEntry(id ident.ID, resources ReusableSeekerResources) (xio.WideEntry, error)
+	SeekWideEntry(
+		id ident.ID,
+		filter schema.WideEntryFilter,
+		resources ReusableSeekerResources,
+	) (xio.WideEntry, error)
 
 	// Range returns the time range associated with data in the volume
 	Range() xtime.Range
@@ -259,7 +264,11 @@ type ConcurrentDataFileSetSeeker interface {
 	SeekIndexEntry(id ident.ID, resources ReusableSeekerResources) (IndexEntry, error)
 
 	// SeekWideEntry is the same as in DataFileSetSeeker.
-	SeekWideEntry(id ident.ID, resources ReusableSeekerResources) (xio.WideEntry, error)
+	SeekWideEntry(
+		id ident.ID,
+		filter schema.WideEntryFilter,
+		resources ReusableSeekerResources,
+	) (xio.WideEntry, error)
 
 	// ConcurrentIDBloomFilter is the same as in DataFileSetSeeker.
 	ConcurrentIDBloomFilter() *ManagedConcurrentBloomFilter

src/dbnode/persist/schema/types.go

@@ -85,6 +85,9 @@ type WideEntry struct {
 	MetadataChecksum int64
 }
 
+// WideEntryFilter provides a filter for wide entries.
+type WideEntryFilter func(entry WideEntry) (bool, error)
+
 // IndexEntryHasher hashes an index entry.
 type IndexEntryHasher interface {
 	// HashIndexEntry computes a hash value for this IndexEntry using its ID, tags,

src/dbnode/storage/block/retriever_manager.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/m3db/m3/src/dbnode/namespace"
+	"github.com/m3db/m3/src/dbnode/persist/schema"
 	"github.com/m3db/m3/src/dbnode/sharding"
 	"github.com/m3db/m3/src/dbnode/x/xio"
 	"github.com/m3db/m3/src/x/context"
@@ -116,10 +117,11 @@ func (r *shardBlockRetriever) StreamWideEntry(
 	ctx context.Context,
 	id ident.ID,
 	blockStart time.Time,
+	filter schema.WideEntryFilter,
 	nsCtx namespace.Context,
 ) (StreamedWideEntry, error) {
 	return r.DatabaseBlockRetriever.StreamWideEntry(ctx, r.shard, id,
-		blockStart, nsCtx)
+		blockStart, filter, nsCtx)
 }
 
 type shardBlockRetrieverManager struct {

src/dbnode/storage/block/types.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/m3db/m3/src/dbnode/encoding"
 	"github.com/m3db/m3/src/dbnode/namespace"
+	"github.com/m3db/m3/src/dbnode/persist/schema"
 	"github.com/m3db/m3/src/dbnode/sharding"
 	"github.com/m3db/m3/src/dbnode/topology"
 	"github.com/m3db/m3/src/dbnode/ts"
@@ -312,6 +313,7 @@ type DatabaseBlockRetriever interface {
 		shard uint32,
 		id ident.ID,
 		blockStart time.Time,
+		filter schema.WideEntryFilter,
 		nsCtx namespace.Context,
 	) (StreamedWideEntry, error)
 
@@ -336,6 +338,7 @@ type DatabaseShardBlockRetriever interface {
 		ctx context.Context,
 		id ident.ID,
 		blockStart time.Time,
+		filter schema.WideEntryFilter,
 		nsCtx namespace.Context,
 	) (StreamedWideEntry, error)
 }

src/dbnode/storage/database.go

@@ -941,12 +941,9 @@ func (d *db) ReadEncoded(
 	return n.ReadEncoded(ctx, id, start, end)
 }
 
-// batchProcessWideQuery runs the given query against the namespace index,
-// iterating in a batchwise fashion across all matching IDs, applying the given
-// IDBatchProcessor batch processing function to each ID discovered.
-func (d *db) batchProcessWideQuery(
+func (d *db) BatchProcessWideQuery(
 	ctx context.Context,
-	n databaseNamespace,
+	n Namespace,
 	query index.Query,
 	batchProcessor IDBatchProcessor,
 	opts index.WideQueryOptions,
@@ -1029,8 +1026,9 @@ func (d *db) WideQuery(
 	streamedWideEntries := make([]block.StreamedWideEntry, 0, batchSize)
 	indexChecksumProcessor := func(batch *ident.IDBatch) error {
 		streamedWideEntries = streamedWideEntries[:0]
+
 		for _, shardID := range batch.ShardIDs {
-			streamedWideEntry, err := n.FetchWideEntry(ctx, shardID.ID, start)
+			streamedWideEntry, err := n.FetchWideEntry(ctx, shardID.ID, start, nil)
 			if err != nil {
 				return err
 			}
@@ -1050,7 +1048,7 @@ func (d *db) WideQuery(
 		return nil
 	}
 
-	err = d.batchProcessWideQuery(ctx, n, query, indexChecksumProcessor, opts)
+	err = d.BatchProcessWideQuery(ctx, n, query, indexChecksumProcessor, opts)
 	if err != nil {
 		return nil, err
 	}

src/dbnode/storage/database_test.go

@@ -902,7 +902,7 @@ func TestWideQuery(t *testing.T) {
 		ns *MockdatabaseNamespace, d *db, q index.Query,
 		now time.Time, shards []uint32, iterOpts index.IterationOptions) {
 		ns.EXPECT().FetchWideEntry(gomock.Any(),
-			ident.StringID("foo"), gomock.Any()).
+			ident.StringID("foo"), gomock.Any(), nil).
 			Return(block.EmptyStreamedWideEntry, nil)
 
 		_, err := d.WideQuery(ctx, ident.StringID("testns"), q, now, shards, iterOpts)

src/dbnode/storage/namespace.go

@@ -31,6 +31,7 @@ import (
 	"github.com/m3db/m3/src/dbnode/namespace"
 	"github.com/m3db/m3/src/dbnode/persist"
 	"github.com/m3db/m3/src/dbnode/persist/fs"
+	"github.com/m3db/m3/src/dbnode/persist/schema"
 	"github.com/m3db/m3/src/dbnode/sharding"
 	"github.com/m3db/m3/src/dbnode/storage/block"
 	"github.com/m3db/m3/src/dbnode/storage/bootstrap"
@@ -922,6 +923,7 @@ func (n *dbNamespace) FetchWideEntry(
 	ctx context.Context,
 	id ident.ID,
 	blockStart time.Time,
+	filter schema.WideEntryFilter,
 ) (block.StreamedWideEntry, error) {
 	callStart := n.nowFn()
 	shard, nsCtx, err := n.readableShardFor(id)
@@ -931,7 +933,7 @@ func (n *dbNamespace) FetchWideEntry(
 		return block.EmptyStreamedWideEntry, err
 	}
 
-	res, err := shard.FetchWideEntry(ctx, id, blockStart, nsCtx)
+	res, err := shard.FetchWideEntry(ctx, id, blockStart, filter, nsCtx)
 	n.metrics.read.ReportSuccessOrError(err, n.nowFn().Sub(callStart))
 
 	return res, err