Configure Jaeger distributed tracing on Kubernetes cluster with Helm charts and Elasticsearch backend

#!/usr/bin/env bash
set -euo pipefail

# Jaeger distributed tracing installation script for Kubernetes
# Supports Ubuntu, Debian, AlmaLinux, Rocky Linux, CentOS, RHEL, Fedora

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

# Default values
NAMESPACE="${1:-jaeger-system}"
ELASTICSEARCH_STORAGE="${2:-50Gi}"

# Usage message
usage() {
    echo "Usage: $0 [namespace] [elasticsearch_storage_size]"
    echo "  namespace: Kubernetes namespace (default: jaeger-system)"
    echo "  elasticsearch_storage_size: Storage size for Elasticsearch (default: 50Gi)"
    exit 1
}

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Cleanup function for rollback
cleanup() {
    log_error "Installation failed. Cleaning up..."
    helm uninstall jaeger -n "$NAMESPACE" 2>/dev/null || true
    helm uninstall elasticsearch -n "$NAMESPACE" 2>/dev/null || true
    kubectl delete namespace "$NAMESPACE" 2>/dev/null || true
    rm -f elasticsearch-values.yaml jaeger-values.yaml get_helm.sh
    exit 1
}

trap cleanup ERR

# Detect distribution
detect_distro() {
    if [ ! -f /etc/os-release ]; then
        log_error "Cannot detect Linux distribution"
        exit 1
    fi
    
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_INSTALL="apt-get install -y"
            PKG_UPDATE="apt-get 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"
            ;;
        *)
            log_error "Unsupported distribution: $ID"
            exit 1
            ;;
    esac
    
    log_info "Detected distribution: $PRETTY_NAME"
    export PKG_MGR PKG_INSTALL PKG_UPDATE
}

# Check prerequisites
check_prerequisites() {
    echo "[1/8] Checking prerequisites..."
    
    if [[ $EUID -eq 0 ]]; then
        log_error "This script should not be run as root for security reasons"
        exit 1
    fi
    
    if ! command -v sudo >/dev/null 2>&1; then
        log_error "sudo is required but not installed"
        exit 1
    fi
    
    if ! kubectl cluster-info >/dev/null 2>&1; then
        log_error "kubectl is not configured or Kubernetes cluster is not accessible"
        exit 1
    fi
    
    log_info "Prerequisites check passed"
}

# Install Helm
install_helm() {
    echo "[2/8] Installing Helm package manager..."
    
    if command -v helm >/dev/null 2>&1; then
        log_info "Helm is already installed"
        return 0
    fi
    
    case "$PKG_MGR" in
        apt)
            curl -fsSL https://baltocdn.com/helm/signing.asc | gpg --dearmor | sudo tee /usr/share/keyrings/helm.gpg > /dev/null
            sudo $PKG_INSTALL apt-transport-https
            echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/helm.gpg] https://baltocdn.com/helm/stable/debian/ all main" | sudo tee /etc/apt/sources.list.d/helm-stable-debian.list
            sudo $PKG_UPDATE
            sudo $PKG_INSTALL helm
            ;;
        dnf|yum)
            curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3
            chmod 755 get_helm.sh
            ./get_helm.sh
            sudo mv /usr/local/bin/helm /usr/bin/ 2>/dev/null || true
            ;;
    esac
    
    log_info "Helm installed successfully"
}

# Add Helm repositories
add_helm_repos() {
    echo "[3/8] Adding Helm repositories..."
    
    helm repo add jaegertracing https://jaegertracing.github.io/helm-charts
    helm repo add elastic https://helm.elastic.co
    helm repo update
    
    log_info "Helm repositories added and updated"
}

# Create namespace
create_namespace() {
    echo "[4/8] Creating Kubernetes namespace..."
    
    if kubectl get namespace "$NAMESPACE" >/dev/null 2>&1; then
        log_warn "Namespace $NAMESPACE already exists"
    else
        kubectl create namespace "$NAMESPACE"
        log_info "Namespace $NAMESPACE created"
    fi
}

# Create Elasticsearch configuration
create_elasticsearch_config() {
    echo "[5/8] Creating Elasticsearch configuration..."
    
    cat > elasticsearch-values.yaml << 'EOF'
replicas: 3
minimumMasterNodes: 2

resources:
  requests:
    cpu: 1000m
    memory: 2Gi
  limits:
    cpu: 2000m
    memory: 4Gi

volumeClaimTemplate:
  storageClassName: "standard"
  resources:
    requests:
      storage: STORAGE_SIZE

esConfig:
  elasticsearch.yml: |
    cluster.name: jaeger-elasticsearch
    network.host: 0.0.0.0
    bootstrap.memory_lock: false
    discovery.seed_hosts: "elasticsearch-master-headless"
    cluster.initial_master_nodes: "elasticsearch-master-0,elasticsearch-master-1,elasticsearch-master-2"
    xpack.security.enabled: false

persistence:
  enabled: true

antiAffinity: "soft"
EOF
    
    sed -i "s/STORAGE_SIZE/$ELASTICSEARCH_STORAGE/" elasticsearch-values.yaml
    chmod 644 elasticsearch-values.yaml
    log_info "Elasticsearch configuration created"
}

# Deploy Elasticsearch
deploy_elasticsearch() {
    echo "[6/8] Deploying Elasticsearch..."
    
    helm install elasticsearch elastic/elasticsearch \
        -n "$NAMESPACE" \
        -f elasticsearch-values.yaml \
        --wait --timeout=10m
    
    log_info "Elasticsearch deployed successfully"
}

# Create Jaeger configuration
create_jaeger_config() {
    echo "[7/8] Creating Jaeger configuration..."
    
    cat > jaeger-values.yaml << 'EOF'
provisionDataStore:
  cassandra: false
  elasticsearch: true

storage:
  type: elasticsearch
  elasticsearch:
    host: elasticsearch-master
    port: 9200
    scheme: http
    nodesWanOnly: false
    extraEnv:
      - name: ES_SERVER_URLS
        value: http://elasticsearch-master:9200
    cmdlineParams:
      es.num-shards: 5
      es.num-replicas: 1
      es.index-prefix: jaeger

query:
  enabled: true
  replicaCount: 2
  resources:
    limits:
      cpu: 500m
      memory: 512Mi
    requests:
      cpu: 256m
      memory: 256Mi
  service:
    type: ClusterIP
    port: 16686

collector:
  enabled: true
  replicaCount: 2
  resources:
    limits:
      cpu: 1
      memory: 1Gi
    requests:
      cpu: 500m
      memory: 512Mi
  service:
    type: ClusterIP
    grpc:
      port: 14250
    http:
      port: 14268
    zipkin:
      port: 9411

agent:
  enabled: true
  daemonset:
    useHostPort: true
  resources:
    limits:
      cpu: 500m
      memory: 512Mi
    requests:
      cpu: 128m
      memory: 128Mi

esIndexCleaner:
  enabled: true
  numberOfDays: 7
  resources:
    limits:
      cpu: 500m
      memory: 512Mi
    requests:
      cpu: 256m
      memory: 256Mi
EOF
    
    chmod 644 jaeger-values.yaml
    log_info "Jaeger configuration created"
}

# Deploy Jaeger
deploy_jaeger() {
    echo "[8/8] Deploying Jaeger..."
    
    helm install jaeger jaegertracing/jaeger \
        -n "$NAMESPACE" \
        -f jaeger-values.yaml \
        --wait --timeout=10m
    
    log_info "Jaeger deployed successfully"
}

# Verify installation
verify_installation() {
    echo "Verifying installation..."
    
    log_info "Waiting for pods to be ready..."
    kubectl wait --for=condition=ready pod -l app=elasticsearch-master -n "$NAMESPACE" --timeout=300s
    kubectl wait --for=condition=ready pod -l app.kubernetes.io/name=jaeger -n "$NAMESPACE" --timeout=300s
    
    echo ""
    log_info "Installation completed successfully!"
    echo ""
    echo "Jaeger components:"
    kubectl get pods -n "$NAMESPACE" -l app.kubernetes.io/instance=jaeger
    echo ""
    echo "To access Jaeger UI, run:"
    echo "kubectl port-forward -n $NAMESPACE svc/jaeger-query 16686:16686"
    echo "Then visit: http://localhost:16686"
    echo ""
    echo "Jaeger collector endpoints:"
    echo "  gRPC: jaeger-collector.$NAMESPACE.svc.cluster.local:14250"
    echo "  HTTP: jaeger-collector.$NAMESPACE.svc.cluster.local:14268"
    echo "  Zipkin: jaeger-collector.$NAMESPACE.svc.cluster.local:9411"
}

# Main execution
main() {
    if [[ "${1:-}" == "-h" ]] || [[ "${1:-}" == "--help" ]]; then
        usage
    fi
    
    detect_distro
    check_prerequisites
    install_helm
    add_helm_repos
    create_namespace
    create_elasticsearch_config
    deploy_elasticsearch
    create_jaeger_config
    deploy_jaeger
    verify_installation
    
    # Cleanup temporary files
    rm -f get_helm.sh
}

main "$@"