Changed thrift_processor_test.go

Author: AdityaS8804

receiver/jaegerreceiver/go.mod

@@ -3,6 +3,7 @@ module github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaeger
 go 1.23.0
 
 require (
+	github.com/HdrHistogram/hdrhistogram-go v1.1.2
 	github.com/apache/thrift v0.21.0
 	github.com/gorilla/mux v1.8.1
 	github.com/jaegertracing/jaeger v1.66.0
@@ -11,7 +12,7 @@ require (
 	github.com/open-telemetry/opentelemetry-collector-contrib/pkg/translator/jaeger v0.120.1
 	github.com/stretchr/testify v1.10.0
 	go.opentelemetry.io/collector/component v0.120.1-0.20250226024140-8099e51f9a77
-	go.opentelemetry.io/collector/component/componentstatus v0.120.1-0.20250226024140-8099e51f9a77
+	go.opentelemetry.io/collector/component/componentstatus v0.119.0
 	go.opentelemetry.io/collector/component/componenttest v0.120.1-0.20250226024140-8099e51f9a77
 	go.opentelemetry.io/collector/config/configgrpc v0.120.1-0.20250226024140-8099e51f9a77
 	go.opentelemetry.io/collector/config/confighttp v0.120.1-0.20250226024140-8099e51f9a77

receiver/jaegerreceiver/go.sum

@@ -0,0 +1,68 @@
+// Copyright The OpenTelemetry Authors
+// Copyright (c) 2019 The Jaeger Authors.
+// Copyright (c) 2017 Uber Technologies, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package customtransport
+
+import (
+	"bytes"
+	"context"
+
+	"github.com/apache/thrift/lib/go/thrift"
+)
+
+// TBufferedReadTransport is a thrift.TTransport that reads from a buffer
+type TBufferedReadTransport struct {
+	readBuf *bytes.Buffer
+}
+
+var _ thrift.TTransport = (*TBufferedReadTransport)(nil)
+
+// NewTBufferedReadTransport creates a buffer backed TTransport
+func NewTBufferedReadTransport(readBuf *bytes.Buffer) (*TBufferedReadTransport, error) {
+	return &TBufferedReadTransport{readBuf: readBuf}, nil
+}
+
+// IsOpen does nothing as transport is not maintaining the connection
+// Required to maintain thrift.TTransport interface
+func (*TBufferedReadTransport) IsOpen() bool {
+	return true
+}
+
+// Open does nothing as transport is not maintaining the connection
+// Required to maintain thrift.TTransport interface
+func (*TBufferedReadTransport) Open() error {
+	return nil
+}
+
+// Close does nothing as transport is not maintaining the connection
+// Required to maintain thrift.TTransport interface
+func (*TBufferedReadTransport) Close() error {
+	return nil
+}
+
+// Read reads bytes from the local buffer and puts them in the specified buf
+func (p *TBufferedReadTransport) Read(buf []byte) (int, error) {
+	in, err := p.readBuf.Read(buf)
+	return in, thrift.NewTTransportExceptionFromError(err)
+}
+
+// RemainingBytes returns the number of bytes left to be read from the readBuf
+func (p *TBufferedReadTransport) RemainingBytes() uint64 {
+	//nolint: gosec // G115
+	return uint64(p.readBuf.Len())
+}
+
+// Write writes bytes into the read buffer
+// Required to maintain thrift.TTransport interface
+func (p *TBufferedReadTransport) Write(buf []byte) (int, error) {
+	p.readBuf = bytes.NewBuffer(buf)
+	return len(buf), nil
+}
+
+// Flush does nothing as udp server does not write responses back
+// Required to maintain thrift.TTransport interface
+func (*TBufferedReadTransport) Flush(_ context.Context) error {
+	return nil
+}

receiver/jaegerreceiver/internal/cmd/customtransport/buffered_read_transport.go

@@ -0,0 +1,67 @@
+// Copyright The OpenTelemetry Authors
+// Copyright (c) 2019 The Jaeger Authors.
+// Copyright (c) 2017 Uber Technologies, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package customtransport
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver/internal/cmd/testutils"
+)
+
+// TestTBufferedReadTransport tests the TBufferedReadTransport
+func TestTBufferedReadTransport(t *testing.T) {
+	buffer := bytes.NewBuffer([]byte("testString"))
+	trans, err := NewTBufferedReadTransport(buffer)
+	require.NotNil(t, trans)
+	require.NoError(t, err)
+	require.Equal(t, uint64(10), trans.RemainingBytes())
+
+	firstRead := make([]byte, 4)
+	n, err := trans.Read(firstRead)
+	require.NoError(t, err)
+	require.Equal(t, 4, n)
+	require.Equal(t, []byte("test"), firstRead)
+	require.Equal(t, uint64(6), trans.RemainingBytes())
+
+	secondRead := make([]byte, 7)
+	n, err = trans.Read(secondRead)
+	require.NoError(t, err)
+	require.Equal(t, 6, n)
+	require.Equal(t, []byte("String"), secondRead[0:6])
+	require.Equal(t, uint64(0), trans.RemainingBytes())
+}
+
+// TestTBufferedReadTransportEmptyFunctions tests the empty functions in TBufferedReadTransport
+func TestTBufferedReadTransportEmptyFunctions(t *testing.T) {
+	byteArr := make([]byte, 1)
+	trans, err := NewTBufferedReadTransport(bytes.NewBuffer(byteArr))
+	require.NotNil(t, trans)
+	require.NoError(t, err)
+
+	err = trans.Open()
+	require.NoError(t, err)
+
+	err = trans.Close()
+	require.NoError(t, err)
+
+	err = trans.Flush(context.Background())
+	require.NoError(t, err)
+
+	n, err := trans.Write(byteArr)
+	require.Equal(t, 1, n)
+	require.NoError(t, err)
+
+	isOpen := trans.IsOpen()
+	require.True(t, isOpen)
+}
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

receiver/jaegerreceiver/internal/cmd/customtransport/buffered_read_transport_test.go

@@ -0,0 +1,15 @@
+// Copyright The OpenTelemetry Authors
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package processors
+
+import (
+	"testing"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver/internal/cmd/testutils"
+)
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

receiver/jaegerreceiver/internal/cmd/processors/package_test.go

@@ -0,0 +1,12 @@
+// Copyright The OpenTelemetry Authors
+// Copyright (c) 2019 The Jaeger Authors.
+// Copyright (c) 2017 Uber Technologies, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package processors
+
+// Processor processes metrics in multiple formats
+type Processor interface {
+	Serve()
+	Stop()
+}

receiver/jaegerreceiver/internal/cmd/processors/processor.go

@@ -0,0 +1,121 @@
+// Copyright The OpenTelemetry Authors
+// Copyright (c) 2019 The Jaeger Authors.
+// Copyright (c) 2017 Uber Technologies, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package processors
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/apache/thrift/lib/go/thrift"
+	"go.uber.org/zap"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver/internal/cmd/customtransport"
+	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver/internal/cmd/servers"
+	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver/internal/pkg/metrics"
+)
+
+// ThriftProcessor is a server that processes spans using a TBuffered Server
+type ThriftProcessor struct {
+	server        servers.Server
+	handler       AgentProcessor
+	protocolPool  *sync.Pool
+	numProcessors int
+	processing    sync.WaitGroup
+	logger        *zap.Logger
+	metrics       struct {
+		// Amount of time taken for processor to close
+		ProcessorCloseTimer metrics.Timer `metric:"thrift.udp.t-processor.close-time"`
+
+		// Number of failed buffer process operations
+		HandlerProcessError metrics.Counter `metric:"thrift.udp.t-processor.handler-errors"`
+	}
+}
+
+// AgentProcessor handler used by the processor to process thrift and call the reporter
+// with the deserialized struct. This interface is implemented directly by Thrift generated
+// code, e.g. jaegerThrift.NewAgentProcessor(handler), where handler implements the Agent
+// Thrift service interface, which is invoked with the deserialized struct.
+type AgentProcessor interface {
+	Process(ctx context.Context, iprot, oprot thrift.TProtocol) (success bool, err thrift.TException)
+}
+
+// NewThriftProcessor creates a TBufferedServer backed ThriftProcessor
+func NewThriftProcessor(
+	server servers.Server,
+	numProcessors int,
+	mFactory metrics.Factory,
+	factory thrift.TProtocolFactory,
+	handler AgentProcessor,
+	logger *zap.Logger,
+) (*ThriftProcessor, error) {
+	if numProcessors <= 0 {
+		return nil, fmt.Errorf(
+			"number of processors must be greater than 0, called with %d", numProcessors)
+	}
+	protocolPool := &sync.Pool{
+		New: func() any {
+			trans := &customtransport.TBufferedReadTransport{}
+			return factory.GetProtocol(trans)
+		},
+	}
+
+	res := &ThriftProcessor{
+		server:        server,
+		handler:       handler,
+		protocolPool:  protocolPool,
+		logger:        logger,
+		numProcessors: numProcessors,
+	}
+	metrics.Init(&res.metrics, mFactory, nil)
+	res.processing.Add(res.numProcessors)
+	for i := 0; i < res.numProcessors; i++ {
+		go func() {
+			res.processBuffer()
+			res.processing.Done()
+		}()
+	}
+	return res, nil
+}
+
+// Serve starts serving traffic
+func (s *ThriftProcessor) Serve() {
+	s.server.Serve()
+}
+
+// IsServing indicates whether the server is currently serving traffic
+func (s *ThriftProcessor) IsServing() bool {
+	return s.server.IsServing()
+}
+
+// Stop stops the serving of traffic and waits until the queue is
+// emptied by the readers
+func (s *ThriftProcessor) Stop() {
+	stopwatch := metrics.StartStopwatch(s.metrics.ProcessorCloseTimer)
+	s.server.Stop()
+	s.processing.Wait()
+	stopwatch.Stop()
+}
+
+// processBuffer reads data off the channel and puts it into a custom transport for
+// the processor to process
+func (s *ThriftProcessor) processBuffer() {
+	for readBuf := range s.server.DataChan() {
+		protocol := s.protocolPool.Get().(thrift.TProtocol)
+		payload := readBuf.GetBytes()
+		protocol.Transport().Write(payload)
+		s.logger.Debug("Span(s) received by the agent", zap.Int("bytes-received", len(payload)))
+
+		// NB: oddly, thrift-gen/agent/agent.go:L156 does this: `return true, thrift.WrapTException(err2)`
+		// So we check for both OK and error.
+		if ok, err := s.handler.Process(context.Background(), protocol, protocol); !ok || err != nil {
+			s.logger.Error("Processor failed", zap.Error(err))
+			s.metrics.HandlerProcessError.Inc(1)
+		}
+		s.protocolPool.Put(protocol)
+		s.server.DataRecd(readBuf) // acknowledge receipt and release the buffer
+	}
+}