feat(search): add search request observation logs and result metrics

metric/ingestor.go

@@ -30,6 +30,12 @@ var (
 		Name:      "errors_total",
 		Help:      "Number of search requests completed with error",
 	})
+	SearchResults = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: "seq_db_ingestor",
+		Subsystem: "search",
+		Name:      "results_total",
+		Help:      "Number of search requests by result and storage tier",
+	}, []string{"result", "tier"})
 	IngestorPanics = promauto.NewCounter(prometheus.CounterOpts{
 		Namespace: "seq_db_ingestor",
 		Subsystem: "common",

proxy/search/docs_iterator.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 	"time"
 
 	"go.uber.org/zap"
@@ -170,3 +171,47 @@ func (e *explainWrapperIterator) Next() (StreamingDoc, error) {
 	}
 	return d, nil
 }
+
+type DocsIteratorWithEvents struct {
+	it       DocsIterator
+	start    bool
+	finish   bool
+	onStart  func()
+	onFinish func()
+}
+
+func SetDocsIteratorEvents(it DocsIterator, onStart, onFinish func()) DocsIterator {
+	return &DocsIteratorWithEvents{
+		it:       it,
+		onStart:  onStart,
+		onFinish: onFinish,
+	}
+}
+
+func (e *DocsIteratorWithEvents) Next() (StreamingDoc, error) {
+	if !e.start {
+		e.onStart()
+		e.start = true
+	}
+	doc, err := e.it.Next()
+	if err != nil && !e.finish {
+		e.onFinish()
+		e.finish = true
+	}
+	return doc, err
+}
+
+func DocsIteratorSeq(it DocsIterator) iter.Seq2[StreamingDoc, error] {
+	return func(yield func(StreamingDoc, error) bool) {
+		doc, err := it.Next()
+		for ; err == nil; doc, err = it.Next() {
+			if !yield(doc, nil) {
+				return
+			}
+		}
+		if err != io.EOF {
+			var empty StreamingDoc
+			yield(empty, err)
+		}
+	}
+}

proxy/search/ingestor.go

@@ -35,6 +35,29 @@ type Config struct {
 	MirrorAddr      string
 }
 
+// StorageTier shows whether a search was served by hot or cold stores.
+type StorageTier string
+
+const (
+	StorageTierHot  StorageTier = "hot"
+	StorageTierCold StorageTier = "cold"
+	// StorageTierNone is used when a request never reached the ingestor
+	// (e.g. it failed validation in the proxy layer).
+	StorageTierNone StorageTier = "none"
+)
+
+// SearchStats carries observability data collected while executing a search.
+type SearchStats struct {
+	Size    int
+	HasHist bool
+	HasAgg  bool
+
+	StorageTier         StorageTier
+	HotSearchDuration   time.Duration
+	ColdSearchDuration  time.Duration
+	TotalSearchDuration time.Duration
+}
+
 type Ingestor struct {
 	config         Config
 	clients        map[string]storeapi.StoreApiClient
@@ -67,7 +90,7 @@ func (si *Ingestor) Search(
 ) (
 	qpr *seq.QPR,
 	docsStream DocsIterator,
-	overallDuration time.Duration,
+	stats *SearchStats,
 	err error,
 ) {
 	if sr.Explain {
@@ -77,8 +100,16 @@ func (si *Ingestor) Search(
 			zap.String("to", sr.To.String()),
 		)
 	}
+
+	stats = &SearchStats{
+		Size:        sr.Size + sr.Offset,
+		HasHist:     sr.Interval > 0,
+		HasAgg:      len(sr.AggQ) > 0,
+		StorageTier: StorageTierNone,
+	}
+
 	if sr.Size < 0 || sr.Offset < 0 {
-		return nil, nil, 0, fmt.Errorf("%w: negative size or offset", consts.ErrInvalidArgument)
+		return nil, nil, stats, fmt.Errorf("%w: negative size or offset", consts.ErrInvalidArgument)
 	}
 
 	startTime := time.Now()
@@ -87,31 +118,40 @@ func (si *Ingestor) Search(
 		searchStores = si.config.HotReadStores
 	}
 	qprs, err := si.searchStores(ctx, sr, searchStores, tr)
+
+	stats.StorageTier = StorageTierHot
+	stats.HotSearchDuration = time.Since(startTime)
+
 	var partialRespErr error
 
 	if err != nil {
 		switch {
 		case errors.Is(err, consts.ErrIngestorQueryWantsOldData):
 			if len(si.config.ReadStores.Shards) == 0 {
 				logger.Error("no cold stores, but hot mode is enabled, bad configuration of stores!")
-				return nil, nil, 0, err
+				return nil, nil, stats, err
 			}
 			metric.SearchColdTotal.Inc()
+
+			coldStart := time.Now()
 			qprs, err = si.searchStores(ctx, sr, si.config.ReadStores, tr)
+			stats.StorageTier = StorageTierCold
+			stats.ColdSearchDuration = time.Since(coldStart)
+
 			if err != nil {
 				metric.SearchColdErrors.Add(1)
 				if errors.Is(err, consts.ErrPartialResponse) {
 					partialRespErr = err // consider partial response from cold stores as a result
 				} else {
 					// errors from both hot and cold stores, return error
-					return nil, nil, 0, err
+					return nil, nil, stats, err
 				}
 			}
 		case errors.Is(err, consts.ErrPartialResponse):
 			partialRespErr = err // consider partial response from hot stores as a result
 		default:
 			// unexpected error on all hot replica sets (usually bad query)
-			return nil, nil, 0, err
+			return nil, nil, stats, err
 		}
 	}
 
@@ -143,19 +183,19 @@ func (si *Ingestor) Search(
 	docsStream = EmptyDocsStream{}
 	if sr.ShouldFetch && size > 0 {
 		if util.IsCancelled(ctx) {
-			return nil, nil, 0, ctx.Err()
+			return nil, nil, stats, ctx.Err()
 		}
 		metric.DocumentsRequested.Observe(float64(len(ids)))
 
 		fieldsFilter := tryParseFieldsFilter(string(sr.Q))
 		docsStream, err = si.FetchDocsStream(ctx, ids, sr.Explain, true, fieldsFilter)
 		if err != nil {
-			return nil, nil, 0, err
+			return nil, nil, stats, err
 		}
 	}
 
 	fetchDuration := time.Since(t)
-	overallDuration = time.Since(startTime)
+	stats.TotalSearchDuration = time.Since(startTime)
 
 	if sr.Explain {
 		logger.Info("data", zap.Any("histogram", qpr.Histogram))
@@ -166,11 +206,11 @@ func (si *Ingestor) Search(
 			util.ZapDurationWithPrec("query_ms", queryDuration, "ms", 2),
 			util.ZapDurationWithPrec("merge_ms", mergeDuration, "ms", 2),
 			util.ZapDurationWithPrec("fetch_ms", fetchDuration, "ms", 2),
-			util.ZapDurationWithPrec("all_ms", overallDuration, "ms", 2),
+			util.ZapDurationWithPrec("all_ms", stats.TotalSearchDuration, "ms", 2),
 		)
 	}
 
-	return qpr, docsStream, overallDuration, partialRespErr
+	return qpr, docsStream, stats, partialRespErr
 }
 
 // tryParseFieldsFilter tries to parse seq-ql query to extract fields/remove pipe.

proxy/search/merged_docs_iterator.go

@@ -86,19 +86,22 @@ func (m *mergedDocStream) Next() (StreamingDoc, error) {
 }
 
 func newNMergedStreams(streams []DocsIterator, less func(seq.IDSource, seq.IDSource) bool) DocsIterator {
-	if len(streams) == 0 {
+	l := len(streams)
+	if l == 0 {
 		return EmptyDocsStream{}
 	}
-
-	if len(streams) == 1 {
+	if l == 1 {
 		return newMergedDocsStream(streams[0], EmptyDocsStream{}, less)
 	}
-
-	merged := newMergedDocsStream(streams[0], streams[1], less)
-	for _, s := range streams[2:] {
-		merged = newMergedDocsStream(merged, s, less)
+	if l == 2 {
+		return newMergedDocsStream(streams[0], streams[1], less)
 	}
-	return merged
+
+	half := l / 2
+	a := newNMergedStreams(streams[:half], less)
+	b := newNMergedStreams(streams[half:], less)
+
+	return newMergedDocsStream(a, b, less)
 }
 
 type mergedStreamIterator struct {

proxyapi/grpc_async_search.go

@@ -225,7 +225,10 @@ func (g *grpcV1) ExportAsyncSearch(req *seqproxyapi.ExportAsyncSearchRequest, st
 		asyncsearcher.ExportSize.WithLabelValues(protocol).Observe(float64(wrapped.size))
 	}()
 
-	for doc, err := docsStream.Next(); err == nil; doc, err = docsStream.Next() {
+	for doc, err := range search.DocsIteratorSeq(docsStream) {
+		if err != nil {
+			return status.Errorf(codes.Internal, "docs reading error: %v", err)
+		}
 		eResp := &seqproxyapi.ExportResponse{
 			Doc: &seqproxyapi.Document{
 				Id:   doc.ID.String(),

proxyapi/grpc_complex_search.go

@@ -13,7 +13,7 @@ import (
 
 func (g *grpcV1) ComplexSearch(
 	ctx context.Context, req *seqproxyapi.ComplexSearchRequest,
-) (*seqproxyapi.ComplexSearchResponse, error) {
+) (_ *seqproxyapi.ComplexSearchResponse, retErr error) {
 	ctx, cancel := context.WithTimeout(ctx, g.config.SearchTimeout)
 	defer cancel()
 
@@ -22,7 +22,8 @@ func (g *grpcV1) ComplexSearch(
 	}
 
 	tr := querytracer.New(req.Query.Explain, "proxy/ComplexSearch")
-	sResp, err := g.doSearch(ctx, req, true, true, tr)
+	sResp, obs, err := g.doSearch(ctx, req, true, true, tr)
+	defer func() { obs.finish("ComplexSearch", retErr) }()
 	if err != nil {
 		return nil, err
 	}

proxyapi/grpc_complex_search_test.go

@@ -153,10 +153,10 @@ func prepareComplexSearchTestData(t *testing.T, cData cSearchTestCaseData) cSear
 		siSearchMock = &siSearchMockData{
 			sr: sr,
 			ret: siSearchRet{
-				qpr:  qpr,
-				docs: docs,
-				took: time.Second,
-				err:  cData.siErr,
+				qpr:   qpr,
+				docs:  docs,
+				stats: &search.SearchStats{},
+				err:   cData.siErr,
 			},
 		}
 	}

proxyapi/grpc_export.go

@@ -12,6 +12,7 @@ import (
 	"github.com/ozontech/seq-db/config"
 	"github.com/ozontech/seq-db/metric"
 	"github.com/ozontech/seq-db/pkg/seqproxyapi/v1"
+	"github.com/ozontech/seq-db/proxy/search"
 )
 
 type metricStream struct {
@@ -24,7 +25,7 @@ func (s *metricStream) Send(resp *seqproxyapi.ExportResponse) error {
 	return s.SeqProxyApi_ExportServer.Send(resp)
 }
 
-func (g *grpcV1) Export(req *seqproxyapi.ExportRequest, stream seqproxyapi.SeqProxyApi_ExportServer) error {
+func (g *grpcV1) Export(req *seqproxyapi.ExportRequest, stream seqproxyapi.SeqProxyApi_ExportServer) (retErr error) {
 	ctx, cancel := context.WithTimeout(stream.Context(), g.config.ExportTimeout)
 	defer cancel()
 
@@ -47,7 +48,8 @@ func (g *grpcV1) Export(req *seqproxyapi.ExportRequest, stream seqproxyapi.SeqPr
 		Offset:    req.Offset,
 		WithTotal: false,
 	}
-	sResp, err := g.doSearch(ctx, proxyReq, true, true, nil)
+	sResp, obs, err := g.doSearch(ctx, proxyReq, true, true, nil)
+	defer func() { obs.finish("Export", retErr) }()
 	if err != nil {
 		return err
 	}
@@ -63,7 +65,10 @@ func (g *grpcV1) Export(req *seqproxyapi.ExportRequest, stream seqproxyapi.SeqPr
 		metric.ExportSize.WithLabelValues(protocol).Observe(float64(wrapped.size))
 	}()
 
-	for doc, err := sResp.docsStream.Next(); err == nil; doc, err = sResp.docsStream.Next() {
+	for doc, err := range search.DocsIteratorSeq(sResp.docsStream) {
+		if err != nil {
+			return status.Errorf(codes.Internal, "docs reading error: %v", err)
+		}
 		eResp := &seqproxyapi.ExportResponse{
 			Doc: &seqproxyapi.Document{
 				Id:   doc.ID.String(),

proxyapi/grpc_export_test.go

@@ -81,10 +81,10 @@ func prepareExportTestData(cData exportTestCaseData) exportTestData {
 				ShouldFetch: true,
 			},
 			ret: siSearchRet{
-				qpr:  qpr,
-				docs: docs,
-				took: time.Second,
-				err:  cData.siErr,
+				qpr:   qpr,
+				docs:  docs,
+				stats: &search.SearchStats{},
+				err:   cData.siErr,
 			},
 		}
 	}

proxyapi/grpc_fetch.go

@@ -60,7 +60,10 @@ func (g *grpcV1) Fetch(req *seqproxyapi.FetchRequest, stream seqproxyapi.SeqProx
 	if err != nil {
 		return status.Errorf(codes.Internal, "can't fetch: %v", err)
 	}
-	for doc, err := docsStream.Next(); err == nil; doc, err = docsStream.Next() {
+	for doc, err := range search.DocsIteratorSeq(docsStream) {
+		if err != nil {
+			return status.Errorf(codes.Internal, "docs reading error: %v", err)
+		}
 		err := stream.Send(&seqproxyapi.Document{
 			Id:   doc.ID.String(),
 			Data: doc.Data,

proxyapi/grpc_get_aggregation.go

@@ -11,7 +11,7 @@ import (
 
 func (g *grpcV1) GetAggregation(
 	ctx context.Context, req *seqproxyapi.GetAggregationRequest,
-) (*seqproxyapi.GetAggregationResponse, error) {
+) (_ *seqproxyapi.GetAggregationResponse, retErr error) {
 	ctx, cancel := context.WithTimeout(ctx, g.config.SearchTimeout)
 	defer cancel()
 
@@ -24,7 +24,8 @@ func (g *grpcV1) GetAggregation(
 		Aggs:  req.Aggs,
 	}
 
-	sResp, err := g.doSearch(ctx, proxyReq, false, true, nil)
+	sResp, obs, err := g.doSearch(ctx, proxyReq, false, true, nil)
+	defer func() { obs.finish("GetAggregation", retErr) }()
 	if err != nil {
 		return nil, err
 	}

proxyapi/grpc_get_aggregation_test.go

@@ -105,9 +105,9 @@ func prepareGetAggregationTestData(t *testing.T, cData getAggregationTestCaseDat
 		siSearchMock = &siSearchMockData{
 			sr: sr,
 			ret: siSearchRet{
-				qpr:  qpr,
-				took: time.Second,
-				err:  cData.siErr,
+				qpr:   qpr,
+				stats: &search.SearchStats{},
+				err:   cData.siErr,
 			},
 		}
 	}