from timeit import timeit
from typing import Any

from django.core.management.base import BaseCommand, CommandParser

from zerver.lib.queue import SimpleQueueClient, queue_json_publish
from zerver.worker.queue_processors import BatchNoopWorker, NoopWorker, QueueProcessingWorker


class Command(BaseCommand):
    help = """Times the overhead of enqueuing and dequeuing messages from RabbitMQ."""

    def add_arguments(self, parser: CommandParser) -> None:
        parser.add_argument(
            "--count", help="Number of messages to enqueue", default=10000, type=int
        )
        parser.add_argument("--reps", help="Iterations of enqueue/dequeue", default=1, type=int)
        parser.add_argument("--batch", help="Enables batch dequeuing", action="store_true")
        parser.add_argument(
            "--prefetch",
            help="Limits the prefetch size; RabbitMQ defaults to unbounded (0)",
            default=0,
            type=int,
        )
        parser.add_argument(
            "--slow",
            help="Which request numbers should take 60s (1-based)",
            action="append",
            type=int,
            default=[],
        )

    def handle(self, *args: Any, **options: Any) -> None:
        print("Purging queue...")
        queue = SimpleQueueClient()
        queue_name = "noop_batch" if options["batch"] else "noop"
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge("noop"))

        count = options["count"]
        reps = options["reps"]

        worker: QueueProcessingWorker = NoopWorker(count, options["slow"])
        if options["batch"]:
            worker = BatchNoopWorker(count, options["slow"])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options["prefetch"])

        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() -> None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(
                lambda: queue_json_publish(queue_name, {}),
                number=count,
            )
            total_dequeue_time += timeit(
                lambda: worker.start(),
                number=1,
            )

        rate = lambda time, iterations: int(iterations / time)

        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f"Total rate per rep: {rate(total_reps_time, reps)} / sec")

        print(f"Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec")
        print(f"Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec")

032fa9c99aa6bae96f2f0bb0a13416ef6c9de244
<mask token>


class Command(BaseCommand):
    <mask token>
    <mask token>

    def handle(self, *args: Any, **options: Any) ->None:
        print('Purging queue...')
        queue = SimpleQueueClient()
        queue_name = 'noop_batch' if options['batch'] else 'noop'
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge(
            'noop'))
        count = options['count']
        reps = options['reps']
        worker: QueueProcessingWorker = NoopWorker(count, options['slow'])
        if options['batch']:
            worker = BatchNoopWorker(count, options['slow'])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options['prefetch'])
        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() ->None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(lambda : queue_json_publish(
                queue_name, {}), number=count)
            total_dequeue_time += timeit(lambda : worker.start(), number=1)
        rate = lambda time, iterations: int(iterations / time)
        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f'Total rate per rep: {rate(total_reps_time, reps)} / sec')
        print(f'Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec')
        print(f'Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec')

<mask token>


class Command(BaseCommand):
    <mask token>

    def add_arguments(self, parser: CommandParser) ->None:
        parser.add_argument('--count', help='Number of messages to enqueue',
            default=10000, type=int)
        parser.add_argument('--reps', help='Iterations of enqueue/dequeue',
            default=1, type=int)
        parser.add_argument('--batch', help='Enables batch dequeuing',
            action='store_true')
        parser.add_argument('--prefetch', help=
            'Limits the prefetch size; RabbitMQ defaults to unbounded (0)',
            default=0, type=int)
        parser.add_argument('--slow', help=
            'Which request numbers should take 60s (1-based)', action=
            'append', type=int, default=[])

    def handle(self, *args: Any, **options: Any) ->None:
        print('Purging queue...')
        queue = SimpleQueueClient()
        queue_name = 'noop_batch' if options['batch'] else 'noop'
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge(
            'noop'))
        count = options['count']
        reps = options['reps']
        worker: QueueProcessingWorker = NoopWorker(count, options['slow'])
        if options['batch']:
            worker = BatchNoopWorker(count, options['slow'])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options['prefetch'])
        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() ->None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(lambda : queue_json_publish(
                queue_name, {}), number=count)
            total_dequeue_time += timeit(lambda : worker.start(), number=1)
        rate = lambda time, iterations: int(iterations / time)
        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f'Total rate per rep: {rate(total_reps_time, reps)} / sec')
        print(f'Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec')
        print(f'Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec')

<mask token>


class Command(BaseCommand):
    help = (
        'Times the overhead of enqueuing and dequeuing messages from RabbitMQ.'
        )

    def add_arguments(self, parser: CommandParser) ->None:
        parser.add_argument('--count', help='Number of messages to enqueue',
            default=10000, type=int)
        parser.add_argument('--reps', help='Iterations of enqueue/dequeue',
            default=1, type=int)
        parser.add_argument('--batch', help='Enables batch dequeuing',
            action='store_true')
        parser.add_argument('--prefetch', help=
            'Limits the prefetch size; RabbitMQ defaults to unbounded (0)',
            default=0, type=int)
        parser.add_argument('--slow', help=
            'Which request numbers should take 60s (1-based)', action=
            'append', type=int, default=[])

    def handle(self, *args: Any, **options: Any) ->None:
        print('Purging queue...')
        queue = SimpleQueueClient()
        queue_name = 'noop_batch' if options['batch'] else 'noop'
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge(
            'noop'))
        count = options['count']
        reps = options['reps']
        worker: QueueProcessingWorker = NoopWorker(count, options['slow'])
        if options['batch']:
            worker = BatchNoopWorker(count, options['slow'])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options['prefetch'])
        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() ->None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(lambda : queue_json_publish(
                queue_name, {}), number=count)
            total_dequeue_time += timeit(lambda : worker.start(), number=1)
        rate = lambda time, iterations: int(iterations / time)
        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f'Total rate per rep: {rate(total_reps_time, reps)} / sec')
        print(f'Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec')
        print(f'Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec')

from timeit import timeit
from typing import Any
from django.core.management.base import BaseCommand, CommandParser
from zerver.lib.queue import SimpleQueueClient, queue_json_publish
from zerver.worker.queue_processors import BatchNoopWorker, NoopWorker, QueueProcessingWorker


class Command(BaseCommand):
    help = (
        'Times the overhead of enqueuing and dequeuing messages from RabbitMQ.'
        )

    def add_arguments(self, parser: CommandParser) ->None:
        parser.add_argument('--count', help='Number of messages to enqueue',
            default=10000, type=int)
        parser.add_argument('--reps', help='Iterations of enqueue/dequeue',
            default=1, type=int)
        parser.add_argument('--batch', help='Enables batch dequeuing',
            action='store_true')
        parser.add_argument('--prefetch', help=
            'Limits the prefetch size; RabbitMQ defaults to unbounded (0)',
            default=0, type=int)
        parser.add_argument('--slow', help=
            'Which request numbers should take 60s (1-based)', action=
            'append', type=int, default=[])

    def handle(self, *args: Any, **options: Any) ->None:
        print('Purging queue...')
        queue = SimpleQueueClient()
        queue_name = 'noop_batch' if options['batch'] else 'noop'
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge(
            'noop'))
        count = options['count']
        reps = options['reps']
        worker: QueueProcessingWorker = NoopWorker(count, options['slow'])
        if options['batch']:
            worker = BatchNoopWorker(count, options['slow'])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options['prefetch'])
        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() ->None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(lambda : queue_json_publish(
                queue_name, {}), number=count)
            total_dequeue_time += timeit(lambda : worker.start(), number=1)
        rate = lambda time, iterations: int(iterations / time)
        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f'Total rate per rep: {rate(total_reps_time, reps)} / sec')
        print(f'Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec')
        print(f'Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec')

from timeit import timeit
from typing import Any

from django.core.management.base import BaseCommand, CommandParser

from zerver.lib.queue import SimpleQueueClient, queue_json_publish
from zerver.worker.queue_processors import BatchNoopWorker, NoopWorker, QueueProcessingWorker


class Command(BaseCommand):
    help = """Times the overhead of enqueuing and dequeuing messages from RabbitMQ."""

    def add_arguments(self, parser: CommandParser) -> None:
        parser.add_argument(
            "--count", help="Number of messages to enqueue", default=10000, type=int
        )
        parser.add_argument("--reps", help="Iterations of enqueue/dequeue", default=1, type=int)
        parser.add_argument("--batch", help="Enables batch dequeuing", action="store_true")
        parser.add_argument(
            "--prefetch",
            help="Limits the prefetch size; RabbitMQ defaults to unbounded (0)",
            default=0,
            type=int,
        )
        parser.add_argument(
            "--slow",
            help="Which request numbers should take 60s (1-based)",
            action="append",
            type=int,
            default=[],
        )

    def handle(self, *args: Any, **options: Any) -> None:
        print("Purging queue...")
        queue = SimpleQueueClient()
        queue_name = "noop_batch" if options["batch"] else "noop"
        queue.ensure_queue(queue_name, lambda channel: channel.queue_purge("noop"))

        count = options["count"]
        reps = options["reps"]

        worker: QueueProcessingWorker = NoopWorker(count, options["slow"])
        if options["batch"]:
            worker = BatchNoopWorker(count, options["slow"])
        worker.ENABLE_TIMEOUTS = True
        worker.setup()
        assert worker.q is not None
        assert worker.q.channel is not None
        worker.q.channel.basic_qos(prefetch_count=options["prefetch"])

        total_enqueue_time = 0.0
        total_dequeue_time = 0.0

        def one_rep() -> None:
            nonlocal total_enqueue_time, total_dequeue_time
            total_enqueue_time += timeit(
                lambda: queue_json_publish(queue_name, {}),
                number=count,
            )
            total_dequeue_time += timeit(
                lambda: worker.start(),
                number=1,
            )

        rate = lambda time, iterations: int(iterations / time)

        total_reps_time = timeit(one_rep, number=reps)
        if reps > 1:
            print(f"Total rate per rep: {rate(total_reps_time, reps)} / sec")

        print(f"Enqueue rate: {rate(total_enqueue_time, count * reps)} / sec")
        print(f"Dequeue rate: {rate(total_dequeue_time, count * reps)} / sec")
