#!/bin/bash

# Author: Benjamin MALYNOVYTCH <bmalynovytch#admantic.fr>
# Requires:
#   - aws-cli
#   - pv (for progress)
# Usage: s3-large-file-restore.sh </path/to/destination/file>

set -o pipefail
set -e

export LC_ALL=en_US.UTF8
export PATH="/usr/local/bin:$PATH"

THREADS=8
BUCKET=my_bucket
DEST=$1
BS=$(( 1024 * 1024 ))

function restore_chunk {
  CHUNK=$1
  SIZE=$2
  echo "Restoring chunk $CHUNK"
  aws s3 cp s3://$BUCKET/$(basename $DEST)/$CHUNK - | dd of=$DEST bs=$BS seek=$(( $SKIP_O / $BS )) 2> /dev/null
}

( losetup -a | grep -q "$DEST" && umount $DEST ) || true

CHUNKS=$(aws s3 ls s3://$BUCKET/$(basename $DEST)/ | grep -v md4sum | awk '{print $4 " " $3}' | sort -n)
ITER=$(echo "$CHUNKS" | wc -l)

(
  trap 'kill $(jobs -p) &> /dev/null' EXIT
  SKIP_O=0
  IFS=$'\n'
  for CHUNK_SIZE in $CHUNKS ; do
    while [ $(jobs -p | wc -l) -ge $THREADS ] ; do
      sleep 1
    done

    CHUNK=${CHUNK_SIZE% *}
    SIZE=${CHUNK_SIZE#* }

    restore_chunk $CHUNK $SIZE &
    let SKIP_O+=$SIZE
  done
  wait
) | pv -w80 -ls $ITER > /dev/null

echo "All done."