code.grnet.gr Git - snf-network/blob - runlocked

   1 #!/usr/bin/env python
   2 #
   3 # Copyright 2013 GRNET S.A. All rights reserved.
   4 #
   5 # Redistribution and use in source and binary forms, with or
   6 # without modification, are permitted provided that the following
   7 # conditions are met:
   8 #
   9 #   1. Redistributions of source code must retain the above
  10 #      copyright notice, this list of conditions and the following
  11 #      disclaimer.
  12 #
  13 #   2. Redistributions in binary form must reproduce the above
  14 #      copyright notice, this list of conditions and the following
  15 #      disclaimer in the documentation and/or other materials
  16 #      provided with the distribution.
  17 #
  18 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  19 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  20 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  21 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  22 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  25 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  26 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  27 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  28 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  29 # POSSIBILITY OF SUCH DAMAGE.
  30 #
  31 # The views and conclusions contained in the software and
  32 # documentation are those of the authors and should not be
  33 # interpreted as representing official policies, either expressed
  34 # or implied, of GRNET S.A.
  35 #
  36
  37 """Run a command with locking.
  38
  39 Idea shamelessly stolen from http://timkay.com/solo/solo,
  40 but nicely implemented in Python. :)
  41
  42 """
  43
  44 import os
  45 import sys
  46 import time
  47 import errno
  48 import socket
  49
  50
  51 DEFAULT_ID = 10001        # TCP port number to use
  52 DEFAULT_RETRY_SEC = 0.5   # Number of seconds between retries on lock failure
  53
  54
  55 def parse_arguments(args):
  56     from argparse import ArgumentParser, RawDescriptionHelpFormatter, REMAINDER
  57
  58     description = \
  59         ("Run a command with proper locking, ensuring only a single instance\n"
  60          "runs per specified value of `id'. Use the same `id' value for\n"
  61          "commands which must never run simultaneously.\n\n"
  62          "Locking works by binding a TCPv4 socket to 127.0.0.1:<id>, so <id>\n"
  63          "must be a valid port number. Values < 1024 are only usable by "
  64          "root.")
  65
  66     parser = ArgumentParser(description=description,
  67                             formatter_class=RawDescriptionHelpFormatter)
  68     parser.add_argument("-i", "--id", action="store", dest="id",
  69                         default=DEFAULT_ID, metavar="ID",
  70                         help=("Run command with id ID, by binding to "
  71                               "127.0.0.1:ID. Default is %d" % DEFAULT_ID))
  72     parser.add_argument("-r", "--retry-sec", action="store", dest="retry",
  73                         default=DEFAULT_RETRY_SEC, metavar="SECONDS_TO_RETRY",
  74                         help=("In case we cannot get the lock ,retry after "
  75                               "SECONDS_TO_TRY_SECONDS. Default is %d" %
  76                               DEFAULT_RETRY_SEC))
  77     parser.add_argument("command", metavar="COMMAND", nargs=REMAINDER)
  78     args = parser.parse_args()
  79     args = vars(args)
  80
  81     # Sanity checking
  82     try:
  83         args['id'] = int(args['id'])
  84         args['retry'] = float(args['retry'])
  85
  86         if not args['command']:
  87             raise ValueError("The COMMAND to run is mandatory.")
  88     except ValueError as ve:
  89         sys.stderr.write("Argument parsing failed: %s\n" % ve)
  90         sys.exit(1)
  91
  92     return args
  93
  94
  95 def main():
  96     args = parse_arguments(sys.argv)
  97     port = args['id']
  98     retry = args['retry']
  99     cmd = args['command']
 100
 101     # Lock!
 102     s = socket.socket(socket.AF_INET)
 103     while True:
 104         try:
 105             s.bind(("127.0.0.1", port))
 106             break
 107         except socket.error as se:
 108             if se.errno != errno.EADDRINUSE:
 109                 raise
 110             sys.stderr.write(("Could not get the lock on TCPv4 "
 111                               "127.0.0.1:%d, retrying in %fs...\n" %
 112                               (port, retry)))
 113             time.sleep(retry)
 114
 115     # Now that we have the lock,
 116     # replace ourselves with the command in args['command'],
 117     # allowing it to inherit our environment.
 118     #
 119     # The lock is freed by the kernel when this process dies.
 120     try:
 121         os.execvpe(cmd[0], cmd, os.environ)
 122     except OSError as oe:
 123         sys.stderr.write("Command execution failed: %s\n" % oe)
 124         sys.exit(2)
 125
 126
 127 if __name__ == "__main__":
 128     sys.exit(main())