Repository: apache/airflow
File: task-sdk/src/airflow/sdk/execution_time/task_runner.py (Python)

Code before:
        try:
            get_listener_manager().hook.on_task_instance_failed(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
            )
        except Exception:
            log.exception("error calling listener")
        if error and task.email_on_failure and task.email:
            _send_error_email_notification(task, ti, context, error, log)

    try:
        get_listener_manager().hook.before_stopping(component=TaskRunnerMarker())
    except Exception:
        log.exception("error calling listener")


def main():
    log = structlog.get_logger(logger_name="task")

    global SUPERVISOR_COMMS
    SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log)

    try:
        try:
            ti, context, log = startup()
        except AirflowRescheduleException as reschedule:
            log.info("Rescheduling task during startup, marking task as UP_FOR_RESCHEDULE")
            SUPERVISOR_COMMS.send(
                msg=RescheduleTask(
                    reschedule_date=reschedule.reschedule_date,
                    end_date=datetime.now(tz=timezone.utc),
                )
            )
            sys.exit(0)
        with BundleVersionLock(
            bundle_name=ti.bundle_instance.name,
            bundle_version=ti.bundle_instance.version,
        ):
            state, _, error = run(ti, context, log)
            context["exception"] = error
            finalize(ti, state, context, log, error)
    except KeyboardInterrupt:
        log.exception("Ctrl-c hit")
        sys.exit(2)
    except Exception:
        log.exception("Top level error")
        sys.exit(1)
    finally:
        # Ensure the request socket is closed on the child side in all circumstances
        # before the process fully terminates.
        if SUPERVISOR_COMMS and SUPERVISOR_COMMS.socket:
            with suppress(Exception):

Review comment (suggestion):
```suggestion
            log.warning("Rescheduling task during startup, marking task as UP_FOR_RESCHEDULE")
```

Code after:
    _xcom_push(ti, BaseXCom.XCOM_RETURN_KEY, result, mapped_length=mapped_length)


def finalize(
    ti: RuntimeTaskInstance,
    state: TaskInstanceState,
    context: Context,
    log: Logger,
    error: BaseException | None = None,
):
    # Record task duration metrics for all terminal states
    if ti.start_date and ti.end_date:
        duration_ms = (ti.end_date - ti.start_date).total_seconds() * 1000
        stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}

        Stats.timing(f"dag.{ti.dag_id}.{ti.task_id}.duration", duration_ms)
        Stats.timing("task.duration", duration_ms, tags=stats_tags)

    task = ti.task
    # Pushing xcom for each operator extra links defined on the operator only.
    for oe in task.operator_extra_links:
        try:
            link, xcom_key = oe.get_link(operator=task, ti_key=ti), oe.xcom_key  # type: ignore[arg-type]
            log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
            _xcom_push_to_db(ti, key=xcom_key, value=link)
        except Exception:
            log.exception(
                "Failed to push an xcom for task operator extra link",
                link_name=oe.name,
                xcom_key=oe.xcom_key,
                ti=ti,
            )

    if getattr(ti.task, "overwrite_rtif_after_execution", False):
        log.debug("Overwriting Rendered template fields.")
        if ti.task.template_fields:
            SUPERVISOR_COMMS.send(SetRenderedFields(rendered_fields=_serialize_rendered_fields(ti.task)))

    log.debug("Running finalizers", ti=ti)
    if state == TaskInstanceState.SUCCESS:
        _run_task_state_change_callbacks(task, "on_success_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_success(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti
            )
        except Exception:
            log.exception("error calling listener")
    elif state == TaskInstanceState.SKIPPED:
        _run_task_state_change_callbacks(task, "on_skipped_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_skipped(