diff --git a/buildscripts/s3_lock.py b/buildscripts/s3_lock.py new file mode 100644 index 00000000000..4ee5e610e03 --- /dev/null +++ b/buildscripts/s3_lock.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +""" +S3-based lock mechanism. + +This script attempts to create a lock file in S3. If the file already exists, +it exits with failure (exit code 1). If the file doesn't exist and is successfully +created, it exits with success (exit code 0). + +Usage: + python s3_lock.py --bucket BUCKET_NAME --key LOCK_KEY [--content CONTENT] + +Example: + python s3_lock.py --bucket my-bucket --key locks/mylock.txt --content "locked by build 123" +""" + +import sys + +import typer +from botocore.exceptions import ClientError + +from buildscripts.util.download_utils import get_s3_client + + +def acquire_s3_lock(bucket, key, content="locked"): + """ + Attempt to acquire a lock by uploading a file to S3 only if it doesn't exist. + + Args: + bucket: S3 bucket name + key: S3 object key (file path) + content: Content to write to the lock file + + Returns: + True if lock was acquired (file didn't exist), False if lock already exists + """ + s3_client = get_s3_client() + + try: + # Try to upload the file only if it doesn't already exist + # using the IfNoneMatch condition with "*" (no existing ETag should match) + s3_client.put_object( + Bucket=bucket, + Key=key, + Body=content.encode("utf-8"), + IfNoneMatch="*", # Only succeed if object doesn't exist + ) + return True + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + + # PreconditionFailed means the object already exists + if error_code == "PreconditionFailed": + return False + + # Any other error should be raised + raise + + +app = typer.Typer() + + +@app.command() +def main( + bucket: str = typer.Option(..., help="S3 bucket name"), + key: str = typer.Option(..., help="S3 object key (path) for the lock file"), + content: str = typer.Option("locked", help="Content to write to the lock file"), +): + """Acquire an S3-based lock by creating a file if it doesn't exist.""" + try: + lock_acquired = acquire_s3_lock(bucket, key, content) + + if lock_acquired: + print(f"Lock acquired: s3://{bucket}/{key}") + sys.exit(0) + else: + print(f"Lock already exists: s3://{bucket}/{key}", file=sys.stderr) + sys.exit(1) + + except ClientError as e: + print(f"S3 error: {e}", file=sys.stderr) + sys.exit(2) + except Exception as e: + print(f"Unexpected error: {e}", file=sys.stderr) + sys.exit(2) + + +if __name__ == "__main__": + app() diff --git a/etc/evergreen_yml_components/tasks/misc_tasks.yml b/etc/evergreen_yml_components/tasks/misc_tasks.yml index 8ca514a3c17..1676af495cd 100644 --- a/etc/evergreen_yml_components/tasks/misc_tasks.yml +++ b/etc/evergreen_yml_components/tasks/misc_tasks.yml @@ -15,6 +15,10 @@ variables: params: exec_timeout_secs: 7200 # 2 hrs timeout_secs: 600 # 10 mins + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: suite_name @@ -77,6 +81,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_stepdown_terminate_kill_primary_with_balancer_and_config_transitions_and_add_remove_shard @@ -94,6 +102,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_replication_with_balancer_and_config_transitions_and_add_remove_shard @@ -111,6 +123,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: jstestfuzz_sharded_with_config_transitions @@ -128,6 +144,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_multi_stmt_txn_stepdown_terminate_kill_primary @@ -145,6 +165,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_stepdown_terminate_kill_primary_with_balancer @@ -162,6 +186,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_initial_sync @@ -179,6 +207,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_initsync_jscore_passthrough @@ -196,6 +228,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_stepdown_terminate_kill_primary_with_balancer @@ -216,6 +252,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: concurrency_sharded_replication @@ -237,6 +277,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: core @@ -254,6 +298,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: fcv_upgrade_downgrade_sharded_collections_jscore_passthrough @@ -273,6 +321,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: jstestfuzz_sharded @@ -290,6 +342,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_jscore_passthrough @@ -307,6 +363,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_lag_oplog_application_jscore_passthrough @@ -324,6 +384,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_kill_primary_jscore_passthrough @@ -341,6 +405,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_kill_secondaries_jscore_passthrough @@ -358,6 +426,10 @@ tasks: commands: - func: "do setup" - func: "do setup for antithesis" + - func: "set up remote credentials" + vars: + aws_key_remote: ${aws_key} + aws_secret_remote: ${aws_secret} - func: "antithesis image build and push" vars: suite: replica_sets_initsync_lagged_sync_source_jscore_passthrough diff --git a/evergreen/antithesis_image_build_and_push.sh b/evergreen/antithesis_image_build_and_push.sh index 032ec663a39..8c90fffd2bc 100644 --- a/evergreen/antithesis_image_build_and_push.sh +++ b/evergreen/antithesis_image_build_and_push.sh @@ -75,15 +75,29 @@ if [ $RET -ne 0 ]; then exit $RET fi -# Push Image +# Push Config Image sudo docker tag "${suite}:$tag" "$antithesis_repo/${task_name}:$tag" sudo docker push "$antithesis_repo/${task_name}:$tag" -sudo docker tag "mongo-binaries:$tag" "$antithesis_repo/mongo-binaries:$tag" -sudo docker push "$antithesis_repo/mongo-binaries:$tag" +# Push workload and binary images with s3 lock to prevent multiple pushes across different tasks +set +o errexit +$python buildscripts/s3_lock.py --bucket mciuploads --key ${project}/${version_id}/${build_variant}/antithesis_lock +RET=$? +set -o errexit -sudo docker tag "workload:$tag" "$antithesis_repo/workload:$tag" -sudo docker push "$antithesis_repo/workload:$tag" +if [ $RET -eq 0 ]; then + echo "Aquired lock for workload and binary images, pushing to antithesis." + sudo docker tag "mongo-binaries:$tag" "$antithesis_repo/mongo-binaries:$tag" + sudo docker push "$antithesis_repo/mongo-binaries:$tag" + + sudo docker tag "workload:$tag" "$antithesis_repo/workload:$tag" + sudo docker push "$antithesis_repo/workload:$tag" +elif [ $RET -eq 1 ]; then + echo "Failed to acquire lock for workload and binary images, skipping push." +else + echo "Error occurred when attempting to acquire s3 lock, exiting." + exit 1 +fi # Logout sudo docker logout https://us-central1-docker.pkg.dev