You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

processmanager.py 3.2 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. #!/usr/bin/python
  2. from nilmdb.utils.printf import *
  3. import threading
  4. import multiprocessing
  5. import cStringIO
  6. import sys
  7. import os
  8. import time
  9. import uuid
  10. class LogReceiver(object):
  11. """Spawn a thread that listens to a pipe for log messages,
  12. and stores them locally."""
  13. def __init__(self, pipe):
  14. self.pipe = pipe
  15. self.log = cStringIO.StringIO()
  16. self.thread = threading.Thread(target = self.run)
  17. self.thread.start()
  18. def run(self):
  19. try:
  20. while True:
  21. data = self.pipe.recv_bytes()
  22. self.log.write(data)
  23. except EOFError:
  24. self.pipe.close()
  25. return
  26. def getvalue(self):
  27. return self.log.getvalue()
  28. def clear(self):
  29. self.log = cStringIO.StringIO()
  30. class LogSender(object): # pragma: no cover (runs in a different process)
  31. """File-like object that writes output to a pipe"""
  32. def __init__(self, pipe):
  33. self.pipe = pipe
  34. def close(self):
  35. if self.pipe:
  36. self.pipe.close()
  37. self.pipe = None
  38. def write(self, data):
  39. if self.pipe:
  40. self.pipe.send_bytes(data)
  41. def flush(self):
  42. pass
  43. def isatty(self):
  44. return False
  45. class Process(object):
  46. """Spawn and manage a running process"""
  47. def __init__(self, name, function, parameters):
  48. self.parameters = parameters
  49. self.start_time = None
  50. self.name = name
  51. pipes = multiprocessing.Pipe(duplex = False)
  52. self._log = LogReceiver(pipes[0])
  53. self._process = multiprocessing.Process(
  54. target = self._tramp, name = name,
  55. args = (function, pipes, parameters))
  56. self._process.daemon = True
  57. self._process.start()
  58. pipes[1].close()
  59. self.start_time = time.time()
  60. self.pid = str(uuid.uuid1(self._process.pid or 0))
  61. def _tramp(self, function, pipes, parameters): # pragma: no cover
  62. # Remap stdio in the child before calling function
  63. pipes[0].close()
  64. logfile = LogSender(pipes[1])
  65. sys.stdin = open(os.devnull, 'r')
  66. sys.stdout = logfile
  67. sys.stderr = logfile
  68. function(parameters)
  69. def terminate(self, timeout = 1):
  70. self._process.join(timeout)
  71. if self.alive:
  72. self._process.terminate()
  73. self._process.join(timeout)
  74. if self.alive:
  75. return False
  76. return True
  77. def clear_log(self):
  78. self._log.clear()
  79. @property
  80. def log(self):
  81. return self._log.getvalue()
  82. @property
  83. def alive(self):
  84. return self._process.is_alive()
  85. @property
  86. def exitcode(self):
  87. return self._process.exitcode
  88. class ProcessManager(object):
  89. def __init__(self):
  90. self.processes = {}
  91. def __iter__(self):
  92. return iter(self.processes.keys())
  93. def __getitem__(self, key):
  94. return self.processes[key]
  95. def run(self, name, function, parameters):
  96. new = Process(name, function, parameters)
  97. self.processes[new.pid] = new
  98. return new.pid
  99. def terminate(self, pid):
  100. return self.processes[pid].terminate()
  101. def remove(self, pid):
  102. del self.processes[pid]