Repository: yarpc/yab
Pull request #414 (2024-05-31 18:10:07)

Add an option to force capture of Jaeger traces
Currently, the existing --jaeger flag adds trace, span and parent span information to outgoing RPCs.

However it does not set the sampling bit, which means analysing the trace is impossible if the backend infrastructure is not recording the distributed trace.

This new option --force-jaeger-sample ensures that every originating yab RPC is captured, most useful in the single yab (non benchmarking) case.

It uses constSampler with a decision value of true to ensure the sampling bit in the Jaeger headers. See:
https://www.jaegertracing.io/docs/1.22/client-libraries/#tracespan-identity


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e250536..aaa1c25 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,9 @@
 Changelog
 =========
 
+# Unreleased
+* Introduce new options --force-jaeger-sample to ensure yab requests are traced.
+
 # 0.22.0 (2023-02-27)
 * Fix: output benchmark errors in JSON format.
 
diff --git a/README.md b/README.md
index f955925..c952e43 100644
--- a/README.md
+++ b/README.md
@@ -96,6 +96,8 @@ Transport Options:
                                  traffic group.
       --jaeger                   Use the Jaeger tracing client to send Uber
                                  style traces and baggage headers
+      --force-jaeger-sample      If Jaeger tracing is enabled with --jaeger, force all requests
+                                 to be sampled.
   -T, --topt=                    Transport options for TChannel, protocol
                                  headers for HTTP
       --http-method=             The HTTP method to use (default: POST)
diff --git a/main.go b/main.go
index 7f0e3b1..49ebd5c 100644
--- a/main.go
+++ b/main.go
@@ -379,7 +379,7 @@ func createJaegerTracer(opts Options, out output) (opentracing.Tracer, io.Closer
 		// throttling threshold. Better to use "always false" sampling and
 		// only enable the span when we have not hit the throttling
 		// threshold.
-		jaeger_config.Sampler(jaeger.NewConstSampler(false)),
+		jaeger_config.Sampler(jaeger.NewConstSampler(opts.TOpts.ForceJaegerSample)),
 		jaeger_config.Reporter(jaeger.NewNullReporter()),
 	)
 	if err != nil {
@@ -389,6 +389,9 @@ func createJaegerTracer(opts Options, out output) (opentracing.Tracer, io.Closer
 }
 
 func getTracer(opts Options, out output) (opentracing.Tracer, io.Closer) {
+	if !opts.TOpts.Jaeger && opts.TOpts.ForceJaegerSample {
+		out.Fatalf("Cannot force Jaeger sampling without enabling Jaeger")
+	}
 	if opts.TOpts.Jaeger && !opts.TOpts.NoJaeger {
 		return createJaegerTracer(opts, out)
 	}
diff --git a/options.go b/options.go
index bb03194..c96cd2e 100644
--- a/options.go
+++ b/options.go
@@ -93,7 +93,7 @@ type TransportOptions struct {
 	TransportHeaders    map[string]string `short:"T" long:"topt" description:"Transport options for TChannel, protocol headers for HTTP"`
 	HTTPMethod          string            `long:"http-method" description:"The HTTP method to use"`
 	GRPCMaxResponseSize int               `long:"grpc-max-response-size" description:"Maximum response size for gRPC requests. Default value is 4MB"`
-
+	ForceJaegerSample   bool              `long:"force-jaeger-sample" description:"Force all requests to be sampled for Jaeger tracing (use with --jaeger)"`
 	// This is a hack to work around go-flags not allowing disabling flags:
 	// https://github.com/jessevdk/go-flags/issues/191
 	// Do not specify this value in a defaults.ini file as it is not possible
diff --git a/template.go b/template.go
index a378865..1a30f09 100644
--- a/template.go
+++ b/template.go
@@ -50,12 +50,13 @@ type template struct {
 	RoutingKey      string `yaml:"routingKey" yaml-aliases:"routingkey,routing-key,rk"`
 	RoutingDelegate string `yaml:"routingDelegate" yaml-aliases:"routingdelegate,routing-delegate,rd"`
 
-	Headers  map[string]string             `yaml:"headers"`
-	Baggage  map[string]string             `yaml:"baggage"`
-	Jaeger   bool                          `yaml:"jaeger"`
-	Request  map[interface{}]interface{}   `yaml:"request"`
-	Requests []map[interface{}]interface{} `yaml:"requests"`
-	Timeout  time.Duration                 `yaml:"timeout"`
+	Headers           map[string]string             `yaml:"headers"`
+	Baggage           map[string]string             `yaml:"baggage"`
+	Jaeger            bool                          `yaml:"jaeger"`
+	ForceJaegerSample bool                          `yaml:"forceJaegerSample" yaml-aliases:"forcejaegersample,force-jaeger-sample"`
+	Request           map[interface{}]interface{}   `yaml:"request"`
+	Requests          []map[interface{}]interface{} `yaml:"requests"`
+	Timeout           time.Duration                 `yaml:"timeout"`
 }
 
 func readYAMLFile(yamlTemplate string, templateArgs map[string]string, opts *Options) error {
@@ -150,6 +151,9 @@ func readYAMLRequest(base string, contents []byte, templateArgs map[string]strin
 	opts.ROpts.Baggage = merge(opts.ROpts.Baggage, t.Baggage)
 	if t.Jaeger {
 		opts.TOpts.Jaeger = true
+		if t.ForceJaegerSample {
+			opts.TOpts.ForceJaegerSample = true
+		}
 	}
 
 	if t.Thrift != "" {


Test patch:
diff --git a/integration_test.go b/integration_test.go
index 0e106f4..574f6fb 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -65,9 +65,10 @@ import (
 //go:generate protoc --go_out=plugins=grpc:. ./testdata/protobuf/simple/simple.proto
 
 var integrationTests = []struct {
-	call    int32
-	wantRes string
-	wantErr string
+	call            int32
+	wantRes         string
+	wantErr         string
+	shouldbesampled bool
 }{
 	{
 		call:    0,
@@ -81,6 +82,11 @@ var integrationTests = []struct {
 		call:    5,
 		wantRes: `"result": 5`,
 	},
+	{
+		call:            5,
+		wantRes:         `"result": 5`,
+		shouldbesampled: true,
+	},
 }
 
 func argHandler(arg int32) (int32, error) {
@@ -100,10 +106,15 @@ func verifyBaggage(ctx context.Context) error {
 		return errors.New("missing span")
 	}
 
-	if _, ok := span.Context().(jaeger.SpanContext); !ok {
+	spanContext, ok := span.Context().(jaeger.SpanContext)
+	if !ok {
 		return errors.New("trace context is not from jaeger")
 	}
 
+	if yarpc.CallFromContext(ctx).Header("shouldbesampled") == "true" && !spanContext.IsSampled() {
+		return errors.New("span is expected to be sampled")
+	}
+
 	val := span.BaggageItem("baggagekey")
 	if val == "" {
 		return errors.New("missing baggage")
@@ -267,11 +278,15 @@ func TestIntegrationProtocols(t *testing.T) {
 					ThriftDisableEnvelopes: c.disableEnvelope,
 				},
 				TOpts: TransportOptions{
-					ServiceName: "foo",
-					Peers:       []string{peer},
-					Jaeger:      true,
+					ServiceName:       "foo",
+					Peers:             []string{peer},
+					Jaeger:            true,
+					ForceJaegerSample: tt.shouldbesampled,
 				},
 			}
+			if tt.shouldbesampled {
+				opts.ROpts.Headers["shouldbesampled"] = "true"
+			}
 
 			gotOut, gotErr := runTestWithOpts(opts)
 			assert.Contains(t, gotOut, tt.wantRes, "%v: Unexpected result for %v", c.desc, tt.call)
diff --git a/main_test.go b/main_test.go
index 25a4a75..cfe4d17 100644
--- a/main_test.go
+++ b/main_test.go
@@ -39,6 +39,7 @@ import (
 	"github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/uber/jaeger-client-go"
 	"github.com/uber/tchannel-go/testutils"
 	"github.com/uber/tchannel-go/thrift"
 	"go.uber.org/thriftrw/protocol"
@@ -1234,6 +1235,24 @@ func TestGetTracer(t *testing.T) {
 				},
 			},
 		},
+		{
+			opts: Options{
+				TOpts: TransportOptions{
+					CallerName:        "test",
+					Jaeger:            true,
+					ForceJaegerSample: false,
+				},
+			},
+		},
+		{
+			opts: Options{
+				TOpts: TransportOptions{
+					CallerName:        "test",
+					Jaeger:            true,
+					ForceJaegerSample: true,
+				},
+			},
+		},
 		{
 			opts: Options{
 				TOpts: TransportOptions{
@@ -1279,6 +1298,12 @@ func TestGetTracer(t *testing.T) {
 			continue
 		}
 
+		if jaegerTracer, ok := tracer.(*jaeger.Tracer); assert.True(t, ok, "expected a Jaeger tracer") {
+			constSampler, ok := jaegerTracer.Sampler().(*jaeger.ConstSampler)
+			if assert.True(t, ok, "expected a ConstSampler") {
+				assert.Equal(t, tt.opts.TOpts.ForceJaegerSample, constSampler.Decision, "expected const sampler to be configured correctly")
+			}
+		}
 		assert.NotEqual(t, opentracing.NoopTracer{}, tracer, "Expected %+v to return real tracer")
 	}
 }
diff --git a/template_test.go b/template_test.go
index 637e6dd..fa11d7a 100644
--- a/template_test.go
+++ b/template_test.go
@@ -76,6 +76,7 @@ func TestTemplate(t *testing.T) {
 	assert.Equal(t, "rd", opts.TOpts.RoutingDelegate)
 	assert.Equal(t, map[string]string{"baggage1": "value1", "baggage2": "value2"}, opts.ROpts.Baggage)
 	assert.Equal(t, true, opts.TOpts.Jaeger)
+	assert.Equal(t, true, opts.TOpts.ForceJaegerSample)
 	assert.Equal(t, "location:\n  cityId: 1\n  latitude: 37.7\n  longitude: -122.4\n  message: true\n", opts.ROpts.RequestJSON)
 	assert.Equal(t, timeMillisFlag(4500*time.Millisecond), opts.ROpts.Timeout)
 	assert.True(t, opts.ROpts.ThriftDisableEnvelopes)
@@ -273,6 +274,7 @@ func TestTemplateAlias(t *testing.T) {
 		wantRoutingKey            string
 		wantRoutingDelegate       string
 		wantDisableThriftEnvelope *bool
+		wantForceJaegerSample     bool
 	}{
 		{
 			templates: []string{
@@ -309,6 +311,14 @@ func TestTemplateAlias(t *testing.T) {
 			},
 			wantDisableThriftEnvelope: new(bool),
 		},
+		{
+			templates: []string{
+				`forceJaegerSample: true`,
+				`forcejaegersample: true`,
+				`force-jaeger-sample: true`,
+			},
+			wantForceJaegerSample: true,
+		},
 	}
 
 	for _, tt := range tests {
@@ -321,6 +331,7 @@ func TestTemplateAlias(t *testing.T) {
 				assert.Equal(t, tt.wantRoutingKey, templ.RoutingKey, "routing key aliases expanded")
 				assert.Equal(t, tt.wantRoutingDelegate, templ.RoutingDelegate, "routing delegate aliases expanded")
 				assert.Equal(t, tt.wantDisableThriftEnvelope, templ.DisableThriftEnvelope, "disable thrift envelope aliases expanded")
+				assert.Equal(t, tt.wantForceJaegerSample, templ.ForceJaegerSample, "force jaeger sample aliases expanded")
 			})
 		}
 	}
diff --git a/testdata/templates/foo.yab b/testdata/templates/foo.yab
index ab97882..8c36de7 100644
--- a/testdata/templates/foo.yab
+++ b/testdata/templates/foo.yab
@@ -14,6 +14,7 @@ baggage:
     baggage1: value1
     baggage2: value2
 jaeger: true
+forceJaegerSample: true
 request:
     location:
         latitude: 37.7