#!/bin/bash

# Copyright (c) 2023 DataDirect Networks, Inc.
# Authors: Patrick Farrell, Andreas Dilger
#
# This is a simple tool which can be run on any Linux
# system to estimate the space usage reduction from the
# Lustre Client Side Data Compression (CSDC) feature with
# particular compression settings (algorithm, chunk size,
# and compression level).
#
# When run in a directory, this tool will recursively
# examine files under that directory, sampling the data in
# those files to estimate how much the files will compress.
#
# This tool will sample all files up to a configured number
# (defaulting to 100 files) and after that, it samples a
# configurable percentage of remaining files.
#
# This tools samples throughout the file, so it should
# avoid problems with poor estimates for files with headers
# which differ from the bulk data in the file.
#
# This tool requires the lz4, lzop, zstd, and gzip utilities to
# be installed in order to test those compression types.
# (lzop is the command line utility for lzo compression)

# Default values
version="1.3"
chunk_size=65536
block_size=4096
sample_count=20
min_files=100
default_path="$(pwd)"
percentage=1
compression_type="gzip -"
compression_level=6
compression_margin=5
whole_file="false"
debug=0
quiet=0
start=$SECONDS

# Display description of script behavior
description()
{
	echo "Version $version. Recursively scan PATH "
	echo "sampling data from the first MIN_FILES "
	echo "then sampling data from PERCENTAGE% of remaining files "
	echo "to estimate the average compression ratio using "
	echo "COMPRESSION_TYPE level COMPRESSION_LEVEL"
	echo "and a chunk size of CHUNK_SIZE bytes."
	echo ""
	echo "This tool takes up to SAMPLE_COUNT samples of CHUNK_SIZE bytes"
	echo "from each sampled file and compresses that data with the selected"
	echo "parameters to generate an estimate of the compression ratio for"
	echo "the full dataset."
	echo ""
	echo "You can trade-off estimation accuracy and scan speed by adjusting"
	echo "the per file sample count and percentage of files to sample."
	echo ""
}

runtime_description()
{
	echo "Version $version. Recursively scan '${path[@]}',"
	if (( percentage < 100 )); then
		(( min_files > 1 )) && echo "sampling data from the first $min_files files"
		echo "then sampling data from $percentage% of remaining files"
	fi
	echo "to estimate the average compression ratio using"
	echo "${compression_type/ -*/} level $compression_level"
	echo "and a chunk size of $chunk_size bytes."
	echo ""
	echo "Run with -h to see options for these parameters."
	echo ""
	echo "This tool takes up to $sample_count samples of $chunk_size bytes"
	echo "from each sampled file and compresses that data with the selected"
	echo "parameters to generate an estimate of the compression ratio for"
	echo "the full dataset."
	echo ""
	echo "You can trade-off estimation accuracy and scan speed by adjusting"
	echo "the per file sample count and percentage of files to sample."
	echo ""
	echo "This tool assumes a relatively uniform distribution of file"
	echo "sizes and contents across the directory tree, and is only"
	echo "intended to provide an approximate estimate of the compression"
	echo "potential of a specific dataset, and does not guarantee a"
	echo "particular compression level."
	echo ""
}

# Function to display script usage
usage() {
cat <<- USAGE
Usage: $(basename $0) [-n MIN_FILES] [-p PERCENTAGE] [-s SAMPLE_COUNT]
	[-c CHUNK_SIZE] [-z COMPRESSION_TYPE] [-l COMPRESSION_LEVEL]
	[-h][-w][-q]
	[PATH ...]

Description:
$(description | fmt)

Arguments:
    -n MIN_FILES: Minimum number of files to scan. Default: $min_files.
    -p PERCENTAGE: Fraction of scanned files to process. Default: ${percentage}%.
    -s SAMPLE_COUNT: Maximum number of chunks to sample per file. Default: $sample_count.
    -c CHUNK_SIZE: Size of data chunk in kibibytes (64-4096). Default: $((chunk_size / 1024))KiB.
    -z COMPRESSION_TYPE: One of gzip, lz4, lz4fast, lzo, zstd, zstdfast. Default: ${compression_type/ -*/}.
    -l COMPRESSION_LEVEL: Compression level to use (1-9). Default: $compression_level.
    -w Sample whole file (override -s). With '-p 100' for a full but slow estimate.
    -q Skip printing of usage header.  -qq to also skip runtime status update.
    -h Print this help message.
USAGE
}

# Parse command-line options
while getopts "c:ds:n:p:z:Z:l:m:wqh" opt; do
	case $opt in
	c)
		if (( OPTARG & (OPTARG - 1) )); then
			echo "Chunk size must be a power-of-two value" 1>&2
			exit 1
		fi
		if (( OPTARG < 64 || OPTARG > 4096)); then
			echo "Chunk size must be between 64 and 4096" 1>&2
			exit 1
		fi
		chunk_size=$((OPTARG *= 1024))
		;;
	d)
		((debug += 1))
		;;
	s)
		sample_count=$OPTARG
		;;
	n)
		min_files=$OPTARG
		;;
	p)
		if (( OPTARG < 1 || OPTARG > 100 )); then
			echo "Scan percentage must be between 1 and 100" 1>&2
			exit 1
		fi
		percentage=$OPTARG
		;;
	q)
		((quiet += 1))
		;;
	z|Z)
		case $OPTARG in
		lzo*)
			compression_type="lzop -"
			;;
		lz4fast*)
			compression_type="lz4 --fast="
			;;
		gzip*|lz4*)
			compression_type="${OPTARG%:*} -"
			;;
		zstdfast*)
			compression_type="zstd --fast="
			;;
		zstd*)
			compression_type="${OPTARG%:*} -"
			;;
		*)
			echo "Unknown compression type: $compression_type" 1>&2
			usage 1>&2
			exit 1
			;;
		esac
		[[ "$OPTARG" =~ ":" ]] && compression_level=${OPTARG#*:}
		;;
	l)
		compression_level=$OPTARG
		;;
	m)
		if (( OPTARG < 0 || OPTARG > 100 )); then
			echo "Compression margin must be between 0 and 100" 1>&2
			exit 1
		fi
		compression_margin=$OPTARG
		;;
	w)
		whole_file="true"
		;;
	h)
		usage
		exit 0
		;;
	*)
		usage 1>&2
		exit 1
		;;
	esac
done

if (( compression_level < 1 || compression_level > 12 )); then
	echo "Compression level must be between 1 and 12" 1>&2
	exit 1
fi
if [[ $compression_level -gt 9 && ! $compression_type =~ "lz4" ]]; then
	echo "Compression level must be between 1 and 9 (10+ for lz4 only)" 1>&2
	exit 2
fi

compress="$compression_type$compression_level -q"
shift $((OPTIND - 1))
if [[ -z "$@" ]]; then
	path=($default_path)
else
	path=("$@")
	shift
fi

# Variables to track overall compression efficiency and additional statistics
export total_file_size=0
export total_uncompressed_size=0
export total_compressed_size=0
export total_files_sampled=0
export total_small_files=0
export total_incompressible_files=0
export total_incompressible_size=0
export total_uncompressed_size_sampled=0
export total_compressed_size_estimated=0

round_to_block_size() {
	local size=$*

	echo $(( ((size - 1) | (block_size - 1)) + 1 ))
}

round_to_chunk_size() {
	local size=$*

	echo $(( ((size - 1) | (chunk_size - 1)) + 1 ))
}

export format="--format=%b*%B"
[[ $(uname) != "Darwin" ]] || format="-f %b*512"
# Function to process a file
process_file() {
	local file="$1"
	local file_size=$(stat $format "$file")
	local sum_uncompressed_chunk=0
	local sum_compressed_chunk=0

	# Round up the file_size to the next block (actual space usage)
	file_size=$(round_to_block_size $file_size)
	# Accumulate total size of files scanned (in block_size multiples)
	total_file_size=$((total_file_size + file_size))
	((total_file_count+= 1))

	# always count incompressible files, in case this is a large fraction
	if [[ -z "$file_size" ]] || (( file_size <= block_size )); then
		((total_small_files+= 1))
		((total_files_sampled+= 1))
		sum_uncompressed_chunk=$file_size
		sum_compressed_chunk=$file_size
		estimated_compressed_file_size=$file_size
	else
		# randomly select $percentage of files after sampling min_files,
		# unless file is larger than average of files checked so far
		local average=$((total_file_size / ${total_files_sampled/#0/1}))
		if (( total_files_sampled > min_files &&
		      file_size < 2 * average )); then
			(( RANDOM % 100 < percentage )) || return
		elif (( total_files_sampled > min_files && debug > 0 )); then
			echo -n "***"
		fi

		((total_files_sampled+= 1))

		local segment_size
		if [[ $whole_file == "true" ]] ||
		   (( file_size < chunk_size * sample_count )); then
			segment_size=$chunk_size
			segment_count=$(($(round_to_chunk_size file_size) /
					 chunk_size))
		else
			# Calculate the segment size for the file
			segment_size=$((file_size / sample_count))
			segment_count=$sample_count
		fi

		(( debug < 1 )) ||
			echo -n "$(basename $file): size: $file_size "
		(( debug < 2 )) ||
			echo -n "segs: $segment_count segsz: $segment_size "

		# Read and process each segment
		for ((i = 0; i < segment_count; i++)); do
			offset=$((i * segment_size / chunk_size))
			compressed_size=$(dd if="$file" bs=$chunk_size count=1 \
				skip=$offset 2>/dev/null | $compress | wc -c)

			# if compressed size is zero, something must have failed
			(( compressed_size > 0 )) || continue

			# Round up compressed size to full block size
			compressed_size=$(round_to_block_size compressed_size)

			# Incompressible chunks will not be compressed
			(( compressed_size <= chunk_size )) ||
				compressed_size=$chunk_size

			# Add sampled chunk bytes, but don't inflate last chunk
			last_chunk=$((file_size - offset * chunk_size ))
			(( last_chunk > chunk_size )) && last_chunk=$chunk_size

			((sum_uncompressed_chunk+= last_chunk ))
			((sum_compressed_chunk+= compressed_size))
		done

		# Get current ratio for this file
		current_ratio=$((sum_uncompressed_chunk * 100 / sum_compressed_chunk))
		# Assume compression ratio will be the same for the entire file
		estimated_compressed_file_size=$((file_size * 100 / current_ratio))

		(( debug < 1 )) ||
			echo "uncompr: $sum_uncompressed_chunk compr: $sum_compressed_chunk est: $estimated_compressed_file_size avg: $average"
	fi

	if ((sum_compressed_chunk >= sum_uncompressed_chunk)); then
		((total_incompressible_files+= 1))
		((total_incompressible_size+= file_size))
	fi

	# Accumulate the total uncompressed and compressed byte counts
	((total_uncompressed_size+= sum_uncompressed_chunk))
	((total_compressed_size+= sum_compressed_chunk))

	# Accumulate the estimated uncompressed and compressed byte counts
	((total_uncompressed_size_sampled+= file_size))
	((total_compressed_size_estimated+= estimated_compressed_file_size))
}

# Calculate compression ratio of real compressed chunks vs original (value >= 1)
calculate_ratio() {
	local ratio=$((total_uncompressed_size * 100 / total_compressed_size))

	printf "%u.%02u" $((ratio / 100)) $((ratio % 100))
}

# add correction factor for estimate safety margin with low sample percentage
(( compression_margin == 0 )) && correction=100 ||
	correction=$((100 + compression_margin + 10 * (100 - percentage) / 100))

# Calculate compression ratio from estimated compressed file size (value >= 1)
calculate_estimated_ratio() {
	local ratio=$((total_uncompressed_size_sampled * 100 * 100 /
                       (total_compressed_size_estimated * correction)))

	printf "%u.%02u" $((ratio / 100)) $((ratio % 100))
}

# Calculate estimated compressed size of all files using the ratio from our
# sample data
calculate_estimated_total_compressed_size()
{
	if (( debug )); then
		echo "(total_file_size=$total_file_size *" 1>&2
		echo "total_compressed_size_estimated=$total_compressed_size_estimated *" 1>&2
		echo "correction=$correction) /" 1>&2
		echo "(total_uncompressed_size_sampled=$total_uncompressed_size_sampled * 100)" 1>&2
	fi

	echo $((total_file_size * correction /
		total_uncompressed_size_sampled *
		total_compressed_size_estimated / 100))
}

print_size() {
	local size=$1
	local frac
	local unit

	if (( size > 4 * 2**50 )); then
		frac=$((size / 2**40))
		unit="PiB"
	elif (( size > 4 * 2**40 )); then
		frac=$((size / 2**30))
		unit="TiB"
	elif (( size > 4 * 2**30 )); then
		frac=$((size / 2**20))
		unit="GiB"
	elif (( size > 4 * 2**20 )); then
		frac=$((size / 2**10))
		unit="MiB"
	else
		frac=$size
		unit="KiB"
	fi

	printf "%u.%03u $unit" $((frac / 1024)) $((frac % 1024))
}

print_summary() {
	trap 0
	echo ""
	echo "Compression type: ${compression_type/ -*/} Level: $compression_level"
	echo "Chunk size: $chunk_size"
	echo "Number of files sampled: $total_files_sampled ($((total_files_sampled * 100 / total_file_count))% of $total_file_count total files)"
	echo "Number of files under $block_size bytes (incompressible): $total_small_files"
	echo "Elapsed scanning time: $((SECONDS - start)) seconds"
	echo "Total number of incompressible files: $total_incompressible_files"
	echo "Total size of incompressible files: $(print_size $total_incompressible_size)"
	echo "Total size of files scanned: $(print_size $total_file_size)"
	echo "Total uncompressed size of sampled data: $(print_size $total_uncompressed_size)"
	echo "Total compressed size of sampled data: $(print_size $total_compressed_size)"
	echo "Compression ratio of sampled data: $(calculate_ratio)"
	echo "Estimated compression ratio of sampled files: $(calculate_estimated_ratio)"
	estimated_total_compressed_size=$(calculate_estimated_total_compressed_size)
	echo "Estimated compressed size of all files: $(print_size $estimated_total_compressed_size)"
	exit 0
}
trap print_summary EXIT

(( quiet == 0 )) && runtime_description | fmt

# if stdout is a tty then make output more interactive
if [[ -t 1 ]]; then
	cr="\r"
	lines=100
	interval=30
else
	lf="\n"
	lines=1000
	interval=300
fi

total_file_count=0
last=$SECONDS

echo ""
if [[ "${path[@]}" != "$default_path"  ]]; then
	echo "Scanning '${path[@]}'."
else
	echo "Scanning current directory, '${path[@]}'."
fi
echo ""
echo ""

while read FILE; do
	process_file "$FILE"

	if (( quiet < 2 &&
	      ((min_files > 1 && total_files_scanned == min_files) ||
	       total_files_sampled % lines == 0 ||
	       last + interval < SECONDS) )); then
		if ((total_files_sampled != total_file_count)); then
			echo -ne "${cr}Sampled $total_files_sampled/$total_file_count files in $((SECONDS - start))s so far, estimated compression ratio $(calculate_estimated_ratio)x...${lf}"
		else
			echo -ne "${cr}Sampled $total_files_sampled files in $((SECONDS - start))s so far, estimated compression ratio $(calculate_estimated_ratio)x...${lf}"
		fi
		last=$SECONDS
	fi
done < <(find "${path[@]}" -type f -print)

(( total_file_count == 0 )) &&
	echo "error: no files found in '${path[@]}' to compress" 1>&2 &&
	exit 10
(( total_uncompressed_size == 0 )) &&
	echo "error: only zero-length files found in '${path[@]}'" 1>&2 &&
	exit 11

echo ""
# Report the additional statistics
if (( quiet == 0 )); then
	echo ""
	echo "Finished sampling."
	echo ""
	echo ""
	echo "---------------------"
	echo "Results"
	echo "---------------------"
fi

