Repository: yarpc/yab
Pull request #415 (2024-06-13 09:17:44)

Preparing release v0.23.0


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e250536..70ab813 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,9 @@
 Changelog
 =========
 
+# 0.23.0 (2024-06-13)
+* 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 != "" {
diff --git a/version.go b/version.go
index 3281137..8aebba2 100644
--- a/version.go
+++ b/version.go
@@ -22,4 +22,4 @@ package main
 
 // versionString is the sem-ver version string for yab.
 // It will be bumped explicitly on releases.
-var versionString = "0.22.0"
+var versionString = "0.23.0"


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