We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Use `select_for_update` in cluster expiration email task
There's a bit of a race condition [here](https://github.com/mozilla/telemetry-analysis-service/blob/e6fecbe12d09b2e2338ae62f5276b3b2f39b0b65/atmo/clusters/tasks.py#L38) where two tasks could ask for expiring clusters at the same time and start sending emails before they can be marked as sent. We should 1. wrap the whole task in a transaction and 2. use `select_for_update` on the query for active clusters to lock those rows.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `atmo/clusters/tasks.py`
Content:
```
1 # This Source Code Form is subject to the terms of the Mozilla Public
2 # License, v. 2.0. If a copy of the MPL was not distributed with this
3 # file, you can obtain one at http://mozilla.org/MPL/2.0/.
4 from datetime import timedelta
5 
6 import mail_builder
7 from botocore.exceptions import ClientError
8 from celery.utils.log import get_task_logger
9 from django.conf import settings
10 from django.db import transaction
11 from django.utils import timezone
12 
13 from ..celery import celery
14 from .models import Cluster
15 from .provisioners import ClusterProvisioner
16 
17 logger = get_task_logger(__name__)
18 
19 
20 @celery.task
21 def deactivate_clusters():
22     now = timezone.now()
23     deactivated_clusters = []
24     for cluster in Cluster.objects.active().filter(end_date__lte=now):
25         with transaction.atomic():
26             deactivated_clusters.append([cluster.identifier, cluster.pk])
27             # The cluster is expired
28             logger.info(
29                 'Cluster %s (%s) is expired, deactivating.',
30                 cluster.pk,
31                 cluster.identifier,
32             )
33             cluster.deactivate()
34     return deactivated_clusters
35 
36 
37 @celery.task
38 def send_expiration_mails():
39     deadline = timezone.now() + timedelta(hours=1)
40     soon_expired = Cluster.objects.active().filter(
41         end_date__lte=deadline,
42         expiration_mail_sent=False,
43     )
44     for cluster in soon_expired:
45         with transaction.atomic():
46             message = mail_builder.build_message(
47                 'atmo/clusters/mails/expiration.mail', {
48                     'cluster': cluster,
49                     'deadline': deadline,
50                     'settings': settings,
51                 },
52             )
53             message.send()
54             cluster.expiration_mail_sent = True
55             cluster.save()
56 
57 
58 @celery.task(max_retries=3)
59 @celery.autoretry(ClientError)
60 def update_master_address(cluster_id, force=False):
61     """
62     Update the public IP address for the cluster with the given cluster ID
63     """
64     cluster = Cluster.objects.get(id=cluster_id)
65     # quick way out in case this job was called accidently
66     if cluster.master_address and not force:
67         return
68     # first get the cluster info from AWS
69     info = cluster.info
70     master_address = info.get('public_dns') or ''
71     # then store the public IP of the cluster if found in response
72     if master_address:
73         cluster.master_address = master_address
74         cluster.save()
75         return master_address
76 
77 
78 @celery.task(max_retries=3)
79 @celery.autoretry(ClientError)
80 def update_clusters():
81     """
82     Update the cluster metadata from AWS for the pending
83     clusters.
84 
85     - To be used periodically.
86     - Won't update state if not needed.
87     - Will queue updating the Cluster's public IP address if needed.
88     """
89     # only update the cluster info for clusters that are pending
90     active_clusters = Cluster.objects.active()
91 
92     # Short-circuit for no active clusters (e.g. on weekends)
93     if not active_clusters.exists():
94         return []
95 
96     # get the start dates of the active clusters, set to the start of the day
97     # to counteract time differences between atmo and AWS and use the oldest
98     # start date to limit the ListCluster API call to AWS
99     oldest_start_date = active_clusters.datetimes('start_date', 'day')
100 
101     # build a mapping between jobflow ID and cluster info
102     cluster_mapping = {}
103     provisioner = ClusterProvisioner()
104     cluster_list = provisioner.list(
105         created_after=oldest_start_date[0]
106     )
107     for cluster_info in cluster_list:
108         cluster_mapping[cluster_info['jobflow_id']] = cluster_info
109 
110     # go through pending clusters and update the state if needed
111     updated_clusters = []
112     for cluster in active_clusters:
113         with transaction.atomic():
114             info = cluster_mapping.get(cluster.jobflow_id)
115             # ignore if no info was found for some reason,
116             # the cluster was deleted in AWS but it wasn't deleted here yet
117             if info is None:
118                 continue
119 
120             # don't update the state if it's equal to the already stored state
121             if info['state'] == cluster.most_recent_status:
122                 continue
123 
124             # run an UPDATE query for the cluster
125             cluster.most_recent_status = info['state']
126             cluster.save()
127 
128             updated_clusters.append(cluster.identifier)
129 
130             # if not given enqueue a job to update the public IP address
131             # but only if the cluster is running or waiting, so the
132             # API call isn't wasted
133             if (not cluster.master_address and
134                     cluster.most_recent_status in cluster.READY_STATUS_LIST):
135                 transaction.on_commit(
136                     lambda: update_master_address.delay(cluster.id)
137                 )
138     return updated_clusters
139
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/atmo/clusters/tasks.py b/atmo/clusters/tasks.py
--- a/atmo/clusters/tasks.py
+++ b/atmo/clusters/tasks.py
@@ -37,22 +37,23 @@
 @celery.task
 def send_expiration_mails():
     deadline = timezone.now() + timedelta(hours=1)
-    soon_expired = Cluster.objects.active().filter(
-        end_date__lte=deadline,
-        expiration_mail_sent=False,
-    )
-    for cluster in soon_expired:
-        with transaction.atomic():
-            message = mail_builder.build_message(
-                'atmo/clusters/mails/expiration.mail', {
-                    'cluster': cluster,
-                    'deadline': deadline,
-                    'settings': settings,
-                },
-            )
-            message.send()
-            cluster.expiration_mail_sent = True
-            cluster.save()
+    with transaction.atomic():
+        soon_expired = Cluster.objects.select_for_update().active().filter(
+            end_date__lte=deadline,
+            expiration_mail_sent=False,
+        )
+        for cluster in soon_expired:
+            with transaction.atomic():
+                message = mail_builder.build_message(
+                    'atmo/clusters/mails/expiration.mail', {
+                        'cluster': cluster,
+                        'deadline': deadline,
+                        'settings': settings,
+                    },
+                )
+                message.send()
+                cluster.expiration_mail_sent = True
+                cluster.save()
 
 
 @celery.task(max_retries=3)

{"golden_diff": "diff --git a/atmo/clusters/tasks.py b/atmo/clusters/tasks.py\n--- a/atmo/clusters/tasks.py\n+++ b/atmo/clusters/tasks.py\n@@ -37,22 +37,23 @@\n @celery.task\n def send_expiration_mails():\n     deadline = timezone.now() + timedelta(hours=1)\n-    soon_expired = Cluster.objects.active().filter(\n-        end_date__lte=deadline,\n-        expiration_mail_sent=False,\n-    )\n-    for cluster in soon_expired:\n-        with transaction.atomic():\n-            message = mail_builder.build_message(\n-                'atmo/clusters/mails/expiration.mail', {\n-                    'cluster': cluster,\n-                    'deadline': deadline,\n-                    'settings': settings,\n-                },\n-            )\n-            message.send()\n-            cluster.expiration_mail_sent = True\n-            cluster.save()\n+    with transaction.atomic():\n+        soon_expired = Cluster.objects.select_for_update().active().filter(\n+            end_date__lte=deadline,\n+            expiration_mail_sent=False,\n+        )\n+        for cluster in soon_expired:\n+            with transaction.atomic():\n+                message = mail_builder.build_message(\n+                    'atmo/clusters/mails/expiration.mail', {\n+                        'cluster': cluster,\n+                        'deadline': deadline,\n+                        'settings': settings,\n+                    },\n+                )\n+                message.send()\n+                cluster.expiration_mail_sent = True\n+                cluster.save()\n \n \n @celery.task(max_retries=3)\n", "issue": "Use `select_for_update` in cluster expiration email task\nThere's a bit of a race condition [here](https://github.com/mozilla/telemetry-analysis-service/blob/e6fecbe12d09b2e2338ae62f5276b3b2f39b0b65/atmo/clusters/tasks.py#L38) where two tasks could ask for expiring clusters at the same time and start sending emails before they can be marked as sent. We should 1. wrap the whole task in a transaction and 2. use `select_for_update` on the query for active clusters to lock those rows.\n", "before_files": [{"content": "# This Source Code Form is subject to the terms of the Mozilla Public\n# License, v. 2.0. If a copy of the MPL was not distributed with this\n# file, you can obtain one at http://mozilla.org/MPL/2.0/.\nfrom datetime import timedelta\n\nimport mail_builder\nfrom botocore.exceptions import ClientError\nfrom celery.utils.log import get_task_logger\nfrom django.conf import settings\nfrom django.db import transaction\nfrom django.utils import timezone\n\nfrom ..celery import celery\nfrom .models import Cluster\nfrom .provisioners import ClusterProvisioner\n\nlogger = get_task_logger(__name__)\n\n\n@celery.task\ndef deactivate_clusters():\n    now = timezone.now()\n    deactivated_clusters = []\n    for cluster in Cluster.objects.active().filter(end_date__lte=now):\n        with transaction.atomic():\n            deactivated_clusters.append([cluster.identifier, cluster.pk])\n            # The cluster is expired\n            logger.info(\n                'Cluster %s (%s) is expired, deactivating.',\n                cluster.pk,\n                cluster.identifier,\n            )\n            cluster.deactivate()\n    return deactivated_clusters\n\n\n@celery.task\ndef send_expiration_mails():\n    deadline = timezone.now() + timedelta(hours=1)\n    soon_expired = Cluster.objects.active().filter(\n        end_date__lte=deadline,\n        expiration_mail_sent=False,\n    )\n    for cluster in soon_expired:\n        with transaction.atomic():\n            message = mail_builder.build_message(\n                'atmo/clusters/mails/expiration.mail', {\n                    'cluster': cluster,\n                    'deadline': deadline,\n                    'settings': settings,\n                },\n            )\n            message.send()\n            cluster.expiration_mail_sent = True\n            cluster.save()\n\n\n@celery.task(max_retries=3)\n@celery.autoretry(ClientError)\ndef update_master_address(cluster_id, force=False):\n    \"\"\"\n    Update the public IP address for the cluster with the given cluster ID\n    \"\"\"\n    cluster = Cluster.objects.get(id=cluster_id)\n    # quick way out in case this job was called accidently\n    if cluster.master_address and not force:\n        return\n    # first get the cluster info from AWS\n    info = cluster.info\n    master_address = info.get('public_dns') or ''\n    # then store the public IP of the cluster if found in response\n    if master_address:\n        cluster.master_address = master_address\n        cluster.save()\n        return master_address\n\n\n@celery.task(max_retries=3)\n@celery.autoretry(ClientError)\ndef update_clusters():\n    \"\"\"\n    Update the cluster metadata from AWS for the pending\n    clusters.\n\n    - To be used periodically.\n    - Won't update state if not needed.\n    - Will queue updating the Cluster's public IP address if needed.\n    \"\"\"\n    # only update the cluster info for clusters that are pending\n    active_clusters = Cluster.objects.active()\n\n    # Short-circuit for no active clusters (e.g. on weekends)\n    if not active_clusters.exists():\n        return []\n\n    # get the start dates of the active clusters, set to the start of the day\n    # to counteract time differences between atmo and AWS and use the oldest\n    # start date to limit the ListCluster API call to AWS\n    oldest_start_date = active_clusters.datetimes('start_date', 'day')\n\n    # build a mapping between jobflow ID and cluster info\n    cluster_mapping = {}\n    provisioner = ClusterProvisioner()\n    cluster_list = provisioner.list(\n        created_after=oldest_start_date[0]\n    )\n    for cluster_info in cluster_list:\n        cluster_mapping[cluster_info['jobflow_id']] = cluster_info\n\n    # go through pending clusters and update the state if needed\n    updated_clusters = []\n    for cluster in active_clusters:\n        with transaction.atomic():\n            info = cluster_mapping.get(cluster.jobflow_id)\n            # ignore if no info was found for some reason,\n            # the cluster was deleted in AWS but it wasn't deleted here yet\n            if info is None:\n                continue\n\n            # don't update the state if it's equal to the already stored state\n            if info['state'] == cluster.most_recent_status:\n                continue\n\n            # run an UPDATE query for the cluster\n            cluster.most_recent_status = info['state']\n            cluster.save()\n\n            updated_clusters.append(cluster.identifier)\n\n            # if not given enqueue a job to update the public IP address\n            # but only if the cluster is running or waiting, so the\n            # API call isn't wasted\n            if (not cluster.master_address and\n                    cluster.most_recent_status in cluster.READY_STATUS_LIST):\n                transaction.on_commit(\n                    lambda: update_master_address.delay(cluster.id)\n                )\n    return updated_clusters\n", "path": "atmo/clusters/tasks.py"}], "after_files": [{"content": "# This Source Code Form is subject to the terms of the Mozilla Public\n# License, v. 2.0. If a copy of the MPL was not distributed with this\n# file, you can obtain one at http://mozilla.org/MPL/2.0/.\nfrom datetime import timedelta\n\nimport mail_builder\nfrom botocore.exceptions import ClientError\nfrom celery.utils.log import get_task_logger\nfrom django.conf import settings\nfrom django.db import transaction\nfrom django.utils import timezone\n\nfrom ..celery import celery\nfrom .models import Cluster\nfrom .provisioners import ClusterProvisioner\n\nlogger = get_task_logger(__name__)\n\n\n@celery.task\ndef deactivate_clusters():\n    now = timezone.now()\n    deactivated_clusters = []\n    for cluster in Cluster.objects.active().filter(end_date__lte=now):\n        with transaction.atomic():\n            deactivated_clusters.append([cluster.identifier, cluster.pk])\n            # The cluster is expired\n            logger.info(\n                'Cluster %s (%s) is expired, deactivating.',\n                cluster.pk,\n                cluster.identifier,\n            )\n            cluster.deactivate()\n    return deactivated_clusters\n\n\n@celery.task\ndef send_expiration_mails():\n    deadline = timezone.now() + timedelta(hours=1)\n    with transaction.atomic():\n        soon_expired = Cluster.objects.select_for_update().active().filter(\n            end_date__lte=deadline,\n            expiration_mail_sent=False,\n        )\n        for cluster in soon_expired:\n            with transaction.atomic():\n                message = mail_builder.build_message(\n                    'atmo/clusters/mails/expiration.mail', {\n                        'cluster': cluster,\n                        'deadline': deadline,\n                        'settings': settings,\n                    },\n                )\n                message.send()\n                cluster.expiration_mail_sent = True\n                cluster.save()\n\n\n@celery.task(max_retries=3)\n@celery.autoretry(ClientError)\ndef update_master_address(cluster_id, force=False):\n    \"\"\"\n    Update the public IP address for the cluster with the given cluster ID\n    \"\"\"\n    cluster = Cluster.objects.get(id=cluster_id)\n    # quick way out in case this job was called accidently\n    if cluster.master_address and not force:\n        return\n    # first get the cluster info from AWS\n    info = cluster.info\n    master_address = info.get('public_dns') or ''\n    # then store the public IP of the cluster if found in response\n    if master_address:\n        cluster.master_address = master_address\n        cluster.save()\n        return master_address\n\n\n@celery.task(max_retries=3)\n@celery.autoretry(ClientError)\ndef update_clusters():\n    \"\"\"\n    Update the cluster metadata from AWS for the pending\n    clusters.\n\n    - To be used periodically.\n    - Won't update state if not needed.\n    - Will queue updating the Cluster's public IP address if needed.\n    \"\"\"\n    # only update the cluster info for clusters that are pending\n    active_clusters = Cluster.objects.active()\n\n    # Short-circuit for no active clusters (e.g. on weekends)\n    if not active_clusters.exists():\n        return []\n\n    # get the start dates of the active clusters, set to the start of the day\n    # to counteract time differences between atmo and AWS and use the oldest\n    # start date to limit the ListCluster API call to AWS\n    oldest_start_date = active_clusters.datetimes('start_date', 'day')\n\n    # build a mapping between jobflow ID and cluster info\n    cluster_mapping = {}\n    provisioner = ClusterProvisioner()\n    cluster_list = provisioner.list(\n        created_after=oldest_start_date[0]\n    )\n    for cluster_info in cluster_list:\n        cluster_mapping[cluster_info['jobflow_id']] = cluster_info\n\n    # go through pending clusters and update the state if needed\n    updated_clusters = []\n    for cluster in active_clusters:\n        with transaction.atomic():\n            info = cluster_mapping.get(cluster.jobflow_id)\n            # ignore if no info was found for some reason,\n            # the cluster was deleted in AWS but it wasn't deleted here yet\n            if info is None:\n                continue\n\n            # don't update the state if it's equal to the already stored state\n            if info['state'] == cluster.most_recent_status:\n                continue\n\n            # run an UPDATE query for the cluster\n            cluster.most_recent_status = info['state']\n            cluster.save()\n\n            updated_clusters.append(cluster.identifier)\n\n            # if not given enqueue a job to update the public IP address\n            # but only if the cluster is running or waiting, so the\n            # API call isn't wasted\n            if (not cluster.master_address and\n                    cluster.most_recent_status in cluster.READY_STATUS_LIST):\n                transaction.on_commit(\n                    lambda: update_master_address.delay(cluster.id)\n                )\n    return updated_clusters\n", "path": "atmo/clusters/tasks.py"}]}