Call span.end() in the happy path in withSpan (#33)

* Test that withSpan() calls end() as expected

* Call span.end() in the happy path in withSpan

* Apply suggestions from code review

* Update TracerTests.swift

Co-authored-by: Konrad `ktoso` Malawski <[email protected]>
Co-authored-by: Konrad `ktoso` Malawski <[email protected]>

Author: JensAyton

Sources/Tracing/Tracer.swift

@@ -98,11 +98,11 @@ extension Tracer {
         _ 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)
-            span.end()
             throw error // rethrow
         }
     }
@@ -125,11 +125,11 @@ extension Tracer {
         _ function: (Span) throws -> T
     ) rethrows -> T {
         let span = self.startSpan(operationName, baggage: baggage, ofKind: kind)
+        defer { span.end() }
         do {
             return try function(span)
         } catch {
             span.recordError(error)
-            span.end()
             throw error // rethrow
         }
     }

Tests/TracingTests/TestTracer.swift

@@ -19,6 +19,7 @@ import Tracing
 
 final class TestTracer: Tracer {
     private(set) var spans = [TestSpan]()
+    var onEndSpan: (Span) -> Void = { _ in }
 
     func startSpan(
         _ operationName: String,
@@ -30,8 +31,9 @@ final class TestTracer: Tracer {
             operationName: operationName,
             startTime: time,
             baggage: baggage,
-            kind: kind
-        ) { _ in }
+            kind: kind,
+            onEnd: onEndSpan
+        )
         self.spans.append(span)
         return span
     }

Tests/TracingTests/TracerTests.swift

@@ -63,11 +63,15 @@ final class TracerTests: XCTestCase {
             InstrumentationSystem.bootstrapInternal(NoOpTracer())
         }
 
+        var spanEnded = false
+        tracer.onEndSpan = { _ in spanEnded = true }
+
         let value = tracer.withSpan("hello", baggage: .topLevel) { _ in
             "yes"
         }
 
         XCTAssertEqual(value, "yes")
+        XCTAssertTrue(spanEnded)
     }
 
     func testWithSpan_throws() {
@@ -77,11 +81,15 @@ final class TracerTests: XCTestCase {
             InstrumentationSystem.bootstrapInternal(NoOpTracer())
         }
 
+        var spanEnded = false
+        tracer.onEndSpan = { _ in spanEnded = true }
+
         do {
             _ = try tracer.withSpan("hello", baggage: .topLevel) { _ in
                 throw ExampleSpanError()
             }
         } catch {
+            XCTAssertTrue(spanEnded)
             XCTAssertEqual(error as? ExampleSpanError, ExampleSpanError())
             return
         }