always force OTEL sampling

Author: ChrisPenner

src/Share/Telemetry/Setup.hs

@@ -6,13 +6,16 @@ import Data.HashMap.Lazy qualified as HM
 import Data.Text qualified as Text
 import Network.URI qualified as URI
 import OpenTelemetry.Attributes qualified as Trace
+import OpenTelemetry.Baggage qualified as Baggage
 import OpenTelemetry.Context (Context)
+import OpenTelemetry.Context qualified as Context
 import OpenTelemetry.Trace qualified as Trace
 import OpenTelemetry.Trace.Id (TraceId)
 import OpenTelemetry.Trace.Sampler qualified as Sampler
 import OpenTelemetry.Trace.TraceState (TraceState)
 import OpenTelemetry.Trace.TraceState qualified as TraceState
 import Share.Prelude
+import Unison.Debug qualified as Debug
 
 initSampler :: IO Sampler.Sampler
 initSampler =
@@ -27,14 +30,27 @@ initSampler =
         _ -> Nothing
     dropSample = (Sampler.Drop, HM.empty, TraceState.empty)
     -- Configure some custom sampling logic.
+    shouldForceSample :: Context -> Bool
+    shouldForceSample ctx = fromMaybe False do
+      baggage <- Context.lookupBaggage ctx
+      let baggageMap = Baggage.values baggage
+      forceTraceToken <- Baggage.mkToken "force-trace"
+      forceTrace <- Baggage.value <$> HM.lookup forceTraceToken baggageMap
+      Debug.debugLogM Debug.Temp ("Force trace token: " <> show forceTrace)
+      case Text.toLower forceTrace of
+        "false" -> Just False
+        "0" -> Just False
+        _ -> Just True
     shouldSample :: Sampler.Sampler -> Context -> TraceId -> Text -> Trace.SpanArguments -> IO (Sampler.SamplingResult, HashMap Text Trace.Attribute, TraceState)
-    shouldSample defaultSampler ctx tid name args = do
-      case (lookupTextAttribute "http.target" args >>= URI.parseURIReference . Text.unpack) <&> URI.pathSegments of
-        Just ("metrics" : _) -> pure dropSample
-        Just ("health" : _) -> pure dropSample
-        -- This is currently used in a health check.
-        Just ("users" : "zarelit" : _) -> pure dropSample
-        _ -> Sampler.shouldSample defaultSampler ctx tid name args
+    shouldSample defaultSampler ctx tid name args
+      | shouldForceSample ctx = Sampler.shouldSample Sampler.alwaysOn ctx tid name args
+      | otherwise = do
+          case (lookupTextAttribute "http.target" args >>= URI.parseURIReference . Text.unpack) <&> URI.pathSegments of
+            Just ("metrics" : _) -> pure dropSample
+            Just ("health" : _) -> pure dropSample
+            -- This is currently used in a health check.
+            Just ("users" : "zarelit" : _) -> pure dropSample
+            _ -> Sampler.shouldSample defaultSampler ctx tid name args
 
 withTracer :: Text -> (Trace.Tracer -> IO c) -> IO c
 withTracer commitHash f =