TUN-6689: Utilize new RegisterUDPSession to begin tracing

tracing/client.go

@@ -24,6 +24,9 @@ type InMemoryClient interface {
 	// Spans returns a copy of the list of in-memory stored spans as a base64
 	// encoded otlp protobuf string.
 	Spans() (string, error)
+	// ProtoSpans returns a copy of the list of in-memory stored spans as otlp
+	// protobuf byte array.
+	ProtoSpans() ([]byte, error)
 }
 
 // InMemoryOtlpClient is a client implementation for otlptrace.Client
@@ -55,21 +58,26 @@ func (mc *InMemoryOtlpClient) UploadTraces(_ context.Context, protoSpans []*trac
 
 // Spans returns the list of in-memory stored spans as a base64 encoded otlp protobuf string.
 func (mc *InMemoryOtlpClient) Spans() (string, error) {
-	mc.mu.Lock()
-	defer mc.mu.Unlock()
-	if len(mc.spans) <= 0 {
-		return "", errNoTraces
-	}
-	pbRequest := &coltracepb.ExportTraceServiceRequest{
-		ResourceSpans: mc.spans,
-	}
-	data, err := proto.Marshal(pbRequest)
+	data, err := mc.ProtoSpans()
 	if err != nil {
 		return "", err
 	}
 	return base64.StdEncoding.EncodeToString(data), nil
 }
 
+// ProtoSpans returns the list of in-memory stored spans as the protobuf byte array.
+func (mc *InMemoryOtlpClient) ProtoSpans() ([]byte, error) {
+	mc.mu.Lock()
+	defer mc.mu.Unlock()
+	if len(mc.spans) <= 0 {
+		return nil, errNoTraces
+	}
+	pbRequest := &coltracepb.ExportTraceServiceRequest{
+		ResourceSpans: mc.spans,
+	}
+	return proto.Marshal(pbRequest)
+}
+
 // NoopOtlpClient is a client implementation for otlptrace.Client that does nothing
 type NoopOtlpClient struct{}
 
@@ -89,3 +97,8 @@ func (mc *NoopOtlpClient) UploadTraces(_ context.Context, _ []*tracepb.ResourceS
 func (mc *NoopOtlpClient) Spans() (string, error) {
 	return "", errNoopTracer
 }
+
+// Spans always returns no traces error
+func (mc *NoopOtlpClient) ProtoSpans() ([]byte, error) {
+	return nil, errNoopTracer
+}

tracing/tracing.go

@@ -93,7 +93,7 @@ type TracedContext struct {
 	*cfdTracer
 }
 
-// NewTracedHTTPRequest creates a new tracer for the current HTTP request context.
+// NewTracedContext creates a new tracer for the current context.
 func NewTracedContext(ctx context.Context, traceContext string, log *zerolog.Logger) *TracedContext {
 	ctx, exists := extractTraceFromString(ctx, traceContext)
 	if !exists {
@@ -155,6 +155,24 @@ func (cft *cfdTracer) GetSpans() (enc string) {
 	return
 }
 
+// GetProtoSpans returns the spans as the otlp traces in protobuf byte array.
+func (cft *cfdTracer) GetProtoSpans() (proto []byte) {
+	proto, err := cft.exporter.ProtoSpans()
+	switch err {
+	case nil:
+		break
+	case errNoTraces:
+		cft.log.Trace().Err(err).Msgf("expected traces to be available")
+		return
+	case errNoopTracer:
+		return // noop tracer has no traces
+	default:
+		cft.log.Debug().Err(err)
+		return
+	}
+	return
+}
+
 // AddSpans assigns spans as base64 encoded protobuf otlp traces to provided
 // HTTP headers.
 func (cft *cfdTracer) AddSpans(headers http.Header) {
@@ -171,6 +189,11 @@ func (cft *cfdTracer) AddSpans(headers http.Header) {
 	headers[CanonicalCloudflaredTracingHeader] = []string{enc}
 }
 
+// End will set the OK status for the span and then end it.
+func End(span trace.Span) {
+	endSpan(span, -1, codes.Ok, nil)
+}
+
 // EndWithErrorStatus will set a status for the span and then end it.
 func EndWithErrorStatus(span trace.Span, err error) {
 	endSpan(span, -1, codes.Error, err)