[Tracing Redesign][2/N] Implement tracing unary Interceptors

api/transport/propagation.go

@@ -119,3 +119,13 @@ func UpdateSpanWithErr(span opentracing.Span, err error) error {
 	}
 	return err
 }
+
+// GetPropagationFormat returns the opentracing propagation depends on transport.
+// For TChannel, the format is opentracing.TextMap
+// For HTTP and gRPC, the format is opentracing.HTTPHeaders
+func GetPropagationFormat(transport string) opentracing.BuiltinFormat {
+	if transport == "tchannel" {
+		return opentracing.TextMap
+	}
+	return opentracing.HTTPHeaders
+}

internal/interceptor/inbound.go

@@ -24,7 +24,9 @@
 package interceptor
 
 import (
+	"context"
 	"go.uber.org/yarpc/api/middleware"
+	"go.uber.org/yarpc/api/transport"
 )
 
 type (
@@ -64,3 +66,21 @@ type (
 	// multiple times for the same request.
 	StreamInbound = middleware.StreamInbound
 )
+
+// UnaryInboundFunc adapts a function into a UnaryInbound middleware.
+type UnaryInboundFunc func(context.Context, *transport.Request, transport.ResponseWriter, transport.UnaryHandler) error
+
+// Handle for UnaryInboundFunc.
+func (f UnaryInboundFunc) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
+	return f(ctx, req, resw, h)
+}
+
+type unaryHandlerWithMiddleware struct {
+	h transport.UnaryHandler
+	i UnaryInbound
+}
+
+// Handle applies the UnaryInbound middleware to the handler's Handle method.
+func (h unaryHandlerWithMiddleware) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
+	return h.i.Handle(ctx, req, resw, h.h)
+}

internal/interceptor/inbound_test.go

@@ -0,0 +1,74 @@
+// Copyright (c) 2024 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package interceptor
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/yarpc/api/transport"
+)
+
+// UnaryHandlerFunc allows a function to be treated as a UnaryHandler for testing purposes.
+type UnaryHandlerFunc func(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error
+
+// Handle calls the underlying function in UnaryHandlerFunc.
+func (f UnaryHandlerFunc) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
+	return f(ctx, req, resw)
+}
+
+// TestUnaryInboundFunc ensures that UnaryInboundFunc works correctly.
+func TestUnaryInboundFunc(t *testing.T) {
+	var called bool
+	handler := UnaryHandlerFunc(func(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
+		called = true
+		return nil
+	})
+
+	middleware := UnaryInboundFunc(func(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
+		assert.False(t, called) // Ensure that the middleware is called before the handler.
+		return h.Handle(ctx, req, resw)
+	})
+
+	err := middleware.Handle(context.Background(), &transport.Request{}, nil, handler)
+	assert.NoError(t, err)
+	assert.True(t, called)
+}
+
+// TestUnaryHandlerWithMiddleware ensures that the unaryHandlerWithMiddleware applies the middleware correctly.
+func TestUnaryHandlerWithMiddleware(t *testing.T) {
+	var called bool
+	handler := UnaryHandlerFunc(func(ctx context.Context, req *transport.Request, resw transport.ResponseWriter) error {
+		called = true
+		return nil
+	})
+
+	middleware := UnaryInboundFunc(func(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, h transport.UnaryHandler) error {
+		assert.False(t, called) // Ensure middleware is called before the handler.
+		return h.Handle(ctx, req, resw)
+	})
+
+	wrappedHandler := unaryHandlerWithMiddleware{h: handler, i: middleware}
+	err := wrappedHandler.Handle(context.Background(), &transport.Request{}, nil)
+	assert.NoError(t, err)
+	assert.True(t, called)
+}

internal/interceptor/outbound.go

@@ -21,7 +21,9 @@
 package interceptor
 
 import (
+	"context"
 	"go.uber.org/yarpc/api/middleware"
+	"go.uber.org/yarpc/api/transport"
 )
 
 type (
@@ -64,3 +66,63 @@ type (
 	// multiple times on the same request.
 	StreamOutbound = middleware.StreamOutbound
 )
+
+var (
+	// NopUnaryOutbound is a no-operation unary outbound middleware.
+	NopUnaryOutbound transport.UnaryOutbound = nopUnaryOutbound{}
+)
+
+type nopUnaryOutbound struct{}
+
+// Call processes a unary request and returns a nil response and no error.
+func (nopUnaryOutbound) Call(ctx context.Context, req *transport.Request) (*transport.Response, error) {
+	return nil, nil
+}
+
+// Start starts the outbound middleware. It is a no-op.
+func (nopUnaryOutbound) Start() error {
+	return nil
+}
+
+// Stop stops the outbound middleware. It is a no-op.
+func (nopUnaryOutbound) Stop() error {
+	return nil
+}
+
+// IsRunning checks if the outbound middleware is running. Always returns false.
+func (nopUnaryOutbound) IsRunning() bool {
+	return false
+}
+
+// Transports returns the transports associated with this middleware. Always returns nil.
+func (nopUnaryOutbound) Transports() []transport.Transport {
+	return nil
+}
+
+// UnaryOutboundFunc adapts a function into a UnaryOutbound middleware.
+type UnaryOutboundFunc func(ctx context.Context, req *transport.Request) (*transport.Response, error)
+
+// Call executes the function as a UnaryOutbound call.
+func (f UnaryOutboundFunc) Call(ctx context.Context, req *transport.Request) (*transport.Response, error) {
+	return f(ctx, req)
+}
+
+// Start starts the UnaryOutboundFunc middleware. It is a no-op.
+func (f UnaryOutboundFunc) Start() error {
+	return nil
+}
+
+// Stop stops the UnaryOutboundFunc middleware. It is a no-op.
+func (f UnaryOutboundFunc) Stop() error {
+	return nil
+}
+
+// IsRunning checks if the UnaryOutboundFunc middleware is running. Always returns false.
+func (f UnaryOutboundFunc) IsRunning() bool {
+	return false
+}
+
+// Transports returns the transports associated with this middleware. Always returns nil.
+func (f UnaryOutboundFunc) Transports() []transport.Transport {
+	return nil
+}

internal/interceptor/outbound_test.go

@@ -0,0 +1,63 @@
+// Copyright (c) 2024 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package interceptor
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/yarpc/api/transport"
+)
+
+// TestNopUnaryOutbound ensures NopUnaryOutbound returns nil responses and no error.
+func TestNopUnaryOutbound(t *testing.T) {
+	outbound := NopUnaryOutbound
+
+	resp, err := outbound.Call(context.Background(), &transport.Request{})
+	assert.NoError(t, err)
+	assert.Nil(t, resp)
+
+	assert.False(t, outbound.IsRunning())
+	assert.Nil(t, outbound.Transports())
+
+	assert.NoError(t, outbound.Start())
+	assert.NoError(t, outbound.Stop())
+}
+
+// TestUnaryOutboundFunc tests if the function gets called correctly.
+func TestUnaryOutboundFunc(t *testing.T) {
+	called := false
+	outbound := UnaryOutboundFunc(func(ctx context.Context, req *transport.Request) (*transport.Response, error) {
+		called = true
+		return &transport.Response{}, nil
+	})
+
+	resp, err := outbound.Call(context.Background(), &transport.Request{})
+	assert.NoError(t, err)
+	assert.NotNil(t, resp)
+	assert.True(t, called)
+
+	assert.NoError(t, outbound.Start())
+	assert.NoError(t, outbound.Stop())
+	assert.False(t, outbound.IsRunning())
+	assert.Nil(t, outbound.Transports())
+}