ingestion-edge/ingestion_edge/flush_manager.py

"""Continuously flush ingestion-edge queue from detached persistent volumes."""

from argparse import ArgumentParser
from dataclasses import dataclass
from datetime import datetime, timedelta
from functools import partial
from multiprocessing.pool import ThreadPool
from time import sleep
from typing import Callable, Dict, List
import json
import os

from kubernetes.config import load_incluster_config
from kubernetes.client import (
    BatchV1Api,
    CoreV1Api,
    V1Container,
    V1DeleteOptions,
    V1EnvVar,
    V1Job,
    V1JobSpec,
    V1ObjectMeta,
    V1ObjectReference,
    V1OwnerReference,
    V1PersistentVolume,
    V1PersistentVolumeClaim,
    V1PersistentVolumeClaimSpec,
    V1PersistentVolumeClaimVolumeSource,
    V1PersistentVolumeSpec,
    V1Pod,
    V1PodSpec,
    V1PodTemplateSpec,
    V1Preconditions,
    V1ResourceRequirements,
    V1Volume,
    V1VolumeMount,
)
from kubernetes.client.rest import ApiException

from .config import get_config_dict, logger

JOB_AND_PVC_PREFIX = "flush-"
ALREADY_EXISTS = "AlreadyExists"
CONFLICT = "Conflict"
NOT_FOUND = "Not Found"

DEFAULT_IMAGE_VERSION = "latest"
try:
    with open("version.json") as fp:
        DEFAULT_IMAGE_VERSION = str(
            json.load(fp).get("version") or DEFAULT_IMAGE_VERSION
        )
except (FileNotFoundError, json.decoder.JSONDecodeError):  # pragma: no cover
    pass  # use default

DEFAULT_ENV = [
    V1EnvVar(name=name, value=os.environ[name])
    for name in get_config_dict()
    if name in os.environ
]
FLUSH_JOB_REQUESTS_CPU = os.environ.get("FLUSH_JOB_REQUESTS_CPU", "100m")
FLUSH_JOB_REQUESTS_MEMORY = os.environ.get("FLUSH_JOB_REQUESTS_MEMORY", "100Mi")
FLUSH_JOB_LIMITS_CPU = os.environ.get("FLUSH_JOB_LIMITS_CPU", None)
FLUSH_JOB_LIMITS_MEMORY = os.environ.get("FLUSH_JOB_LIMITS_MEMORY", "500Mi")

parser = ArgumentParser(description=__doc__)
parser.add_argument(
    "--command",
    default=["python", "-m", "ingestion_edge.flush"],
    type=json.loads,
    help="Docker command for flush jobs; must drain queue until empty or exit non-zero",
)
parser.add_argument(
    "--env",
    default=DEFAULT_ENV,
    type=json.loads,
    help="env to use for flush jobs",
)
parser.add_argument(
    "--image",
    default="mozilla/ingestion-edge:" + DEFAULT_IMAGE_VERSION,
    help="Docker image to use for flush jobs",
)
parser.add_argument(
    "--namespace",
    default="default",
    help="Kubernetes namespace to use",
)
parser.add_argument(
    "--service-account-name",
    default="default",
    help="Kubernetes service account name to use",
)
parser.add_argument(
    "--claim-prefix",
    default="queue-",
    help="Prefix for the names of persistent volume claims to delete",
)
parser.add_argument(
    "--pvc-cleanup-delay-seconds",
    default=60,
    type=int,
    help="Number of seconds to wait for persistent volume claims to remain detached "
    "before deleting",
)


@dataclass(frozen=True)
class PvcCacheEntry:
    """Cache entry for tracking detached persistent volume claims."""

    pv: str
    time: datetime


def _job_and_pvc_name_from_pv(pv: V1PersistentVolume) -> str:
    return JOB_AND_PVC_PREFIX + pv.metadata.name


def _pv_name_from_job(job: V1Job) -> str:
    return job.metadata.name.replace(JOB_AND_PVC_PREFIX, "", 1)


def _is_flush_job(job: V1Job) -> bool:
    return job.metadata.name.startswith(JOB_AND_PVC_PREFIX)


def _create_pvc(
    api: CoreV1Api, name: str, namespace: str, pv: V1PersistentVolume
) -> V1PersistentVolumeClaim:
    logger.info(f"creating pvc: {name}")
    try:
        return api.create_namespaced_persistent_volume_claim(
            namespace=namespace,
            body=V1PersistentVolumeClaim(
                api_version="v1",
                kind="PersistentVolumeClaim",
                metadata=V1ObjectMeta(name=name, namespace=namespace),
                spec=V1PersistentVolumeClaimSpec(
                    access_modes=["ReadWriteOnce"],
                    resources=V1ResourceRequirements(requests=pv.spec.capacity),
                    storage_class_name=pv.spec.storage_class_name,
                    volume_name=pv.metadata.name,
                ),
            ),
        )
    except ApiException as e:
        if e.reason == CONFLICT and json.loads(e.body)["reason"] == ALREADY_EXISTS:
            logger.info(f"using existing pvc: {name}")
            return api.read_namespaced_persistent_volume_claim(name, namespace)
        raise


def _bind_pvc(
    api: CoreV1Api, pv: V1PersistentVolume, pvc: V1PersistentVolumeClaim
) -> V1PersistentVolume:
    logger.info(f"binding pv to pvc: {pv.metadata.name}, {pvc.metadata.name}")
    return api.patch_persistent_volume(
        name=pv.metadata.name,
        body=V1PersistentVolume(
            spec=V1PersistentVolumeSpec(
                claim_ref=V1ObjectReference(
                    api_version="v1",
                    kind="PersistentVolumeClaim",
                    name=pvc.metadata.name,
                    namespace=pvc.metadata.namespace,
                    resource_version=pvc.metadata.resource_version,
                    uid=pvc.metadata.uid,
                )
            )
        ),
    )


def _create_flush_job(
    batch_api: BatchV1Api,
    command: List[str],
    env: List[V1EnvVar],
    image: str,
    name: str,
    namespace: str,
    service_account_name: str,
) -> V1Job:
    logger.info(f"creating job: {name}")
    try:
        return batch_api.create_namespaced_job(
            namespace=namespace,
            body=V1Job(
                api_version="batch/v1",
                kind="Job",
                metadata=V1ObjectMeta(name=name, namespace=namespace),
                spec=V1JobSpec(
                    template=V1PodTemplateSpec(
                        spec=V1PodSpec(
                            containers=[
                                V1Container(
                                    image=image,
                                    command=command,
                                    name="flush",
                                    volume_mounts=[
                                        V1VolumeMount(mount_path="/data", name="queue")
                                    ],
                                    env=env,
                                    resources=V1ResourceRequirements(
                                        requests={
                                            key: value
                                            for key, value in (
                                                ("cpu", FLUSH_JOB_REQUESTS_CPU),
                                                ("memory", FLUSH_JOB_REQUESTS_MEMORY),
                                            )
                                            if value
                                        },
                                        limits={
                                            key: value
                                            for key, value in (
                                                ("cpu", FLUSH_JOB_LIMITS_CPU),
                                                ("memory", FLUSH_JOB_LIMITS_MEMORY),
                                            )
                                            if value
                                        },
                                    ),
                                )
                            ],
                            restart_policy="OnFailure",
                            volumes=[
                                V1Volume(
                                    name="queue",
                                    persistent_volume_claim=(
                                        V1PersistentVolumeClaimVolumeSource(
                                            claim_name=name
                                        )
                                    ),
                                )
                            ],
                            service_account_name=service_account_name,
                        )
                    )
                ),
            ),
        )
    except ApiException as e:
        if e.reason == CONFLICT and json.loads(e.body)["reason"] == ALREADY_EXISTS:
            logger.info(f"using existing job: {name}")
            return batch_api.read_namespaced_job(name, namespace)
        raise


def flush_released_pvs(
    api: CoreV1Api,
    batch_api: BatchV1Api,
    command: List[str],
    env: List[V1EnvVar],
    image: str,
    namespace: str,
    service_account_name: str,
):
    """
    Flush persistent volumes.

    Gracefully handle resuming after an interruption, because this is not atomic.
    """
    existing_jobs = {
        job.metadata.name for job in batch_api.list_namespaced_job(namespace).items
    }
    for pv in api.list_persistent_volume().items:
        name = _job_and_pvc_name_from_pv(pv)
        if (
            name not in existing_jobs
            and pv.spec.claim_ref
            and pv.spec.claim_ref.namespace == namespace
            and pv.spec.persistent_volume_reclaim_policy != "Delete"
            and pv.status
            and (pv.status.phase == "Released" or pv.spec.claim_ref.name == name)
        ):
            logger.info(f"flushing unbound pv: {pv.metadata.name}")
            if pv.status.phase != "Bound":
                pvc = _create_pvc(api, name, namespace, pv)
                _bind_pvc(api, pv, pvc)
            _create_flush_job(
                batch_api, command, env, image, name, namespace, service_account_name
            )


def delete_complete_jobs(api: CoreV1Api, batch_api: BatchV1Api, namespace: str):
    """Delete complete jobs."""
    for job in batch_api.list_namespaced_job(namespace).items:
        if (
            job.status.conditions
            and job.status.conditions[0].type == "Complete"
            and not job.metadata.deletion_timestamp
            and _is_flush_job(job)
        ):
            logger.info(f"deleting complete job: {job.metadata.name}")
            # configure persistent volume claims to be deleted with the job
            pv_name = _pv_name_from_job(job)
            logger.info(f"including pv in pvc delete: {pv_name}")
            api.patch_persistent_volume(
                name=pv_name,
                body=V1PersistentVolume(
                    spec=V1PersistentVolumeSpec(
                        persistent_volume_reclaim_policy="Delete",
                    )
                ),
            )
            logger.info(f"including pvc in job delete: {job.metadata.name}")
            api.patch_namespaced_persistent_volume_claim(
                name=job.metadata.name,
                namespace=namespace,
                body=V1PersistentVolumeClaim(
                    metadata=V1ObjectMeta(
                        owner_references=[
                            V1OwnerReference(
                                api_version="batch/v1",
                                kind="Job",
                                name=job.metadata.name,
                                uid=job.metadata.uid,
                                block_owner_deletion=True,
                            )
                        ]
                    )
                ),
            )
            try:
                batch_api.delete_namespaced_job(
                    name=job.metadata.name,
                    namespace=namespace,
                    body=V1DeleteOptions(
                        grace_period_seconds=0,
                        propagation_policy="Foreground",
                        preconditions=V1Preconditions(
                            resource_version=job.metadata.resource_version,
                            uid=job.metadata.uid,
                        ),
                    ),
                )
            except ApiException as e:
                if e.reason not in (CONFLICT, NOT_FOUND):
                    raise
                logger.info(f"job already deleted or updated: {job.metadata.name}")


def flush_released_pvs_and_delete_complete_jobs(
    api: CoreV1Api,
    batch_api: BatchV1Api,
    command: List[str],
    env: List[V1EnvVar],
    image: str,
    namespace: str,
    service_account_name: str,
):
    """Flush released persistent volumes then delete complete jobs.

    Run sequentially to avoid race conditions.
    """
    flush_released_pvs(
        api, batch_api, command, env, image, namespace, service_account_name
    )
    delete_complete_jobs(api, batch_api, namespace)


def _unschedulable_due_to_pvc(pod: V1Pod):
    return (
        pod.status
        and pod.status.phase == "Pending"
        and any(
            condition.reason == "Unschedulable"
            and condition.message
            and condition.message.startswith('persistentvolumeclaim "')
            and condition.message.endswith('" not found')
            for condition in (pod.status.conditions or [])
            if condition
        )
        and pod.metadata.owner_references
        and any(ref.kind == "StatefulSet" for ref in pod.metadata.owner_references)
    )


def delete_detached_pvcs(
    api: CoreV1Api,
    namespace: str,
    claim_prefix: str,
    pvc_cleanup_delay: timedelta,
    detached_pvc_cache: Dict[str, PvcCacheEntry],
):
    """
    Delete persistent volume claims that are not attached to any pods.

    If a persistent volume claim is deleted while attached to a pod, then the
    underlying persistent volume will remain bound until the delete is
    finalized, and the delete will not be finalized until the pod is also
    deleted.

    If a stateful set immediately recreates a pod (e.g. via `kubectl rollout
    restart`) that was attached to a persistent volume claim that was deleted,
    then the stateful set may still try to reuse the persistent volume claim
    after the delete is finalized. Delete the pod again to cause the stateful
    set to recreate the persistent volume claim when it next recreates the pod.
    """
    attached_pvcs = {
        volume.persistent_volume_claim.claim_name
        for pod in api.list_namespaced_pod(namespace).items
        if not _unschedulable_due_to_pvc(pod) and pod.spec and pod.spec.volumes
        for volume in pod.spec.volumes
        if volume.persistent_volume_claim
    }
    for pvc in api.list_namespaced_persistent_volume_claim(namespace).items:
        if (
            pvc.metadata.name.startswith(claim_prefix)
            and pvc.metadata.name not in attached_pvcs
            and not pvc.metadata.deletion_timestamp
        ):
            name, pv, now = pvc.metadata.name, pvc.spec.volume_name, datetime.utcnow()
            if name not in detached_pvc_cache or detached_pvc_cache[name].pv != pv:
                logger.info(f"found newly detached pvc: {pvc.metadata.name}")
                detached_pvc_cache[name] = PvcCacheEntry(pv, now)
            if (now - detached_pvc_cache[name].time) < pvc_cleanup_delay:
                # wait for pvc to remain detached for pvc_cleanup_delay before deleting
                continue
            logger.info(f"deleting detached pvc: {pvc.metadata.name}")
            try:
                api.delete_namespaced_persistent_volume_claim(
                    name=pvc.metadata.name,
                    namespace=namespace,
                    body=V1DeleteOptions(
                        grace_period_seconds=0,
                        propagation_policy="Background",
                        preconditions=V1Preconditions(
                            resource_version=pvc.metadata.resource_version,
                            uid=pvc.metadata.uid,
                        ),
                    ),
                )
            except ApiException as e:
                if e.reason not in (CONFLICT, NOT_FOUND):
                    raise
                logger.info(f"pvc already deleted or updated: {pvc.metadata.name}")
        else:
            # pvc is not detached, drop from cache if present
            detached_pvc_cache.pop(pvc.metadata.name, None)


def delete_unschedulable_pods(api: CoreV1Api, namespace: str):
    """
    Delete pods that are unschedulable due to a missing persistent volume claim.

    A stateful set may create a pod attached to a missing persistent volume
    claim if the pod is recreated while the persistent volume claim is pending
    delete.

    When this happens, delete the pod so that the stateful set will create a
    new persistent volume claim when it next creates the pod.
    """
    for pod in api.list_namespaced_pod(namespace).items:
        if _unschedulable_due_to_pvc(pod):
            logger.info(f"deleting unschedulable pod: {pod.metadata.name}")
            try:
                api.delete_namespaced_pod(
                    name=pod.metadata.name,
                    namespace=namespace,
                    body=V1DeleteOptions(
                        grace_period_seconds=0,
                        propagation_policy="Background",
                        preconditions=V1Preconditions(
                            resource_version=pod.metadata.resource_version,
                            uid=pod.metadata.uid,
                        ),
                    ),
                )
            except ApiException as e:
                if e.reason not in (CONFLICT, NOT_FOUND):
                    raise
                logger.info(f"pod already deleted or updated: {pod.metadata.name}")


def run_task(func: Callable[[], None]):
    """Continuously run func and print exceptions."""
    while True:
        try:
            func()
        except Exception:
            logger.exception("unhandled exception")
        else:
            sleep(1)


def main():
    """Continuously flush and delete detached persistent volumes."""
    args = parser.parse_args()
    load_incluster_config()
    api = CoreV1Api()
    batch_api = BatchV1Api()
    tasks = [
        partial(
            flush_released_pvs_and_delete_complete_jobs,
            api,
            batch_api,
            args.command,
            args.env,
            args.image,
            args.namespace,
            args.service_account_name,
        ),
        partial(
            delete_detached_pvcs,
            api,
            args.namespace,
            args.claim_prefix,
            timedelta(seconds=args.pvc_cleanup_delay_seconds),
            {},  # detached_pvc_cache
        ),
        partial(delete_unschedulable_pods, api, args.namespace),
    ]
    with ThreadPool(len(tasks)) as pool:
        pool.map(run_task, tasks, chunksize=1)


if __name__ == "__main__":  # pragma: no cover
    main()