Configure OpenTelemetry sampling strategies for high-traffic applications

#!/usr/bin/env bash

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Default configuration
OTEL_VERSION="0.91.0"
SAMPLING_TYPE="probabilistic"
SAMPLING_PERCENTAGE="10.0"

# Usage function
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Options:"
    echo "  -t, --type        Sampling type: probabilistic or tail (default: probabilistic)"
    echo "  -p, --percentage  Sampling percentage for probabilistic sampling (default: 10.0)"
    echo "  -v, --version     OpenTelemetry Collector version (default: 0.91.0)"
    echo "  -h, --help        Show this help message"
    exit 1
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -t|--type)
            SAMPLING_TYPE="$2"
            shift 2
            ;;
        -p|--percentage)
            SAMPLING_PERCENTAGE="$2"
            shift 2
            ;;
        -v|--version)
            OTEL_VERSION="$2"
            shift 2
            ;;
        -h|--help)
            usage
            ;;
        *)
            echo -e "${RED}Unknown option: $1${NC}"
            usage
            ;;
    esac
done

# Validate sampling type
if [[ "$SAMPLING_TYPE" != "probabilistic" && "$SAMPLING_TYPE" != "tail" ]]; then
    echo -e "${RED}Error: Sampling type must be 'probabilistic' or 'tail'${NC}"
    exit 1
fi

# Error handling
cleanup() {
    echo -e "${RED}Installation failed. Cleaning up...${NC}"
    sudo rm -f /tmp/otelcol_${OTEL_VERSION}_linux_amd64.tar.gz
    sudo systemctl stop otelcol 2>/dev/null || true
    sudo systemctl disable otelcol 2>/dev/null || true
}

trap cleanup ERR

# Check if running as root or with sudo
if [[ $EUID -eq 0 ]]; then
    echo -e "${YELLOW}Warning: Running as root${NC}"
elif ! sudo -n true 2>/dev/null; then
    echo -e "${RED}Error: This script requires sudo privileges${NC}"
    exit 1
fi

# Detect distribution
echo -e "${BLUE}[1/8] Detecting distribution...${NC}"
if [ -f /etc/os-release ]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_INSTALL="apt install -y"
            PKG_UPDATE="apt update"
            ;;
        almalinux|rocky|centos|rhel|ol|fedora)
            PKG_MGR="dnf"
            PKG_INSTALL="dnf install -y"
            PKG_UPDATE="dnf check-update || true"
            ;;
        amzn)
            PKG_MGR="yum"
            PKG_INSTALL="yum install -y"
            PKG_UPDATE="yum check-update || true"
            ;;
        *)
            echo -e "${RED}Unsupported distribution: $ID${NC}"
            exit 1
            ;;
    esac
    echo -e "${GREEN}Detected: $PRETTY_NAME${NC}"
else
    echo -e "${RED}Error: Cannot detect distribution${NC}"
    exit 1
fi

# Update package manager
echo -e "${BLUE}[2/8] Updating package manager...${NC}"
sudo $PKG_UPDATE

# Install required packages
echo -e "${BLUE}[3/8] Installing required packages...${NC}"
sudo $PKG_INSTALL curl tar

# Download and install OpenTelemetry Collector
echo -e "${BLUE}[4/8] Downloading OpenTelemetry Collector v${OTEL_VERSION}...${NC}"
cd /tmp
curl -LO "https://github.com/open-telemetry/opentelemetry-collector-releases/releases/download/v${OTEL_VERSION}/otelcol_${OTEL_VERSION}_linux_amd64.tar.gz"

echo -e "${BLUE}[5/8] Installing OpenTelemetry Collector...${NC}"
tar -xzf "otelcol_${OTEL_VERSION}_linux_amd64.tar.gz"
sudo mv otelcol /usr/local/bin/
sudo chmod 755 /usr/local/bin/otelcol
sudo chown root:root /usr/local/bin/otelcol

# Create OpenTelemetry user and directories
echo -e "${BLUE}[6/8] Creating OpenTelemetry user and directories...${NC}"
sudo useradd --system --shell /bin/false --home-dir /var/lib/otelcol otel 2>/dev/null || echo "User otel already exists"
sudo mkdir -p /etc/otelcol /var/log/otelcol /var/lib/otelcol
sudo chown -R otel:otel /etc/otelcol /var/log/otelcol /var/lib/otelcol
sudo chmod 755 /etc/otelcol /var/log/otelcol /var/lib/otelcol

# Create configuration file based on sampling type
echo -e "${BLUE}[7/8] Creating OpenTelemetry configuration...${NC}"
if [[ "$SAMPLING_TYPE" == "probabilistic" ]]; then
    sudo tee /etc/otelcol/config.yaml > /dev/null <<EOF
receivers:
  otlp:
    protocols:
      grpc:
        endpoint: 0.0.0.0:4317
      http:
        endpoint: 0.0.0.0:4318
  jaeger:
    protocols:
      grpc:
        endpoint: 0.0.0.0:14250
      thrift_http:
        endpoint: 0.0.0.0:14268

processors:
  probabilistic_sampler:
    sampling_percentage: ${SAMPLING_PERCENTAGE}
  batch:
    timeout: 1s
    send_batch_size: 1024
    send_batch_max_size: 2048
  memory_limiter:
    limit_mib: 512
    spike_limit_mib: 128
    check_interval: 5s

exporters:
  logging:
    loglevel: info

service:
  pipelines:
    traces:
      receivers: [otlp, jaeger]
      processors: [memory_limiter, probabilistic_sampler, batch]
      exporters: [logging]
  telemetry:
    logs:
      level: "info"
      development: false
      output_paths:
        - "/var/log/otelcol/collector.log"
      error_output_paths:
        - "/var/log/otelcol/collector-error.log"
EOF
else
    sudo tee /etc/otelcol/config.yaml > /dev/null <<EOF
receivers:
  otlp:
    protocols:
      grpc:
        endpoint: 0.0.0.0:4317
      http:
        endpoint: 0.0.0.0:4318

processors:
  tail_sampling:
    decision_wait: 30s
    num_traces: 100000
    expected_new_traces_per_sec: 1000
    policies:
      - name: errors_policy
        type: status_code
        status_code:
          status_codes: [ERROR]
      - name: latency_policy
        type: latency
        latency:
          threshold_ms: 5000
      - name: probabilistic_policy
        type: probabilistic
        probabilistic:
          sampling_percentage: ${SAMPLING_PERCENTAGE}
      - name: rate_limiting_policy
        type: rate_limiting
        rate_limiting:
          spans_per_second: 100
  batch:
    timeout: 1s
    send_batch_size: 1024
    send_batch_max_size: 2048
  memory_limiter:
    limit_mib: 1024
    spike_limit_mib: 256
    check_interval: 5s

exporters:
  logging:
    loglevel: info

service:
  pipelines:
    traces:
      receivers: [otlp]
      processors: [memory_limiter, tail_sampling, batch]
      exporters: [logging]
  telemetry:
    logs:
      level: "info"
      development: false
      output_paths:
        - "/var/log/otelcol/collector.log"
      error_output_paths:
        - "/var/log/otelcol/collector-error.log"
EOF
fi

sudo chown otel:otel /etc/otelcol/config.yaml
sudo chmod 644 /etc/otelcol/config.yaml

# Create systemd service
sudo tee /etc/systemd/system/otelcol.service > /dev/null <<EOF
[Unit]
Description=OpenTelemetry Collector
After=network.target

[Service]
Type=simple
User=otel
Group=otel
ExecStart=/usr/local/bin/otelcol --config=/etc/otelcol/config.yaml
Restart=on-failure
RestartSec=5
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF

# Configure firewall based on distribution
if command -v ufw >/dev/null 2>&1; then
    sudo ufw allow 4317/tcp comment "OpenTelemetry GRPC"
    sudo ufw allow 4318/tcp comment "OpenTelemetry HTTP"
elif command -v firewall-cmd >/dev/null 2>&1; then
    sudo firewall-cmd --permanent --add-port=4317/tcp
    sudo firewall-cmd --permanent --add-port=4318/tcp
    sudo firewall-cmd --reload
fi

# Start and enable the service
echo -e "${BLUE}[8/8] Starting OpenTelemetry Collector service...${NC}"
sudo systemctl daemon-reload
sudo systemctl enable otelcol
sudo systemctl start otelcol

# Cleanup
sudo rm -f /tmp/otelcol_${OTEL_VERSION}_linux_amd64.tar.gz

# Verification
echo -e "${BLUE}Verifying installation...${NC}"
sleep 3

if sudo systemctl is-active --quiet otelcol; then
    echo -e "${GREEN}✓ OpenTelemetry Collector is running${NC}"
else
    echo -e "${RED}✗ OpenTelemetry Collector failed to start${NC}"
    sudo systemctl status otelcol
    exit 1
fi

if netstat -tuln 2>/dev/null | grep -q ":4317\|:4318" || ss -tuln 2>/dev/null | grep -q ":4317\|:4318"; then
    echo -e "${GREEN}✓ OpenTelemetry Collector is listening on configured ports${NC}"
else
    echo -e "${YELLOW}⚠ Warning: Cannot verify port status${NC}"
fi

echo -e "${GREEN}OpenTelemetry Collector installation completed successfully!${NC}"
echo -e "${BLUE}Configuration:${NC}"
echo "  - Type: $SAMPLING_TYPE sampling"
echo "  - Percentage: $SAMPLING_PERCENTAGE%"
echo "  - Config file: /etc/otelcol/config.yaml"
echo "  - Log files: /var/log/otelcol/"
echo "  - Service: systemctl status otelcol"