Ensure span cancellation is done properly with concurrency protection

.gitignore

@@ -101,3 +101,4 @@ logs/*
 
 # Other
 .DS_Store
+features/fixtures/ios/Fixture/Info.plist

Sources/BugsnagPerformance/Private/BugsnagPerformanceImpl.mm

@@ -21,13 +21,6 @@
     return NSSearchPathForDirectoriesInDomains(NSCachesDirectory, NSUserDomainMask, YES).firstObject;
 }
 
-void (^generateOnSpanStarted(BugsnagPerformanceImpl *impl))(void) {
-    __block auto blockImpl = impl;
-  return ^{
-      blockImpl->onSpanStarted();
-  };
-}
-
 BugsnagPerformanceImpl::BugsnagPerformanceImpl(std::shared_ptr<Reachability> reachability,
                                                AppStateTracker *appStateTracker) noexcept
 : persistence_(std::make_shared<Persistence>(getPersistenceDir()))
@@ -36,7 +29,7 @@
 , reachability_(reachability)
 , batch_(std::make_shared<Batch>())
 , sampler_(std::make_shared<Sampler>())
-, tracer_(std::make_shared<Tracer>(spanStackingHandler_, sampler_, batch_, generateOnSpanStarted(this)))
+, tracer_(std::make_shared<Tracer>(spanStackingHandler_, sampler_, batch_, ^{this->onSpanStarted();}))
 , retryQueue_(std::make_unique<RetryQueue>([persistence_->bugsnagPerformanceDir() stringByAppendingPathComponent:@"retry-queue"]))
 , appStateTracker_(appStateTracker)
 , viewControllersToSpans_([NSMapTable mapTableWithKeyOptions:NSMapTableWeakMemory | NSMapTableObjectPointerPersonality

Sources/BugsnagPerformance/Private/BugsnagPerformanceLibrary.mm

@@ -51,18 +51,17 @@
 : appStateTracker_([[AppStateTracker alloc] init])
 , reachability_(std::make_shared<Reachability>())
 , bugsnagPerformanceImpl_(std::make_shared<BugsnagPerformanceImpl>(reachability_, appStateTracker_))
-{
-    auto impl = bugsnagPerformanceImpl_;
-    bugsnagPerformanceImpl_->setOnViewLoadSpanStarted([=](NSString *className) {
-        impl->didStartViewLoadSpan(className);
-    });
-}
+{}
 
 void BugsnagPerformanceLibrary::earlyConfigure(BSGEarlyConfiguration *config) noexcept {
     bugsnagPerformanceImpl_->earlyConfigure(config);
 }
 
 void BugsnagPerformanceLibrary::earlySetup() noexcept {
+    auto impl = bugsnagPerformanceImpl_;
+    bugsnagPerformanceImpl_->setOnViewLoadSpanStarted([=](NSString *className) {
+        impl->didStartViewLoadSpan(className);
+    });
     bugsnagPerformanceImpl_->earlySetup();
 }
 

Sources/BugsnagPerformance/Private/Instrumentation/AppStartupInstrumentation.h

@@ -23,7 +23,7 @@ class AppStartupInstrumentation: public PhasedStartup {
                               std::shared_ptr<SpanAttributesProvider> spanAttributesProvider) noexcept;
 
     void earlyConfigure(BSGEarlyConfiguration *) noexcept {}
-    void earlySetup() noexcept {}
+    void earlySetup() noexcept;
     void configure(BugsnagPerformanceConfiguration *config) noexcept;
     void start() noexcept {}
 

Sources/BugsnagPerformance/Private/Instrumentation/AppStartupInstrumentation.mm

@@ -64,9 +64,10 @@ static inline bool isActivePrewarm(void) {
 , tracer_(tracer)
 , spanAttributesProvider_(spanAttributesProvider)
 , didStartProcessAtTime_(getProcessStartTime())
-, didCallMainFunctionAtTime_(CFAbsoluteTimeGetCurrent())
 , isColdLaunch_(isColdLaunch())
-{
+{}
+
+void AppStartupInstrumentation::earlySetup() noexcept {
     if (!canInstallInstrumentation()) {
         disable();
     }
@@ -86,6 +87,7 @@ static inline bool isActivePrewarm(void) {
 
     beginAppStartSpan();
     beginPreMainSpan();
+    didCallMainFunctionAtTime_ = CFAbsoluteTimeGetCurrent();
     [preMainSpan_ endWithAbsoluteTime:didCallMainFunctionAtTime_];
     beginPostMainSpan();
 

Sources/BugsnagPerformance/Private/Tracer.h

@@ -65,12 +65,16 @@ class Tracer: public PhasedStartup {
     NSMutableArray<BugsnagPerformanceSpan *> *earlyNetworkSpans_;
 
     std::shared_ptr<Batch> batch_;
-    void (^onSpanStarted_)(){nil};
-    std::function<void(NSString *)> onViewLoadSpanStarted_{};
-    BugsnagPerformanceNetworkRequestCallback networkRequestCallback_;
+    void (^onSpanStarted_)(){ ^(){} };
+    std::function<void(NSString *)> onViewLoadSpanStarted_{ [](NSString *){} };
+    BugsnagPerformanceNetworkRequestCallback networkRequestCallback_ {
+        ^BugsnagPerformanceNetworkRequestInfo * _Nonnull(BugsnagPerformanceNetworkRequestInfo * _Nonnull info) {
+            return info;
+        }
+    };
 
     BugsnagPerformanceSpan *startSpan(NSString *name, SpanOptions options, BSGFirstClass defaultFirstClass) noexcept;
-    void tryAddSpanToBatch(std::shared_ptr<SpanData> spanData);
+    void trySampleAndAddSpanToBatch(std::shared_ptr<SpanData> spanData);
     void markEarlyNetworkSpan(BugsnagPerformanceSpan *span) noexcept;
     void endEarlySpansPhase() noexcept;
 };

Sources/BugsnagPerformance/Private/Tracer.mm

@@ -41,7 +41,7 @@
     // Now that the sampler has been configured, re-sample everything.
     auto unsampledBatch = batch_->drain(true);
     for (auto spanData: *unsampledBatch) {
-        tryAddSpanToBatch(spanData);
+        trySampleAndAddSpanToBatch(spanData);
     }
 }
 
@@ -69,19 +69,17 @@
                                                               firstClass,
                                        ^void(std::shared_ptr<SpanData> spanData) {
         blockThis->spanStackingHandler_->didEnd(spanData->spanId);
-        blockThis->tryAddSpanToBatch(spanData);
+        blockThis->trySampleAndAddSpanToBatch(spanData);
     })];
     if (options.makeCurrentContext) {
         spanStackingHandler_->push(span);
     }
     [span addAttributes:SpanAttributes::get()];
-    if (onSpanStarted_) {
-        onSpanStarted_();
-    }
+    onSpanStarted_();
     return span;
 }
 
-void Tracer::tryAddSpanToBatch(std::shared_ptr<SpanData> spanData) {
+void Tracer::trySampleAndAddSpanToBatch(std::shared_ptr<SpanData> spanData) {
     if (sampler_->sampled(*spanData)) {
         batch_->add(spanData);
     }
@@ -116,21 +114,18 @@
 
 BugsnagPerformanceSpan *
 Tracer::startNetworkSpan(NSURL *url, NSString *httpMethod, SpanOptions options) noexcept {
-    if (networkRequestCallback_ != nil) {
-        auto info = [BugsnagPerformanceNetworkRequestInfo new];
-        info.url = url;
-        info = networkRequestCallback_(info);
-        url = info.url;
-    }
+    auto info = [BugsnagPerformanceNetworkRequestInfo new];
+    info.url = url;
+    info = networkRequestCallback_(info);
+    url = info.url;
     if (url == nil) {
         return nil;
     }
 
     auto name = [NSString stringWithFormat:@"[HTTP/%@]", httpMethod];
     auto span = startSpan(name, options, BSGFirstClassUnset);
     [span addAttribute:httpUrlAttributeKey withValue:(NSString *_Nonnull)url.absoluteString];
-    // Check both isEarlySpansPhase_ and networkRequestCallback_ to counter potential race condition
-    if (isEarlySpansPhase_ && !networkRequestCallback_) {
+    if (isEarlySpansPhase_) {
         markEarlyNetworkSpan(span);
     }
     return span;
@@ -150,23 +145,24 @@
 
 void Tracer::markEarlyNetworkSpan(BugsnagPerformanceSpan *span) noexcept {
     std::lock_guard<std::mutex> guard(earlySpansMutex_);
-    [earlyNetworkSpans_ addObject:span];
+    if (isEarlySpansPhase_) {
+        [earlyNetworkSpans_ addObject:span];
+    }
 }
 
 void Tracer::endEarlySpansPhase() noexcept {
-    isEarlySpansPhase_ = false;
     std::lock_guard<std::mutex> guard(earlySpansMutex_);
-    if (networkRequestCallback_ != nil) {
-        for (BugsnagPerformanceSpan *span: earlyNetworkSpans_) {
-            auto info = [BugsnagPerformanceNetworkRequestInfo new];
-            NSString *urlString = [span getAttribute:httpUrlAttributeKey];
-            info.url = [NSURL URLWithString:urlString];
-            info = networkRequestCallback_(info);
-            if (info.url == nil) {
-                batch_->removeSpan(span.traceId, span.spanId);
-            } else {
-                [span addAttribute:httpUrlAttributeKey withValue:(NSString *_Nonnull)info.url.absoluteString];
-            }
+    isEarlySpansPhase_ = false;
+    for (BugsnagPerformanceSpan *span: earlyNetworkSpans_) {
+        auto info = [BugsnagPerformanceNetworkRequestInfo new];
+        NSString *urlString = [span getAttribute:httpUrlAttributeKey];
+        info.url = [NSURL URLWithString:urlString];
+        // We have to check again because the real callback might not have been set initially.
+        info = networkRequestCallback_(info);
+        if (info.url != nil) {
+            [span addAttribute:httpUrlAttributeKey withValue:(NSString *_Nonnull)info.url.absoluteString];
+        } else {
+            cancelQueuedSpan(span);
         }
     }
     [earlyNetworkSpans_ removeAllObjects];