Golang heap search tool

search-heap

#!/usr/bin/env bash

# search-heap - Search for functions in Go heap profile files
# Install: chmod +x search-heap && mv search-heap /usr/local/bin/

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

usage() {
    cat << EOF
Usage: search-heap [OPTIONS] <directory> <function_name>

Search for a specific function in Go heap profile files.

Arguments:
    directory       Directory containing heap profile files
    function_name   Function name to search for (can be partial)
                    Examples: "Writer" "Write" "(*Writer).Write"

Options:
    -h, --help      Show this help message
    -v, --verbose   Show verbose output including file sizes
    --list          Just list which files contain the pattern

Examples:
    search-heap ./heaps Write
    search-heap ./heaps "(*Writer).Write"
    search-heap -v /tmp/profiles Writer
    search-heap --list . Function

Notes:
    - Searches for literal string matches (special chars like () * . work as-is)
    - Use quotes for patterns with special characters or spaces
    - Searches both the raw profile data and go tool pprof output
EOF
    exit 0
}

# Default options
VERBOSE=0
LIST_ONLY=0

# Parse options
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            usage
            ;;
        -v|--verbose)
            VERBOSE=1
            shift
            ;;
        --list)
            LIST_ONLY=1
            shift
            ;;
        -*)
            echo -e "${RED}Error: Unknown option $1${NC}" >&2
            echo "Use --help for usage information" >&2
            exit 1
            ;;
        *)
            break
            ;;
    esac
done

# Check arguments
if [ $# -ne 2 ]; then
    echo -e "${RED}Error: Missing required arguments${NC}" >&2
    echo "Usage: search-heap <directory> <function_name>" >&2
    echo "Use --help for more information" >&2
    exit 1
fi

HEAP_DIR="$1"
SEARCH_PATTERN="$2"

# Validate directory
if [ ! -d "$HEAP_DIR" ]; then
    echo -e "${RED}Error: Directory '$HEAP_DIR' does not exist${NC}" >&2
    exit 1
fi

echo -e "${BLUE}Searching for pattern: ${NC}${SEARCH_PATTERN}"
echo -e "${BLUE}In directory: ${NC}${HEAP_DIR}"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""

# Find all heap/profile files
mapfile -t FILES < <(find "$HEAP_DIR" -type f \( -name "*.heap" -o -name "*.prof" -o -name "*heap*.pb.gz" -o -name "*profile*" -o -name "*.pprof" \) 2>/dev/null)

if [ ${#FILES[@]} -eq 0 ]; then
    echo -e "${YELLOW}No heap profile files found in $HEAP_DIR${NC}"
    echo ""
    echo "Looking for files with extensions: .heap, .prof, .pprof, or containing 'heap' or 'profile'"
    exit 1
fi

FOUND_COUNT=0
TOTAL_COUNT=${#FILES[@]}
current_file=0

search_file() {
    local file="$1"
    local pattern="$2"
    local found=0
    local all_matches=""

    # Strategy 1: Try go tool pprof -text first (most reliable for go profiles)
    if command -v go &> /dev/null; then
        local pprof_output
        # Try without timeout first, but with stderr redirected
        pprof_output=$(go tool pprof -text "$file" 2>&1) || true
        # Check if output looks valid (not an error message)
        if echo "$pprof_output" | head -n1 | grep -q -E '(flat|cum|File:)' 2>/dev/null; then
            # Use grep -F for fixed string matching (no regex interpretation)
            local pprof_matches
            pprof_matches=$(echo "$pprof_output" | grep -F "$pattern") || true
            if [ -n "$pprof_matches" ]; then
                all_matches="$pprof_matches"
                found=1
            fi
        fi
    fi

    # Strategy 2: Use strings on the raw file
    if [ $found -eq 0 ] && command -v strings &> /dev/null; then
        local string_matches
        string_matches=$(strings "$file" 2>/dev/null | grep -F "$pattern") || true
        if [ -n "$string_matches" ]; then
            all_matches="$string_matches"
            found=1
        fi
    fi

    # Strategy 3: Direct binary grep as last resort
    if [ $found -eq 0 ]; then
        local grep_matches
        grep_matches=$(grep -a -F "$pattern" "$file" 2>/dev/null | head -20) || true
        if [ -n "$grep_matches" ]; then
            all_matches="$grep_matches"
            found=1
        fi
    fi

    if [ $found -eq 1 ]; then
        echo "$all_matches"
        return 0
    fi

    return 1
}

for file in "${FILES[@]}"; do
    current_file=$((current_file + 1))

    if [ $VERBOSE -eq 1 ]; then
        size=$(du -h "$file" | cut -f1)
        echo -e "${BLUE}[${current_file}/${TOTAL_COUNT}]${NC} $(basename "$file") (${size})"
    else
        # Show progress in non-verbose mode
        printf "\r${BLUE}Checking:${NC} %d/%d files" "$current_file" "$TOTAL_COUNT"
    fi

    if MATCHES=$(search_file "$file" "$SEARCH_PATTERN"); then
        FOUND_COUNT=$((FOUND_COUNT + 1))

        if [ $LIST_ONLY -eq 1 ]; then
            echo "$file"
        else
            echo -e "${GREEN}✓ FOUND${NC} in: $(basename "$file")"

            if [ $VERBOSE -eq 1 ]; then
                echo -e "  ${BLUE}Full path:${NC} $file"
            fi

            echo -e "  ${BLUE}Matching lines:${NC}"

            # Show matches with limited output
            line_count=0
            while IFS= read -r line; do
                if [ -n "$line" ] && [ $line_count -lt 15 ]; then
                    echo "    $line"
                    line_count=$((line_count + 1))
                fi
            done <<< "$MATCHES"

            total_matches=$(echo "$MATCHES" | grep -c '^' || echo 0)
            if [ $total_matches -gt 15 ]; then
                remaining=$((total_matches - 15))
                echo -e "    ${YELLOW}... and $remaining more matches${NC}"
            fi

            echo ""
        fi
    fi
done

# Clear progress line in non-verbose mode
if [ $VERBOSE -eq 0 ]; then
    printf "\r%80s\r" ""  # Clear the progress line
fi

if [ $LIST_ONLY -eq 0 ]; then
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo -e "${BLUE}Summary:${NC}"
    echo "  Total files searched: $TOTAL_COUNT"
    echo "  Files containing '$SEARCH_PATTERN': $FOUND_COUNT"
fi

if [ $FOUND_COUNT -eq 0 ]; then
    if [ $LIST_ONLY -eq 0 ]; then
        echo ""
        echo -e "${YELLOW}❌ Pattern not found in any heap files.${NC}"
        echo ""
        echo "Troubleshooting:"
        echo "  • Try a simpler pattern (e.g., just 'Write' instead of '(*Writer).Write')"
        echo "  • Manually check with: go tool pprof -text <file> | grep -F 'pattern'"
        echo "  • List all functions with: go tool pprof -text <file> | less"
    fi
    exit 1
else
    if [ $LIST_ONLY -eq 0 ]; then
        echo ""
        echo -e "${GREEN}✓ Pattern found in $FOUND_COUNT file(s)${NC}"
    fi
    exit 0
fi