Implement traceID-ratio and parent-based samplers in eBPF (#1150)

* Implements traceID-ration and parent base samplers in eBPF * Add changelog, and run ebpf tests on x86 linux in CI * fix changelog entry * code review adjustments * Update internal/pkg/instrumentation/probe/sampling/sampling_linux_test.go Co-authored-by: Tyler Yahn <[email protected]> * update makefile trarget for test_ebpf --------- Co-authored-by: Tyler Yahn <[email protected]>
open-telemetry · Oct 19, 2024 · 83f7338 · 83f7338
1 parent 756cc23
commit 83f7338
Show file tree

Hide file tree

Showing 29 changed files with 654 additions and 106 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -27,6 +27,7 @@ jobs:
           uname -p;
           clang --version;
       - run: make test
+      - run: sudo make test_ebpf
       - run: make check-clean-work-tree
   generate-and-test-arm64:
     runs-on: macos-latest

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 - Support `google.golang.org/grpc` `1.65.1`. ([#1174](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/1174))
 - Support `go.opentelemetry.io/[email protected]`. ([#1178](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/1178))
 - Support `google.golang.org/grpc` `1.69.0-dev`. ([#1203](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/1203))
+- Implement traceID ratio and parent-based samplers. ([#1150](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/1150))
 
 ## [v0.15.0-alpha] - 2024-10-01
 

diff --git a/Makefile b/Makefile
@@ -51,13 +51,21 @@ tools: $(GOLICENSES) $(MULTIMOD) $(GOLANGCI_LINT) $(DBOTCONF) $(OFFSETGEN)
 
 ALL_GO_MODS := $(shell find . -type f -name 'go.mod' ! -path '$(TOOLS_MOD_DIR)/*' ! -path './LICENSES/*' | sort)
 GO_MODS_TO_TEST := $(ALL_GO_MODS:%=test/%)
+GO_MODS_TO_EBPF_TEST := $(ALL_GO_MODS:%=test_ebpf/%)
 
 .PHONY: test
 test: generate $(GO_MODS_TO_TEST)
 test/%: GO_MOD=$*
 test/%:
 	cd $(shell dirname $(GO_MOD)) && $(GOCMD) test -v ./...
 
+# These tests need to be run as privileged user/with sudo
+.PHONY: test_ebpf
+test_ebpf: generate $(GO_MODS_TO_EBPF_TEST)
+test_ebpf/%: GO_MOD=$*
+test_ebpf/%:
+	cd $(shell dirname $(GO_MOD)) && $(GOCMD) test -v -tags=ebpf_test -run ^TestEBPF ./...
+
 .PHONY: generate
 generate: export CFLAGS := $(BPF_INCLUDE)
 generate: go-mod-tidy

diff --git a/internal/include/trace/sampling.h b/internal/include/trace/sampling.h
@@ -7,15 +7,163 @@
 #include "common.h"
 #include "span_context.h"
 
+#define MAX_SAMPLER_CONFIG_SIZE 256
+#define MAX_SAMPLERS 32
+
+typedef u32 sampler_id_t;
+
+struct parent_based_config {
+    sampler_id_t root;
+    sampler_id_t remote_parent_sampled;
+    sampler_id_t remote_parent_not_sampled;
+    sampler_id_t local_parent_sampled;
+    sampler_id_t local_parent_not_sampled;
+};
+
+enum sampler_type {
+    // OpenTelemetry spec defined samplers
+    ALWAYS_ON = 0,
+    ALWAYS_OFF = 1,
+    TRACE_ID_RATIO = 2,
+    PARENT_BASED = 3,
+    // Custom samplers
+
+};
+
+struct sampling_config {
+    enum sampler_type type;
+    union {
+        u64 sampling_rate_numerator;
+        struct parent_based_config parent_based;
+        char buf[MAX_SAMPLER_CONFIG_SIZE];
+    } config_data;
+};
+
 typedef struct sampling_parameters {
     struct span_context *psc;
     u8 *trace_id;
     // TODO: add more fields
 } sampling_parameters_t;
 
-static __always_inline bool should_sample(sampling_parameters_t *params) {
-    // TODO
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(key_size, sizeof(sampler_id_t));
+	__uint(value_size, sizeof(struct sampling_config));
+	__uint(max_entries, MAX_SAMPLERS);
+} samplers_config_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(key_size, sizeof(u32));
+	__uint(value_size, sizeof(sampler_id_t));
+	__uint(max_entries, 1);
+} probe_active_sampler_map SEC(".maps");
+
+static const u8 FLAG_SAMPLED = 1;
+
+static __always_inline bool trace_flags_is_sampled(u8 flags)
+{
+    return ((flags & FLAG_SAMPLED) == FLAG_SAMPLED);
+}
+
+static __always_inline bool is_sampled(struct span_context *ctx)
+{
+    return trace_flags_is_sampled(ctx->TraceFlags);
+}
+
+// This value should be in sync with user-space code which configures the sampler
+static const u64 sampling_rate_denominator = ((1ULL<<32) - 1);
+
+static __always_inline bool _traceIDRatioSampler_should_sample(u64 sampling_rate_numerator, u8 *trace_id) {
+    if (sampling_rate_numerator == 0) {
+        return false;
+    }
+
+    if (sampling_rate_numerator >= sampling_rate_denominator) {
+        return true;
+    }
+
+    u64 trace_id_num = 0;
+    __builtin_memcpy(&trace_id_num, &trace_id[8], 8);
+    u64 trace_id_upper_bound = ((1ULL << 63) / sampling_rate_denominator) * sampling_rate_numerator;
+    return (trace_id_num >> 1) < trace_id_upper_bound;
+}
+
+static __always_inline bool traceIDRatioSampler_should_sample(struct sampling_config* config, sampling_parameters_t *params) {
+    return _traceIDRatioSampler_should_sample(config->config_data.sampling_rate_numerator, params->trace_id);
+}
+
+static __always_inline bool alwaysOnSampler_should_sample(struct sampling_config* config, sampling_parameters_t *params) {
     return true;
 }
 
+static __always_inline bool alwaysOffSampler_should_sample(struct sampling_config* config, sampling_parameters_t *params) {
+    return false;
+}
+
+static __always_inline bool parentBasedSampler_should_sample(struct sampling_config* config, sampling_parameters_t *params) {
+    sampler_id_t sampler_id; 
+    if (params->psc == NULL) {
+        sampler_id = config->config_data.parent_based.root;
+    } else {
+        // TODO: once we add remote parent field to span context, we should check if it's remote or local
+        // currently assuming local parent
+        if (trace_flags_is_sampled(params->psc->TraceFlags)) {
+            sampler_id = config->config_data.parent_based.local_parent_sampled;
+        } else {
+            sampler_id = config->config_data.parent_based.local_parent_not_sampled;
+        }
+    }
+
+    struct sampling_config *base_config = bpf_map_lookup_elem(&samplers_config_map, &sampler_id);
+    if (base_config == NULL) {
+        bpf_printk("No sampler config found for parent based sampler\n");
+        return false;
+    }
+
+    if (base_config->type == PARENT_BASED) {
+        bpf_printk("Parent based sampler can't have parent based sampler as base\n");
+        return false;
+    }
+
+    switch (base_config->type) {
+        case ALWAYS_ON:
+            return alwaysOnSampler_should_sample(base_config, params);
+        case ALWAYS_OFF:
+            return alwaysOffSampler_should_sample(base_config, params);
+        case TRACE_ID_RATIO:
+            return traceIDRatioSampler_should_sample(base_config, params);
+        default:
+            return false;
+    }
+}
+
+static __always_inline bool should_sample(sampling_parameters_t *params) {
+    u32 active_sampler_map_key = 0;
+    sampler_id_t *active_sampler_id = bpf_map_lookup_elem(&probe_active_sampler_map, &active_sampler_map_key);
+    if (active_sampler_id == NULL) {
+        bpf_printk("No active sampler found\n");
+        return false;
+    }
+
+    struct sampling_config *config = bpf_map_lookup_elem(&samplers_config_map, active_sampler_id);
+    if (config == NULL) {
+        bpf_printk("No sampler config found\n");
+        return false;
+    }
+
+    switch (config->type) {
+        case ALWAYS_ON:
+            return alwaysOnSampler_should_sample(config, params);
+        case ALWAYS_OFF:
+            return alwaysOffSampler_should_sample(config, params);
+        case TRACE_ID_RATIO:
+            return traceIDRatioSampler_should_sample(config, params);
+        case PARENT_BASED:
+            return parentBasedSampler_should_sample(config, params);
+        default:
+            return false;
+    }
+}
+
 #endif
diff --git a/internal/include/trace/span_context.h b/internal/include/trace/span_context.h
@@ -16,7 +16,6 @@
 #define TRACE_FLAGS_SIZE 1
 #define TRACE_FLAGS_STRING_SIZE 2
 
-static const u8 FLAG_SAMPLED = 1;
 
 struct span_context
 {
@@ -73,14 +72,4 @@ static __always_inline void w3c_string_to_span_context(char *str, struct span_co
     hex_string_to_bytes(str + trace_flags_start_pos, TRACE_FLAGS_STRING_SIZE, &ctx->TraceFlags);
 }
 
-static __always_inline bool trace_flags_is_sampled(u8 flags)
-{
-    return ((flags & FLAG_SAMPLED) == FLAG_SAMPLED);
-}
-
-static __always_inline bool is_sampled(struct span_context *ctx)
-{
-    return trace_flags_is_sampled(ctx->TraceFlags);
-}
-
 #endif
diff --git a/internal/include/trace/span_output.h b/internal/include/trace/span_output.h
@@ -1,9 +1,8 @@
 // Copyright The OpenTelemetry Authors
 // SPDX-License-Identifier: Apache-2.0
 
-#include "trace/span_context.h"
 #include "common.h"
-#include "trace/span_context.h"
+#include "trace/sampling.h"
 
 #ifndef _SPAN_OUTPUT_H_
 #define _SPAN_OUTPUT_H_

diff --git a/internal/pkg/instrumentation/bpf/database/sql/bpf_arm64_bpfel.go b/internal/pkg/instrumentation/bpf/database/sql/bpf_arm64_bpfel.go
diff --git a/internal/pkg/instrumentation/bpf/database/sql/bpf_x86_bpfel.go b/internal/pkg/instrumentation/bpf/database/sql/bpf_x86_bpfel.go
diff --git a/internal/pkg/instrumentation/bpf/github.com/segmentio/kafka-go/consumer/bpf_arm64_bpfel.go b/internal/pkg/instrumentation/bpf/github.com/segmentio/kafka-go/consumer/bpf_arm64_bpfel.go
diff --git a/internal/pkg/instrumentation/bpf/github.com/segmentio/kafka-go/consumer/bpf_x86_bpfel.go b/internal/pkg/instrumentation/bpf/github.com/segmentio/kafka-go/consumer/bpf_x86_bpfel.go