From 5bbfc22cbc5bdcb65eb5c354c264b33f9faa9d76 Mon Sep 17 00:00:00 2001
From: humivo <hvoam@amazon.com>
Date: Mon, 5 Apr 2021 16:33:16 -0400
Subject: [PATCH] Make ExportSpans for Jaeger Exporter honor deadline (#1773)

* Make ExportSpans for Jaeger honor deadline

* Make variable name more descriptive

* new commit

* Revert "new commit"

This reverts commit 06e24cc38d52482eb2187df31a1181e3fddb72f7.

* Change PR number in changelog

* Take out separate goroutine and add back TODO

* Check error string

* Fix error assert

Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com>
---
 CHANGELOG.md                          |  1 +
 exporters/trace/jaeger/jaeger.go      |  8 +++-
 exporters/trace/jaeger/jaeger_test.go | 61 +++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f4a44b60..ec9f9c49c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -54,6 +54,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
   This changes it to make `SamplingParameters` conform with the OpenTelemetry specification. (#1749)
 - Modify `BatchSpanProcessor.ForceFlush` to abort after timeout/cancellation. (#1757)
 - Improve OTLP/gRPC exporter connection errors. (#1737)
+- Make `ExportSpans` in Jaeger Exporter honor context deadline. (#1773)
 
 ### Removed
 
diff --git a/exporters/trace/jaeger/jaeger.go b/exporters/trace/jaeger/jaeger.go
index e4d45afce..7eb6b2823 100644
--- a/exporters/trace/jaeger/jaeger.go
+++ b/exporters/trace/jaeger/jaeger.go
@@ -217,11 +217,17 @@ func (e *Exporter) ExportSpans(ctx context.Context, ss []*export.SpanSnapshot) e
 
 	for _, span := range ss {
 		// TODO(jbd): Handle oversized bundlers.
-		err := e.bundler.Add(span, 1)
+		err := e.bundler.AddWait(ctx, span, 1)
 		if err != nil {
 			return fmt.Errorf("failed to bundle %q: %w", span.Name, err)
 		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
 	}
+
 	return nil
 }
 
diff --git a/exporters/trace/jaeger/jaeger_test.go b/exporters/trace/jaeger/jaeger_test.go
index e7cbd5466..0666f17af 100644
--- a/exporters/trace/jaeger/jaeger_test.go
+++ b/exporters/trace/jaeger/jaeger_test.go
@@ -649,6 +649,67 @@ func TestErrorOnExportShutdownExporter(t *testing.T) {
 	assert.NoError(t, e.ExportSpans(context.Background(), nil))
 }
 
+func TestExporterExportSpansHonorsCancel(t *testing.T) {
+	e, err := NewRawExporter(withTestCollectorEndpoint())
+	require.NoError(t, err)
+	now := time.Now()
+	ss := []*export.SpanSnapshot{
+		{
+			Name: "s1",
+			Resource: resource.NewWithAttributes(
+				semconv.ServiceNameKey.String("name"),
+				attribute.Key("r1").String("v1"),
+			),
+			StartTime: now,
+			EndTime:   now,
+		},
+		{
+			Name: "s2",
+			Resource: resource.NewWithAttributes(
+				semconv.ServiceNameKey.String("name"),
+				attribute.Key("r2").String("v2"),
+			),
+			StartTime: now,
+			EndTime:   now,
+		},
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	assert.EqualError(t, e.ExportSpans(ctx, ss), context.Canceled.Error())
+}
+
+func TestExporterExportSpansHonorsTimeout(t *testing.T) {
+	e, err := NewRawExporter(withTestCollectorEndpoint())
+	require.NoError(t, err)
+	now := time.Now()
+	ss := []*export.SpanSnapshot{
+		{
+			Name: "s1",
+			Resource: resource.NewWithAttributes(
+				semconv.ServiceNameKey.String("name"),
+				attribute.Key("r1").String("v1"),
+			),
+			StartTime: now,
+			EndTime:   now,
+		},
+		{
+			Name: "s2",
+			Resource: resource.NewWithAttributes(
+				semconv.ServiceNameKey.String("name"),
+				attribute.Key("r2").String("v2"),
+			),
+			StartTime: now,
+			EndTime:   now,
+		},
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
+	defer cancel()
+	<-ctx.Done()
+
+	assert.EqualError(t, e.ExportSpans(ctx, ss), context.DeadlineExceeded.Error())
+}
+
 func TestJaegerBatchList(t *testing.T) {
 	newString := func(value string) *string {
 		return &value