import os
import sys
import zmq
import time
import shutil
import logging
import tempfile
import argparse
import threading

from ami import LogConfig
from ami.comm import Ports


logger = logging.getLogger(__name__)


class SimSync:
    """Class for simulating timestamp synchronization for AMI data sources.

    This class provides a timestamp synchronization service for AMI data
    sources that whose timestamps are not generated by an external source (like
    the DRP, which gets it timestamps from the accelerators timing system).

    This is primarily useful when the simulated data sources provided by AMI on
    a cluster of machines. Each time the simulated data source generates an
    event it can request a timestamp for that event from the timestamp sync
    service via zeromq.

    Args:
        addr (str): the zmq address of the timestamp request socket
                    (e.g. tcp://localhost:5555).
        comm_addr (str): the zmq address of the communication socket of the
                         service.
        start (int): the starting value for use for the generated timestamps.
        interval (int): the reporting interval (in units of timestamps) for
                        logging statistics.
        ctx (zmq.Context): optional zmq context for the node to use. If none is
                           passed it creates one.
    """

    def __init__(self, addr, comm_addr, start, interval, ctx=None):
        self.running = True
        if ctx is None:
            self.ctx = zmq.Context()
            self.owner = True
        else:
            self.ctx = ctx
            self.owner = False
        self.comm = self.ctx.socket(zmq.REP)
        self.comm.bind(comm_addr)
        self.sock = self.ctx.socket(zmq.REP)
        self.sock.bind(addr)
        self.poller = zmq.Poller()
        self.poller.register(self.sock, zmq.POLLIN)
        self.poller.register(self.comm, zmq.POLLIN)
        self.handlers = {
            self.sock: self.timestamp_request,
            self.comm: self.comm_request,
        }
        self.ts = start
        self.tlast = None
        self.interval = interval

    def comm_request(self):
        """
        Called whenever data is available on the the communication socket.
        If a value of 'exit' is received then the receive loop is stopped.
        """
        request = self.comm.recv_string()
        if request == 'exit':
            self.running = False
            self.comm.send_pyobj(0)
        else:
            self.comm.send_pyobj(1)

    def timestamp_request(self):
        """
        Called whenever data is available on the the timestamp request socket.
        The only request it can handle is 'ts'. If a 'ts' request is made the
        reply is the next sequential timestamp value. If an invalid request is
        made then None is the reply.
        """
        request = self.sock.recv_string()
        if request == 'ts':
            self.sock.send_pyobj((self.ts, time.time()))
            self.ts += 1
            if self.ts % self.interval == 0:
                tcurrent = time.time()
                if self.tlast is not None:
                    tdelta = tcurrent - self.tlast
                    logger.info("Processing %f events per second", (self.interval/tdelta))
                self.tlast = tcurrent
        else:
            self.sock.send_pyobj(None)

    def run(self):
        """
        Handles receiving of requests over zmq. Runs until told stop via the
        communication socket.
        """
        while self.running:
            for sock, flag in self.poller.poll():
                if flag != zmq.POLLIN:
                    continue
                if sock in self.handlers:
                    self.handlers[sock]()

    def close(self):
        """
        Function should be called to clean up zmq resources. All sockets are
        closed and the zmq.Context is destroyed if it is owned by this
        instance.
        """
        if self.owner:
            self.ctx.destroy()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


def run_syncer(addr, comm_addr, start, interval, ctx=None):
    with SimSync(addr, comm_addr, start, interval, ctx=ctx) as syncer:
        syncer.run()


def main():
    parser = argparse.ArgumentParser(description='AMII timestamp syncer for simulated data')

    parser.add_argument(
        '-H',
        '--host',
        default='*',
        help='interface the AMII timestamp syncer listens on (default: all)'
    )

    addr_group = parser.add_mutually_exclusive_group()

    addr_group.add_argument(
        '-p',
        '--port',
        type=int,
        default=Ports.Sync,
        help='port for the AMII timestamp syncer (default: %d)' % Ports.Sync
    )

    addr_group.add_argument(
        '--ipc',
        action='store_true',
        help='use ipc for the AMII timestamp syncer'
    )

    parser.add_argument(
        '-s',
        '--start',
        type=int,
        default=0,
        help='starting value to use for the timestamps (default: 0)'
    )

    parser.add_argument(
        '-i',
        '--interval',
        type=int,
        default=10,
        help='the interval (in timestamp increments) to use for reporting (default: 10)'
    )

    parser.add_argument(
        '--log-level',
        default=LogConfig.Level,
        help='the logging level of the application (default %s)' % LogConfig.Level
    )

    parser.add_argument(
        '--log-file',
        help='an optional file to write the log output to'
    )

    args = parser.parse_args()

    ctx = None
    ipcdir = None
    comm_addr = "inproc://sync-comm"
    if args.ipc:
        ipcdir = tempfile.mkdtemp()
        addr = "ipc://%s/sync" % ipcdir
    else:
        addr = "tcp://%s:%d" % (args.host, args.port)

    log_handlers = [logging.StreamHandler()]
    if args.log_file is not None:
        file_fmt = logging.Formatter(LogConfig.Format)
        file_handler = logging.FileHandler(args.log_file)
        file_handler.setFormatter(file_fmt)
        log_handlers.append(file_handler)
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    logging.basicConfig(format=LogConfig.BasicFormat, level=log_level, handlers=log_handlers)

    logger.info("Starting timestamp syncer at %s", addr)

    try:
        ctx = zmq.Context()
        sync_thread = threading.Thread(target=run_syncer,
                                       args=(addr, comm_addr, args.start, args.interval, ctx),
                                       name='syncer',
                                       daemon=True)
        sync_thread.start()

        # create a socket for messaging the syncer thread
        comm = ctx.socket(zmq.REQ)
        comm.connect(comm_addr)

        input("Press any key to exit...\n")
        comm.send_string('exit')
        exit_code = comm.recv_pyobj()
        logger.info("Waiting for timestamp syncer to terminate...")
        sync_thread.join()

        return exit_code
    except KeyboardInterrupt:
        logger.info("Simulated timestamp syncer killed by user...")
        return 0
    finally:
        # cleanup the zmq context
        if ctx is not None:
            ctx.destroy()
        # remove the ipcdir if it exists
        if ipcdir is not None and os.path.exists(ipcdir):
            shutil.rmtree(ipcdir)


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

import os
import sys
import zmq
import time
import shutil
import logging
import tempfile
import argparse
import threading

from ami import LogConfig
from ami.comm import Ports


logger = logging.getLogger(__name__)


class SimSync:
    """Class for simulating timestamp synchronization for AMI data sources.

    This class provides a timestamp synchronization service for AMI data
    sources that whose timestamps are not generated by an external source (like
    the DRP, which gets it timestamps from the accelerators timing system).

    This is primarily useful when the simulated data sources provided by AMI on
    a cluster of machines. Each time the simulated data source generates an
    event it can request a timestamp for that event from the timestamp sync
    service via zeromq.

    Args:
        addr (str): the zmq address of the timestamp request socket
                    (e.g. tcp://localhost:5555).
        comm_addr (str): the zmq address of the communication socket of the
                         service.
        start (int): the starting value for use for the generated timestamps.
        interval (int): the reporting interval (in units of timestamps) for
                        logging statistics.
        ctx (zmq.Context): optional zmq context for the node to use. If none is
                           passed it creates one.
    """

    def __init__(self, addr, comm_addr, start, interval, ctx=None):
        self.running = True
        if ctx is None:
            self.ctx = zmq.Context()
            self.owner = True
        else:
            self.ctx = ctx
            self.owner = False
        self.comm = self.ctx.socket(zmq.REP)
        self.comm.bind(comm_addr)
        self.sock = self.ctx.socket(zmq.REP)
        self.sock.bind(addr)
        self.poller = zmq.Poller()
        self.poller.register(self.sock, zmq.POLLIN)
        self.poller.register(self.comm, zmq.POLLIN)
        self.handlers = {
            self.sock: self.timestamp_request,
            self.comm: self.comm_request,
        }
        self.ts = start
        self.tlast = None
        self.interval = interval

    def comm_request(self):
        """
        Called whenever data is available on the the communication socket.
        If a value of 'exit' is received then the receive loop is stopped.
        """
        request = self.comm.recv_string()
        if request == 'exit':
            self.running = False
            self.comm.send_pyobj(0)
        else:
            self.comm.send_pyobj(1)

    def timestamp_request(self):
        """
        Called whenever data is available on the the timestamp request socket.
        The only request it can handle is 'ts'. If a 'ts' request is made the
        reply is the next sequential timestamp value. If an invalid request is
        made then None is the reply.
        """
        request = self.sock.recv_string()
        if request == 'ts':
            self.sock.send_pyobj((self.ts, time.time()))
            self.ts += 1
            if self.ts % self.interval == 0:
                tcurrent = time.time()
                if self.tlast is not None:
                    tdelta = tcurrent - self.tlast
                    logger.info("Processing %f events per second", (self.interval/tdelta))
                self.tlast = tcurrent
        else:
            self.sock.send_pyobj(None)

    def run(self):
        """
        Handles receiving of requests over zmq. Runs until told stop via the
        communication socket.
        """
        while self.running:
            for sock, flag in self.poller.poll():
                if flag != zmq.POLLIN:
                    continue
                if sock in self.handlers:
                    self.handlers[sock]()

    def close(self):
        """
        Function should be called to clean up zmq resources. All sockets are
        closed and the zmq.Context is destroyed if it is owned by this
        instance.
        """
        if self.owner:
            self.ctx.destroy()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


def run_syncer(addr, comm_addr, start, interval, ctx=None):
    with SimSync(addr, comm_addr, start, interval, ctx=ctx) as syncer:
        syncer.run()


def main():
    parser = argparse.ArgumentParser(description='AMII timestamp syncer for simulated data')

    parser.add_argument(
        '-H',
        '--host',
        default='*',
        help='interface the AMII timestamp syncer listens on (default: all)'
    )

    addr_group = parser.add_mutually_exclusive_group()

    addr_group.add_argument(
        '-p',
        '--port',
        type=int,
        default=Ports.Sync,
        help='port for the AMII timestamp syncer (default: %d)' % Ports.Sync
    )

    addr_group.add_argument(
        '--ipc',
        action='store_true',
        help='use ipc for the AMII timestamp syncer'
    )

    parser.add_argument(
        '-s',
        '--start',
        type=int,
        default=0,
        help='starting value to use for the timestamps (default: 0)'
    )

    parser.add_argument(
        '-i',
        '--interval',
        type=int,
        default=10,
        help='the interval (in timestamp increments) to use for reporting (default: 10)'
    )

    parser.add_argument(
        '--log-level',
        default=LogConfig.Level,
        help='the logging level of the application (default %s)' % LogConfig.Level
    )

    parser.add_argument(
        '--log-file',
        help='an optional file to write the log output to'
    )

    args = parser.parse_args()

    ctx = None
    ipcdir = None
    comm_addr = "inproc://sync-comm"
    if args.ipc:
        ipcdir = tempfile.mkdtemp()
        addr = "ipc://%s/sync" % ipcdir
    else:
        addr = "tcp://%s:%d" % (args.host, args.port)

    log_handlers = [logging.StreamHandler()]
    if args.log_file is not None:
        file_fmt = logging.Formatter(LogConfig.Format)
        file_handler = logging.FileHandler(args.log_file)
        file_handler.setFormatter(file_fmt)
        log_handlers.append(file_handler)
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    logging.basicConfig(format=LogConfig.BasicFormat, level=log_level, handlers=log_handlers)

    logger.info("Starting timestamp syncer at %s", addr)

    try:
        ctx = zmq.Context()
        sync_thread = threading.Thread(target=run_syncer,
                                       args=(addr, comm_addr, args.start, args.interval, ctx),
                                       name='syncer',
                                       daemon=True)
        sync_thread.start()

        # create a socket for messaging the syncer thread
        comm = ctx.socket(zmq.REQ)
        comm.connect(comm_addr)

        input("Press any key to exit...\n")
        comm.send_string('exit')
        exit_code = comm.recv_pyobj()
        logger.info("Waiting for timestamp syncer to terminate...")
        sync_thread.join()

        return exit_code
    except KeyboardInterrupt:
        logger.info("Simulated timestamp syncer killed by user...")
        return 0
    finally:
        # cleanup the zmq context
        if ctx is not None:
            ctx.destroy()
        # remove the ipcdir if it exists
        if ipcdir is not None and os.path.exists(ipcdir):
            shutil.rmtree(ipcdir)


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

Class for simulating timestamp synchronization for AMI data sources.      This class provides a timestamp synchronization service for AMI data     sources that whose timestamps are not generated by an external source (like     the DRP, which gets it timestamps from the accelerators timing system).      This is primarily useful when the simulated data sources provided by AMI on     a cluster of machines. Each time the simulated data source generates an     event it can request a timestamp for that event from the timestamp sync     service via zeromq.      Args:         addr (str): the zmq address of the timestamp request socket                     (e.g. tcp://localhost:5555).         comm_addr (str): the zmq address of the communication socket of the                          service.         start (int): the starting value for use for the generated timestamps.         interval (int): the reporting interval (in units of timestamps) for                         logging statistics.         ctx (zmq.Context): optional zmq context for the node to use. If none is                            passed it creates one. Called whenever data is available on the the communication socket.         If a value of 'exit' is received then the receive loop is stopped. Called whenever data is available on the the timestamp request socket.         The only request it can handle is 'ts'. If a 'ts' request is made the         reply is the next sequential timestamp value. If an invalid request is         made then None is the reply. Handles receiving of requests over zmq. Runs until told stop via the         communication socket. Function should be called to clean up zmq resources. All sockets are         closed and the zmq.Context is destroyed if it is owned by this         instance. # create a socket for messaging the syncer thread # cleanup the zmq context # remove the ipcdir if it exists