|
- #!/usr/bin/python
-
- import nilmdb.client
- from nilmdb.utils.printf import *
- import nilmdb.utils.lock
- import nilmtools
-
- import time
- import sys
- import os
- import argparse
- import subprocess
- import tempfile
- import threading
- import select
- import signal
- import Queue
- import daemon
-
- def parse_args(argv = None):
- parser = argparse.ArgumentParser(
- formatter_class = argparse.ArgumentDefaultsHelpFormatter,
- version = nilmtools.__version__,
- description = """\
- Pipe data from 'generator' to 'consumer'. This is intended to be
- executed frequently from cron, and will exit if another copy is
- already running. If 'generator' or 'consumer' returns an error,
- or if 'generator' stops sending data for a while, it will exit.
-
- Intended for use with ethstream (generator) and nilm-insert
- (consumer). Commands are executed through the shell.
- """)
- parser.add_argument("-d", "--daemon", action="store_true",
- help="Run in background")
- parser.add_argument("-l", "--lock", metavar="FILENAME", action="store",
- default=tempfile.gettempdir() +
- "/nilm-pipewatch.lock",
- help="Lock file for detecting running instance")
- parser.add_argument("-t", "--timeout", metavar="SECONDS", action="store",
- type=float, default=30,
- help="Restart if no output from " +
- "generator for this long")
- group = parser.add_argument_group("commands to execute")
- group.add_argument("generator", action="store",
- help="Data generator (e.g. \"ethstream -r 8000\")")
- group.add_argument("consumer", action="store",
- help="Data consumer (e.g. \"nilm-insert /foo/bar\")")
- args = parser.parse_args(argv)
-
- return args
-
- def reader_thread(queue, fd):
- # Read from a file descriptor, write to queue.
- try:
- while True:
- (r, w, x) = select.select([fd], [], [fd], 0.25)
- if x:
- raise Exception # generator died?
- if not r:
- # short timeout -- just try again. This is to catch the
- # fd being closed elsewhere, which is only detected
- # when select restarts.
- continue
- data = os.read(fd, 65536)
- if data == "": # generator EOF
- raise Exception
- queue.put(data)
- except Exception:
- queue.put(None)
-
- def watcher_thread(queue, procs):
- # Put None in the queue if either process dies
- while True:
- for p in procs:
- if p.poll() is not None:
- queue.put(None)
- return
- time.sleep(0.25)
-
- def pipewatch(args):
- # Run the processes, etc
- with open(os.devnull, "r") as devnull:
- generator = subprocess.Popen(args.generator, shell = True,
- bufsize = -1, close_fds = True,
- stdin = devnull,
- stdout = subprocess.PIPE,
- stderr = None,
- preexec_fn = os.setpgrp)
- consumer = subprocess.Popen(args.consumer, shell = True,
- bufsize = -11, close_fds = True,
- stdin = subprocess.PIPE,
- stdout = None,
- stderr = None,
- preexec_fn = os.setpgrp)
-
- queue = Queue.Queue(maxsize = 4)
- reader = threading.Thread(target = reader_thread,
- args = (queue, generator.stdout.fileno()))
- reader.start()
- watcher = threading.Thread(target = watcher_thread,
- args = (queue, [generator, consumer]))
- watcher.start()
- try:
- while True:
- try:
- data = queue.get(True, args.timeout)
- if data is None:
- break
- consumer.stdin.write(data)
- except Queue.Empty:
- # Timeout: kill the generator
- fprintf(sys.stderr, "pipewatch: timeout\n")
- generator.terminate()
- break
-
- generator.stdout.close()
- consumer.stdin.close()
- except IOError:
- fprintf(sys.stderr, "pipewatch: I/O error\n")
-
- def kill(proc):
- # Wait for a process to end, or kill it
- def poll_timeout(proc, timeout):
- for x in range(1+int(timeout / 0.1)):
- if proc.poll() is not None:
- break
- time.sleep(0.1)
- return proc.poll()
- try:
- if poll_timeout(proc, 0.5) is None:
- os.killpg(proc.pid, signal.SIGTERM)
- if poll_timeout(proc, 0.5) is None:
- os.killpg(proc.pid, signal.SIGKILL)
- except OSError:
- pass
- return poll_timeout(proc, 0.5)
-
- # Wait for them to die, or kill them
- cret = kill(consumer)
- gret = kill(generator)
-
- # Consume all remaining data in the queue until the reader
- # and watcher threads are done
- while reader.is_alive() or watcher.is_alive():
- queue.get(True, 0.1)
-
- fprintf(sys.stderr, "pipewatch: generator returned %d, " +
- "consumer returned %d\n", gret, cret)
- if gret == 0 and cret == 0:
- sys.exit(0)
- sys.exit(1)
-
- def main(argv = None):
- args = parse_args(argv)
-
- lockfile = open(args.lock, "w")
- if not nilmdb.utils.lock.exclusive_lock(lockfile):
- printf("pipewatch process already running (according to %s)\n",
- args.lock)
- sys.exit(0)
- try:
- # Run as a daemon if requested, otherwise run directly.
- if args.daemon:
- with daemon.DaemonContext(files_preserve = [ lockfile ]):
- pipewatch(args)
- else:
- pipewatch(args)
- finally:
- # Clean up lockfile
- try:
- os.unlink(args.lock)
- except OSError:
- pass
-
- if __name__ == "__main__":
- main()
|