Implement ProcessCdnLogQueue background job

Author: Turbo87

.env.sample

@@ -45,6 +45,14 @@ export TEST_DATABASE_URL=
 # Uses AWS credentials.
 # export CLOUDFRONT_DISTRIBUTION=
 
+
+# Configuration for the CDN log queue. You can leave these commented out if
+# you're not using the CDN log queue.
+# export CDN_LOG_QUEUE_ACCESS_KEY=
+# export CDN_LOG_QUEUE_SECRET_KEY=
+# export CDN_LOG_QUEUE_URL=
+# export CDN_LOG_QUEUE_REGION=
+
 # Upstream location of the registry index. Background jobs will push to
 # this URL. The default points to a local index for development.
 # Run `./script/init-local-index.sh` to initialize this repo.

Cargo.lock

@@ -52,6 +52,7 @@ async-trait = "=0.1.77"
 aws-credential-types = { version = "=1.1.4", features = ["hardcoded-credentials"] }
 aws-ip-ranges = "=0.90.0"
 aws-sdk-cloudfront = "=1.12.0"
+aws-sdk-sqs = "=1.12.0"
 axum = { version = "=0.7.4", features = ["macros", "matched-path"] }
 axum-extra = { version = "=0.9.2", features = ["cookie-signed", "typed-header"] }
 base64 = "=0.21.7"

Cargo.toml

@@ -31,6 +31,7 @@ pub enum Command {
         #[arg()]
         name: String,
     },
+    ProcessCdnLogQueue(jobs::ProcessCdnLogQueue),
     SyncAdmins {
         /// Force a sync even if one is already in progress
         #[arg(long)]
@@ -89,6 +90,9 @@ pub fn run(command: Command) -> Result<()> {
         Command::DailyDbMaintenance => {
             jobs::DailyDbMaintenance.enqueue(conn)?;
         }
+        Command::ProcessCdnLogQueue(job) => {
+            job.enqueue(conn)?;
+        }
         Command::SquashIndex => {
             jobs::SquashIndex.enqueue(conn)?;
         }

src/admin/enqueue_job.rs

@@ -1,7 +1,13 @@
+mod message;
+
 use crate::config::CdnLogStorageConfig;
+use crate::tasks::spawn_blocking;
 use crate::worker::Environment;
 use anyhow::Context;
+use aws_credential_types::Credentials;
+use aws_sdk_sqs::config::{BehaviorVersion, Region};
 use crates_io_cdn_logs::{count_downloads, Decompressor};
+use crates_io_env_vars::required_var;
 use crates_io_worker::BackgroundJob;
 use object_store::aws::AmazonS3Builder;
 use object_store::local::LocalFileSystem;
@@ -130,6 +136,142 @@ impl ProcessCdnLog {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize, clap::Parser)]
+pub struct ProcessCdnLogQueue {
+    /// The maximum number of messages to receive from the queue and process.
+    #[clap(long, default_value = "1")]
+    max_messages: usize,
+}
+
+impl BackgroundJob for ProcessCdnLogQueue {
+    const JOB_NAME: &'static str = "process_cdn_log_queue";
+
+    type Context = Arc<Environment>;
+
+    async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
+        const MAX_BATCH_SIZE: usize = 10;
+
+        let access_key = required_var("CDN_LOG_QUEUE_ACCESS_KEY")?;
+        let secret_key = required_var("CDN_LOG_QUEUE_SECRET_KEY")?;
+        let queue_url = required_var("CDN_LOG_QUEUE_URL")?;
+        let region = required_var("CDN_LOG_QUEUE_REGION")?;
+
+        let credentials = Credentials::from_keys(access_key, secret_key, None);
+
+        let config = aws_sdk_sqs::Config::builder()
+            .credentials_provider(credentials)
+            .region(Region::new(region))
+            .behavior_version(BehaviorVersion::v2023_11_09())
+            .build();
+
+        let client = aws_sdk_sqs::Client::from_conf(config);
+
+        info!("Receiving messages from the CDN log queue…");
+        let mut num_remaining = self.max_messages;
+        while num_remaining > 0 {
+            let batch_size = num_remaining.min(MAX_BATCH_SIZE);
+            num_remaining -= batch_size;
+
+            debug!("Receiving next {batch_size} messages from the CDN log queue…");
+            let response = client
+                .receive_message()
+                .queue_url(&queue_url)
+                .max_number_of_messages(batch_size as i32)
+                .send()
+                .await?;
+
+            let messages = response.messages();
+            debug!(
+                "Received {num_messages} messages from the CDN log queue",
+                num_messages = messages.len()
+            );
+            if messages.is_empty() {
+                info!("No more messages to receive from the CDN log queue");
+                break;
+            }
+
+            for message in messages {
+                let message_id = message.message_id().unwrap_or("<unknown>");
+                debug!("Processing message: {message_id}");
+
+                let Some(receipt_handle) = message.receipt_handle() else {
+                    warn!("Message {message_id} has no receipt handle; skipping");
+                    continue;
+                };
+
+                debug!("Deleting message {message_id} from the CDN log queue…");
+                client
+                    .delete_message()
+                    .queue_url(&queue_url)
+                    .receipt_handle(receipt_handle)
+                    .send()
+                    .await
+                    .with_context(|| {
+                        format!("Failed to delete message {message_id} from the CDN log queue")
+                    })?;
+
+                let Some(body) = message.body() else {
+                    warn!("Message {message_id} has no body; skipping");
+                    continue;
+                };
+
+                let message = match serde_json::from_str::<message::Message>(body) {
+                    Ok(message) => message,
+                    Err(err) => {
+                        warn!("Failed to parse message {message_id}: {err}");
+                        continue;
+                    }
+                };
+
+                if message.records.is_empty() {
+                    warn!("Message {message_id} has no records; skipping");
+                    continue;
+                }
+
+                let pool = ctx.connection_pool.clone();
+                spawn_blocking({
+                    let message_id = message_id.to_owned();
+                    move || {
+                        let mut conn = pool
+                            .get()
+                            .context("Failed to acquire database connection")?;
+
+                        for record in message.records {
+                            let region = record.aws_region;
+                            let bucket = record.s3.bucket.name;
+                            let path = record.s3.object.key;
+
+                            let path = match object_store::path::Path::from_url_path(&path) {
+                                Ok(path) => path,
+                                Err(err) => {
+                                    warn!("Failed to parse path ({path}): {err}");
+                                    continue;
+                                }
+                            };
+
+                            info!("Enqueuing processing job for message {message_id}… ({path})");
+                            let job = ProcessCdnLog::new(region, bucket, path.as_ref().to_owned());
+
+                            job.enqueue(&mut conn).with_context(|| {
+                                format!("Failed to enqueue processing job for message {message_id}")
+                            })?;
+
+                            debug!("Enqueued processing job for message {message_id}");
+                        }
+
+                        Ok::<_, anyhow::Error>(())
+                    }
+                })
+                .await?;
+
+                debug!("Processed message: {message_id}");
+            }
+        }
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/worker/jobs/downloads.rs

@@ -0,0 +1,3 @@
+{
+    "Records": []
+}

src/worker/jobs/downloads/fixtures/empty-event.json

@@ -0,0 +1,72 @@
+{
+    "Records": [
+        {
+            "eventVersion": "2.1",
+            "eventSource": "aws:s3",
+            "awsRegion": "us-west-1",
+            "eventTime": "2024-01-16T16:24:06.554Z",
+            "eventName": "ObjectCreated:Put",
+            "userIdentity": {
+                "principalId": "AWS:AROASLZSKMR5WNLFCYRXL:prod.iad.dbs.datafeeds.aws.internal"
+            },
+            "requestParameters": {
+                "sourceIPAddress": "72.21.217.31"
+            },
+            "responseElements": {
+                "x-amz-request-id": "WK0MBXBCXJ4Y1KDX",
+                "x-amz-id-2": "Rwojcx6M4ZoUCE8gETkqTmOPO2zKbF7AFeZshk1/nYxdtng47pBPRwsSxbscaQTZu2PVbT9w+u9LTq2KVsoleencxqJ3vCex"
+            },
+            "s3": {
+                "s3SchemaVersion": "1.0",
+                "configurationId": "cloudfront",
+                "bucket": {
+                    "name": "rust-staging-crates-io-logs",
+                    "ownerIdentity": {
+                        "principalId": "A266IOH8VA1FN7"
+                    },
+                    "arn": "arn:aws:s3:::rust-staging-crates-io-logs"
+                },
+                "object": {
+                    "key": "cloudfront/index.staging.crates.io/E35K556QRQDZXW.2024-01-16-16.d01d5f13.gz",
+                    "size": 710,
+                    "eTag": "4c8266e632a8f1ed489c19821680da6e",
+                    "sequencer": "0065A6ADA6765BF0F0"
+                }
+            }
+        },
+        {
+            "eventVersion": "2.1",
+            "eventSource": "aws:s3",
+            "awsRegion": "us-west-1",
+            "eventTime": "2024-01-16T16:24:06.554Z",
+            "eventName": "ObjectCreated:Put",
+            "userIdentity": {
+                "principalId": "AWS:AROASLZSKMR5WNLFCYRXL:prod.iad.dbs.datafeeds.aws.internal"
+            },
+            "requestParameters": {
+                "sourceIPAddress": "72.21.217.31"
+            },
+            "responseElements": {
+                "x-amz-request-id": "WK0MBXBCXJ4Y1KDX",
+                "x-amz-id-2": "Rwojcx6M4ZoUCE8gETkqTmOPO2zKbF7AFeZshk1/nYxdtng47pBPRwsSxbscaQTZu2PVbT9w+u9LTq2KVsoleencxqJ3vCex"
+            },
+            "s3": {
+                "s3SchemaVersion": "1.0",
+                "configurationId": "cloudfront",
+                "bucket": {
+                    "name": "rust-staging-crates-io-logs",
+                    "ownerIdentity": {
+                        "principalId": "A266IOH8VA1FN7"
+                    },
+                    "arn": "arn:aws:s3:::rust-staging-crates-io-logs"
+                },
+                "object": {
+                    "key": "cloudfront/index.staging.crates.io/E35K556QRQDZXW.2024-01-16-16.d01d5f13.gz",
+                    "size": 710,
+                    "eTag": "4c8266e632a8f1ed489c19821680da6e",
+                    "sequencer": "0065A6ADA6765BF0F0"
+                }
+            }
+        }
+    ]
+}

src/worker/jobs/downloads/fixtures/multi-event.json

@@ -0,0 +1,38 @@
+{
+    "Records": [
+        {
+            "eventVersion": "2.1",
+            "eventSource": "aws:s3",
+            "awsRegion": "us-west-1",
+            "eventTime": "2024-01-16T16:24:06.554Z",
+            "eventName": "ObjectCreated:Put",
+            "userIdentity": {
+                "principalId": "AWS:AROASLZSKMR5WNLFCYRXL:prod.iad.dbs.datafeeds.aws.internal"
+            },
+            "requestParameters": {
+                "sourceIPAddress": "72.21.217.31"
+            },
+            "responseElements": {
+                "x-amz-request-id": "WK0MBXBCXJ4Y1KDX",
+                "x-amz-id-2": "Rwojcx6M4ZoUCE8gETkqTmOPO2zKbF7AFeZshk1/nYxdtng47pBPRwsSxbscaQTZu2PVbT9w+u9LTq2KVsoleencxqJ3vCex"
+            },
+            "s3": {
+                "s3SchemaVersion": "1.0",
+                "configurationId": "cloudfront",
+                "bucket": {
+                    "name": "rust-staging-crates-io-logs",
+                    "ownerIdentity": {
+                        "principalId": "A266IOH8VA1FN7"
+                    },
+                    "arn": "arn:aws:s3:::rust-staging-crates-io-logs"
+                },
+                "object": {
+                    "key": "cloudfront/index.staging.crates.io/E35K556QRQDZXW.2024-01-16-16.d01d5f13.gz",
+                    "size": 710,
+                    "eTag": "4c8266e632a8f1ed489c19821680da6e",
+                    "sequencer": "0065A6ADA6765BF0F0"
+                }
+            }
+        }
+    ]
+}