prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <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.
  
  </issue>
  <code>
  [start of atmo/clusters/tasks.py]
  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 of atmo/clusters/tasks.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: 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)
  
verification_info: {"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"}]}