notifier

#!/usr/bin/env python
#
# Copyright 2015 Kenneth A. Giusti
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import sys
import threading
import time
import uuid

from oslo_config import cfg
from oslo_log import log as logging
import oslo_messaging


_options = [
    cfg.StrOpt("topic",
               default="my-topic",
               help="target topic, default 'my-topic'"),
    cfg.StrOpt("url",
               required=True,
               default="rabbit://localhost:5672",
               help="transport address, default 'rabbit://localhost:5672'"),
    cfg.BoolOpt("quiet",
                default=False,
                help="Suppress all stdout output"),
    cfg.StrOpt("type",
               default="test.event",
               help="Type of event, default 'test.event'"),
    cfg.StrOpt("severity",
               default="debug",
               help="Severity of event, default 'debug'"),
    cfg.IntOpt("count",
               default=1,
               help="Send COUNT notifications per thread, default 1"),
    cfg.StrOpt("payload",
               help="Text for message body."),
    cfg.BoolOpt("stats",
                default=False,
                help="Calculate throughput"),
    cfg.IntOpt("threads",
               default=1,
               help="Number of threads (default 1)"),
    cfg.StrOpt("log_levels",
               help="Set module specific log levels, e.g."
               " 'amqp=WARN,oslo.messaging=INFO,...'")
]


class ClientThread(threading.Thread):
    def __init__(self, conf, name, transport, topic):
        super(ClientThread, self).__init__()
        self.transport = transport
        self.name = name
        self.topic = topic
        self.conf = conf
        self.daemon = True
        self.start()

    def run(self):
        conf = self.conf
        cls = oslo_messaging.notify.notifier.Notifier
        client = cls(self.transport,
                     self.name,
                     driver='messaging',
                     topics=[self.topic])
        try:
            f = getattr(client, conf.severity, 'debug')
            for t in range(conf.count):
                f({}, conf.type, {"payload": conf.payload})
        except Exception as e:
            logging.error("Unexpected exception occured: %s" % str(e))


def main(argv=None):

    conf = cfg.CONF
    logging.register_options(conf)
    conf.register_cli_opts(_options)
    conf(sys.argv[1:])
    if cfg.CONF.log_levels:
        logging.set_defaults(
            default_log_levels=cfg.CONF.log_levels.split(',')
        )
    logging.setup(conf, "notifier")

    topic = conf.topic
    quiet = conf.quiet
    url = conf.url
    severity = conf.severity

    transport = oslo_messaging.get_notification_transport(cfg.CONF, url=url)

    threads = []
    start_time = time.time()
    for t in range(conf.threads):
        name = "notifier-" + uuid.uuid4().hex,
        threads.append(ClientThread(conf, name, transport, topic))
        if not quiet:
            print("notifier %s: url=%s, topic=%s severity=%s"
                  % (name, url, topic, severity))
    try:
        for t in threads:
            t.join()
    except KeyboardInterrupt:
        pass

    if conf.stats:
        delta = time.time() - start_time
        msg_count = conf.count * conf.threads
        stats = float(msg_count)/float(delta) if delta else 0
        print("Messages per second: %6.4f" % stats)

    transport.cleanup()
    return 0

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