run-queue

#!/usr/bin/env python3

# This file is part of Cockpit.
#
# Copyright (C) 2018 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import argparse
import json
import logging
import os
import platform
import random
import smtplib
import subprocess
import sys
import time
from collections.abc import Sequence

import pika

from lib.directories import get_images_data_dir
from lib.network import redhat_network
from lib.stores import LOG_STORE
from task import distributed_queue

logging.basicConfig(level=logging.INFO)

statistics_queue = os.environ.get("RUN_STATISTICS_QUEUE")

# as per pika docs
DeliveryTag = int

ConsumeResult = tuple[Sequence[str] | str | None, DeliveryTag | None]


# Returns a command argv to execute and the delivery tag needed to ack the message
def consume_webhook_queue(dq: distributed_queue.DistributedQueue) -> ConsumeResult:
    # interpret payload
    # call tests-scan or issue-scan appropriately
    method_frame, _header_frame, message = dq.channel.basic_get(queue='webhook')
    if not method_frame or not message:
        return None, None

    body = json.loads(message)
    event = body['event']
    request = body['request']
    repo = None
    cmd = None
    if event == 'pull_request':
        pull_request = request['pull_request']
        repo = pull_request['base']['repo']['full_name']
        action = request['action']
        # scan for body changes (edited) and the bots label; we only use that for image refreshes in the bots repo
        if repo.endswith('/bots') and (
                action == 'labeled' or (action == 'edited' and 'body' in request.get('changes', {}))):
            cmd = ['./issue-scan', '--issues-data', json.dumps(request), '--amqp', dq.address]
        elif action in ['opened', 'synchronize']:
            cmd = ['./tests-scan', '--pull-data', json.dumps(request), '--amqp', dq.address, '--repo', repo]
        # When PR was merged, generate task for storing tests
        elif action == 'closed' and pull_request.get('merged', False):
            sha = pull_request['head']['sha']
            db = os.path.join(get_images_data_dir(), "test-results.db")
            body = {
                "command": (
                    f"./store-tests --db {db} --repo {repo} {sha} && "
                    # prune the db if it gets too big; 14 days weighs about 15 MB
                    f"if [ $(stat -c %s {db}) -gt 25000000 ]; then ./prometheus-stats --db {db} --prune 14; fi &&"
                    f"./prometheus-stats --db {db} --s3 {os.path.join(LOG_STORE, 'prometheus')}"
                )
            }
            dq.channel.basic_publish('', 'statistics', json.dumps(body),
                                     properties=pika.BasicProperties(priority=distributed_queue.MAX_PRIORITY))
            cmd = None

    elif event == 'status':
        repo = request['repository']['full_name']
        sha = request['sha']
        context = request['context']
        cmd = ['./tests-scan', '--sha', sha, '--amqp', dq.address, '--context', context, '--repo', repo]
    elif event == 'issues':
        action = request['action']
        # scan for opened, body changes (edited), and the bots label
        if action in ['opened', 'labeled'] or (action == 'edited' and 'body' in request.get('changes', {})):
            cmd = ['./issue-scan', '--issues-data', json.dumps(request), '--amqp', dq.address]
    else:
        logging.error('Unkown event type in the webhook queue')
        return None, None

    return cmd, method_frame.delivery_tag


# Returns a command to execute and the delivery tag needed to ack the message
def consume_task_queue(dq: distributed_queue.DistributedQueue) -> ConsumeResult:
    if statistics_queue and dq.queue_counts['statistics'] > 0:
        # statistics queue is quick to process, always do that first
        queue = 'statistics'
    elif os.path.exists('/dev/kvm') or os.getenv('JOB_RUNNER_CONFIG'):
        # only process test queues in capable environments: with KVM or job-runner support
        queue = 'public'
        if redhat_network():
            # Try the rhel queue if the public queue is empty
            if dq.queue_counts['public'] == 0:
                queue = 'rhel'
                # If both are non-empty, shuffle
            elif dq.queue_counts['rhel'] > 0:
                queue = ['public', 'rhel'][random.randrange(2)]
    else:
        # nothing to do
        return None, None

    method_frame, _header_frame, message = dq.channel.basic_get(queue=queue)
    if not method_frame or not message:
        return None, None

    body = json.loads(message)
    if job := body.get('job'):
        command = ['./job-runner', 'json', json.dumps(job)]
    else:
        command = body['command']
    return command, method_frame.delivery_tag


def mail_notification(body: str) -> None:
    mx = os.environ.get("TEST_NOTIFICATION_MX")
    if not mx:
        return

    sender = "noreply@redhat.com"
    receivers = os.environ.get("TEST_NOTIFICATION_TO", "").split(',')
    # if sending the notification fails, then fail the pod -- we do *not* want to ack messages in this case,
    # otherwise we would silently miss failure notifications; thus no exception handling here
    s = smtplib.SMTP(mx)
    s.sendmail(sender, receivers, f"""From: {sender}
To: {', '.join(receivers)}
Subject: cockpituous: run-queue crash on {platform.node()}

{body}""")


# Consume from the webhook queue and republish to the task queue via *-scan
# Consume from the task queue (endpoint)
def main() -> int:
    parser = argparse.ArgumentParser(description='Bot: read a single test command from the queue and execute it')
    parser.add_argument('--amqp', default=distributed_queue.DEFAULT_AMQP_SERVER,
                        help='The host:port of the AMQP server to consume from (default: %(default)s)')
    opts = parser.parse_args()

    with distributed_queue.DistributedQueue(opts.amqp, ['webhook', 'rhel', 'public', 'statistics']) as dq:
        cmd, delivery_tag = consume_webhook_queue(dq)
        if not cmd and delivery_tag:
            logging.info("Webhook message interpretation generated no command")
            dq.channel.basic_ack(delivery_tag)
            return 0

        if not cmd:
            cmd, delivery_tag = consume_task_queue(dq)
        if not cmd:
            logging.info("All queues are empty")
            return 1

        if isinstance(cmd, Sequence):
            cmd_str = ' '.join(cmd)
        else:
            cmd_str = cmd
        logging.info("Consuming message with command: %s", cmd_str)
        p = subprocess.Popen(cmd, shell=isinstance(cmd, str))
        while p.poll() is None:
            dq.connection.process_data_events()
            time.sleep(3)
        if p.returncode != 0:
            logging.error("%s failed with exit code %i", cmd_str, p.returncode)
            mail_notification("""The queue command

  %s

failed with exit code %i. Please check the container logs for details.""" % (cmd_str, p.returncode))

        if delivery_tag is not None:
            dq.channel.basic_ack(delivery_tag)

    return 0


if __name__ == '__main__':
    sys.exit(main())