Skip to content

Instantly share code, notes, and snippets.

@phobos182
Created November 20, 2012 22:48
Show Gist options
  • Select an option

  • Save phobos182/4121804 to your computer and use it in GitHub Desktop.

Select an option

Save phobos182/4121804 to your computer and use it in GitHub Desktop.

Revisions

  1. phobos182 created this gist Nov 20, 2012.
    255 changes: 255 additions & 0 deletions diamond
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,255 @@
    #!/usr/bin/python

    import os
    import sys
    import configobj
    import pwd
    import grp

    for path in [
    os.path.join('opt', 'diamond', 'lib'),
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src'))
    ]:
    if os.path.exists(os.path.join(path, 'diamond', '__init__.py')):
    sys.path.append(path)
    break

    from diamond import *
    from diamond.server import Server
    from diamond.util import get_diamond_version

    import optparse
    import signal
    import logging.config

    def main():
    try:
    # Initialize Options
    parser = optparse.OptionParser()
    parser.add_option("-c", "--configfile", dest="configfile", default="/etc/diamond/diamond.conf", help="config file")
    parser.add_option("-f", "--foreground", dest="foreground", default=False, action="store_true", help="run in foreground")
    parser.add_option("-l", "--log-stdout", dest="log_stdout", default=False, action="store_true", help="log to stdout")
    parser.add_option("-p", "--pidfile", dest="pidfile", default=None, help="pid file")
    parser.add_option("-r", "--run", dest="collector", default=None, help="run a given collector once and exit")
    parser.add_option("-v", "--version", dest="version", default=False, action="store_true", help="display the version and exit")
    parser.add_option("--skip-pidfile", dest="skip_pidfile", default=False, action="store_true", help="Skip creating PID file")
    parser.add_option("--skip-change-user", dest="skip_change_user", default=False, action="store_true", help="Skip changing to an unprivilegd user")
    parser.add_option("--skip-fork", dest="skip_fork", default=False, action="store_true", help="Skip forking (damonizing) process")

    # Parse Command Line Args
    (options, args) = parser.parse_args()

    if options.version:
    print "Diamond version %s" % (get_diamond_version())
    sys.exit(0)

    # Initialize Config
    if os.path.exists(options.configfile):
    config = configobj.ConfigObj(os.path.abspath(options.configfile))
    config['configfile'] = options.configfile
    else:
    print >> sys.stderr, "ERROR: Config file: %s does not exist." % (options.configfile)
    parser.print_help(sys.stderr)
    sys.exit(1)

    # Initialize Logging
    log = logging.getLogger('diamond')

    if options.log_stdout:
    log.setLevel(logging.DEBUG)
    # Configure Logging Format
    formatter = logging.Formatter('[%(asctime)s] [%(threadName)s] %(message)s')
    # handler
    streamHandler = logging.StreamHandler(sys.stdout)
    streamHandler.setFormatter(formatter)
    streamHandler.setLevel(logging.DEBUG)
    log.addHandler(streamHandler)
    else:
    try:
    if sys.version_info >= (2,6):
    logging.config.fileConfig(options.configfile,
    disable_existing_loggers=False)
    else:
    # python <= 2.5 does not have disable_existing_loggers
    # default was to always disable them, in our case we want to
    # keep any logger created by handlers
    logging.config.fileConfig(options.configfile)
    for logger in logging.root.manager.loggerDict.values():
    logger.disabled = 0
    except Exception, e:
    sys.stderr.write("Error occurs when initialize logging: ")
    sys.stderr.write(str(e))
    sys.stderr.write(os.linesep)

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
    raise SystemExit

    except Exception, e:
    import traceback
    sys.stderr.write("Unhandled exception: %s" % str(e))
    sys.stderr.write("traceback: %s" % traceback.format_exc())
    sys.exit(1)

    # Switch to using the logging system
    try:
    # PID MANAGEMENT
    if not options.skip_pidfile:
    # Initialize Pid file
    if not options.pidfile:
    options.pidfile = str(config['server']['pid_file'])

    # Read existing pid file
    try:
    pf = file(options.pidfile,'r')
    pid = int(pf.read().strip())
    pf.close()
    except IOError:
    pid = None

    # Check existing pid file
    if pid:
    # Check if pid is real
    if not os.path.exists("/".join(["/proc", str(pid), "cmdline"])):
    # Pid is not real
    os.unlink(options.pidfile)
    pid = None
    print >> sys.stderr, "WARN: Bogus pid file was found. I deleted it."
    else:
    print >> sys.stderr, "ERROR: Pidfile already exists. Server already running?"
    sys.exit(1)

    # Get final GIDs
    gid = -1
    if len(config['server']['group']) :
    gid = grp.getgrnam(config['server']['group']).gr_gid

    # Get final UID
    uid = -1
    if len(config['server']['user']) :
    uid = pwd.getpwnam(config['server']['user']).pw_uid

    # Fix up pid permissions
    if not options.foreground and not options.collector:
    # Write pid file
    pid = str(os.getpid())
    try:
    pf = file(options.pidfile,'w+')
    except IOError, e:
    print >> sys.stderr, "Failed to write PID file: %s" % (e)
    sys.exit(1)
    pf.write("%s\n" % pid)
    pf.close()
    os.chown(options.pidfile, uid, gid)
    # Log
    log.debug("Wrote First PID file: %s" % (options.pidfile))

    # USER MANAGEMENT
    if not options.skip_change_user:
    # Switch user to specified user/group if required
    try:
    if gid != -1 and os.getgid() != gid:
    # Set GID
    os.setgid(gid)

    if uid != -1 and os.getuid() != uid:
    # Set UID
    os.setuid(uid)

    except Exception, e:
    print >> sys.stderr, "ERROR: Failed to set UID/GID. %s" % (e)
    sys.exit(1)

    # Log
    log.info('Changed UID: %d (%s) GID: %d (%s).' % (os.getuid(), config['server']['user'], os.getgid(), config['server']['group']))

    # DAEMONIZE MANAGEMENT
    if not options.skip_fork:
    # Detatch Process
    if not options.foreground and not options.collector:

    # Double fork to serverize process
    log.info('Detaching Process.')

    # Fork 1
    try:
    pid = os.fork()
    if pid > 0:
    # Exit first paren
    sys.exit(0)
    except OSError, e:
    print >> sys.stderr, "Failed to fork process." % (e)
    sys.exit(1)
    # Decouple from parent environmen
    os.setsid()
    os.umask(0)
    # Fork 2
    try:
    pid = os.fork()
    if pid > 0:
    # Exit second paren
    sys.exit(0)
    except OSError, e:
    print >> sys.stderr, "Failed to fork process." % (e)
    sys.exit(1)
    # Close file descriptors so that we can detach
    sys.stdout.close()
    sys.stderr.close()
    sys.stdin.close()
    os.close(0)
    os.close(1)
    os.close(2)

    # PID MANAGEMENT
    if not options.skip_pidfile:
    # Finish Initialize PID file
    if not options.foreground and not options.collector:
    # Write pid file
    pid = str(os.getpid())
    try:
    pf = file(options.pidfile,'w+')
    except IOError, e:
    log.error("Failed to write child PID file: %s" % (e))
    sys.exit(1)
    pf.write("%s\n" % pid)
    pf.close()
    # Log
    log.debug("Wrote child PID file: %s" % (options.pidfile))

    # Initialize Server
    server = Server(config)

    def sigint_handler(signum, frame):
    # Log
    log.debug("Signal Received: %d" % (signum))
    # Stop Server
    server.stop()
    # Delete Pidfile
    if os.path.exists(options.pidfile):
    os.remove(options.pidfile)
    # Log
    log.debug("Removed PID file: %s" % (options.pidfile))

    # Set the signal handlers
    signal.signal(signal.SIGINT, sigint_handler)
    signal.signal(signal.SIGTERM, sigint_handler)

    if options.collector:
    # Run Server with one collector
    server.run_one(options.collector)
    else:
    # Run Server
    server.run()

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
    raise SystemExit

    except Exception, e:
    import traceback
    log.error("Unhandled exception: %s" % str(e))
    log.error("traceback: %s" % traceback.format_exc())
    sys.exit(1)

    if __name__ == "__main__":
    main()