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

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
+}