Synnefo » snf-image

#!/usr/bin/env python

# Copyright (C) 2011, 2012 GRNET S.A.

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful, but

# WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

# General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

# 02110-1301, USA.

"""Utility to monitor the progress of image deployment

A small utility to monitor the progress of image deployment

by watching the contents of /proc/<pid>/io and producing

notifications.

"""

import os

import sys

import time

import json

import prctl

import signal

import socket

def parse_arguments(args):

    from optparse import OptionParser

    kw = {}

    kw['usage'] = "%prog [options] command [args...]"

    kw['description'] = \

        "%prog runs 'command' with the specified arguments, monitoring the " \

        "number of bytes read by it. 'command' is assumed to be " \

        "A program used to install the OS for a Ganeti instance. %prog " \

        "periodically issues notifications of type 'copy-progress'."

    parser = OptionParser(**kw)

    parser.disable_interspersed_args()

    parser.add_option("-r", "--read-bytes",

                      action="store", type="int", dest="read_bytes",

                      metavar="BYTES_TO_READ",

                      help="The expected number of bytes to be read, " \

                           "used to compute input progress",

                      default=None)

    parser.add_option("-o", "--output", dest="output", default=None,

                    metavar="FILE",

                    help="Write output notifications to this file")

    (opts, args) = parser.parse_args(args)

    if opts.read_bytes is None:

        sys.stderr.write("Fatal: Option '-r' is mandatory.\n")

        parser.print_help()

        sys.exit(1)

    if opts.output is None:

        sys.stderr.write("Fatal: Option '-o' is mandatory.\n")

        parser.print_help()

        sys.exit(1)

    if len(args) == 0:

        sys.stderr.write("Fatal: You need to specify the command to run.\n")

        parser.print_help()

        sys.exit(1)

    return (opts, args)

def report_wait_status(pid, status):

    if os.WIFEXITED(status):

        sys.stderr.write("Child PID = %d exited, status = %d\n" %

                         (pid, os.WEXITSTATUS(status)))

    elif os.WIFSIGNALED(status):

        sys.stderr.write("Child PID = %d died by signal, signal = %d\n" %

                         (pid, os.WTERMSIG(status)))

    elif os.WIFSTOPPED(status):

        sys.stderr.write("Child PID = %d stopped by signal, signal = %d\n" %

                         (pid, os.WSTOPSIG(status)))

    else:

        sys.stderr.write("Internal error: Unhandled case, " \

                         "PID = %d, status = %d\n" % (pid, status))

        sys.exit(1)

    sys.stderr.flush()

def send_message(to, message):

    message['timestamp'] = time.time()

    to.write("%s\n" % json.dumps(message))

    to.flush()

def main():

    (opts, args) = parse_arguments(sys.argv[1:])

    with open(opts.output, 'w') as out:

        pid = os.fork()

        if pid == 0:

            # In child process:

            # Make sure we die with the parent and are not left behind

            # WARNING: This uses the prctl(2) call and is Linux-specific.

            prctl.set_pdeathsig(signal.SIGHUP)

            # exec command specified in arguments,

            # searching the $PATH, keeping all environment

            os.execvpe(args[0], args, os.environ)

            sys.stderr.write("execvpe failed, exiting with non-zero status")

            os.exit(1)

        # In parent process:

        iofname = "/proc/%d/io" % pid

        iof = open(iofname, "r", 0)   # 0: unbuffered open

        sys.stderr.write("%s: created child PID = %d, monitoring file %s\n" %

                         (sys.argv[0], pid, iofname))

        message = {}

        message['type'] = 'copy-progress'

        message['total'] = opts.read_bytes

        while True:

            # check if the child process is still alive

            (wpid, status) = os.waitpid(pid, os.WNOHANG)

            if wpid == pid:

                report_wait_status(pid, status)

                if (os.WIFEXITED(status) or os.WIFSIGNALED(status)):

                    if not (os.WIFEXITED(status) and

                                                os.WEXITSTATUS(status) == 0):

                        return 1

                    else:

                        message['position'] = message['total']

                        message['progress'] = float(100)

                        send_message(out, message)

                        return 0

            iof.seek(0)

            for l in iof.readlines():

                if l.startswith("rchar:"):

                    message['position'] = int(l.split(': ')[1])

                    message['progress'] = float(100) if opts.read_bytes == 0 \

                        else float("%2.2f" % (

                            message['position'] * 100.0 / message['total']))

                    send_message(out, message)

                    break

            # Sleep for a while

            time.sleep(3)

if __name__ == "__main__":

    sys.exit(main())

# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :