Use FOR KEY SHARE UPDATE instead of FOR UPDATE (apache#42082)

Generally speaking, the goal of our SELECT FOR UPDATE statements is concurrency control.  We want to make sure that nothing else is handling the dag run (or ti) at the same time.  In no case that I'm aware of, are we actually going to update the PK value for the record.  So we can take the weaker lock, and still have that concurrency control.  But this only works on postgres.  Mysql has no such weaker for update lock.

One negative consequence of taking the stronger lock, is it blocks any insert into a table of a row that references the locked row in a FK.  This can cause slowdowns and pileups in some circumstances.  In one case we do an extra FOR UPDATE lock because that FK wait issue was associated with a deadlock. This change makes that extra unnecessary in the postgres case so we remove it here.  And with the mysql case in mind, we add a commit to avoid the deadlocking.

Author: dstandish

airflow/models/taskinstance.py

@@ -1673,8 +1673,6 @@ def _handle_reschedule(
 
     ti = _coalesce_to_orm_ti(ti=ti, session=session)
 
-    from airflow.models.dagrun import DagRun  # Avoid circular import
-
     ti.refresh_from_db(session)
 
     if TYPE_CHECKING:
@@ -1683,16 +1681,16 @@ def _handle_reschedule(
     ti.end_date = timezone.utcnow()
     ti.set_duration()
 
-    # Lock DAG run to be sure not to get into a deadlock situation when trying to insert
-    # TaskReschedule which apparently also creates lock on corresponding DagRun entity
-    with_row_locks(
-        session.query(DagRun).filter_by(
-            dag_id=ti.dag_id,
-            run_id=ti.run_id,
-        ),
-        session=session,
-    ).one()
-    # Log reschedule request
+    # set state
+    ti.state = TaskInstanceState.UP_FOR_RESCHEDULE
+
+    ti.clear_next_method_args()
+
+    session.merge(ti)
+    session.commit()
+
+    # we add this in separate commit to reduce likelihood of deadlock
+    # see https://github.com/apache/airflow/pull/21362 for more info
     session.add(
         TaskReschedule(
             ti.task_id,
@@ -1705,13 +1703,6 @@ def _handle_reschedule(
             ti.map_index,
         )
     )
-
-    # set state
-    ti.state = TaskInstanceState.UP_FOR_RESCHEDULE
-
-    ti.clear_next_method_args()
-
-    session.merge(ti)
     session.commit()
     return ti
 

airflow/utils/sqlalchemy.py

@@ -312,6 +312,7 @@ def with_row_locks(
     *,
     nowait: bool = False,
     skip_locked: bool = False,
+    key_share: bool = True,
     **kwargs,
 ) -> Query:
     """
@@ -329,6 +330,7 @@ def with_row_locks(
     :param session: ORM Session
     :param nowait: If set to True, will pass NOWAIT to supported database backends.
     :param skip_locked: If set to True, will pass SKIP LOCKED to supported database backends.
+    :param key_share: If true, will lock with FOR KEY SHARE UPDATE (at least on postgres).
     :param kwargs: Extra kwargs to pass to with_for_update (of, nowait, skip_locked, etc)
     :return: updated query
     """
@@ -343,6 +345,8 @@ def with_row_locks(
         kwargs["nowait"] = True
     if skip_locked:
         kwargs["skip_locked"] = True
+    if key_share:
+        kwargs["key_share"] = True
     return query.with_for_update(**kwargs)
 
 

tests/utils/test_sqlalchemy.py

@@ -147,7 +147,7 @@ def test_with_row_locks(
             returned_value = with_row_locks(query=query, session=session, nowait=True)
 
         if expected_use_row_level_lock:
-            query.with_for_update.assert_called_once_with(nowait=True)
+            query.with_for_update.assert_called_once_with(nowait=True, key_share=True)
         else:
             assert returned_value == query
             query.with_for_update.assert_not_called()