[v2][storage] Create v2 query service to operate on otlp data model (#…

…6343)

## Which problem is this PR solving?
- Towards #6337 

## Description of the changes
- Implement a v2 version of the query service that operates on the OTLP
data model. This PR will be followed up by a series of PRs where this
this new query service will be updated with the existing handlers. Once
all the handlers have been migrated to use this query service, we can
remove the old one.

## How was this change tested?
- Added unit tests

## Checklist
- [x] I have read
https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md
- [x] I have signed all commits
- [x] I have added unit tests for the new functionality
- [x] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `yarn lint` and `yarn test`

---------

Signed-off-by: Mahad Zaryab <[email protected]>

cmd/jaeger/internal/integration/trace_reader.go

@@ -90,7 +90,7 @@ func (r *traceReader) GetServices(ctx context.Context) ([]string, error) {
 	return res.Services, nil
 }
 
-func (r *traceReader) GetOperations(ctx context.Context, query tracestore.OperationQueryParameters) ([]tracestore.Operation, error) {
+func (r *traceReader) GetOperations(ctx context.Context, query tracestore.OperationQueryParams) ([]tracestore.Operation, error) {
 	var operations []tracestore.Operation
 	res, err := r.client.GetOperations(ctx, &api_v3.GetOperationsRequest{
 		Service:  query.ServiceName,

cmd/query/app/querysvc/v2/querysvc/package_test.go

@@ -0,0 +1,14 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package querysvc
+
+import (
+	"testing"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+)
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

cmd/query/app/querysvc/v2/querysvc/service.go

@@ -0,0 +1,203 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package querysvc
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"go.opentelemetry.io/collector/pdata/pcommon"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+
+	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc/v2/adjuster"
+	"github.com/jaegertracing/jaeger/internal/jptrace"
+	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/pkg/iter"
+	"github.com/jaegertracing/jaeger/storage_v2/depstore"
+	"github.com/jaegertracing/jaeger/storage_v2/tracestore"
+)
+
+var errNoArchiveSpanStorage = errors.New("archive span storage was not configured")
+
+const (
+	defaultMaxClockSkewAdjust = time.Second
+)
+
+// QueryServiceOptions holds the configuration options for the query service.
+type QueryServiceOptions struct {
+	// ArchiveTraceReader is used to read archived traces from the storage.
+	ArchiveTraceReader tracestore.Reader
+	// ArchiveTraceWriter is used to write traces to the archive storage.
+	ArchiveTraceWriter tracestore.Writer
+	// Adjuster is used to adjust traces before they are returned to the client.
+	// If not set, the default adjuster will be used.
+	Adjuster adjuster.Adjuster
+}
+
+// StorageCapabilities is a feature flag for query service
+type StorageCapabilities struct {
+	ArchiveStorage bool `json:"archiveStorage"`
+	// TODO: Maybe add metrics Storage here
+	// SupportRegex     bool
+	// SupportTagFilter bool
+}
+
+// QueryService provides methods to query data from the storage.
+type QueryService struct {
+	traceReader      tracestore.Reader
+	dependencyReader depstore.Reader
+	options          QueryServiceOptions
+}
+
+// GetTraceParams defines the parameters for retrieving traces using the GetTraces function.
+type GetTraceParams struct {
+	// TraceIDs is a slice of trace identifiers to fetch.
+	TraceIDs []tracestore.GetTraceParams
+	// RawTraces indicates whether to retrieve raw traces.
+	// If set to false, the traces will be adjusted using QueryServiceOptions.Adjuster.
+	RawTraces bool
+}
+
+// TraceQueryParams represents the parameters for querying a batch of traces.
+type TraceQueryParams struct {
+	tracestore.TraceQueryParams
+	// RawTraces indicates whether to retrieve raw traces.
+	// If set to false, the traces will be adjusted using QueryServiceOptions.Adjuster.
+	RawTraces bool
+}
+
+func NewQueryService(
+	traceReader tracestore.Reader,
+	dependencyReader depstore.Reader,
+	options QueryServiceOptions,
+) *QueryService {
+	qsvc := &QueryService{
+		traceReader:      traceReader,
+		dependencyReader: dependencyReader,
+		options:          options,
+	}
+
+	if qsvc.options.Adjuster == nil {
+		qsvc.options.Adjuster = adjuster.Sequence(
+			adjuster.StandardAdjusters(defaultMaxClockSkewAdjust)...)
+	}
+	return qsvc
+}
+
+// GetTraces retrieves traces with given trace IDs from the primary reader,
+// and if any of them are not found it then queries the archive reader.
+// The iterator is single-use: once consumed, it cannot be used again.
+func (qs QueryService) GetTraces(
+	ctx context.Context,
+	params GetTraceParams,
+) iter.Seq2[[]ptrace.Traces, error] {
+	getTracesIter := qs.traceReader.GetTraces(ctx, params.TraceIDs...)
+	return func(yield func([]ptrace.Traces, error) bool) {
+		foundTraceIDs, proceed := qs.receiveTraces(getTracesIter, yield, params.RawTraces)
+		if proceed && qs.options.ArchiveTraceReader != nil {
+			var missingTraceIDs []tracestore.GetTraceParams
+			for _, id := range params.TraceIDs {
+				if _, found := foundTraceIDs[id.TraceID]; !found {
+					missingTraceIDs = append(missingTraceIDs, id)
+				}
+			}
+			if len(missingTraceIDs) > 0 {
+				getArchiveTracesIter := qs.options.ArchiveTraceReader.GetTraces(ctx, missingTraceIDs...)
+				qs.receiveTraces(getArchiveTracesIter, yield, params.RawTraces)
+			}
+		}
+	}
+}
+
+func (qs QueryService) GetServices(ctx context.Context) ([]string, error) {
+	return qs.traceReader.GetServices(ctx)
+}
+
+func (qs QueryService) GetOperations(
+	ctx context.Context,
+	query tracestore.OperationQueryParams,
+) ([]tracestore.Operation, error) {
+	return qs.traceReader.GetOperations(ctx, query)
+}
+
+func (qs QueryService) FindTraces(
+	ctx context.Context,
+	query TraceQueryParams,
+) iter.Seq2[[]ptrace.Traces, error] {
+	return func(yield func([]ptrace.Traces, error) bool) {
+		tracesIter := qs.traceReader.FindTraces(ctx, query.TraceQueryParams)
+		qs.receiveTraces(tracesIter, yield, query.RawTraces)
+	}
+}
+
+// ArchiveTrace archives a trace specified by the given query parameters.
+// If the ArchiveTraceWriter is not configured, it returns
+// an error indicating that there is no archive span storage available.
+func (qs QueryService) ArchiveTrace(ctx context.Context, query tracestore.GetTraceParams) error {
+	if qs.options.ArchiveTraceWriter == nil {
+		return errNoArchiveSpanStorage
+	}
+	getTracesIter := qs.GetTraces(
+		ctx, GetTraceParams{TraceIDs: []tracestore.GetTraceParams{query}},
+	)
+	var archiveErr error
+	getTracesIter(func(traces []ptrace.Traces, err error) bool {
+		if err != nil {
+			archiveErr = err
+			return false
+		}
+		for _, trace := range traces {
+			err = qs.options.ArchiveTraceWriter.WriteTraces(ctx, trace)
+			if err != nil {
+				archiveErr = errors.Join(archiveErr, err)
+			}
+		}
+		return true
+	})
+	return archiveErr
+}
+
+func (qs QueryService) GetDependencies(ctx context.Context, endTs time.Time, lookback time.Duration) ([]model.DependencyLink, error) {
+	return qs.dependencyReader.GetDependencies(ctx, depstore.QueryParameters{
+		StartTime: endTs.Add(-lookback),
+		EndTime:   endTs,
+	})
+}
+
+func (qs QueryService) GetCapabilities() StorageCapabilities {
+	return StorageCapabilities{
+		ArchiveStorage: qs.options.hasArchiveStorage(),
+	}
+}
+
+func (opts *QueryServiceOptions) hasArchiveStorage() bool {
+	return opts.ArchiveTraceReader != nil && opts.ArchiveTraceWriter != nil
+}
+
+func (qs QueryService) receiveTraces(
+	seq iter.Seq2[[]ptrace.Traces, error],
+	yield func([]ptrace.Traces, error) bool,
+	rawTraces bool,
+) (map[pcommon.TraceID]struct{}, bool) {
+	aggregatedTraces := jptrace.AggregateTraces(seq)
+	foundTraceIDs := make(map[pcommon.TraceID]struct{})
+	proceed := true
+	aggregatedTraces(func(trace ptrace.Traces, err error) bool {
+		if err != nil {
+			proceed = yield(nil, err)
+			return proceed
+		}
+		if !rawTraces {
+			qs.options.Adjuster.Adjust(trace)
+		}
+		jptrace.SpanIter(trace)(func(_ jptrace.SpanIterPos, span ptrace.Span) bool {
+			foundTraceIDs[span.TraceID()] = struct{}{}
+			return true
+		})
+		proceed = yield([]ptrace.Traces{trace}, nil)
+		return proceed
+	})
+	return foundTraceIDs, proceed
+}