Discard broken S3 downloads (#65)

nienbo · Oct 30, 2023 · 475f937 · 475f937
1 parent 9df04eb
commit 475f937
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 7 deletions.
diff --git a/lib/backends/s3.bash b/lib/backends/s3.bash
@@ -112,7 +112,15 @@ function restore() {
  fi
 
  if [[ ! "${BK_AWS_FOUND}" =~ (false) ]]; then
- aws s3 cp ${BK_CUSTOM_AWS_ARGS} "s3://${BUCKET}/${TAR_FILE}" .
+ TMP_FILE="$(mktemp)"
+ aws s3 cp ${BK_CUSTOM_AWS_ARGS} "s3://${BUCKET}/${TAR_FILE}" "${TMP_FILE}" || s3_download_failed=true
+ if ${s3_download_failed:-false}; then
+ echo -e "S3 download failed, soft failing and skipping cache restore..."
+ return 0
+ else
+ mv -f "${TMP_FILE}" "${TAR_FILE}"
+ fi
+
  [ "${BK_CACHE_SAVE_CACHE}" == "true" ] && cp "${TAR_FILE}" "${BK_CACHE_LOCAL_PATH}/${TAR_FILE}"
  tar ${BK_TAR_EXTRACT_ARGS} "${TAR_FILE}" -C .
  else

diff --git a/tests/command.bats b/tests/command.bats
@@ -10,10 +10,15 @@ setup() {
 }
 
 @test "Pre-command restores cache with basic key" {
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-key.tar : true"
 
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-key.tar -C . : echo Extracted tar archive"
@@ -40,6 +45,8 @@ setup() {
 
  unstub aws
  unstub tar
+ unstub mv
+ unstub mktemp
 }
 
 @test "Pre-command restores S3 backed cache using local file" {
@@ -118,9 +125,15 @@ setup() {
 @test "Cache key template evaluation on file" {
  CHECKSUM=355831032f586e782b45744f2ed79316cc830244
 
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-key-${CHECKSUM}.tar : true"
+
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-key-${CHECKSUM}.tar -C . : echo Extracted tar archive"
@@ -144,16 +157,24 @@ setup() {
  unset BUILDKITE_PIPELINE_SLUG
  unset BUILDKITE_ORGANIZATION_SLUG
 
+ unstub mktemp
+ unstub mv
  unstub aws
  unstub tar
 }
 
 @test "Cache key template evaluation on dir" {
  CHECKSUM=4cfa4e590847976f26d761074e355e4d95fa8107
 
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-key-${CHECKSUM}.tar : true"
+
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-key-${CHECKSUM}.tar -C . : echo Extracted tar archive"
@@ -177,16 +198,24 @@ setup() {
  unset BUILDKITE_PIPELINE_SLUG
  unset BUILDKITE_ORGANIZATION_SLUG
 
+ unstub mktemp
+ unstub mv
  unstub aws
  unstub tar
 }
 
 @test "Cache key multi-template evaluation" {
  CHECKSUMS=355831032f586e782b45744f2ed79316cc830244-241bc31c8ddc004c48e6d88d7fa51ee981b8ce51
 
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-key-${CHECKSUMS}.tar : true"
+
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-key-${CHECKSUMS}.tar -C . : echo Extracted tar archive"
@@ -210,16 +239,24 @@ setup() {
  unset BUILDKITE_PIPELINE_SLUG
  unset BUILDKITE_ORGANIZATION_SLUG
 
+ unstub mktemp
+ unstub mv
  unstub aws
  unstub tar
 }
 
 @test "Cache key template evaluation in middle of key" {
  CHECKSUM=355831032f586e782b45744f2ed79316cc830244
 
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-${CHECKSUM}-key.tar : true"
+
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-$CHECKSUM-key.tar -C . : echo Extracted tar archive"
@@ -243,6 +280,8 @@ setup() {
  unset BUILDKITE_PIPELINE_SLUG
  unset BUILDKITE_ORGANIZATION_SLUG
 
+ unstub mktemp
+ unstub mv
  unstub aws
  unstub tar
 }
@@ -271,9 +310,15 @@ setup() {
 
 @test "S3 arguments are passed through to copy command" {
 
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub mv \
+ "-f /tmp/tempfile v1-cache-key.tar : true"
+
  stub aws \
  "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
- "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar . : echo Copied from S3"
+ "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : echo Copied from S3"
 
  stub tar \
  "-xf v1-cache-key.tar -C . : echo Extracted tar archive"
@@ -299,6 +344,41 @@ setup() {
  unset BUILDKITE_PIPELINE_SLUG
  unset BUILDKITE_ORGANIZATION_SLUG
 
+ unstub mktemp
+ unstub mv
  unstub aws
  unstub tar
 }
+
+@test "S3 download errors cause soft failure" {
+
+ stub mktemp \
+ " : echo '/tmp/tempfile'"
+
+ stub aws \
+ "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
+ "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : false"
+
+ export BUILDKITE_ORGANIZATION_SLUG="my-org"
+ export BUILDKITE_PIPELINE_SLUG="my-pipeline"
+ export BUILDKITE_PLUGIN_CACHE_S3_BUCKET="my-bucket"
+ export BUILDKITE_PLUGIN_CACHE_S3_PROFILE="my-profile"
+ export BUILDKITE_PLUGIN_CACHE_S3_ARGS="--acl bucket-owner-full-control"
+ export BUILDKITE_PLUGIN_CACHE_BACKEND="s3"
+ export BUILDKITE_PLUGIN_CACHE_KEY="v1-cache-key"
+
+ run "$PWD/hooks/pre-command"
+ assert_success
+ assert_output --partial "S3 download failed"
+
+ unset BUILDKITE_PLUGIN_CACHE_KEY
+ unset BUILDKITE_PLUGIN_CACHE_BACKEND
+ unset BUILDKITE_PLUGIN_CACHE_S3_PROFILE
+ unset BUILDKITE_PLUGIN_CACHE_S3_ARGS
+ unset BUILDKITE_PLUGIN_CACHE_S3_BUCKET
+ unset BUILDKITE_PIPELINE_SLUG
+ unset BUILDKITE_ORGANIZATION_SLUG
+
+ unstub mktemp
+ unstub aws
+}