From dd1c66783ae5fd2cd8b416687e465d9f1b558282 Mon Sep 17 00:00:00 2001
From: "Moritz E. Beber" <midnighter@posteo.net>
Date: Fri, 6 May 2022 17:09:01 +0200
Subject: [PATCH] refactor: use a template to add retrying

---
 modules/sratools/prefetch/main.nf             | 27 ++-------
 .../prefetch/templates/retry_with_backoff.sh  | 59 +++++++++++++++++++
 2 files changed, 65 insertions(+), 21 deletions(-)
 create mode 100755 modules/sratools/prefetch/templates/retry_with_backoff.sh

diff --git a/modules/sratools/prefetch/main.nf b/modules/sratools/prefetch/main.nf
index 75fa17a1..92a66616 100644
--- a/modules/sratools/prefetch/main.nf
+++ b/modules/sratools/prefetch/main.nf
@@ -12,31 +12,16 @@ process SRATOOLS_PREFETCH {
     tuple val(meta), val(id)
 
     output:
-    tuple val(meta), path("$id"), emit: sra
+    tuple val(meta), path(id), emit: sra
     path "versions.yml"         , emit: versions
 
     when:
     task.ext.when == null || task.ext.when
 
-    script:
-    def args = task.ext.args ?: ''
-    def config = "/LIBS/GUID = \"${UUID.randomUUID().toString()}\"\\n/libs/cloud/report_instance_identity = \"true\"\\n"
-    """
-    eval "\$(vdb-config -o n NCBI_SETTINGS | sed 's/[" ]//g')"
-    if [[ ! -f "\${NCBI_SETTINGS}" ]]; then
-        mkdir -p "\$(dirname "\${NCBI_SETTINGS}")"
-        printf '${config}' > "\${NCBI_SETTINGS}"
-    fi
+    shell:
+    args = task.ext.args ?: ''
+    args2 = task.ext.args2 ?: '5 1 100'  // <num retries> <base delay in seconds> <max delay in seconds>
+    config = "/LIBS/GUID = \"${UUID.randomUUID().toString()}\"\\n/libs/cloud/report_instance_identity = \"true\"\\n"
 
-    prefetch \\
-        $args \\
-        $id
-
-    vdb-validate $id
-
-    cat <<-END_VERSIONS > versions.yml
-    "${task.process}":
-        sratools: \$(prefetch --version 2>&1 | grep -Eo '[0-9.]+')
-    END_VERSIONS
-    """
+    template 'retry_with_backoff.sh'
 }
diff --git a/modules/sratools/prefetch/templates/retry_with_backoff.sh b/modules/sratools/prefetch/templates/retry_with_backoff.sh
new file mode 100755
index 00000000..fbcb6532
--- /dev/null
+++ b/modules/sratools/prefetch/templates/retry_with_backoff.sh
@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+
+set -u
+
+retry_with_backoff() {
+    local max_attempts=${1}
+    local delay=${2}
+    local max_time=${3}
+    local attempt=1
+    local output=
+    local status=
+
+    # Remove the first three arguments to this function in order to access
+    # the 'real' command with `${@}`.
+    shift 3
+
+    while [ ${attempt} -le ${max_attempts} ]; do
+        output=$("${@}")
+        status=${?}
+
+        if [ ${status} -eq 0 ]; then
+            break
+        fi
+
+        if [ ${attempt} -lt ${max_attempts} ]; then
+            echo "Failed attempt ${attempt} of ${max_attempts}. Retrying in ${delay} s." >&2
+            sleep ${delay}
+        elif [ ${attempt} -eq ${max_attempts} ]; then
+            echo "Failed after ${attempt} attempts." >&2
+            return ${status}
+        fi
+
+        attempt=$(( ${attempt} + 1 ))
+        delay=$(( ${delay} * 2 ))
+        if [ ${delay} -ge ${max_time} ]; then
+            delay=${max_time}
+        fi
+    done
+
+    echo "${output}"
+}
+
+eval "$(vdb-config -o n NCBI_SETTINGS | sed 's/[" ]//g')"
+if [[ ! -f "${NCBI_SETTINGS}" ]]; then
+    mkdir -p "$(dirname "${NCBI_SETTINGS}")"
+    printf '!{config}' > "${NCBI_SETTINGS}"
+fi
+
+retry_with_backoff !{args2} \
+    prefetch \
+    !{args} \
+    !{id}
+
+vdb-validate !{id}
+
+cat <<-END_VERSIONS > versions.yml
+"!{task.process}":
+    sratools: $(prefetch --version 2>&1 | grep -Eo '[0-9.]+')
+END_VERSIONS