#!/bin/bash
#####################################################################################
# Copyright (c) Members of the EGEE Collaboration. 2004.
# See http://www.eu-egee.org/partners/ for details on the copyright
# holders.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS
# OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#####################################################################################
#
# AUTHORS: Dimitar Shiyachki <Dimitar.Shiyachki@cern.ch>
#
#####################################################################################
#
# Helper function for filesystem selection tests
#
#####################################################################################
# Change status for all filesystems
#    @Param 1: status
# meta: preconfig=../../DPM-config

function change_allfs_status
{
   func=${python_change_allfs_status//NEW_STATUS/$1}

   python -c "$func"
   return $?
}

python_change_allfs_status="

import sys
import dpm2 as dpm

pools = dpm.dpm_getpools ()
for pool in pools :
   filesystems = dpm.dpm_getpoolfs (pool.poolname)
   for filesystem in filesystems :
      dpm.dpm_modifyfs (filesystem.server, filesystem.fs, NEW_STATUS, filesystem.weight)

"
#####################################################################################
# Shell stub for creating a dpm filesystem using a loopback device
#
#    FILESYSTEM_NAME: string
#    FILESYSTEM_SIZE: integer

shell_create_filesystem="

   fsname=FILESYSTEM_NAME;
   fssize=FILESYSTEM_SIZE;

   current_loops=\$(losetup -a | wc -l)
   if [ \$current_loops -eq 0 ]; then
      rmmod loop;
      modprobe loop max_loop=32;
      sleep 1
   fi

   mkdir /tmp/\${fsname};
   dd if=/dev/zero of=/tmp/loop_\${fsname} bs=1048576 count=\$fssize;
   DEVICE=\$(losetup -f);
   if [ \$? -ne 0 ]; then
      max_loop=\$(ls /dev/loop* | sed -e 's/\\/dev\\/loop//' | sort -n | tail -1);
      
      rmdir /tmp/\${fsname};
      rm -f /tmp/loop_\${fsname};
      exit -1;
   fi

   losetup \$DEVICE /tmp/loop_\${fsname};
   mkfs.ext3 \$DEVICE;
   mount \$DEVICE /tmp/\${fsname};
   chown dpmmgr.dpmmgr /tmp/\${fsname};

   exit 0;
"

#####################################################################################
# Shell stub for destroying a dpm filesystem and releasing the resources
# 

shell_destroy_filesystem="

   fsdir=FILESYSTEM_DIR;

   DEVICE=\$(mount | grep \"\${fsdir} \" | cut -f 1 -d \" \");
   umount \${fsdir};
   losetup -d \$DEVICE;
   rmdir \${fsdir};
   fsname=\$(echo \$fsdir | sed -e 's/\/tmp\///');
   rm -f /tmp/loop_\${fsname};
"

#####################################################################################
# Creates an input file for dpm-buildfsv
#    @Param 1: file name
#

function write_fsv_input
{
   func=${python_write_fsv_input//FILENAME/$1}

   python -c "$func"
   return $?
}

python_write_fsv_input="

import sys
import dpm2 as dpm

filename=\"FILENAME\";
fp = open(filename, \"w\");

pools = dpm.dpm_getpools ()
for pool in pools :
   filesystems = dpm.dpm_getpoolfs (pool.poolname)
   for filesystem in filesystems :
      if filesystem.status == 0 :
         fp.write (\"%s %s %d\n\" % (filesystem.server, filesystem.fs, filesystem.weight))
      else :
         fp.write (\"%s %s %d\n\" % (filesystem.server, filesystem.fs, 0))

fp.close()

"

#####################################################################################
# Get replica location for a file ( poolname:hostname:filesystem )
#    $Param 1: filename

function get_single_replica_location
{
   func=${python_get_single_replica_location//FILENAME/$1}

   python -c "$func"
   return $?
}

python_get_single_replica_location="

import sys
import dpm2 as dpm

replicas = dpm.dpns_getreplica (\"FILENAME\", None, None);
if ( len (replicas) != 1 ) :
   print len (replicas)
   sys.exit(-1)

print \"%s:%s:%s\" % (replicas[0].poolname, replicas[0].host, replicas[0].fs);
sys.exit (0)

"

#####################################################################################
# Return all filesystems in a pool ( hostname:filesystem )
#    $Param 1: poolname

function list_pool_filesystems
{
   func=${python_list_pool_filesystems//POOLNAME/$1}

   python -c "$func"
   return $?
}

python_list_pool_filesystems="

import sys
import dpm2 as dpm

filesystems = dpm.dpm_getpoolfs (\"POOLNAME\");
for fs in filesystems :
   print \"%s:%s\" % (fs.server, fs.fs);

sys.exit (0)

"
#####################################################################################
# File selection tests cleanup
#    @Param 1: file count to cleanup

function cleanup
{
   echo "Cleaning up..."
   remove_files $1

   rm -f /tmp/fsgroupselcheck
   rm -f /tmp/fsv_input

   # Reenable preexisting filesystems
   change_allfs_status 0
}

#####################################################################################
# Remove files
#    @Param 1: file count
#    @Param 2: file prefix

function remove_files
{
   local FILE_COUNT=$1
   local FILE_PREFIX=$2

   for i in `seq $FILE_COUNT`; do
      rfrm $DPNS_HOME/${FILE_PREFIX}_${i}_ > /dev/null 2>&1 &
      if [ $[$i % 10] -eq 0 ]; then
         wait
      fi
   done
   wait
}

#####################################################################################
# Check the replicas for a file sequence 
#    (fs selection order should be the same as in dpm-buildfsv)
#
#   @Param1: poolname
#   @Param2: file prefix
#   @Param3: file count
#   @Param4: fs weight sum

function check_pool_sequence
{
   local i=0
   local POOLNAME=$1
   local FILE_PREFIX=$2
   local FILE_COUNT=$3
   local WEIGHT_SUM=$4

   write_fsv_input /tmp/fsv_input
   for s in $(dpm-buildfsv /tmp/fsv_input | grep "/tmp/${POOLNAME}_fs" | cut -f 3- -d " " | sed -e 's/ /:/'); do
     SELECTION[$i]=$s
      i=$[$i + 1]
   done

   if [ $i -ne $WEIGHT_SUM ]; then 
      ERROR="Unexpected vector lenght"
      return -1
   fi

   echo "Printing pool configuration with testing weights configured..."
   dpm-qryconf --weight
   if [ $? -ne 0 ]; then
      ERROR="Cannot qrery pool configuration with --weight"
      return -1
   fi

   echo "The sum of all weights is $WEIGHT_SUM"
   for i in `seq $FILE_COUNT`; do

      rl=`get_single_replica_location ${DPNS_HOME}/${FILE_PREFIX}_${i}_`
      if [ $? -ne 0 ]; then
         ERROR="Unexpected replica count for ${DPNS_HOME}/${FILE_PREFIX}_${i}_"
         return -1
      fi

      el="wpool:${SELECTION[$[($i-1) % $WEIGHT_SUM]]}"
      if [ "x$rl" == "x$el" ]; then
         echo "${DPNS_HOME}/${FILE_PREFIX}_${i}_: replica OK"
      else
         ERROR="${DPNS_HOME}/${FILE_PREFIX}_${i}_: expected ($el), actual ($rl)"
         return -1
      fi
   done
   return 0
}

#####################################################################################
# Check the replicas for a file group
#    @Param 1: poolname
#    @Param 2: file prefix
#    @Param 3: file count
#    @Param 4: fs weight sum
#    @Param 5: fs count

function check_pool_group
{
   local POOLNAME=$1
   local FILE_PREFIX=$2
   local FILE_COUNT=$3
   local WEIGHT_SUM=$4
   local FS_COUNT=$5

   rm -f /tmp/fsgroupselcheck

   for i in `seq $FILE_COUNT`; do
      echo "Checking replica location for ${DPNS_HOME}/${FILE_PREFIX}_${i}_"
      rl=`get_single_replica_location ${DPNS_HOME}/${FILE_PREFIX}_${i}_`
      if [ $? -ne 0 ]; then
         ERROR="${DPNS_HOME}/${FILE_PREFIX}_${i}_:${rl}"
         return -1
      else
         echo $rl >> /tmp/fsgroupselcheck
      fi
   done
   for i in `seq 0 $[$FS_COUNT-1]`; do

      next_server=$(get_pool_server $POOLNAME $i)
      next_filesystem=$(get_pool_directory $POOLNAME $i)
      next_weight=$(get_fs_weight $POOLNAME $i)

      ccount=$(grep ":${next_server}:${next_filesystem}$" /tmp/fsgroupselcheck | wc -l)
      if [ $ccount -ne $[${next_weight} * $FILE_COUNT / $WEIGHT_SUM] ]; then
         ERROR="Filesystem ${next_server}:${next_filesystem}: $ccount replicas; UNEXPECTED"
         return -1
      else
         echo "Filesystem ${next_server}:${next_filesystem} OK"
      fi
   done
   return 0
}

#####################################################################################
# Copy a set of files to the DPM server 
#    @Param 1: file count
#    @Param 2: file prefix
#    @Param 3: parallel count (default 1)

function create_files
{
   local FILE_COUNT=$1
   local FILE_PREFIX=$2
   local PARALLEL_COUNT=$3

   if [ -z "$PARALLEL_COUNT" ]; then
      parallelc=1
   else
      parallelc=$3
   fi
   for i in `seq $FILE_COUNT`; do
      echo "Copying ${DPNS_HOME}/${FILE_PREFIX}_${i}_ ..."
      rfcp /bin/ls ${DPNS_HOME}/${FILE_PREFIX}_${i}_ &
      if [ $[$i % $parallelc] -eq 0 ]; then
         echo "Waiting to perform..."
         wait
      fi
   done
   wait
}

#####################################################################################
# Creates a pool and adds filesystems to it
#    @Param1: poolname
#    @Param2: filesystem count
#    @Param3: filesystem size

function create_pool
{
   local POOLNAME=$1
   local FS_PREFIX=$1
   local FS_COUNT=$2
   local FS_SIZE=$3
   
   dpm-addpool --poolname $POOLNAME
   if [ $? -ne 0 ]; then
      echo "Cannot create pool $POOLNAME"
      TEST_FAILED "Pool creation error"
   fi

   local current_count=0
   while [ $current_count -ne $FS_COUNT ]; do
      for server in $DISK_SERVERS; do

         local FS_NAME="${FS_PREFIX}_fs${current_count}"
         echo "Adding filesystem /tmp/${FS_NAME} on $server"

         set_pool_server $POOLNAME $current_count $server
         set_pool_directory $POOLNAME $current_count "/tmp/${FS_PREFIX}_fs${current_count}"

         c_shell_create_filesystem=${shell_create_filesystem//FILESYSTEM_NAME/$FS_NAME}
         c_shell_create_filesystem=${c_shell_create_filesystem//FILESYSTEM_SIZE/$FS_SIZE}

         DPM_HOST_EXEC "$c_shell_create_filesystem" $server
         NO_PROXY_NEEDED
         if [ $? -ne 0 ]; then
            cleanup 0;
            echo "Cannot create filesystem /tmp/${FS_NAME} on host $server"
            TEST_FAILED "FS creation failure"
         fi

         dpm-addfs --poolname $POOLNAME --server $server --fs /tmp/${FS_NAME} --weight 2813
         if [ $? -ne 0 ]; then
           cleanup 0;
            echo "Cannot add filesystem /tmp/${FS_NAME} on host $server"
            TEST_FAILED "dpm-addfs failure"
         fi

         current_count=$[$current_count + 1]
         if [ $current_count -eq $FS_COUNT ]; then
            break
         fi
      done
   done

   echo "Printing pool configurations after fs creation"
   dpm-qryconf --weight
   if [ $? -ne 0 ]; then
      cleanup 0
      TEST_FAILED "Cannot initial qrery pool configuration with --weight"
   else
      fsc_mlines=$(dpm-qryconf --weight | grep "/tmp/${FS_PREFIX}_fs" | grep "ONLINE 2813" | wc -l)
      if [ $fsc_mlines -ne $FS_COUNT ]; then
         cleanup 0
         TEST_FAILED "Unexpected initial filesystem configuration (wrong initial weight reported by dpm-qryconf)"
      fi
   fi
}

#####################################################################################
# Removes a pool and all filesystems associated with it
#    @Param1: pool name

function destroy_pool
{
   local POOLNAME=$1
   
   for fs in `list_pool_filesystems $POOLNAME`; do
      local server=$(echo $fs | cut -f 1 -d ":")
      local directory=$(echo $fs | cut -f 2 -d ":")
      
      echo "Removing filesystem ${directory} from ${server}"
      c_shell_destroy_filesystem=${shell_destroy_filesystem//FILESYSTEM_DIR/$directory}
      
      DPM_HOST_EXEC "$c_shell_destroy_filesystem" $server
      NO_PROXY_NEEDED
      dpm-rmfs --server $server --fs $directory
   done
   
   # Remove the pool from DPM
   dpm-rmpool --poolname $POOLNAME
}

#####################################################################################
#
# Use /etc/shift.conf to look up all trusted nodes
#DISK_SERVERS=$(grep "DPM TRUST" /etc/shift.conf | sed -e 's/DPM TRUST //')
#if [ -z "$DISK_SERVERS" ]; then
#   TEST_FAILED "Disk servers lookup has failed"
#fi
#
# Disable all available filesystems
#change_allfs_status 1
#if [ $? -ne 0 ]; then
#   TEST_FAILED "Cannot disable existing filesystems to prepare the test environment"
#fi
#
#####################################################################################
# Get/Set the weight of the specified filesystem for a pool
#    @Param1: pool name (getter/setter)
#    @Param2: fs index (getter/setter)
#    @Param3: fs weight (setter only)

function get_fs_weight 
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local FS_INDEX=$2
   
   eval echo '$'{`echo ${POOLNAME}`_WEIGHTS[`echo $FS_INDEX`]}
}

function set_fs_weight
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local FS_INDEX=$2
   local FS_WEIGHT=$3

   eval `echo ${POOLNAME}`_WEIGHTS[`echo $FS_INDEX`]=$FS_WEIGHT
}

#####################################################################################
# Get/Set the servers for a pool
#    @Param1: pool name (getter/setter)
#    @Param2: server index (getter/setter)
#    @Param3: server name (setter only)

function get_pool_server
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local SERVER_INDEX=$2
   
   eval echo '$'{`echo ${POOLNAME}`_SERVERS[`echo $SERVER_INDEX`]}
}

function set_pool_server
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local SERVER_INDEX=$2
   local SERVER_NAME=$3

   eval `echo ${POOLNAME}`_SERVERS[`echo $SERVER_INDEX`]=$SERVER_NAME
}
        
#####################################################################################
# Get/Set the directories for a pool
#    @Param1: pool name (getter/setter)
#    @Param2: directory index (getter/setter)
#    @Param3: directory name (setter only)

function get_pool_directory
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local DIRECTORY_INDEX=$2
   
   eval echo '$'{`echo ${POOLNAME}`_DIRECTORIES[`echo $DIRECTORY_INDEX`]}
}

function set_pool_directory
{
   local POOLNAME=$(echo $1 | tr [:lower:] [:upper:])
   local DIRECTORY_INDEX=$2
   local DIRECTORY_NAME=$3

   eval `echo ${POOLNAME}`_DIRECTORIES[`echo $DIRECTORY_INDEX`]=$DIRECTORY_NAME
}
        
#####################################################################################
# Configure fs weights of a pool
#    @Param1: pool name
#    @Param2: pool fs count

function configure_pool
{
   local POOLNAME=$1
   local FS_COUNT=$2
   
   local SUM=0
   for i in `seq 0 $[$FS_COUNT - 1]`; do

      next_server=$(get_pool_server $POOLNAME $i)
      next_filesystem=$(get_pool_directory $POOLNAME $i)
      next_weight=$(get_fs_weight $POOLNAME $i)

      dpm-modifyfs --server ${next_server} --fs ${next_filesystem} --weight ${next_weight}
      if [ $? -ne 0 ]; then
         return -1
      fi
 
      SUM=$[$SUM + ${next_weight}]
   done

   return $SUM
}

#####################################################################################
# Create permanent files in spacetoken 
#    @Param 1: file count
#    @Param 2: file prefix
#    @Param 3: parallel count (default 1)
#    @Param 4: spacetoken

function create_permanent_files
{
   local FILE_COUNT=$1
   local FILE_PREFIX=$2
   local PARALLEL_COUNT=$3
   local SPACETOKEN=$4

   if [ -z "$PARALLEL_COUNT" ]; then
      parallelc=1
   else
      parallelc=$3
   fi
   for i in `seq $FILE_COUNT`; do
      echo "Copying ${DPNS_HOME}/${FILE_PREFIX}_${i}_ ..."
      if [ -z "$SPACETOKEN" ]; then
         lcg-cp -b -D srmv2 /bin/ls srm://${DPM_HOST}:8446/srm/managerv2?SFN=${DPNS_HOME}/${FILE_PREFIX}_${i}_ &
      else
         lcg-cp -b -S $SPACETOKEN -D srmv2 /bin/ls srm://${DPM_HOST}:8446/srm/managerv2?SFN=${DPNS_HOME}/${FILE_PREFIX}_${i}_ &
      fi
      if [ $[$i % $parallelc] -eq 0 ]; then
         echo "Waiting to perform..."
         wait
      fi
   done
   wait
}
