Reload search blocks and replay search WAL

modules/ingester/ingester.go

@@ -6,6 +6,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/grafana/tempo/tempodb/search"
+
 	"github.com/cortexproject/cortex/pkg/ring"
 	"github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/go-kit/kit/log/level"
@@ -106,9 +108,6 @@ func (i *Ingester) starting(ctx context.Context) error {
 		return fmt.Errorf("failed to rediscover local blocks %w", err)
 	}
 
-	// Search data is considered experimental and removed on every startup.
-	i.clearSearchData()
-
 	// Now that user states have been created, we can start the lifecycler.
 	// Important: we want to keep lifecycler running until we ask it to stop, so we need to give it independent context
 	if err := i.lifecycler.StartAsync(context.Background()); err != nil {
@@ -335,6 +334,30 @@ func (i *Ingester) replayWal() error {
 		return fmt.Errorf("fatal error replaying wal %w", err)
 	}
 
+	searchBlocks, err := search.RescanBlocks(i.store.WAL().GetFilepath())
+	if err != nil {
+		return fmt.Errorf("fatal error replaying search wal %w", err)
+	}
+
+	// clear any searchBlock that does not have a corresponding entry in tracesBlocks
+	for j := len(searchBlocks) - 1; j >= 0; j-- {
+		clear := true
+		for _, tracesBlock := range blocks {
+			if searchBlocks[j].BlockID == tracesBlock.BlockID() {
+				clear = false
+				break
+			}
+		}
+
+		if clear {
+			err := searchBlocks[j].Clear()
+			if err != nil { // just log the error
+				level.Warn(log.Logger).Log("msg", "error clearing search WAL file", "blockID", searchBlocks[j].BlockID, "err", err)
+			}
+			searchBlocks = append(searchBlocks[:j], searchBlocks[j+1:]...)
+		}
+	}
+
 	for _, b := range blocks {
 		tenantID := b.Meta().TenantID
 
@@ -354,7 +377,13 @@ func (i *Ingester) replayWal() error {
 			return err
 		}
 
-		instance.AddCompletingBlock(b)
+		var searchWALBlock *search.StreamingSearchBlock
+		for _, s := range searchBlocks {
+			if b.BlockID() == s.BlockID {
+				searchWALBlock = s
+			}
+		}
+		instance.AddCompletingBlock(b, searchWALBlock)
 
 		i.enqueue(&flushOp{
 			kind:    opKindComplete,

modules/ingester/ingester_search.go

@@ -3,8 +3,6 @@ package ingester
 import (
 	"context"
 
-	"github.com/cortexproject/cortex/pkg/util/log"
-	"github.com/go-kit/kit/log/level"
 	"github.com/grafana/tempo/pkg/tempopb"
 	"github.com/weaveworks/common/user"
 )
@@ -66,11 +64,3 @@ func (i *Ingester) SearchTagValues(ctx context.Context, req *tempopb.SearchTagVa
 
 	return resp, nil
 }
-
-func (i *Ingester) clearSearchData() {
-	// clear wal
-	err := i.store.WAL().ClearFolder(searchDir)
-	if err != nil {
-		level.Error(log.Logger).Log("msg", "error clearing search data from wal")
-	}
-}

modules/ingester/ingester_test.go

@@ -212,6 +212,42 @@ func TestWal(t *testing.T) {
 	}
 }
 
+//func TestSearchWAL(t *testing.T) {
+//	tmpDir := t.TempDir()
+//
+//	i, _, _ := defaultIngester(t, tmpDir)
+//	inst, ok := i.getInstanceByID("test")
+//	require.True(t, ok)
+//
+//	// Write wal
+//	err := inst.CutCompleteTraces(0, true)
+//	require.NoError(t, err)
+//	_, err = inst.CutBlockIfReady(0, 0, true)
+//	require.NoError(t, err)
+//
+//	// assert that search WAL is being searched
+//	ctx := user.InjectOrgID(context.Background(), "test")
+//	searchReq := &tempopb.SearchRequest{Tags: map[string]string{
+//		search.SecretExhaustiveSearchTag: "",
+//	}}
+//	results, err := inst.Search(ctx, searchReq)
+//	assert.NoError(t, err)
+//	assert.Greater(t, results.Metrics.InspectedTraces, 0)
+//
+//	// Shutdown
+//	err = i.stopping(nil)
+//	require.NoError(t, err)
+//
+//	// replay wal
+//	i, _, _ = defaultIngester(t, tmpDir)
+//	inst, ok = i.getInstanceByID("test")
+//	require.True(t, ok)
+//
+//	results, err = inst.Search(ctx, searchReq)
+//	assert.NoError(t, err)
+//	assert.Greater(t, results.Metrics.InspectedTraces, 0)
+//}
+
 // TestWalReplayDeletesLocalBlocks simulates the condition where an ingester restarts after a wal is completed
 // to the local disk, but before the wal is deleted. On startup both blocks exist, and the ingester now errs
 // on the side of caution and chooses to replay the wal instead of rediscovering the local block.

modules/ingester/instance.go

@@ -447,11 +447,17 @@ func (i *instance) FindTraceByID(ctx context.Context, id []byte) (*tempopb.Trace
 // AddCompletingBlock adds an AppendBlock directly to the slice of completing blocks.
 // This is used during wal replay. It is expected that calling code will add the appropriate
 // jobs to the queue to eventually flush these.
-func (i *instance) AddCompletingBlock(b *wal.AppendBlock) {
+func (i *instance) AddCompletingBlock(b *wal.AppendBlock, s *search.StreamingSearchBlock) {
 	i.blocksMtx.Lock()
 	defer i.blocksMtx.Unlock()
 
 	i.completingBlocks = append(i.completingBlocks, b)
+
+	// search WAL
+	if s == nil {
+		return
+	}
+	i.searchAppendBlocks[b] = &searchStreamingBlockEntry{b: s}
 }
 
 // getOrCreateTrace will return a new trace object for the given request
@@ -623,7 +629,7 @@ func (i *instance) rediscoverLocalBlocks(ctx context.Context) error {
 
 		i.blocksMtx.Lock()
 		i.completeBlocks = append(i.completeBlocks, ib)
-		//i.searchCompleteBlocks[ib] = sb
+		//i.searchCompleteBlocks[ib] = &searchLocalBlockEntry{b: sb}
 		i.blocksMtx.Unlock()
 
 		level.Info(log.Logger).Log("msg", "reloaded local block", "tenantID", i.instanceID, "block", id.String(), "flushed", ib.FlushedTime())

tempodb/search/rescan_blocks.go

@@ -0,0 +1,120 @@
+package search
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	cortex_util "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/go-kit/log/level"
+	"github.com/google/uuid"
+
+	"github.com/grafana/tempo/pkg/tempofb"
+	"github.com/grafana/tempo/tempodb/backend"
+	"github.com/grafana/tempo/tempodb/encoding"
+	"github.com/grafana/tempo/tempodb/wal"
+)
+
+// RescanBlocks scans through the search directory in the WAL folder and replays files
+// todo: copied from wal.RescanBlocks(), see if we can reduce duplication?
+func RescanBlocks(walPath string) ([]*StreamingSearchBlock, error) {
+	searchFilepath := filepath.Join(walPath, "search")
+	files, err := ioutil.ReadDir(searchFilepath)
+	if err != nil {
+		return nil, err
+	}
+
+	blocks := make([]*StreamingSearchBlock, 0, len(files))
+	for _, f := range files {
+		if f.IsDir() {
+			continue
+		}
+		start := time.Now()
+		level.Info(cortex_util.Logger).Log("msg", "beginning replay", "file", f.Name(), "size", f.Size())
+
+		b, warning, err := newStreamingSearchBlockFromWALReplay(f.Name())
+
+		remove := false
+		if err != nil {
+			// wal replay failed, clear and warn
+			level.Warn(cortex_util.Logger).Log("msg", "failed to replay block. removing.", "file", f.Name(), "err", err)
+			remove = true
+		}
+
+		if b != nil && b.appender.Length() == 0 {
+			level.Warn(cortex_util.Logger).Log("msg", "empty wal file. ignoring.", "file", f.Name(), "err", err)
+			remove = true
+		}
+
+		if warning != nil {
+			level.Warn(cortex_util.Logger).Log("msg", "received warning while replaying block. partial replay likely.", "file", f.Name(), "warning", warning, "records", b.appender.Length())
+		}
+
+		if remove {
+			err = os.Remove(filepath.Join(searchFilepath, f.Name()))
+			if err != nil {
+				return nil, err
+			}
+			continue
+		}
+
+		level.Info(cortex_util.Logger).Log("msg", "replay complete", "file", f.Name(), "duration", time.Since(start))
+
+		blocks = append(blocks, b)
+	}
+
+	return blocks, nil
+}
+
+// newStreamingSearchBlockFromWALReplay creates a StreamingSearchBlock with in-memory records from a search WAL file
+func newStreamingSearchBlockFromWALReplay(filename string) (*StreamingSearchBlock, error, error) {
+	f, err := os.OpenFile(filename, os.O_RDONLY, 0644)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	blockID, _, _, _, _, err := parseFilename(filename)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// version is pinned to v2 for now
+	v, err := encoding.FromVersion("v2")
+	if err != nil {
+		return nil, nil, err
+	}
+
+	blockHeader := tempofb.NewSearchBlockHeaderMutable()
+	records, warning, err := wal.ReplayWALAndGetRecords(f, v, backend.EncNone, func(bytes []byte) error {
+		entry := tempofb.SearchEntryFromBytes(bytes)
+		blockHeader.AddEntry(entry)
+		return nil
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+	return &StreamingSearchBlock{
+		BlockID:  blockID,
+		file:     f,
+		appender: encoding.NewRecordAppender(records),
+		header:   blockHeader,
+		encoding: v,
+	}, warning, nil
+}
+
+func parseFilename(filename string) (blockID uuid.UUID, tenantID string, version string, encoding string, dataEncoding string, err error) {
+	splits := strings.Split(filename, ":")
+
+	if len(splits) < 6 {
+		return uuid.Nil, "", "", "", "", err
+	}
+
+	id, err := uuid.Parse(splits[0])
+	if err != nil {
+		return uuid.Nil, "", "", "", "", err
+	}
+	// todo: any other validation?
+	return id, splits[1], splits[2], splits[3], splits[4], nil
+}