[refactor] Move internal TracesData type to package jptrace

cmd/jaeger/internal/integration/trace_reader.go

@@ -143,7 +143,7 @@ func (*traceReader) FindTraceIDs(
 }
 
 type traceStream interface {
-	Recv() (*api_v3.TracesData, error)
+	Recv() (*jptrace.TracesData, error)
 }
 
 // consumeTraces reads the stream and calls yield for each chunk.

cmd/query/app/apiv3/grpc_handler.go

@@ -59,7 +59,7 @@ func (h *Handler) FindTraces(request *api_v3.FindTracesRequest, stream api_v3.Qu
 func (h *Handler) internalFindTraces(
 	ctx context.Context,
 	request *api_v3.FindTracesRequest,
-	streamSend func(*api_v3.TracesData) error,
+	streamSend func(*jptrace.TracesData) error,
 ) error {
 	query := request.GetQuery()
 	if query == nil {
@@ -130,14 +130,14 @@ func (h *Handler) GetOperations(ctx context.Context, request *api_v3.GetOperatio
 
 func receiveTraces(
 	seq iter.Seq2[[]ptrace.Traces, error],
-	sendFn func(*api_v3.TracesData) error,
+	sendFn func(*jptrace.TracesData) error,
 ) error {
 	for traces, err := range seq {
 		if err != nil {
 			return err
 		}
 		for _, trace := range traces {
-			tracesData := api_v3.TracesData(trace)
+			tracesData := jptrace.TracesData(trace)
 			if err := sendFn(&tracesData); err != nil {
 				return status.Error(codes.Internal,
 					fmt.Sprintf("failed to send response stream chunk to client: %v", err))

cmd/query/app/apiv3/grpc_handler_test.go

@@ -18,6 +18,7 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc/v2/querysvc"
+	"github.com/jaegertracing/jaeger/internal/jptrace"
 	"github.com/jaegertracing/jaeger/internal/proto/api_v3"
 	dependencyStoreMocks "github.com/jaegertracing/jaeger/internal/storage/v2/api/depstore/mocks"
 	"github.com/jaegertracing/jaeger/internal/storage/v2/api/tracestore"
@@ -204,7 +205,7 @@ func TestFindTracesSendError(t *testing.T) {
 				StartTimeMax: time.Now(),
 			},
 		},
-		/* streamSend= */ func(*api_v3.TracesData) error {
+		/* streamSend= */ func(*jptrace.TracesData) error {
 			return assert.AnError
 		},
 	)

cmd/query/app/apiv3/http_gateway.go

@@ -24,6 +24,7 @@ import (
 	"github.com/jaegertracing/jaeger-idl/model/v1"
 	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc/v2/querysvc"
 	"github.com/jaegertracing/jaeger/internal/jiter"
+	"github.com/jaegertracing/jaeger/internal/jptrace"
 	"github.com/jaegertracing/jaeger/internal/proto/api_v3"
 	"github.com/jaegertracing/jaeger/internal/storage/v1/api/spanstore"
 	"github.com/jaegertracing/jaeger/internal/storage/v2/api/tracestore"
@@ -113,7 +114,7 @@ func (h *HTTPGateway) tryParamError(w http.ResponseWriter, err error, paramName
 }
 
 func (h *HTTPGateway) returnTrace(td ptrace.Traces, w http.ResponseWriter) {
-	tracesData := api_v3.TracesData(td)
+	tracesData := jptrace.TracesData(td)
 	response := &api_v3.GRPCGatewayWrapper{
 		Result: &tracesData,
 	}

internal/jptrace/traces.go

@@ -0,0 +1,93 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jptrace
+
+import (
+	"github.com/gogo/protobuf/jsonpb"
+	"github.com/gogo/protobuf/proto"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+
+	"github.com/jaegertracing/jaeger/pkg/gogocodec"
+)
+
+// TracesData is an alias to ptrace.Traces that supports Gogo marshaling.
+// Our .proto APIs may refer to otlp.TraceData type, but its corresponding
+// protoc-generated struct is internal in OTel Collector, so we substitute
+// it for this TracesData type that implements marshaling methods by
+// delegating to public functions in the OTel Collector's ptrace module.
+type TracesData ptrace.Traces
+
+var (
+	_ gogocodec.CustomType = (*TracesData)(nil)
+	_ proto.Message        = (*TracesData)(nil)
+)
+
+func (td TracesData) ToTraces() ptrace.Traces {
+	return ptrace.Traces(td)
+}
+
+// Marshal implements gogocodec.CustomType.
+func (td *TracesData) Marshal() ([]byte, error) {
+	return new(ptrace.ProtoMarshaler).MarshalTraces(td.ToTraces())
+}
+
+// MarshalTo implements gogocodec.CustomType.
+func (td *TracesData) MarshalTo(buf []byte) (n int, err error) {
+	return td.MarshalToSizedBuffer(buf)
+}
+
+// MarshalToSizedBuffer is used by Gogo.
+func (td *TracesData) MarshalToSizedBuffer(buf []byte) (int, error) {
+	data, err := td.Marshal()
+	if err != nil {
+		return 0, err
+	}
+	n := copy(buf, data)
+	return n, nil
+}
+
+// MarshalJSONPB implements gogocodec.CustomType.
+func (td *TracesData) MarshalJSONPB(*jsonpb.Marshaler) ([]byte, error) {
+	return new(ptrace.JSONMarshaler).MarshalTraces(td.ToTraces())
+}
+
+// UnmarshalJSONPB implements gogocodec.CustomType.
+func (td *TracesData) UnmarshalJSONPB(_ *jsonpb.Unmarshaler, data []byte) error {
+	t, err := new(ptrace.JSONUnmarshaler).UnmarshalTraces(data)
+	if err != nil {
+		return err
+	}
+	*td = TracesData(t)
+	return nil
+}
+
+// Size implements gogocodec.CustomType.
+func (td *TracesData) Size() int {
+	return new(ptrace.ProtoMarshaler).TracesSize(td.ToTraces())
+}
+
+// Unmarshal implements gogocodec.CustomType.
+func (td *TracesData) Unmarshal(data []byte) error {
+	t, err := new(ptrace.ProtoUnmarshaler).UnmarshalTraces(data)
+	if err != nil {
+		return err
+	}
+	*td = TracesData(t)
+	return nil
+}
+
+// ProtoMessage implements proto.Message.
+func (*TracesData) ProtoMessage() {
+	// nothing to do here
+}
+
+// Reset implements proto.Message.
+func (td *TracesData) Reset() {
+	*td = TracesData(ptrace.NewTraces())
+}
+
+// String implements proto.Message.
+func (*TracesData) String() string {
+	return "*TracesData"
+}

internal/jptrace/traces_test.go

@@ -0,0 +1,57 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jptrace
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+)
+
+func TestTracesData(t *testing.T) {
+	td := TracesData(ptrace.NewTraces())
+
+	// Test ToTraces
+	assert.Equal(t, ptrace.Traces(td), td.ToTraces())
+
+	// Test Marshal
+	_, err := td.Marshal()
+	require.NoError(t, err)
+
+	// Test MarshalTo
+	_, err = td.MarshalTo(make([]byte, td.Size()))
+	require.NoError(t, err)
+
+	// Test MarshalJSONPB
+	_, err = td.MarshalJSONPB(nil)
+	require.NoError(t, err)
+
+	// Test UnmarshalJSONPB
+	err = td.UnmarshalJSONPB(nil, []byte(`{"resourceSpans":[]}`))
+	require.NoError(t, err)
+
+	err = td.UnmarshalJSONPB(nil, []byte(`{"resourceSpans":123}`))
+	require.Error(t, err)
+
+	// Test Size
+	assert.Equal(t, 0, td.Size())
+
+	// Test Unmarshal
+	err = td.Unmarshal([]byte{})
+	require.NoError(t, err)
+	err = td.Unmarshal([]byte{1})
+	require.Error(t, err)
+
+	// Test ProtoMessage
+	td.ProtoMessage()
+
+	// Test Reset
+	td.Reset()
+	assert.Equal(t, TracesData(ptrace.NewTraces()), td)
+
+	// Test String
+	assert.Equal(t, "*TracesData", td.String())
+}

scripts/makefiles/Protobuf.mk

@@ -141,7 +141,7 @@ patch-api-v3:
 .PHONY: proto-api-v3
 proto-api-v3: patch-api-v3
 	$(call proto_compile, $(API_V3_PATH), $(API_V3_PATCHED), -I$(API_V3_PATCHED_DIR) -Iidl/opentelemetry-proto)
-	@echo "🏗️  replace TracesData with internal custom type"
-	$(SED) -i 's/v1.TracesData/TracesData/g' $(API_V3_PATH)/query_service.pb.go
-	@echo "🏗️  remove OTEL import because we're not using any other OTLP types"
+	@echo "🏗️  replace first instance of OTEL import with internal type"
+	$(SED) -i '0,/go.opentelemetry.io\/proto\/otlp\/trace\/v1/s/go.opentelemetry.io\/proto\/otlp\/trace\/v1/github.com\/jaegertracing\/jaeger\/internal\/jptrace/' $(API_V3_PATH)/query_service.pb.go
+	@echo "🏗️  remove all remaining OTEL imports because we're not using any other OTLP types"
 	$(SED) -i 's+^.*v1 "go.opentelemetry.io/proto/otlp/trace/v1".*$$++' $(API_V3_PATH)/query_service.pb.go