From 73f2b523f0f1b04b0c58979a10fafc7860819889 Mon Sep 17 00:00:00 2001
From: Moritz Lang <hi@slashmo.codes>
Date: Mon, 21 Dec 2020 15:18:12 +0100
Subject: [PATCH] Delegate to withSpan(_ baggage:) from withSpan(_ context:)
 (#36)

Co-authored-by: Konrad `ktoso` Malawski <ktoso@apple.com>
---
 Sources/Tracing/Tracer.swift                |  9 +--------
 Tests/TracingTests/TracerTests+XCTest.swift |  1 +
 Tests/TracingTests/TracerTests.swift        | 15 +++++++++++++++
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/Sources/Tracing/Tracer.swift b/Sources/Tracing/Tracer.swift
index 21adec8..af68022 100644
--- a/Sources/Tracing/Tracer.swift
+++ b/Sources/Tracing/Tracer.swift
@@ -97,14 +97,7 @@ extension Tracer {
         ofKind kind: SpanKind = .internal,
         _ function: (Span) throws -> T
     ) rethrows -> T {
-        let span = self.startSpan(operationName, context: context, ofKind: kind)
-        defer { span.end() }
-        do {
-            return try function(span)
-        } catch {
-            span.recordError(error)
-            throw error // rethrow
-        }
+        return try self.withSpan(operationName, baggage: context.baggage, function)
     }
 
     /// Execute a specific task within a newly created `Span`.
diff --git a/Tests/TracingTests/TracerTests+XCTest.swift b/Tests/TracingTests/TracerTests+XCTest.swift
index 3744864..55231cd 100644
--- a/Tests/TracingTests/TracerTests+XCTest.swift
+++ b/Tests/TracingTests/TracerTests+XCTest.swift
@@ -29,6 +29,7 @@ extension TracerTests {
                 ("testContextPropagationWithNoOpSpan", testContextPropagationWithNoOpSpan),
                 ("testWithSpan_success", testWithSpan_success),
                 ("testWithSpan_throws", testWithSpan_throws),
+                ("testWithSpan_context", testWithSpan_context),
            ]
    }
 }
diff --git a/Tests/TracingTests/TracerTests.swift b/Tests/TracingTests/TracerTests.swift
index 54dd779..82cf8cc 100644
--- a/Tests/TracingTests/TracerTests.swift
+++ b/Tests/TracingTests/TracerTests.swift
@@ -95,6 +95,21 @@ final class TracerTests: XCTestCase {
         }
         XCTFail("Should have throw")
     }
+
+    func testWithSpan_context() {
+        let tracer = TestTracer()
+        InstrumentationSystem.bootstrapInternal(tracer)
+        defer {
+            InstrumentationSystem.bootstrapInternal(NoOpTracer())
+        }
+
+        var spanEnded = false
+        tracer.onEndSpan = { _ in spanEnded = true }
+
+        tracer.withSpan("", context: DefaultLoggingContext.topLevel(logger: Logger(label: "test"))) { _ in }
+
+        XCTAssertTrue(spanEnded)
+    }
 }
 
 struct ExampleSpanError: Error, Equatable {}