flowchart TD
    subgraph K[KERNEL]
        direction TB
        k[kprobes / tracepoints]
        p[PER-CPU Ring Buff]
    end

    subgraph U[User Space (Rust)]
        direction LR
        L[eBPF Loader libbpf-rs] --> S[Span Correlation]
        S --> E[OTLP Exporter]
    end

    subgraph O[OTEL Collector]
        T[Jaeger / Tempo]
    end

    k --> p
    p -->|streams events| L
    E -->|OTLP/gRPC| T

// tracesend.bpf.c
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

struct span_event {
    __u64 timestamp_ns;   // boot‑relative ns, converted to epoch in userspace
    __u32 pid;
    __u32 syscall_nr;
    char comm[TASK_COMM_LEN];
    __u64 sock_addr;      // packed IPv4 address + port
    __s32 ret;            // return value of syscall
    __u32 bytes;          // bytes transferred
    __u32 direction;      // 0 = send, 1 = recv
};

struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024 * 1024); // 256 MiB ring buffer
} events SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 65536);
    __type(key, __u32);   // pid
    __type(value, __u32); // parent pid
} pid_tree SEC(".maps");

SEC("kprobe/tcp_sendmsg")
int trace_tcp_sendmsg(struct sock *sk, struct msghdr *msg)
{
    struct span_event ev = {};
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 pid = pid_tgid >> 32;

    ev.timestamp_ns = bpf_ktime_get_ns();
    ev.pid = pid;
    ev.syscall_nr = 46;                 // __NR_sendmsg on x86_64
    bpf_get_current_comm(&ev.comm, sizeof(ev.comm));
    ev.direction = 0;
    ev.ret = 0;

    // Pack IPv4 address + port (simplified for illustration)
    if (sk) {
        __u32 ip = 0, port = 0;
        // Use BTF‑aware macro for portability
        ip = BPF_CORE_READ(sk, __sk_common.sk_rcv_saddr);
        port = BPF_CORE_READ(sk, __sk_common.sk_dport);
        ev.sock_addr = ((__u64)port << 32) | ip;
    }

    bpf_ringbuf_output(&events, &ev, sizeof(ev), 0);
    return 0;
}

SEC("kretprobe/tcp_sendmsg")
int trace_tcp_sendmsg_ret(struct pt_regs *ctx)
{
    struct span_event ev = {};
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 pid = pid_tgid >> 32;
    __s32 ret = PT_REGS_RC(ctx);

    ev.timestamp_ns = bpf_ktime_get_ns();
    ev.pid = pid;
    ev.syscall_nr = 46;
    ev.ret = ret;
    ev.direction = 0;
    bpf_get_current_comm(&ev.comm, sizeof(ev.comm));

    // Extract bytes from return value (only positive returns)
    ev.bytes = (ret > 0) ? ( __u32)ret : 0;
    bpf_ringbuf_output(&events, &ev, sizeof(ev), 0);
    return 0;
}

SEC("tracepoint/sched/sched_process_fork")
int trace_fork(struct trace_event_raw_sched_process_template *ctx)
{
    __u32 child = ctx->pid;
    __u32 parent = ctx->parent_pid;
    bpf_map_update_elem(&pid_tree, &child, &parent, BPF_ANY);
    return 0;
}

char LICENSE[] SEC("license") = "GPL";

// src/main.rs
use libbpf_rs::{Map, program::Kprobe};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::runtime::Builder;

mod span_reconstruction;
mod otlp_exporter;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Load the compiled eBPF object (generated by clang/llvm with BTF)
    let mut prog = bpfs::load()?;
    prog.set_auto_raise_rlimit(true);

    // Attach kprobes
    let send_probe = prog.attach_kprobe("tcp_sendmsg")?;
    let ret_probe  = prog.attach_kretprobe("tcp_sendmsg")?;

    // Open the ring‑buffer map
    let map: &Map = prog.map("events").unwrap();
    let mut reader = map.open_ring_buffer()?;

    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();

    // Background task that drives the async OTLP exporter
    let otlp = otlp_exporter::OtlpBatchExporter::new("http://localhost:4317").await?;
    let mut span_recon = span_reconstruction::SpanReconstructor::new();

    // Spawn a Tokio runtime for the reconstruction + export loop
    let rt = Builder::new_multi_thread()
        .enable_all()
        .build()?;

    rt.spawn(async move {
        span_reconstruction_loop(reader, &mut span_recon, otlp, r).await;
    });

    // Ctrl‑C handling
    tokio::signal::ctrl_c().await?;
    running.store(false, Ordering::SeqCst);
    Ok(())
}

// src/span_reconstruction.rs
use std::collections::{HashMap, VecDeque};
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Clone, Debug)]
pub struct RawEvent {
    pub timestamp_ns: u64,   // boot‑relative; converted to epoch in userspace
    pub pid: u32,
    pub comm: String,
    pub direction: Direction,
    pub bytes: u32,
    pub ret: i32,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Direction {
    Send,
    Receive,
}

struct ActiveSpan {
    start_ns: u64,
    operation: String,
    span_id: [u8; 8],
    bytes_accumulated: u64,
}

pub struct SpanReconstructor {
    pending: HashMap<u32, VecDeque<ActiveSpan>>,
    pid_tree: HashMap<u32, u32>, // child → parent
    completed: Vec<Span>,
}

impl SpanReconstructor {
    pub fn new() -> Self {
        Self {
            pending: HashMap::new(),
            pid_tree: HashMap::new(),
            completed: Vec::new(),
        }
    }

    pub fn process_event(&mut self, ev: RawEvent) {
        match ev.direction {
            Direction::Send if ev.ret >= 0 => {
                // Entry: push a pending span
                let span_id = self.generate_span_id(ev.pid);
                let op_name = format!("{} (bytes={})", ev.comm, ev.bytes);
                self.pending
                    .entry(ev.pid)
                    .or_default()
                    .push_back(ActiveSpan {
                        start_ns: ev.timestamp_ns,
                        operation: op_name,
                        span_id,
                        bytes_accumulated: ev.bytes as u64,
                    });
            }
            Direction::Send => {
                // Return: pop and create a finished span
                if let Some(mut stack) = self.pending.remove(&ev.pid) {
                    if let Some(active) = stack.pop_back() {
                        let duration = ev.timestamp_ns.saturating_sub(active.start_ns);
                        let peer_port = self.extract_peer_port(&ev.pid);
                        let span = Span {
                            trace_id: self.trace_id(&ev.pid, self.parent_of(ev.pid)),
                            span_id,
                            parent_span_id: self.parent_span_id(&ev.pid),
                            operation_name: active.operation,
                            start_time: active.start_ns,
                            duration_ns: duration.max(1),
                            tags: self.build_tags(&ev, active, peer_port),
                        };
                        self.completed.push(span);
                    }
                }
            }
            Direction::Receive => {
                // Similar logic for receive syscalls (omitted for brevity)
                self.process_receive(ev);
            }
        }
    }

    fn generate_span_id(&self, pid: u32) -> [u8; 8] {
        // Simple deterministic ID: timestamp low 8 bytes + pid high 4 bytes
        let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
        let mut id = [0u8; 8];
        id[0..4].copy_from_slice(&(ts & 0xFFFFFFFF) as u64);
        id[4..8].copy_from_slice(&(pid as u32).to_be_bytes());
        id
    }

    fn parent_of(&self, pid: u32) -> u32 {
        self.pid_tree.get(&pid).copied().unwrap_or(0)
    }

    fn trace_id(&self, pid: u32, parent: u32) -> [u8; 16] {
        // Deterministic 16‑byte trace ID from PID + parent PID
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(pid.to_be_bytes());
        hasher.update(parent.to_be_bytes());
        let hash = hasher.finalize();
        let mut out = [0u8; 16];
        out[..16].copy_from_slice(&hash[..16]);
        out
    }

    fn build_tags(&self, ev: &RawEvent, active: &ActiveSpan, peer_port: u16) -> HashMap<String, String> {
        let mut tags = HashMap::new();
        tags.insert("net.peer.port".to_string(), peer_port.to_string());
        tags.insert("net.bytes".to_string(), active.bytes_accumulated.to_string());
        tags.insert("proc.name".to_string(), ev.comm.clone());
        tags.insert("pid".to_string(), ev.pid.to_string());
        tags.insert("parent.pid".to_string(), self.parent_of(ev.pid).to_string());
        tags
    }

    fn extract_peer_port(&self, pid: u32) -> u16 {
        // In a real implementation this would read the socket from userspace
        // or use a helper eBPF map; here we mock a constant.
        8080
    }

    fn process_receive(&mut self, ev: RawEvent) {
        // For brevity, receive handling mirrors send logic.
        // It creates a receive span and pairs entry/exit events similarly.
    }
}

// src/otlp_exporter.rs
use opentelemetry::sdk::trace::{Span, SpanData, SpanExporter};
use opentelemetry::trace::SpanKind;
use opentelemetry::sdk::export::traces::ExportResult;
use opentelemetry_proto::tonic::collector::trace::v1::{
    ExportTraceServiceRequest, ResourceSpans, ScopeSpans, Span as ProtoSpan,
};
use tonic::transport::Channel;
use std::time::Duration;

pub struct OtlpBatchExporter {
    client: opentelemetry::sdk::trace::export::ExportSpanService,
    buffer: Vec<ProtoSpan>,
    max_batch: usize,
}

impl OtlpBatchExporter {
    pub async fn new(endpoint: &str) -> anyhow::Result<Self> {
        let channel = Channel::from_static(endpoint)
            .connect()
            .await?;
        let client = opentelemetry::sdk::trace::export::SpanExporter::new(channel);
        Ok(Self {
            client,
            buffer: Vec::with_capacity(1024),
            max_batch: 1024,
        })
    }

    pub fn buffer_span(&mut self, span: Span) {
        let proto = ProtoSpan {
            trace_id: span.trace_id().to_vec(),
            span_id: span.span_id().to_vec(),
            parent_span_id: span.parent_span_id().map(|id| id.to_vec()),
            name: span.name().to_string(),
            kind: SpanKind::Client as i32, // or Server depending on direction
            start_time_unix_nano: span.start_time_unix_nano(),
            end_time_unix_nano: span.end_time_unix_nano(),
            attributes: span
                .attributes()
                .into_iter()
                .map(|(k, v)| {
                    opentelemetry_proto::tonic::common::v1::KeyValue {
                        key: k,
                        value: Some(opentelemetry_proto::tonic::common::v1::key_value::Value::StringValue(v)),
                    }
                })
                .collect(),
            ..Default::default()
        };
        self.buffer.push(proto);
        if self.buffer.len() >= self.max_batch {
            self.flush().await;
        }
    }

    pub async fn flush(&mut self) {
        if self.buffer.is_empty() {
            return;
        }
        let request = ExportTraceServiceRequest {
            resource_spans: vec![ResourceSpans {
                resource: opentelemetry::sdk::Resource::new(),
                scope_spans: vec![ScopeSpans {
                    scope: opentelemetry::sdk::Resource::new(),
                    spans: std::mem::take(&mut self.buffer),
                }],
            }],
        };
        match self.client.export(request).await {
            Ok(_) => tracing::info!("Exported {} spans", self.buffer.len()),
            Err(e) => tracing::error!("OTLP export failed: {}", e),
        }
    }
}

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: ebpf-tracer
  namespace: observability
spec:
  selector:
    matchLabels:
      app: ebpf-tracer
  template:
    metadata:
      labels:
        app: ebpf-tracer
    spec:
      hostNetwork: true
      hostPID: true
      containers:
        - name: tracer
          image: registry.internal/ebpf-tracer:sha-2025-03-12
          securityContext:
            privileged: false
            runAsNonRoot: true
            capabilities:
              add:
                - SYS_ADMIN
                - BPF
                - PERFMON
          volumeMounts:
            - name: bpffs
              mountPath: /sys/fs/bpf
            - name: debugfs
              mountPath: /sys/kernel/debug
            - name: lib-modules
              mountPath: /lib/modules
              readOnly: true
            - name: usr-src
              mountPath: /usr/src
              readOnly: true
          resources:
            requests:
              cpu: 100m
              memory: 128Mi
            limits:
              cpu: 2000m
              memory: 512Mi
      volumes:
        - name: bpffs
          hostPath:
            path: /sys/fs/bpf
        - name: debugfs
          hostPath:
            path: /sys/kernel/debug
        - name: lib-modules
          hostPath:
            path: /lib/modules
        - name: usr-src
          hostPath:
            path: /usr/src