bench_allocators.sh

#!/bin/bash
# Benchmark different memory allocators for the EVTX parser
# Compares: jemalloc, mimalloc, mimalloc-secure, and system allocator
# Tests both single-threaded and multi-threaded performance

# Exit on error, undefined variables, and propagate errors in pipelines
set -euo pipefail

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

# Define benchmark parameters
WARMUP_RUNS=3
BENCHMARK_RUNS=10
OUTPUT_DIR="benchmark_results"

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Check for hyperfine
echo -e "${BLUE}Checking dependencies...${NC}"
if ! command -v hyperfine &>/dev/null; then
    echo -e "${BLUE}Installing hyperfine...${NC}"
    brew install hyperfine
else
    echo -e "${GREEN}Hyperfine already installed.${NC}"
fi

# Function to determine the default target
get_default_target() {
    case "$(uname -sm)" in
    "Darwin x86_64") echo "x86_64-apple-darwin" ;;
    "Darwin arm64") echo "aarch64-apple-darwin" ;;
    "Linux x86_64") echo "x86_64-unknown-linux-gnu" ;;
    "Linux aarch64") echo "aarch64-unknown-linux-gnu" ;;
    *) echo "unknown" ;;
    esac
}

# Use the provided TARGET or default to the current machine's target
TARGET=${TARGET:-$(get_default_target)}

if [ "$TARGET" = "unknown" ]; then
    echo -e "${RED}Error: Unable to determine default target. Please specify TARGET explicitly.${NC}"
    exit 1
fi

echo -e "${BLUE}Using target: $TARGET${NC}"

# Find the largest sample file to benchmark with
echo -e "${BLUE}Finding largest EVTX sample file...${NC}"
SAMPLE_FILE=$(find samples -name "*.evtx" -type f -exec du -k {} \; | sort -nr | head -1 | cut -f2)

if [ -z "$SAMPLE_FILE" ]; then
    echo -e "${RED}No sample EVTX files found. Please place at least one .evtx file in the samples directory.${NC}"
    exit 1
fi

# Get file size in KB for display
FILE_SIZE_KB=$(du -k "$SAMPLE_FILE" | cut -f1)
echo -e "${GREEN}Using largest sample file: $SAMPLE_FILE (${FILE_SIZE_KB} KB)${NC}"

# Build all variants of evtx_dump with different allocators
echo -e "${BLUE}Building evtx_dump with different allocators...${NC}"

# System allocator (no special feature)
echo -e "${BLUE}Building with system allocator...${NC}"
cargo build --release --bin evtx_dump --target $TARGET || {
    echo -e "${RED}Failed to build with system allocator${NC}"
    exit 1
}
cp "target/$TARGET/release/evtx_dump" "target/$TARGET/release/evtx_dump_system"

# Jemalloc
echo -e "${BLUE}Building with jemalloc...${NC}"
cargo build --release --bin evtx_dump --target $TARGET --features fast-alloc-jemalloc || {
    echo -e "${RED}Failed to build with jemalloc${NC}"
    exit 1
}
cp "target/$TARGET/release/evtx_dump" "target/$TARGET/release/evtx_dump_jemalloc"

# Mimalloc
echo -e "${BLUE}Building with mimalloc...${NC}"
cargo build --release --bin evtx_dump --target $TARGET --features fast-alloc-mimalloc || {
    echo -e "${RED}Failed to build with mimalloc${NC}"
    exit 1
}
cp "target/$TARGET/release/evtx_dump" "target/$TARGET/release/evtx_dump_mimalloc"

# Mimalloc secure
echo -e "${BLUE}Building with mimalloc secure...${NC}"
cargo build --release --bin evtx_dump --target $TARGET --features fast-alloc-mimalloc-secure || {
    echo -e "${RED}Failed to build with mimalloc secure${NC}"
    exit 1
}
cp "target/$TARGET/release/evtx_dump" "target/$TARGET/release/evtx_dump_mimalloc_secure"

# Create result directories by test type
mkdir -p "$OUTPUT_DIR/single_thread"
mkdir -p "$OUTPUT_DIR/multi_thread"

# Run the single-threaded benchmarks with hyperfine
echo -e "\n${YELLOW}Running single-threaded benchmarks (-t 1)...${NC}"
hyperfine --warmup $WARMUP_RUNS \
    --export-markdown "$OUTPUT_DIR/single_thread/benchmarks.md" \
    --export-json "$OUTPUT_DIR/single_thread/benchmarks.json" \
    --export-csv "$OUTPUT_DIR/single_thread/benchmarks.csv" \
    --setup "sync && sudo purge" \
    --prepare "sleep 1" \
    --runs $BENCHMARK_RUNS \
    --command-name "jemalloc" "target/$TARGET/release/evtx_dump_jemalloc -t 1 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "mimalloc" "target/$TARGET/release/evtx_dump_mimalloc -t 1 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "mimalloc-secure" "target/$TARGET/release/evtx_dump_mimalloc_secure -t 1 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "system" "target/$TARGET/release/evtx_dump_system -t 1 -o json $SAMPLE_FILE > /dev/null"

# Run the multi-threaded benchmarks with hyperfine
echo -e "\n${YELLOW}Running multi-threaded benchmarks (-t 8)...${NC}"
hyperfine --warmup $WARMUP_RUNS \
    --export-markdown "$OUTPUT_DIR/multi_thread/benchmarks.md" \
    --export-json "$OUTPUT_DIR/multi_thread/benchmarks.json" \
    --export-csv "$OUTPUT_DIR/multi_thread/benchmarks.csv" \
    --setup "sync && sudo purge" \
    --prepare "sleep 1" \
    --runs $BENCHMARK_RUNS \
    --command-name "jemalloc" "target/$TARGET/release/evtx_dump_jemalloc -t 8 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "mimalloc" "target/$TARGET/release/evtx_dump_mimalloc -t 8 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "mimalloc-secure" "target/$TARGET/release/evtx_dump_mimalloc_secure -t 8 -o json $SAMPLE_FILE > /dev/null" \
    --command-name "system" "target/$TARGET/release/evtx_dump_system -t 8 -o json $SAMPLE_FILE > /dev/null"

# If hyperfine succeeded, display the results location
if [ $? -eq 0 ]; then
    echo -e "\n${GREEN}Benchmarks complete!${NC}"

    # Single-threaded results
    echo -e "\n${YELLOW}Single-threaded (-t 1) results:${NC}"
    cat "$OUTPUT_DIR/single_thread/benchmarks.md"

    # Multi-threaded results
    echo -e "\n${YELLOW}Multi-threaded (-t 8) results:${NC}"
    cat "$OUTPUT_DIR/multi_thread/benchmarks.md"

    echo -e "\n${GREEN}Detailed results saved to:${NC}"
    echo -e "  - ${BLUE}$OUTPUT_DIR/single_thread/${NC} (Single-threaded)"
    echo -e "  - ${BLUE}$OUTPUT_DIR/multi_thread/${NC} (Multi-threaded)"
else
    echo -e "${RED}Benchmark failed.${NC}"
    exit 1
fi