| 
 | 1 | +/*  | 
 | 2 | + * Copyright The OpenTelemetry Authors  | 
 | 3 | + * SPDX-License-Identifier: Apache-2.0  | 
 | 4 | + */  | 
 | 5 | + | 
 | 6 | +package io.opentelemetry.contrib.sampler;  | 
 | 7 | + | 
 | 8 | +import static java.util.Objects.requireNonNull;  | 
 | 9 | + | 
 | 10 | +import dev.cel.common.types.CelProtoTypes;  | 
 | 11 | +import dev.cel.common.types.SimpleType;  | 
 | 12 | +import dev.cel.compiler.CelCompiler;  | 
 | 13 | +import dev.cel.compiler.CelCompilerFactory;  | 
 | 14 | +import dev.cel.runtime.CelEvaluationException;  | 
 | 15 | +import dev.cel.runtime.CelRuntime;  | 
 | 16 | +import dev.cel.runtime.CelRuntimeFactory;  | 
 | 17 | +import io.opentelemetry.api.common.Attributes;  | 
 | 18 | +import io.opentelemetry.api.trace.SpanBuilder;  | 
 | 19 | +import io.opentelemetry.api.trace.SpanKind;  | 
 | 20 | +import io.opentelemetry.context.Context;  | 
 | 21 | +import io.opentelemetry.sdk.trace.data.LinkData;  | 
 | 22 | +import io.opentelemetry.sdk.trace.samplers.Sampler;  | 
 | 23 | +import io.opentelemetry.sdk.trace.samplers.SamplingResult;  | 
 | 24 | +import java.util.HashMap;  | 
 | 25 | +import java.util.List;  | 
 | 26 | +import java.util.Map;  | 
 | 27 | +import java.util.logging.Level;  | 
 | 28 | +import java.util.logging.Logger;  | 
 | 29 | +import java.util.regex.Pattern;  | 
 | 30 | + | 
 | 31 | +/**  | 
 | 32 | + * This sampler accepts a list of {@link CelBasedSamplingExpression}s and tries to match every  | 
 | 33 | + * proposed span against those rules. Every rule describes a span's attribute, a pattern against  | 
 | 34 | + * which to match attribute's value, and a sampler that will make a decision about given span if  | 
 | 35 | + * match was successful.  | 
 | 36 | + *  | 
 | 37 | + * <p>Matching is performed by {@link Pattern}.  | 
 | 38 | + *  | 
 | 39 | + * <p>Provided span kind is checked first and if differs from the one given to {@link  | 
 | 40 | + * #builder(Sampler)}, the default fallback sampler will make a decision.  | 
 | 41 | + *  | 
 | 42 | + * <p>Note that only attributes that were set on {@link SpanBuilder} will be taken into account,  | 
 | 43 | + * attributes set after the span has been started are not used  | 
 | 44 | + *  | 
 | 45 | + * <p>If none of the rules matched, the default fallback sampler will make a decision.  | 
 | 46 | + */  | 
 | 47 | +public final class CelBasedSampler implements Sampler {  | 
 | 48 | + | 
 | 49 | +  private static final Logger logger = Logger.getLogger(CelBasedSampler.class.getName());  | 
 | 50 | + | 
 | 51 | +  public static final CelCompiler celCompiler =  | 
 | 52 | +      CelCompilerFactory.standardCelCompilerBuilder()  | 
 | 53 | +          .addVar("name", SimpleType.STRING)  | 
 | 54 | +          .addVar("traceId", SimpleType.STRING)  | 
 | 55 | +          .addVar("spanKind", SimpleType.STRING)  | 
 | 56 | +          .addVar("attribute", CelProtoTypes.createMap(CelProtoTypes.STRING, CelProtoTypes.DYN))  | 
 | 57 | +          .setResultType(SimpleType.BOOL)  | 
 | 58 | +          .build();  | 
 | 59 | + | 
 | 60 | +  final CelRuntime celRuntime;  | 
 | 61 | + | 
 | 62 | +  private final List<CelBasedSamplingExpression> expressions;  | 
 | 63 | +  private final Sampler fallback;  | 
 | 64 | + | 
 | 65 | +  public CelBasedSampler(List<CelBasedSamplingExpression> expressions, Sampler fallback) {  | 
 | 66 | +    this.expressions = requireNonNull(expressions, "expressions must not be null");  | 
 | 67 | +    this.expressions.forEach(  | 
 | 68 | +        expr -> {  | 
 | 69 | +          if (!expr.abstractSyntaxTree.isChecked()) {  | 
 | 70 | +            throw new IllegalArgumentException(  | 
 | 71 | +                "Expression and its AST is not checked: " + expr.expression);  | 
 | 72 | +          }  | 
 | 73 | +        });  | 
 | 74 | +    this.fallback = requireNonNull(fallback);  | 
 | 75 | +    this.celRuntime = CelRuntimeFactory.standardCelRuntimeBuilder().build();  | 
 | 76 | +  }  | 
 | 77 | + | 
 | 78 | +  public static CelBasedSamplerBuilder builder(Sampler fallback) {  | 
 | 79 | +    return new CelBasedSamplerBuilder(  | 
 | 80 | +        requireNonNull(fallback, "fallback sampler must not be null"), celCompiler);  | 
 | 81 | +  }  | 
 | 82 | + | 
 | 83 | +  @Override  | 
 | 84 | +  public SamplingResult shouldSample(  | 
 | 85 | +      Context parentContext,  | 
 | 86 | +      String traceId,  | 
 | 87 | +      String name,  | 
 | 88 | +      SpanKind spanKind,  | 
 | 89 | +      Attributes attributes,  | 
 | 90 | +      List<LinkData> parentLinks) {  | 
 | 91 | + | 
 | 92 | +    // Prepare the evaluation context with span data  | 
 | 93 | +    Map<String, Object> evaluationContext = new HashMap<>();  | 
 | 94 | +    evaluationContext.put("name", name);  | 
 | 95 | +    evaluationContext.put("traceId", traceId);  | 
 | 96 | +    evaluationContext.put("spanKind", spanKind.name());  | 
 | 97 | +    evaluationContext.put("attribute", convertAttributesToMap(attributes));  | 
 | 98 | + | 
 | 99 | +    for (CelBasedSamplingExpression expression : expressions) {  | 
 | 100 | +      try {  | 
 | 101 | +        CelRuntime.Program program = celRuntime.createProgram(expression.abstractSyntaxTree);  | 
 | 102 | +        Object result = program.eval(evaluationContext);  | 
 | 103 | +        // Happy path: Perform sampling based on the boolean result  | 
 | 104 | +        if (result instanceof Boolean && ((Boolean) result)) {  | 
 | 105 | +          return expression.delegate.shouldSample(  | 
 | 106 | +              parentContext, traceId, name, spanKind, attributes, parentLinks);  | 
 | 107 | +        }  | 
 | 108 | +        // If result is not boolean, treat as false  | 
 | 109 | +        logger.log(  | 
 | 110 | +            Level.FINE,  | 
 | 111 | +            "Expression '" + expression.expression + "' returned non-boolean result: " + result);  | 
 | 112 | +      } catch (CelEvaluationException e) {  | 
 | 113 | +        logger.log(  | 
 | 114 | +            Level.FINE,  | 
 | 115 | +            "Expression '" + expression.expression + "' evaluation error: " + e.getMessage());  | 
 | 116 | +      }  | 
 | 117 | +    }  | 
 | 118 | + | 
 | 119 | +    return fallback.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);  | 
 | 120 | +  }  | 
 | 121 | + | 
 | 122 | +  /** Convert OpenTelemetry Attributes to a Map that CEL can work with */  | 
 | 123 | +  private static Map<String, Object> convertAttributesToMap(Attributes attributes) {  | 
 | 124 | +    Map<String, Object> map = new HashMap<>();  | 
 | 125 | +    attributes.forEach(  | 
 | 126 | +        (key, value) -> {  | 
 | 127 | +          map.put(key.getKey(), value);  | 
 | 128 | +        });  | 
 | 129 | +    return map;  | 
 | 130 | +  }  | 
 | 131 | + | 
 | 132 | +  @Override  | 
 | 133 | +  public String getDescription() {  | 
 | 134 | +    return "CelBasedSampler{" + "fallback=" + fallback + ", expressions=" + expressions + '}';  | 
 | 135 | +  }  | 
 | 136 | + | 
 | 137 | +  @Override  | 
 | 138 | +  public String toString() {  | 
 | 139 | +    return getDescription();  | 
 | 140 | +  }  | 
 | 141 | +}  | 
0 commit comments