rest otel support (#943)

rest/handler/otelhandler.go

@@ -0,0 +1,36 @@
+package handler
+
+import (
+	"net/http"
+
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+func OtelHandler(path string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		if !opentelemetry.Enabled() {
+			return next
+		}
+
+		propagator := otel.GetTextMapPropagator()
+		tracer := otel.GetTracerProvider().Tracer(opentelemetry.TraceName)
+
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+			spanCtx, span := tracer.Start(
+				ctx,
+				path,
+				oteltrace.WithSpanKind(oteltrace.SpanKindServer),
+				oteltrace.WithAttributes(semconv.HTTPServerAttributesFromHTTPRequest("", path, r)...),
+			)
+			defer span.End()
+
+			next.ServeHTTP(w, r.WithContext(spanCtx))
+		})
+	}
+}

rest/handler/otelhandler_test.go

@@ -0,0 +1,49 @@
+package handler
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/opentelemetry"
+	"go.opentelemetry.io/otel/propagation"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func TestOtelHandler(t *testing.T) {
+	opentelemetry.StartAgent(opentelemetry.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+
+	ts := httptest.NewServer(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+			spanCtx := trace.SpanContextFromContext(ctx)
+			assert.Equal(t, true, spanCtx.IsValid())
+		}),
+	)
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
+		defer span.End()
+
+		req, _ := http.NewRequest("GET", ts.URL, nil)
+		otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
+
+		res, err := client.Do(req)
+		assert.Equal(t, err, nil)
+		_ = res.Body.Close()
+		return nil
+	}(context.Background())
+
+	assert.Equal(t, err, nil)
+}