benchmarks/benchmark2.py

import time
import sys
import json
import requests
from datetime import datetime
from dataclasses import dataclass
import argparse
from multiprocessing import Process, Queue
from queue import Empty
import paho.mqtt.client
from threading import Thread, current_thread


@dataclass(frozen=False, init=False)
class UploadRequest:
    id: int
    upload_started: float = None
    upload_finished: float = None
    upload_status_code: bool = False
    alert_arrived: float = None

    def __init__(self, _id: int):
        self.id = _id


@dataclass(frozen=False, init=True)
class ReturnedRequest:
    id: int
    arrived: float


def print_progress_meme(char: str):
    print(char, end='')
    sys.stdout.flush()


def sender_worker(input_queue: Queue, sent_queue: Queue, target_url: str, file_contents: bytes, request_timeout: float):
    content_len = len(file_contents)
    session = requests.Session()
    try:
        while True:
            try:
                job = input_queue.get_nowait()
            except Empty:
                return

            files = {
                "file": (
                    f"birbbenchmark2_testfile_{job.id}.wav",
                    file_contents,
                    'audio/wave',
                    {'Content-length': content_len}
                ),
                "description": (
                    None,
                    json.dumps({'date': datetime.now().isoformat(), 'device_id': str(job.id)}),
                    "application/json"
                )
            }

            # SUPER HIGH PRECISION MEASUREMENT COMPARTMENT DEPARTMENT ADMIRAL GENERAL ALADEEN
            job.upload_started = time.time()
            try:
                r = session.post(target_url, files=files, timeout=request_timeout)
                job.upload_finished = time.time()
                if r.status_code == 200:
                    print_progress_meme('.')
                else:
                    print_progress_meme('!')

                job.upload_status_code = r.status_code
            except requests.Timeout:
                print_progress_meme('?')
            # /'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\./'\

            sent_queue.put(job)  # Will have none as status_code and upload_finished if the request timed out

    except KeyboardInterrupt:
        return


def mqtt_on_connect(client, userdata, flags, rc):
    client.subscribe(f"command/#")


def mqtt_on_command(client, userdata, message):
    # msg = json.loads(message.payload.decode())
    # if msg['command'] == 'doAlert':
    _id = int(message.topic[8:])  # len('command/') = 8
    userdata.put(ReturnedRequest(id=_id, arrived=time.time()))
    print_progress_meme(',')


def infinite_job_generator(input_queue: Queue, target_ready_jobs: int):
    i = 0
    target_ready_jobs_local = target_ready_jobs
    first = True
    while True:
        jobs_ready = input_queue.qsize()
        if jobs_ready < target_ready_jobs_local:

            if not first:  # In the first run there are exactly zero jobs ready, so
                if jobs_ready < (target_ready_jobs_local / 2):
                    target_ready_jobs_local *= 2
            else:
                first = False

            jobs_to_generate = target_ready_jobs_local - jobs_ready
            for _ in range(jobs_to_generate):
                i += 1
                input_queue.put(UploadRequest(i))
        time.sleep(0.02)

        if not current_thread().active:
            return


def run_benchmark(
        num_requests: int,
        num_workers: int,
        sound_file_path: str,
        target_url: str,
        timeout: float,
        inflight_timeout: float,
        request_timeout: float,
        mqtt_client: paho.mqtt.client.Client):
    input_queue = Queue()
    sent_queue = Queue()
    returned_queue = Queue()
    sender_workers = []

    input_generator_thread = None

    if num_requests == 0:
        print(f"Unlimited amount of requests will be generated...")
        input_generator_thread = Thread(target=infinite_job_generator, args=(input_queue, num_workers * 100))
        input_generator_thread.active = True
        input_generator_thread.start()
    else:
        print(f"Preparing {num_requests} requests...")
        for i in range(num_requests):
            input_queue.put(UploadRequest(i + 1))

    mqtt_client.user_data_set(returned_queue)
    mqtt_client.loop_start()  # Puts mqtt stuff to a background thread

    print(f"Preparing {num_workers} workers...")

    with open(sound_file_path, "rb") as f:
        file_contents = f.read()

    for i in range(num_workers):
        sender_workers.append(
            Process(target=sender_worker,
                    args=(
                        input_queue,
                        sent_queue,
                        target_url,
                        file_contents,
                        request_timeout
                    ))
        )

    print("Baszatás!")
    # Starting all workers
    for worker in sender_workers:
        worker.start()

    # Waiting for workers to finish
    try:
        for worker in sender_workers:
            worker.join()
    except KeyboardInterrupt:
        # Azt jelenti: jovan, megállítottam a küldést, most már csak az elveszett válaszokat várom
        print_progress_meme('|')
        # Interrupt forwarded for workers, so we just wait for them to quit
        for worker in sender_workers:
            exitcode = worker.join(timeout=2)
            if not exitcode:  # = it did not exit
                worker.terminate()  # Not the nicest solution

    # Stop the input generator
    input_generator_thread.active = False

    # read out all send requests to an array for handling later:
    results = {}
    total_successful_uploads = 0
    while True:
        try:
            job = sent_queue.get_nowait()
        except Empty:
            break

        if job.upload_status_code == 200:
            total_successful_uploads += 1

        results[job.id] = job

    # total_successful_uploads = len([req for req in results.values() if req.upload_status_code == 200])

    # Wait for inflight messages for a little
    for _ in range(int(inflight_timeout)):
        time.sleep(1)
        if returned_queue.qsize() >= total_successful_uploads:
            break

    mqtt_client.loop_stop()  # This stops further recieving MQTT messages
    print()  # Everything that might print from other threads are now stopped, so we put a newline

    if returned_queue.qsize() < total_successful_uploads:
        print("Waiting for inflight MQTT messages timed out!")

    # Fill out the missing fields in the reuslts
    total_answered = 0
    while True:
        try:
            returned_job = returned_queue.get_nowait()
        except Empty:
            break

        try:
            results[returned_job.id].alert_arrived = returned_job.arrived
        except KeyError:
            print("Bruh moment: Alert arrived with an id that's not sent", returned_job.id)

        total_answered += 1

    # mini statistics
    print(f"{len(results)} requests completed: {total_successful_uploads} successfully uploaded and {total_answered} answered")

    return results


def main():
    parser = argparse.ArgumentParser(description='Birbnetes Benchmarker 2')
    parser.add_argument('--workers', type=int, required=False, default=1, help='Number of workers')
    parser.add_argument('--requests', type=int, required=False, default=100, help='Number of requests (0 for infinity)')
    parser.add_argument('--timeout', type=int, required=False, default=0,
                        help='Maximum time for request sending (0 for infinity)')

    parser.add_argument('--inflight-timeout', type=int, required=False, default=30,
                        help='Number of seconds to wait for MQTT messages to arrive after uploading finished (not very precise)')

    parser.add_argument('--request-timeout', type=float, required=False, default=5,
                        help='Timeout for HTTP requests')

    parser.add_argument('--file', type=str, required=False, default="oIbBz.wav",
                        help='Name of the sound file to upload')

    parser.add_argument('--target-url', type=str, required=False, default="https://birb.k8s.kmlabz.com/sample",
                        help='The target endpoint')

    parser.add_argument('--mqtt-username', type=str, required=False, default="birbnetes",
                        help="Username for the MQTT server")
    parser.add_argument('--mqtt-port', type=int, required=False, default=1883, help="Port for the MQTT server")
    parser.add_argument('--mqtt-host', type=str, required=False, default="mqtt.kmlabz.com",
                        help="Hostname for the MQTT server")
    parser.add_argument('--mqtt-password', type=str, required=False, default=None, help="Username for the MQTT server")

    args = parser.parse_args()

    print(f"Preparing MQTT (host: {args.mqtt_host}:{args.mqtt_port}, login: {bool(args.mqtt_password)})")
    client = paho.mqtt.client.Client(client_id="10")
    client.on_connect = mqtt_on_connect
    client.on_message = mqtt_on_command
    if args.mqtt_password:
        client.username_pw_set(args.mqtt_username, args.mqtt_password)
    client.connect(args.mqtt_host, args.mqtt_port, 60)

    results = run_benchmark(args.requests, args.workers, args.file, args.target_url, args.timeout,
                            args.inflight_timeout, args.request_timeout, client)

    print("Eposzi kiértékelés idő...")
    print("id  Upload Dur         Return Dur            Total Dur")
    # read out all recieved requests to an array for easier handling:
    for id_, req in results.items():

        if not req.alert_arrived:
            if req.upload_status_code == 200:
                print(id_, req.upload_finished - req.upload_started, '???', '???')
            else:
                print(id_, req.upload_finished - req.upload_started, '!', '???', '???')
        else:
            print(id_, req.upload_finished - req.upload_started, req.alert_arrived - req.upload_finished,
                  req.alert_arrived - req.upload_started)


if __name__ == '__main__':
    main()