[extension/jaegerencodingextension] Enhance jaegerEncodingExtension

.chloggen/jaeger_encoding_extension_enhance.yaml

@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: jaegerencodingextension
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Make jaegerEncodingExtension impl ptrace.Unmarshal.
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [6272]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext:
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []

extension/encoding/jaegerencodingextension/buffer.go

@@ -0,0 +1,125 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerencodingextension // import "github.com/open-telemetry/opentelemetry-collector-contrib/extension/encoding/jaegerencodingextension"
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+)
+
+var (
+	errInvalidBuff = errors.New("invalid buffer")
+	header         = []byte{'o', 't', 'l', 'p'}
+)
+
+type readonlyLengthFieldBuffer struct {
+	buff      *bytes.Buffer
+	hasHeader bool
+}
+
+func newReadableBuffer(bts []byte) *readonlyLengthFieldBuffer {
+	hasHeader := checkHeader(bts)
+	buffer := bytes.NewBuffer(bts)
+	// Skip the first 4 bytes.
+	if hasHeader {
+		buffer.Next(4)
+	}
+	return &readonlyLengthFieldBuffer{buff: buffer, hasHeader: hasHeader}
+}
+
+func checkHeader(bts []byte) bool {
+	return len(bts) > 4 && bts[0] == header[0] && bts[1] == header[1] && bts[2] == header[2] && bts[3] == header[3]
+}
+
+// slices Return slices
+func (b *readonlyLengthFieldBuffer) slices() ([][]byte, error) {
+	arr := make([][]byte, 0)
+	if !b.hasHeader {
+		arr = append(arr, b.buff.Bytes())
+		return arr, nil
+	}
+
+	for {
+		if b.buff.Len() == 0 {
+			break
+		}
+		size, err := b.readInt()
+		if err != nil {
+			return nil, errInvalidBuff
+		}
+		bts := make([]byte, size)
+		n, err := b.buff.Read(bts)
+		if n != size {
+			return nil, errInvalidBuff
+		}
+		if err != nil {
+			return nil, err
+		}
+		arr = append(arr, bts)
+	}
+	return arr, nil
+}
+
+// readInt Read an int number from the buffer.
+func (b *readonlyLengthFieldBuffer) readInt() (int, error) {
+	var v int32
+	err := binary.Read(b.buff, binary.BigEndian, &v)
+	if err != nil {
+		return 0, err
+	}
+	return int(v), nil
+}
+
+// Write only length filed byte buffer.
+type writeOnlyLengthFieldBuffer struct {
+	buff *bytes.Buffer
+}
+
+func newWritableBuffer(initCap int) (*writeOnlyLengthFieldBuffer, error) {
+	arr := make([]byte, 0, initCap)
+	buff := bytes.NewBuffer(arr)
+	// write head to buff first.
+	n, err := buff.Write(header)
+	if n != len(header) || err != nil {
+		return nil, errors.New("unable to create buffer")
+	}
+	return &writeOnlyLengthFieldBuffer{buff: buff}, nil
+}
+
+func calculateCap(bts [][]byte) int {
+	c := 4
+	for _, bt := range bts {
+		c = len(bt) + 4
+	}
+	return c
+}
+
+// writeBytes write a byte array to the buffer.
+func (b *writeOnlyLengthFieldBuffer) writeBytes(bts []byte) error {
+	length := len(bts)
+	err := b.writeInt(length)
+	if err != nil {
+		return err
+	}
+	n, err0 := b.buff.Write(bts)
+	if n != length {
+		return errInvalidBuff
+	}
+	if err0 != nil {
+		return err0
+	}
+	return nil
+}
+
+// writeInt Write an int number to the buffer.
+func (b *writeOnlyLengthFieldBuffer) writeInt(value int) error {
+	v := int32(value)
+	return binary.Write(b.buff, binary.BigEndian, v)
+}
+
+// bytes Return the byte array.
+func (b *writeOnlyLengthFieldBuffer) bytes() []byte {
+	return b.buff.Bytes()
+}

extension/encoding/jaegerencodingextension/buffer_test.go

@@ -0,0 +1,59 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerencodingextension // import "github.com/open-telemetry/opentelemetry-collector-contrib/extension/encoding/jaegerencodingextension"
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBufferInt(t *testing.T) {
+	for n := 0; n < 1024*1024; n++ {
+		buff, err := newWritableBuffer(0)
+		require.NoError(t, err)
+		require.NoError(t, buff.writeInt(n))
+		bts := buff.bytes()
+
+		readBuff := newReadableBuffer(bts)
+		n0, err0 := readBuff.readInt()
+		require.NoError(t, err0)
+		require.Equal(t, n, n0)
+	}
+}
+
+func TestBuffSlices(t *testing.T) {
+	arr := make([]string, 0)
+	for n := 0; n < 100; n++ {
+		v := fmt.Sprintf("OpenTelemetry Collector %d", n)
+		arr = append(arr, v)
+	}
+	writableBuff, err := newWritableBuffer(0)
+	require.NoError(t, err)
+	for _, s := range arr {
+		require.NoError(t, writableBuff.writeBytes([]byte(s)))
+	}
+	bts := writableBuff.bytes()
+
+	readBuff := newReadableBuffer(bts)
+	arrs, err := readBuff.slices()
+	require.NoError(t, err)
+	require.Equal(t, len(arrs), 100)
+
+	for i, e := range arrs {
+		v := string(e)
+		v0 := arr[i]
+		require.Equal(t, v, v0)
+	}
+}
+
+func TestReadOnlyBuffCompatibility(t *testing.T) {
+	str := "OpenTelemetry Collector"
+	buff := newReadableBuffer([]byte(str))
+	slices, err := buff.slices()
+	require.NoError(t, err)
+	require.Equal(t, len(slices), 1)
+	require.Equal(t, string(slices[0]), str)
+}

extension/encoding/jaegerencodingextension/extension.go

@@ -13,21 +13,31 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/encoding"
 )
 
-var _ encoding.TracesUnmarshalerExtension = &jaegerExtension{}
+var (
+	_ encoding.TracesMarshalerExtension   = (*jaegerExtension)(nil)
+	_ encoding.TracesUnmarshalerExtension = (*jaegerExtension)(nil)
+)
 
 type jaegerExtension struct {
 	config      *Config
+	marshaler   ptrace.Marshaler
 	unmarshaler ptrace.Unmarshaler
 }
 
+func (e *jaegerExtension) MarshalTraces(traces ptrace.Traces) ([]byte, error) {
+	return e.marshaler.MarshalTraces(traces)
+}
+
 func (e *jaegerExtension) UnmarshalTraces(buf []byte) (ptrace.Traces, error) {
 	return e.unmarshaler.UnmarshalTraces(buf)
 }
 
 func (e *jaegerExtension) Start(_ context.Context, _ component.Host) error {
 	switch e.config.Protocol {
 	case JaegerProtocolProtobuf:
-		e.unmarshaler = jaegerProtobufTrace{}
+		jaegerProto := jaegerProtobufTrace{}
+		e.marshaler = jaegerProto
+		e.unmarshaler = jaegerProto
 	default:
 		return fmt.Errorf("unsupported protocol: %q", e.config.Protocol)
 	}

extension/encoding/jaegerencodingextension/jaeger.go

@@ -14,14 +14,51 @@ type jaegerProtobufTrace struct {
 }
 
 func (j jaegerProtobufTrace) UnmarshalTraces(buf []byte) (ptrace.Traces, error) {
-	span := &jaegerproto.Span{}
-	err := span.Unmarshal(buf)
+	buff := newReadableBuffer(buf)
+	slices, err := buff.slices()
 	if err != nil {
 		return ptrace.NewTraces(), err
 	}
-	batch := jaegerproto.Batch{
-		Spans:   []*jaegerproto.Span{span},
-		Process: span.Process,
+
+	batches := make([]*jaegerproto.Batch, 0)
+	for _, slice := range slices {
+		batch := jaegerproto.Batch{}
+		e := batch.Unmarshal(slice)
+		if e != nil {
+			return ptrace.NewTraces(), e
+		}
+		batches = append(batches, &batch)
+	}
+
+	return jaeger.ProtoToTraces(batches)
+}
+
+func (j jaegerProtobufTrace) MarshalTraces(traces ptrace.Traces) ([]byte, error) {
+	batches, err := jaeger.ProtoFromTraces(traces)
+	if err != nil {
+		return nil, err
+	}
+
+	arr := make([][]byte, 0)
+	for _, batch := range batches {
+		bts, e := batch.Marshal()
+		if e != nil {
+			return nil, e
+		}
+		arr = append(arr, bts)
+	}
+
+	buff, err := newWritableBuffer(calculateCap(arr))
+	if err != nil {
+		return nil, err
 	}
-	return jaeger.ProtoToTraces([]*jaegerproto.Batch{&batch})
+
+	for _, bytes := range arr {
+		err := buff.writeBytes(bytes)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return buff.bytes(), nil
 }

extension/encoding/jaegerencodingextension/jaeger_test.go

@@ -0,0 +1,78 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerencodingextension // import "github.com/open-telemetry/opentelemetry-collector-contrib/extension/encoding/jaegerencodingextension"
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/extension"
+	"go.opentelemetry.io/collector/pdata/pcommon"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/translator/jaeger"
+)
+
+func buildTraces(num int) ptrace.Traces {
+	td := ptrace.NewTraces()
+	for n := 0; n < num; n++ {
+		span := td.ResourceSpans().AppendEmpty().ScopeSpans().AppendEmpty().Spans().AppendEmpty()
+		span.SetName("foo")
+		span.SetStartTimestamp(pcommon.Timestamp(10))
+		span.SetEndTimestamp(pcommon.Timestamp(20))
+		span.SetTraceID([16]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
+		span.SetSpanID([8]byte{1, 2, 3, 4, 5, 6, 7, 8})
+	}
+	return td
+}
+
+func TestJaegerProto(t *testing.T) {
+	config := &Config{Protocol: JaegerProtocolProtobuf}
+	ex, err := createExtension(context.Background(), extension.CreateSettings{}, config)
+	require.NoError(t, err)
+
+	j := ex.(*jaegerExtension)
+	err = j.Start(context.Background(), componenttest.NewNopHost())
+	require.NoError(t, err)
+
+	require.NotNil(t, j.marshaler)
+	require.NotNil(t, j.unmarshaler)
+
+	trace := buildTraces(5)
+	require.NotNil(t, trace)
+
+	bts, err := j.MarshalTraces(trace)
+	require.NoError(t, err)
+
+	trace1, err := j.UnmarshalTraces(bts)
+	require.NoError(t, err)
+	require.Equal(t, trace.SpanCount(), trace1.SpanCount())
+}
+
+func TestJaegerProtoCompatible(t *testing.T) {
+	config := &Config{Protocol: JaegerProtocolProtobuf}
+	ex, err := createExtension(context.Background(), extension.CreateSettings{}, config)
+	require.NoError(t, err)
+
+	j := ex.(*jaegerExtension)
+	err = j.Start(context.Background(), componenttest.NewNopHost())
+	require.NoError(t, err)
+
+	require.NotNil(t, j.marshaler)
+	require.NotNil(t, j.unmarshaler)
+
+	trace := buildTraces(1)
+
+	batch, err := jaeger.ProtoFromTraces(trace)
+	require.NoError(t, err)
+
+	bts, err := batch[0].Marshal()
+	require.NoError(t, err)
+
+	trace1, err := j.UnmarshalTraces(bts)
+	require.NoError(t, err)
+	require.Equal(t, trace.SpanCount(), trace1.SpanCount())
+}