add option for public endpoint to gRPC

CHANGELOG.md

@@ -8,6 +8,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Added
+
+- Add the `WithPublicEndpoint` and `WithPublicEndpointFn` options to `go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc`. (#7407)
+
 ### Removed
 
 - The deprecated `StreamServerInterceptor` function from `go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc` is removed. (#7362)

instrumentation/google.golang.org/grpc/otelgrpc/config.go

@@ -4,6 +4,8 @@
 package otelgrpc // import "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 
 import (
+	"context"
+
 	"google.golang.org/grpc/stats"
 
 	"go.opentelemetry.io/otel"
@@ -39,6 +41,9 @@ type config struct {
 	SpanAttributes    []attribute.KeyValue
 	MetricAttributes  []attribute.KeyValue
 
+	PublicEndpoint   bool
+	PublicEndpointFn func(ctx context.Context, info *stats.RPCTagInfo) bool
+
 	ReceivedEvent bool
 	SentEvent     bool
 }
@@ -61,6 +66,38 @@ func newConfig(opts []Option) *config {
 	return c
 }
 
+type publicEndpointOption struct{ p bool }
+
+func (o publicEndpointOption) apply(c *config) {
+	c.PublicEndpoint = o.p
+}
+
+// WithPublicEndpoint configures the Handler to link the span with an incoming
+// span context. If this option is not provided, then the association is a child
+// association instead of a link.
+func WithPublicEndpoint() Option {
+	return publicEndpointOption{p: true}
+}
+
+type publicEndpointFnOption struct {
+	fn func(context.Context, *stats.RPCTagInfo) bool
+}
+
+func (o publicEndpointFnOption) apply(c *config) {
+	if o.fn != nil {
+		c.PublicEndpointFn = o.fn
+	}
+}
+
+// WithPublicEndpointFn runs with every request, and allows conditionally
+// configuring the Handler to link the span with an incoming span context. If
+// this option is not provided or returns false, then the association is a
+// child association instead of a link.
+// Note: WithPublicEndpoint takes precedence over WithPublicEndpointFn.
+func WithPublicEndpointFn(fn func(context.Context, *stats.RPCTagInfo) bool) Option {
+	return publicEndpointFnOption{fn: fn}
+}
+
 type propagatorsOption struct{ p propagation.TextMapPropagator }
 
 func (o propagatorsOption) apply(c *config) {

instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go

@@ -112,11 +112,21 @@ func (h *serverHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) cont
 	}
 
 	if record {
+		opts := []trace.SpanStartOption{
+			trace.WithSpanKind(trace.SpanKindServer),
+			trace.WithAttributes(append(attrs, h.SpanAttributes...)...),
+		}
+		if h.PublicEndpoint || (h.PublicEndpointFn != nil && h.PublicEndpointFn(ctx, info)) {
+			opts = append(opts, trace.WithNewRoot())
+			// Linking incoming span context if any for public endpoint.
+			if s := trace.SpanContextFromContext(ctx); s.IsValid() && s.IsRemote() {
+				opts = append(opts, trace.WithLinks(trace.Link{SpanContext: s}))
+			}
+		}
 		ctx, _ = h.tracer.Start(
 			trace.ContextWithRemoteSpanContext(ctx, trace.SpanContextFromContext(ctx)),
 			name,
-			trace.WithSpanKind(trace.SpanKindServer),
-			trace.WithAttributes(append(attrs, h.SpanAttributes...)...),
+			opts...,
 		)
 	}
 

instrumentation/google.golang.org/grpc/otelgrpc/stats_handler_test.go

@@ -6,14 +6,165 @@ package otelgrpc
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
 
 	"go.opentelemetry.io/otel/metric"
 	"go.opentelemetry.io/otel/metric/embedded"
+	"go.opentelemetry.io/otel/propagation"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	"go.opentelemetry.io/otel/trace"
 )
 
+func TestWithPublicEndpoint(t *testing.T) {
+	spanRecorder := tracetest.NewSpanRecorder()
+	provider := sdktrace.NewTracerProvider(
+		sdktrace.WithSpanProcessor(spanRecorder),
+	)
+	remoteSpan := trace.SpanContextConfig{
+		TraceID: trace.TraceID{0x01},
+		SpanID:  trace.SpanID{0x01},
+		Remote:  true,
+	}
+	prop := propagation.TraceContext{}
+	h := NewServerHandler(
+		WithPublicEndpoint(),
+		WithPropagators(prop),
+		WithTracerProvider(provider),
+	)
+
+	sc := trace.NewSpanContext(remoteSpan)
+	ctx := trace.ContextWithSpanContext(context.Background(), sc)
+
+	ctx = h.TagRPC(ctx, &stats.RPCTagInfo{
+		FullMethodName: "some.package/Method",
+		FailFast:       true,
+	})
+
+	h.HandleRPC(ctx, &stats.Begin{
+		Client:                    false,
+		BeginTime:                 time.Time{},
+		FailFast:                  true,
+		IsClientStream:            false,
+		IsServerStream:            false,
+		IsTransparentRetryAttempt: false,
+	})
+
+	h.HandleRPC(ctx, &stats.End{
+		Client:    false,
+		BeginTime: time.Time{},
+		EndTime:   time.Time{},
+		Trailer:   metadata.MD{},
+		Error:     nil,
+	})
+
+	// Recorded span should be linked with an incoming span context.
+	assert.NoError(t, spanRecorder.ForceFlush(ctx))
+	spans := spanRecorder.Ended()
+	require.Len(t, spans, 1)
+	require.Len(t, spans[0].Links(), 1, "should contain link")
+	require.True(t, sc.Equal(spans[0].Links()[0].SpanContext), "should link incoming span context")
+}
+
+func TestWithPublicEndpointFn(t *testing.T) {
+	remoteSpan := trace.SpanContextConfig{
+		TraceID:    trace.TraceID{0x01},
+		SpanID:     trace.SpanID{0x01},
+		TraceFlags: trace.FlagsSampled,
+		Remote:     true,
+	}
+	prop := propagation.TraceContext{}
+
+	for _, tt := range []struct {
+		name          string
+		fn            func(context.Context, *stats.RPCTagInfo) bool
+		handlerAssert func(*testing.T, trace.SpanContext)
+		spansAssert   func(*testing.T, trace.SpanContext, []sdktrace.ReadOnlySpan)
+	}{
+		{
+			name: "with the method returning true",
+			fn: func(context.Context, *stats.RPCTagInfo) bool {
+				return true
+			},
+			handlerAssert: func(t *testing.T, sc trace.SpanContext) {
+				// Should be with new root trace.
+				assert.True(t, sc.IsValid())
+				assert.False(t, sc.IsRemote())
+				assert.NotEqual(t, remoteSpan.TraceID, sc.TraceID())
+			},
+			spansAssert: func(t *testing.T, sc trace.SpanContext, spans []sdktrace.ReadOnlySpan) {
+				require.Len(t, spans, 1)
+				require.Len(t, spans[0].Links(), 1, "should contain link")
+				require.True(t, sc.Equal(spans[0].Links()[0].SpanContext), "should link incoming span context")
+			},
+		},
+		{
+			name: "with the method returning false",
+			fn: func(context.Context, *stats.RPCTagInfo) bool {
+				return false
+			},
+			handlerAssert: func(t *testing.T, sc trace.SpanContext) {
+				// Should have remote span as parent
+				assert.True(t, sc.IsValid())
+				assert.False(t, sc.IsRemote())
+				assert.Equal(t, remoteSpan.TraceID, sc.TraceID())
+			},
+			spansAssert: func(t *testing.T, _ trace.SpanContext, spans []sdktrace.ReadOnlySpan) {
+				require.Len(t, spans, 1)
+				require.Empty(t, spans[0].Links(), "should not contain link")
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			spanRecorder := tracetest.NewSpanRecorder()
+			provider := sdktrace.NewTracerProvider(
+				sdktrace.WithSpanProcessor(spanRecorder),
+			)
+
+			h := NewServerHandler(
+				WithPublicEndpointFn(tt.fn),
+				WithPropagators(prop),
+				WithTracerProvider(provider),
+			)
+
+			sc := trace.NewSpanContext(remoteSpan)
+			ctx := trace.ContextWithSpanContext(context.Background(), sc)
+
+			ctx = h.TagRPC(ctx, &stats.RPCTagInfo{
+				FullMethodName: "some.package/Method",
+				FailFast:       true,
+			})
+
+			h.HandleRPC(ctx, &stats.Begin{
+				Client:                    false,
+				BeginTime:                 time.Time{},
+				FailFast:                  true,
+				IsClientStream:            false,
+				IsServerStream:            false,
+				IsTransparentRetryAttempt: false,
+			})
+
+			h.HandleRPC(ctx, &stats.End{
+				Client:    false,
+				BeginTime: time.Time{},
+				EndTime:   time.Time{},
+				Trailer:   metadata.MD{},
+				Error:     nil,
+			})
+
+			// Recorded span should be linked with an incoming span context.
+			assert.NoError(t, spanRecorder.ForceFlush(ctx))
+			spans := spanRecorder.Ended()
+			tt.spansAssert(t, sc, spans)
+		})
+	}
+}
+
 func TestNilInstruments(t *testing.T) {
 	mp := meterProvider{}
 	opts := []Option{WithMeterProvider(mp)}