#!/bin/bash
#
# bb_start_stage_in
# Perform stage-in of files for a specific job
*
# NOTE: Sample script used for burst buffer management. Used for development
# purposes and is not intended for production use.
#
# File system formats:
# I. BB_DIRECTORY (environment variable): Defines head of burst buffer file
#    system
#    A. .slurm_state, child of BB_DIRECTORY, contains state information, file
#       format is "TotalSize=#"
#    B. Slurm job ID number, child directories of BB_DIRECTORY, one per job
#       1. .slurm_state, child of Slurm job ID directory, contains state
#          information, file format is
#          "UserID=# JobID=# Name=name State=name Size=#"
#          (NOTE: "UserID" must be first)
#          (NOTE: Either the JobID or Name should be specified, but not both)
#

# Validate execute line
valid=0
if [[ $# -ge 4 && $1 = "start_stage_in" ]]; then
	job_id=$2
	user_id=$3
	buf_size=$4
	valid=1
fi
if [[ $valid -eq 0 ]]; then
	echo "Usage: $0 start_stage_in <jobid> <userid> <size> [script]"
	exit 1
fi

# Default BB_DIRECTORY for testing is $HOME/slurm_bb
if [[ -z "$BB_DIRECTORY" ]] ; then
	BB_DIRECTORY=$HOME/slurm_bb
fi

if [[ ! -d $BB_DIRECTORY/$job_id ]]; then
	mkdir $BB_DIRECTORY/$job_id
fi

# Set "State=staging-in" in burst buffer's state file
if [[ -e $BB_DIRECTORY/$job_id/.slurm_state ]]; then
	cp $BB_DIRECTORY/$job_id/.slurm_state $BB_DIRECTORY/$job_id/.slurm_state2
	cat $BB_DIRECTORY/$job_id/.slurm_state2 | awk 'BEGIN {FS=" "}{print $1, $2, "State=staging-in", $4}' >$BB_DIRECTORY/$job_id/.slurm_state
	rm $BB_DIRECTORY/$job_id/.slurm_state2
else
	echo "UserID=$user_id JobId=$job_id State=staging-in Size=$buf_size" >$BB_DIRECTORY/$job_id/.slurm_state
fi

################################################################################
# Move files here
#
# For testing purposes, we just store the script here, then briefly sleep
shift
shift
rm -f $BB_DIRECTORY/$job_id/stage_out_script
while (($#)); do
	echo "$1" >>$BB_DIRECTORY/$job_id/stage_in_script
	shift
done
sleep 10
################################################################################

# Set "State=staged-in" in burst buffer's state file
if [[ -e $BB_DIRECTORY/$job_id/.slurm_state ]]; then
	cp $BB_DIRECTORY/$job_id/.slurm_state $BB_DIRECTORY/$job_id/.slurm_state2
	cat $BB_DIRECTORY/$job_id/.slurm_state2 | awk 'BEGIN {FS=" "}{print $1, $2, "State=staged-in", $4}' >$BB_DIRECTORY/$job_id/.slurm_state
	rm $BB_DIRECTORY/$job_id/.slurm_state2
	exit 0
else
	echo "UserID=$user_id JobId=$job_id State=staged-in Size=$buf_size" >$BB_DIRECTORY/$job_id/.slurm_state
	exit 1
fi
