getFromQueue.py

__author__ = 'Manuela Kuhn <manuela.kuhn@desy.de>'

import zmq
import os
import sys
import logging
import traceback
import cPickle
import shutil


def setup(log, prop):

    # Create zmq socket
    socket        = prop["context"].socket(zmq.PULL)
    connectionStr = "tcp://{ip}:{port}".format( ip=prop["extIp"], port=prop["port"] )
    socket.bind(connectionStr)
    log.info("Start socket (bind): '" + str(connectionStr) + "'")

    # register socket
    prop["socket"] = socket


def getMetadata (log, metadata, chunkSize, localTarget = None):

    #extract fileEvent metadata
    try:
        #TODO validate metadata dict
        sourceFile = metadata["filename"]
    except:
        log.error("Invalid fileEvent message received.", exc_info=True)
        log.debug("metadata=" + str(metadata))
        #skip all further instructions and continue with next iteration
        raise

    #TODO combine better with sourceFile... (for efficiency)
    if localTarget:
        targetFile     = os.path.join(localTarget, filename)
    else:
        targetFile     = None

    try:
        # For quick testing set filesize of file as chunksize
        log.debug("get filesize for '" + str(sourceFile) + "'...")
#        filesize    = os.path.getsize(sourceFile)
#        fileModTime = os.stat(sourceFile).st_mtime
#        chunksize   = filesize    #can be used later on to split multipart message
#        log.debug("filesize(%s) = %s" % (sourceFile, str(filesize)))
#        log.debug("fileModTime(%s) = %s" % (sourceFile, str(fileModTime)))

    except:
        log.error("Unable to create metadata dictionary.", exc_info=True)
        raise

    try:
        log.debug("create metadata for source file...")
        #metadata = {
        #        "filename"     : filename,
        #        "filesize"     : filesize,
        #        "fileModTime"  : fileModTime,
        #        "chunkSize"    : self.zmqMessageChunkSize
        #        }
#        metadata[ "filesize"    ] = filesize
#        metadata[ "fileModTime" ] = fileModTime
        metadata[ "chunkSize"   ] = chunkSize

        log.debug("metadata = " + str(metadata))
    except:
        log.error("Unable to assemble multi-part message.", exc_info=True)
        raise

    return sourceFile, targetFile, metadata


def sendData (log, targets, sourceFile, metadata, openConnections, context, prop):
    #reading source file into memory
    try:
        log.debug("Getting '" + str(sourceFile) + "'...")
        data = prop["socket"].recv()
    except:
        log.error("Unable to get source file '" + str(sourceFile) + "'.", exc_info=True)
        raise

    chunkSize = metadata[ "chunkSize" ]

    #send message
    try:
        log.debug("Passing multipart-message for file " + str(sourceFile) + "...")
        chunkNumber = 0

        #assemble metadata for zmq-message
        metadataExtended = metadata.copy()
        metadataExtended["chunkNumber"] = chunkNumber
        metadataExtended = cPickle.dumps(metadata)

        payload = []
        payload.append(metadataExtended)
        payload.append(data)

        # streaming data
        for target, prio in targets:

            # send data to the data stream to store it in the storage system
            if prio == 0:
                # socket already known
                if target in openConnections:
                    tracker = openConnections[target].send_multipart(payload, copy=False, track=True)
                    log.info("Sending message part from file " + str(sourceFile) + " to '" + target + "' with priority " + str(prio) )
                else:
                    # open socket
                    socket        = context.socket(zmq.PUSH)
                    connectionStr = "tcp://" + str(target)

                    socket.connect(connectionStr)
                    log.info("Start socket (connect): '" + str(connectionStr) + "'")

                    # register socket
                    openConnections[target] = socket

                    # send data
                    tracker = openConnections[target].send_multipart(payload, copy=False, track=True)
                    log.info("Sending message part from file " + str(sourceFile) + " to '" + target + "' with priority " + str(prio) )

                # socket not known
                if not tracker.done:
                    log.info("Message part from file " + str(sourceFile) + " has not been sent yet, waiting...")
                    tracker.wait()
                    log.info("Message part from file " + str(sourceFile) + " has not been sent yet, waiting...done")

            else:
                # socket already known
                if target in openConnections:
                    # send data
                    openConnections[target].send_multipart(payload, zmq.NOBLOCK)
                    log.info("Sending message part from file " + str(sourceFile) + " to " + target)
                # socket not known
                else:
                    # open socket
                    socket        = context.socket(zmq.PUSH)
                    connectionStr = "tcp://" + str(target)

                    socket.connect(connectionStr)
                    log.info("Start socket (connect): '" + str(connectionStr) + "'")

                    # register socket
                    openConnections[target] = socket

                    # send data
                    openConnections[target].send_multipart(payload, zmq.NOBLOCK)
                    log.info("Sending message part from file " + str(sourceFile) + " to " + target)

        log.debug("Passing multipart-message for file " + str(sourceFile) + "...done.")

    except:
        log.error("Unable to send multipart-message for file " + str(sourceFile), exc_info=True)


def finishDataHandling (log, sourceFile, targetFile, removeFlag = False):
    pass


def clean(prop):
    # Close zmq socket
    if prop["socket"]:
        prop["socket"].close(0)
        prop["socket"] = None


if __name__ == '__main__':
    import time
    from shutil import copyfile

    try:
        BASE_PATH = os.path.dirname ( os.path.dirname ( os.path.dirname ( os.path.dirname ( os.path.realpath ( __file__ ) ))))
    except:
        BASE_PATH = os.path.dirname ( os.path.dirname ( os.path.dirname ( os.path.dirname ( os.path.abspath ( sys.argv[0] ) ))))
    print "BASE_PATH", BASE_PATH
    SHARED_PATH  = BASE_PATH + os.sep + "src" + os.sep + "shared"

    if not SHARED_PATH in sys.path:
        sys.path.append ( SHARED_PATH )
    del SHARED_PATH

    import helpers

    logfile = BASE_PATH + os.sep + "logs" + os.sep + "getFromFile.log"

    # Get the log Configuration for the lisener
    h1, h2 = helpers.getLogHandlers(logfile, verbose=True, onScreenLogLevel="debug")

    # Create log and set handler to queue handle
    root = logging.getLogger()
    root.setLevel(logging.DEBUG) # Log level = DEBUG
    root.addHandler(h1)
    root.addHandler(h2)

    receivingPort    = "6005"
    receivingPort2   = "6006"
    extIp            = "0.0.0.0"
    dataFwPort       = "6050"

    context          = zmq.Context.instance()

    dataFwSocket     = context.socket(zmq.PUSH)
    connectionStr    = "tcp://{ip}:{port}".format( ip=extIp, port=dataFwPort )
    dataFwSocket.connect(connectionStr)
    logging.info("=== Start dataFwsocket (connect): '" + str(connectionStr) + "'")

    receivingSocket  = context.socket(zmq.PULL)
    connectionStr    = "tcp://{ip}:{port}".format( ip=extIp, port=receivingPort )
    receivingSocket.bind(connectionStr)
    logging.info("=== receivingSocket connected to " + connectionStr)

    receivingSocket2 = context.socket(zmq.PULL)
    connectionStr    = "tcp://{ip}:{port}".format( ip=extIp, port=receivingPort2 )
    receivingSocket2.bind(connectionStr)
    logging.info("=== receivingSocket2 connected to " + connectionStr)


    prework_sourceFile = BASE_PATH + os.sep + "test_file.cbf"

    #read file to send it in data pipe
    fileDescriptor = open(prework_sourceFile, "rb")
    fileContent = fileDescriptor.read()
    logging.debug("=== File read")
    fileDescriptor.close()

    dataFwSocket.send(fileContent)
    logging.debug("=== File send")

    workload = {
            "sourcePath"  : BASE_PATH + os.sep +"data" + os.sep + "source",
            "relativePath": os.sep + "local" + os.sep + "raw",
            "filename"    : "100.cbf"
            }
    targets = [['localhost:' + receivingPort, 1], ['localhost:' + receivingPort2, 0]]

    chunkSize       = 10485760 ; # = 1024*1024*10 = 10 MiB
    localTarget     = BASE_PATH + os.sep + "data" + os.sep + "target"
    openConnections = dict()

    dataFetcherProp = {
            "type"       : "getFromQueue",
            "context"    : context,
            "extIp"      : extIp,
            "port"       : dataFwPort
            }

    logging.debug("openConnections before function call: " + str(openConnections))

    setup(logging, dataFetcherProp)

    sourceFile, targetFile, metadata = getMetadata (logging, workload, chunkSize, localTarget = None)
    sendData(logging, targets, sourceFile, metadata, openConnections, context, dataFetcherProp)

    finishDataHandling(logging, sourceFile, targetFile, dataFetcherProp)

    logging.debug("openConnections after function call: " + str(openConnections))


    try:
        recv_message = receivingSocket.recv_multipart()
        logging.info("=== received: " + str(cPickle.loads(recv_message[0])))
        recv_message = receivingSocket2.recv_multipart()
        logging.info("=== received 2: " + str(cPickle.loads(recv_message[0])))
    except KeyboardInterrupt:
        pass
    finally:
        receivingSocket.close(0)
        receivingSocket2.close(0)
        clean(dataFetcherProp)
        context.destroy()