scripts: Add an "auto-size" feature

Changed image-builder/image_builder.sh to
automatically calculate the size of the root disk

Fixes #25.

Signed-off-by: Harshal Patil <harshal.patil@in.ibm.com>

image-builder/image_builder.sh

@@ -36,18 +36,26 @@ info()
 	echo "INFO: ${msg}"
 }
 
+warning()
+{
+        local msg="$*"
+        echo "WARNING: ${msg}"
+}
+
 usage()
 {
 	error="${1:-0}"
 	cat <<EOT
 Usage: ${SCRIPT_NAME} [options] <rootfs-dir>
-	This script will create a Kata Containers image file based on the 
+	This script will create a Kata Containers image file of
-	<rootfs-dir> directory.
+        an adequate size based on the <rootfs-dir> directory.
+        The size of the image can be also be specified manually
+        by '-s' flag.
 
 Options:
 	-h Show this help
 	-o path to generate image file ENV: IMAGE
-	-s Image size in MB (default $IMG_SIZE) ENV: IMG_SIZE
+	-s Image size in MB ENV: IMG_SIZE
 
 Extra environment variables:
 	AGENT_BIN:  use it to change the expected agent binary name"
@@ -57,12 +65,32 @@ EOT
 exit "${error}"
 }
 
-while getopts "ho:s:" opt
+# Maximum allowed size in MB for root disk
+MAX_IMG_SIZE_MB=2048
+
+FS_TYPE=${FS_TYPE:-"ext4"}
+
+# In order to support memory hotplug, image must be aligned to 128M
+MEM_BOUNDARY=128
+
+# Maximum no of attempts to create a root disk before giving up
+MAX_ATTEMPTS=5
+
+ATTEMPT_NUM=0
+while getopts "ho:s:f:" opt
 do
 	case "$opt" in
 		h)	usage ;;
 		o)	IMAGE="${OPTARG}" ;;
-		s)	IMG_SIZE="${OPTARG}" ;;
+		s)	IMG_SIZE=${OPTARG}
+                        if [ ${IMG_SIZE} -lt 0 ]; then
+                                die "Image size has to be greater than 0 MB."
+                        fi
+                        if [ ${IMG_SIZE} -gt ${MAX_IMG_SIZE_MB} ]; then
+                                die "Image size should not be greater than ${MAX_IMG_SIZE_MB} MB."
+                        fi
+                        ;;
+                f)      FS_TYPE="${OPTARG}" ;;
 	esac
 done
 
@@ -113,8 +141,48 @@ OK "init is installed"
 OK "Agent installed"
 [ "$(id -u)" -eq 0 ] || die "$0: must be run as root"
 
+ROOTFS_SIZE=$(du -B 1MB -s "${ROOTFS}" | awk '{print $1}')
 BLOCK_SIZE=${BLOCK_SIZE:-4096}
-IMG_SIZE=${IMG_SIZE:-80}
+OLD_IMG_SIZE=0
+
+align_memory()
+{
+    remaining=$(($IMG_SIZE % $MEM_BOUNDARY))
+    if [ "$remaining" != "0" ];then
+        warning "image size '$IMG_SIZE' is not aligned to memory boundary '$MEM_BOUNDARY', aligning it"
+        IMG_SIZE=$(($IMG_SIZE + $MEM_BOUNDARY - $remaining))
+    fi
+}
+
+# Calculate image size based on the rootfs
+calculate_img_size()
+{
+    IMG_SIZE=${IMG_SIZE:-$MEM_BOUNDARY}
+    align_memory
+}
+
+# Cleanup
+cleanup()
+{
+    sync
+    umount -l ${MOUNT_DIR}
+    rmdir ${MOUNT_DIR}
+    fsck -D -y "${DEVICE}p1"
+    losetup -d "${DEVICE}"
+}
+
+create_rootfs_disk()
+{
+    ATTEMPT_NUM=$(($ATTEMPT_NUM+1))
+    info "Create root disk image. Attempt ${ATTEMPT_NUM} out of ${MAX_ATTEMPTS}."
+    if [ ${ATTEMPT_NUM} -gt ${MAX_ATTEMPTS} ]; then
+        die "Unable to create root disk image."
+    fi
+
+    calculate_img_size
+    if [ ${OLD_IMG_SIZE} -ne 0 ]; then
+        info "Image size ${OLD_IMG_SIZE}MB too small, trying again with size ${IMG_SIZE}MB"
+    fi
 
     info "Creating raw disk with size ${IMG_SIZE}M"
     qemu-img create -q -f raw "${IMAGE}" "${IMG_SIZE}M"
@@ -125,7 +193,7 @@ OK "Image file created"
 
     info "Creating partitions"
     parted "${IMAGE}" --script "mklabel gpt" \
-"mkpart ext4 1M -1M"
+    "mkpart ${FS_TYPE} 1M -1M"
     OK "Partitions created"
 
     # Get the loop device bound to the image file (requires /dev mounted in the
@@ -143,25 +211,34 @@ OK "Image formated"
     info "Mounting root paratition"
     mount "${DEVICE}p1" "${MOUNT_DIR}"
     OK "root paratition mounted"
-
     RESERVED_BLOCKS_PERCENTAGE=3
     info "Set filesystem reserved blocks percentage to ${RESERVED_BLOCKS_PERCENTAGE}%"
     tune2fs -m "${RESERVED_BLOCKS_PERCENTAGE}" "${DEVICE}p1"
 
-#TODO: Calculate disk size based on rootfs
-#FIXME: https://github.com/kata-containers/osbuilder/issues/2
-ROOTFS_SIZE=$(du -B 1MB -s "${ROOTFS}" | awk '{print $1}')
     AVAIL_DISK=$(df -B M --output=avail "${DEVICE}p1" | tail -1)
     AVAIL_DISK=${AVAIL_DISK/M}
     info "Free space root partition ${AVAIL_DISK} MB"
+
+    # if the available disk space is less than rootfs size, repeat the process
+    # of disk creation by adding 5% in the inital assumed value $ROOTFS_SIZE
+    if [ $ROOTFS_SIZE -gt $AVAIL_DISK ]; then
+        # Increase the size but remain aligned to 128
+        MEM_BOUNDARY=$(($MEM_BOUNDARY+128))
+        rm -f ${IMAGE}
+        OLD_IMG_SIZE=${IMG_SIZE}
+        unset IMG_SIZE
+        cleanup
+        create_rootfs_disk
+    fi
+}
+
+create_rootfs_disk
+
 info "rootfs size ${ROOTFS_SIZE} MB"
 info "Copying content from rootfs to root partition"
 cp -a "${ROOTFS}"/* ${MOUNT_DIR}
 OK "rootfs copied"
 
-# Cleanup
+cleanup
-sync
+
-umount -l ${MOUNT_DIR}
+info "Image created. Size: ${IMG_SIZE}MB."
-fsck -D -y "${DEVICE}p1"
-losetup -d "${DEVICE}"
-info "Image created"