class RepeatTimer(Thread):
def __init__(self, interval, function, iterations=0, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.iterations = iterations
self.args = args
self.kwargs = kwargs
self.finished = Event()
def run(self):
count = 0
offset = 0
while not self.finished.is_set() and (self.iterations <= 0 or count < self.iterations):
if count == 0:
dt = datetime.now()
secs = (ceil(dt) - dt).total_seconds()
else:
secs = self.interval - offset
self.finished.wait(secs)
if not self.finished.is_set():
t = time.time()
self.function(*self.args, **self.kwargs)
offset = time.time() - t
count += 1
def cancel(self):
self.finished.set()
class PerpetualTimer(Thread):
"""Call a function after a specified number of seconds:
t = PerpetualTimer(30.0, f, args=[], kwargs={})
t.start()
t.cancel() # stop the timer's action if it's still waiting
"""
def __init__(self, interval, function, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
def run(self):
self.function(*self.args, **self.kwargs)
while not self.finished.is_set():
try:
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
except (SystemExit, KeyboardInterrupt):
print('Cancel timer')
self.cancel()
class PeriodicTimer(Thread):
"""Call a function every /interval/ seconds:
t = PeriodicTimer(30.0, f, args=[], kwargs={})
t.start()
t.cancel()
"""
def __init__(self, interval, function, limit=None, finishCallback=None, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.limit = limit
self.finishCallback = finishCallback
self.args = args
self.kwargs = kwargs
self.finished = Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
def run(self):
i = 0
while (not self.finished.is_set()) and (not self.limit or i < self.limit):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
i += 1
if not self.finished.is_set() and self.limit and self.finishCallback:
self.finishCallback()
class RepeatTimer(Thread):
## Class constructor.
#
# Executes the given function each interval time, a certain number of times (iterations).
#
# \param interval Interval, in seconds
# \param function Function or callable object to execute
# \param iterations Maximum number of iterations
# \param args Function arguments
# \param kwargs Arguments cictionary
def __init__(self, interval, function, iterations=0, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.iterations = iterations
self.args = args
self.kwargs = kwargs
self.finished = Event()
## Statrs thread
#
def run(self):
count = 0
while not self.finished.is_set() and (self.iterations <= 0 or count < self.iterations):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
count += 1
## Cancel execution
#
def cancel(self):
self.finished.set()
def test_mock_pin_edges():
pin = Device.pin_factory.pin(2)
assert pin.when_changed is None
fired = Event()
pin.function = 'input'
pin.edges = 'both'
assert pin.edges == 'both'
pin.drive_low()
assert not pin.state
def changed():
fired.set()
pin.when_changed = changed
pin.drive_high()
assert pin.state
assert fired.is_set()
fired.clear()
pin.edges = 'falling'
pin.drive_low()
assert not pin.state
assert fired.is_set()
fired.clear()
pin.drive_high()
assert pin.state
assert not fired.is_set()
assert pin.edges == 'falling'
class RepeatTimer(Thread):
'''Timer with custom number of iterations'''
def __init__(self, interval, function, iterations=0, args=None, kwargs=None):
Thread.__init__(self)
self.interval = interval
self.function = function
self.iterations = iterations
if args is None:
args = []
self.args = args
if kwargs is None:
kwargs = {}
self.kwargs = kwargs
self.finished = Event()
def run(self):
count = 0
while not self.finished.is_set() and (self.iterations <= 0 or count < self.iterations):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
count += 1
def cancel(self):
self.finished.set()
class Behavior(ParralelClass):
def __init__(self, robot, sensation, callback):
ParralelClass.__init__(self)
self.robot = robot
self.sensation = sensation
self.condition = sensation.condition
self.callback = callback
self._running = Event()
self._running.clear()
self._to_detach = Event()
self._to_detach.clear()
def run(self):
while True:
self.condition.acquire()
# print "acquired"
if self._to_detach.is_set():
self.condition.release()
break
# print "released for detach"
break
if self._running.is_set():
if self.sensation is None:
# print "behave from None"
self.callback(self.robot)
else:
# print "waiting for sensation"
self.condition.wait()
# print "Received sensation"
self.callback(self.robot)
self.condition.release()
class RepeatTimer(Thread):
"""This class provides functionality to call a function at regular
intervals or a specified number of iterations. The timer can be cancelled
at any time."""
def __init__(self, interval, function, iterations=0, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.iterations = iterations
self.args = args
self.kwargs = kwargs
self.finished = Event()
self.daemon = True
def run(self):
count = 0
while not self.finished.is_set() and (self.iterations <= 0
or count < self.iterations):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
count += 1
def cancel(self):
self.finished.set()
class WorkerPoolManager(Thread):
""" Thread managing the worker pool
Parameters:
worker_pool - worker pool
semaphore - worker semaphore
job_queue - job queue
timeout - semaphore time-out in seconds
"""
def __init__(self, worker_pool, semaphore, job_queue, timeout=1.0):
Thread.__init__(self)
self._pool = worker_pool
self._semaphore = semaphore
self._job_queue = job_queue
self._stop_event = Event()
self.timeout = timeout
def stop(self):
""" Tell the thread to stop. """
self._stop_event.set()
def run(self):
""" Thread connection handler. """
logger = getLogger(LOGGER_NAME)
logger.debug("WPM: START")
def callback(exception):
""" Worker callback. """
self._semaphore.release()
logger.debug("WPM: SEMAPHORE RELEASED")
if exception:
logger.error(
"Job execution failed! Reason: %s %s",
type(exception).__name__, exception
)
while not self._stop_event.is_set():
if self._semaphore.acquire(True, self.timeout):
logger.debug("WPM: SEMAPHORE ACQUIRED")
while not self._stop_event.is_set():
try:
_, job_id = self._job_queue.get()
except self._job_queue.Empty:
continue
# load the pickled job
try:
with open(get_task_path(job_id), "rb") as fobj:
_, job = pickle.load(fobj)
except: # pylint: disable=bare-except
logger.error(
"Failed to unpickle job %s! The job is ignored!"
)
self._pool.apply_async(
execute_job, job, callback=callback
)
logger.debug("WPM: APPLIED JOB")
break
logger.debug("WPM: STOP")
class RepeatTimer(Thread):
def __init__(self, interval, function, iterations=0, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.function = function
self.iterations = iterations
self.args = args
self.kwargs = kwargs
self.finished = Event()
def run(self):
count = 0
global DOWNLOADED
while not self.finished.is_set() and (self.iterations <= 0 or count < self.iterations):
if not self.finished.is_set():
time1 = time.time()
self.function(*self.args, **self.kwargs)
time2 = time.time()
count += 1
diff = time2 - time1
if self.interval - diff > 0:
self.finished.wait(self.interval - diff)
def cancel(self):
self.finished.set()
def test_integration():
service_added = Event()
service_removed = Event()
type_ = "_http._tcp.local."
registration_name = "xxxyyy.%s" % type_
def on_service_state_change(zeroconf, service_type, state_change, name):
if name == registration_name:
if state_change is ServiceStateChange.Added:
service_added.set()
elif state_change is ServiceStateChange.Removed:
service_removed.set()
zeroconf_browser = Zeroconf()
browser = ServiceBrowser(zeroconf_browser, type_, [on_service_state_change])
zeroconf_registrar = Zeroconf()
desc = {'path': '/~paulsm/'}
info = ServiceInfo(
type_, registration_name,
socket.inet_aton("10.0.1.2"), 80, 0, 0,
desc, "ash-2.local.")
zeroconf_registrar.register_service(info)
try:
service_added.wait(1)
assert service_added.is_set()
zeroconf_registrar.unregister_service(info)
service_removed.wait(1)
assert service_removed.is_set()
finally:
zeroconf_registrar.close()
browser.cancel()
zeroconf_browser.close()
class MarkerThread(Thread):
def __init__(self, da, zl, marker, coord, conf, pixDim):
Thread.__init__(self)
self.da = da
self.update = Event()
self.update.set()
self.__stop = Event()
self.zl = zl
self.marker = marker
self.coord = coord
self.conf = conf
self.pixDim = pixDim
self.img = self.marker.get_marker_pixbuf(zl)
def run(self):
while not self.__stop.is_set():
self.update.wait()
self.update.clear()
self.draw_markers()
def stop(self):
self.__stop.set()
self.update.set()
def draw_markers(self):
for string in self.marker.positions.keys():
if self.update.is_set() or self.__stop.is_set():
break
mpos = self.marker.positions[string]
if (self.zl <= mpos[2]) and (mpos[0], mpos[1]) != (self.coord[0], self.coord[1]):
gtk.threads_enter()
try:
self.da.draw_marker(self.conf, mpos, self.zl, self.img, self.pixDim, string)
finally:
gtk.threads_leave()
class myObj ():
def __init__(self, name, counter):
self._stopper = Event()
self.name = name
self.counter = counter
self.t = None
self.count = 0
def start(self):
self.t = Thread(name = self.name, target = self._doWork, args=(self.name, ) )
self.t.start()
def _doWork(self,threadName):
self._stopper.clear()
while not self._stopper.is_set():
print "%s: %s" % ( current_thread().getName(),time.ctime(time.time()) )
self.count += 1
self._stopper.wait(10)
print "Done"
print ("Count: %d"%self.count)
def stopit(self):
self._stopper.set()
self.t.join()
def stopped(self):
return self._stopper.is_set()
class TimedFunct(Thread):
def __init__(self, interval, function, repetitions=0, args=[], kwargs={}):
Thread.__init__(self)
self.interval = interval
self.repetitions = repetitions
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
@property
def completed(self):
return self.finished.is_set()
def run(self):
counter = self.repetitions
while not self.finished.is_set():
self.function(*self.args, **self.kwargs)
self.finished.wait(self.interval)
if counter != 0:
counter -= 1
elif self.repetitions != 0:
self.cancel()
class InfiniteTimer(Thread):
'''
A class, extending Thread, which behaves just like threading.Timer,
but instead of returning after the call is completed, it starts
the timer again.
'''
def __init__(self, interval, function, args=[], kwargs={},
immediate=False):
Thread.__init__(self)
self.interval = interval
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = Event()
self.immediate = immediate
def run(self):
if self.immediate:
self.function(*self.args, **self.kwargs)
while not self.finished.is_set():
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
def cancel(self):
'''
Call this method to stop the timer prematurely
'''
self.finished.set()
class ProcessCLIInput(Thread):
instance = None
def __init__(self, interval):
# make sure we actually have a qapp
if QApplication.instance() is None:
QApplication(sys.argv)
Thread.__init__(self)
self.interval = interval
self.event = Event()
self.qtapp = QApplication.instance()
def run(self):
while not self.event.is_set():
try:
self.event.wait(self.interval)
if not self.event.is_set():
if os.name == 'posix' or os.name == 'mac' :
import select
try:
i,o,e = select.select([sys.stdin],[],[],0.0001)
for s in i:
if s == sys.stdin: self.qtapp.exit(0)
except:
if sys.stdin.closed : break
pass
else:
import msvcrt
if msvcrt.kbhit(): self.qtapp.exit(0)
except:
if sys.stdin.closed : break
pass
class Worker(Thread):
def __init__(self, port=5557, local=False):
self._address = 'tcp://{}:{}'.format('127.0.0.1' if local else '0.0.0.0', port)
self._closed = Event()
super(Worker, self).__init__(target=self._serve)
self._close_msg = 'CLOSE'
def _serve(self):
context = zmq.Context()
socket = context.socket(zmq.REP)
socket.bind(self._address)
while not self._closed.is_set():
i = socket.recv_pyobj()
if self._closed.is_set() and i == self._close_msg:
break
socket.send_pyobj(produce(i))
socket.close()
def close(self):
self._closed.set()
context = zmq.Context()
with closing(context.socket(zmq.REQ)) as socket:
socket.connect(self._address)
socket.send_pyobj(self._close_msg)
class BCWorkerMonitor(Thread):
def __init__(self, job, workunit, prefs):
self.sigquit = Event()
self.workunit = workunit
self.prefs = prefs
self.job = job
self.started = datetime.datetime.now()
def run(self):
while not self.sigquit.is_set():
#monitor resource consumption of running container
datapoint = {
'timepoint':datetime.datetime.now() - self.started,
'cpu':0,
'memory':0,
'disk_usage':0,
'message':''
}
requests.post(self.prefs['api'] + self.job['datapoints'], datapoint)
if not self.sigquit.is_set():
time.sleep(100)
def quit(self):
self.sigquit.set()
class PWCounter(object): # PowerCounter class
def __init__(self, start_count, end_count, function, step=1):
self.start_count = start_count
self.end_count = end_count
self.function = function
self.finished = Event()
if self.start_count <= self.end_count:
self.step = -step
self.remaining_counts = self.end_count - self.start_count
elif self.start_count > self.end_count:
self.step = step
self.remaining_counts = self.start_count - self.end_count
self.function_triggered = Event()
def update(self):
if not self.finished.is_set():
self.remaining_counts += self.step
if self.remaining_counts <= 0:
self.finished.set()
elif self.finished.is_set() and not self.function_triggered.is_set():
self.function()
self.function_triggered.set()
def restart(self):
self.function_triggered.clear()
self.finished.clear()
if self.start_count <= self.end_count:
self.remaining_counts = self.end_count - self.start_count
elif self.start_count > self.end_count:
self.remaining_counts = self.start_count - self.end_count
class StoppableTimerThread(Thread):
"""Thread class with a stop() method. The thread itself has to check
regularly for the stopped() condition."""
def __init__(self, interval=None, group=None, target=None, name=None,
args=(), kwargs=None, verbose=None):
Thread.__init__( self, group=None, target=target, name=name,
verbose=verbose)
self.__args = args
self.target = self.__dict__["_Thread__target"]
self.stop_event = Event()
self.interval = int(interval)
def stop(self):
self.stop_event.set()
def stopped(self):
return self.stop_event.is_set()
def run(self):
while not self.stop_event.is_set():
if self.stopped():
return 0
self.target(*self.__args)
self.stop_event.wait(self.interval)
class StubCollector(AbstractCollector):
"""Stub collector to inspect server behavior"""
def __init__(self, data=None, frequency=1.0):
super(StubCollector, self).__init__()
self.data = data
self.max_items = None
self.frequency = float(frequency)
self._shutdown = Event()
self._running = Event()
self.queued_data = []
def start(self):
self._running.set()
while not self._shutdown.is_set():
if self.max_items is None or self.max_items > len(
self.queued_data):
self.queue.put(self.data)
self.queued_data.append(self.data)
sleep(random() / self.frequency)
self._running.clear()
def wait_until_queuing_requests(self, timeout=None):
pass
def shutdown(self):
self._shutdown.set()
def is_queuing_requests(self):
return self._running.is_set()
def wait_until_shutdown(self, timeout=None):
self._shutdown.wait(timeout)
class MainInterfaceAPI(API):
def __init__(self, name, *di_args, **di_kwargs):
API.__init__(self, name, *di_args, **di_kwargs)
self._need_socket_event = Event()
self._running = Event()
def run(self):
while not self._running.is_set():
if not self._need_socket_event.is_set():
output = os.popen('ip route').read()
ip = self._find(output, "src (\d+\.\d+\.\d+\.\d+)", None)
if_name = self._find(output, "dev ([a-zA-Z]{3,4}\d)", None)
gateway = self._find(output, "default via (\d+\.\d+\.\d+\.\d+)", None)
if ip is not None:
self.interface_came_up(ip, if_name, if_name[0:-1] if not if_name is None else None, gateway)
self._running.wait(1)
def _find(self, _input, _format, default):
m = re.search(_format, _input, re.IGNORECASE)
if m:
return m.groups()[0]
return default
def socket_state_callback(self, socket, state):
if state == 0 or state == 11:
self._need_socket_event.set()
else:
self._need_socket_event.clear()
class BehaviorMixer(ParralelClass):
def __init__(self, robot):
ParralelClass.__init__(self)
self.period = 1. / robot.freq
self.robot = robot
self.behaviors = robot._behaviors
self.condition = robot.condition
self._running = Event()
self._running.clear()
self._to_terminate = Event()
self._to_terminate.clear()
self.mode = "average"
def run(self):
while True:
if self._to_terminate.is_set():
break
start_time = self.robot.io.get_simulation_current_time()
if self._running.is_set():
self.condition.acquire()
if not self.behaviors:
activations = 0., 0.
else:
if self.mode == "random":
activations = self._random()
elif self.mode == "average":
activations = self._average()
else:
print self.mode, "is not a valid mode. Choices are \"random\" or \"average\""
self.robot.left_wheel, self.robot.right_wheel = activations
self.condition.release()
self.robot.wait(self.period + start_time - self.robot.io.get_simulation_current_time())
def _average(self):
activations = array([(b.left_wheel, b.right_wheel, b.activation) for b in self.behaviors.itervalues()])
activations = average(activations[:, :2], weights=activations[:, 2] + 1e-10, axis=0)
return activations
def _random(self):
beh = choice(list(self.behaviors.values()))
return beh.left_wheel, beh.right_wheel
def set_mode(self, mode):
self.condition.acquire()
self.mode = mode
self.condition.release()
def execute(self):
self._running.set()
def is_executed(self):
return self._running.is_set()
def stop(self):
self._running.clear()
def _terminate(self):
self.stop()
self._to_terminate.set()
class CallFunctionThread(Thread):
"""
Call function in separate thread.
In case the queue has nothing to, the timeout will determine how long it waits for something to come up.
"""
def __init__(self, daemon=True, timeout=1.0, name=""):
Thread.__init__(self, name="CallFunctionThread_" + name)
self.timeout = timeout
self._run_event = Event() # Set when it is time to stop
self.queue = Queue.PriorityQueue()
self.setDaemon(daemon)
self.count = 0 # Number of Empty exceptions in a row
self._task_available = Event() # Set if tasks are available
self._done_event = Event() # Set when run is done
self._wait_for_tasks = True # Allow new tasks and looping
self._pause_event = Event()
self._pause_event.set() # Default is not pausing
def run(self):
while not self._run_event.is_set() or (self._wait_for_tasks and not self.empty()):
if self._task_available.is_set():
try:
self._pause_event.wait()
_, (f, a, d) = self.queue.get()
f(*a, **d)
self.queue.task_done()
except Queue.Empty:
self._task_available.clear()
except Exception:
logger.exception("")
else:
self._task_available.wait(self.timeout)
self._done_event.set()
def put(self, func, *args, **kargs):
if not self._wait_for_tasks:
return False
self.queue.put((kargs.pop("queue_priority", None), (func, args, kargs)))
self._task_available.set()
return True
def pause(self):
self._pause_event.clear()
def unpause(self):
self._pause_event.set()
def empty(self):
return self.queue.empty()
def queued(self):
return self.queue.qsize()
def stop(self, wait_for_tasks=False, timeout=1.0):
self._wait_for_tasks = wait_for_tasks
self._run_event.set() # Stop looping
self._task_available.set() # Stop sleeping immediately
if wait_for_tasks:
self._done_event.wait(timeout) # Wait till nothing is left to do
def test_input_event_activated():
event = Event()
pin = MockPin(2)
with DigitalInputDevice(pin) as device:
device.when_activated = lambda: event.set()
assert not event.is_set()
pin.drive_high()
assert event.is_set()
def test_input_event_activated():
event = Event()
pin = Device.pin_factory.pin(4)
with DigitalInputDevice(4) as device:
device.when_activated = lambda: event.set()
assert not event.is_set()
pin.drive_high()
assert event.is_set()
def main():
"""
main entry point
returns 0 for normal termination (usually SIGTERM)
"""
return_value = 0
initialize_logging(_log_path)
log = logging.getLogger("main")
log.info("program starts")
halt_event = Event()
set_signal_handler(halt_event)
memcached_client = memcache.Client(_memcached_nodes)
zeromq_context = zmq.Context()
sub_socket = _create_sub_socket(zeromq_context)
expected_sequence = {
_cache_update_channel : None,
}
while not halt_event.is_set():
try:
topic = sub_socket.recv()
assert sub_socket.rcvmore
meta = sub_socket.recv()
if sub_socket.rcvmore:
data = sub_socket.recv()
else:
data = ""
_process_one_event(memcached_client,
expected_sequence,
topic,
meta,
data)
except KeyboardInterrupt: # convenience for testing
log.info("keyboard interrupt: terminating normally")
halt_event.set()
except zmq.ZMQError as zmq_error:
if is_interrupted_system_call(zmq_error) and halt_event.is_set():
log.info("interrupted system call - ok at shutdown")
else:
log.exception("zeromq error processing request")
return_value = 1
halt_event.set()
except Exception:
log.exception("error processing request")
return_value = 1
halt_event.set()
sub_socket.close()
zeromq_context.term()
log.info("program teminates: return value = {0}".format(return_value))
return return_value
class SchedulerFactory:
""" Shares data among classes in Scheduler.
This is needed since there are many threads running in Scheduler.
"""
def __init__(self, docker):
super(SchedulerFactory, self).__init__()
self.docker = docker
self.addrs = {} # ip & port of workers
self.has_worker = Event()
self.all_done = Event()
self.no_new_job = Event()
self.no_future = Event()
self.wait_no_future = Event()
self.job_q = Queue()
self.futures = {}
def add_addr(self, name, addr):
self.addrs[name] = addr
if not self.has_worker.is_set() and name in self.docker.containers.keys():
self.has_worker.set()
log.info("{} workers are connected".format(len(self.addrs)))
def remove_addr(self, name):
try:
addr = self.addrs.pop(name)
except KeyError:
log.warn("addr `{}` doesn't exist".format(name))
return
else:
return addr
def add_future(self, fid, future):
self.futures[fid] = future
if self.no_future.is_set():
self.no_future.clear()
glob_executor.submit(self._wait_future_finish, future)
def get_future(self, fid):
if fid in self.futures:
return self.futures[fid]
else:
return None
def remove_future(self, fid):
if fid in self.futures:
return self.futures.pop(fid)
else:
return None
def _wait_future_finish(self, future):
future.finished.wait()
self.remove_future(future.fid)
if len(self.futures) == 0:
if self.no_new_job.is_set():
self.all_done.set()
elif self.wait_no_future.is_set():
self.no_future.set()
def main():
"""
main entry point
"""
return_code = 0
args = _parse_command_line()
halt_event = Event()
set_signal_handler(halt_event)
zeromq_context = zmq.Context()
reporting_socket = zeromq_context.socket(zmq.PUSH)
reporting_socket.setsockopt(zmq.LINGER, 5000)
reporting_socket.setsockopt(zmq.SNDHWM, args.hwm)
reporting_socket.connect(args.reporting_url)
req_socket = None
ping_message = {"message-type": "ping", "message-id": uuid.uuid1().hex}
result_message = {
"message-type": "ping-result",
"url": args.ping_url,
"result": None,
"socket-reconnection-number": 0,
"check-number": 0,
}
while not halt_event.is_set():
try:
req_socket = _process_one_ping(
halt_event, args, zeromq_context, ping_message, result_message, req_socket, reporting_socket
)
except zmq.ZMQError as zmq_error:
if is_interrupted_system_call(zmq_error) and halt_event.is_set():
pass
else:
logging.exception(str(zmq_error))
halt_event.set()
return_code = 1
except Exception as instance:
logging.exception(str(instance))
halt_event.set()
return_code = 1
else:
halt_event.wait(args.interval)
if req_socket is not None:
req_socket.close()
reporting_socket.close()
zeromq_context.term()
return return_code
class Job(Thread):
def __init__(self, pid, cpu_steps, mem_per_step, net_sched, usb_sched, os):
super(Job, self).__init__()
self.pid = pid
self.cpu_steps = cpu_steps
self.total_steps = 0
self.mem_per_step = mem_per_step
self.net_sched = net_sched
self.usb_sched = usb_sched
self.os = os
self.locked_resources = []
self.blocked = False
self.release_all_resources = Event()
self.done = Event()
self.daemon = True
def run(self):
while not self.done.is_set():
if self.release_all_resources.is_set():
self.free_resources()
else:
self.take_step()
self.os.kill(self)
def take_step(self):
resources = self.required_resources()
if resources:
print("Resources required are %s" % resources)
if not self.os.acquire(self, resources):
return
self.total_steps += self.os.run_process(self)
self.free_resources()
if self.cpu_steps <= self.total_steps:
self.done.set()
def free_resources(self):
for resource in self.locked_resources:
resource.release()
self.locked_resources = []
self.blocked = False
self.release_all_resources.clear()
def required_resources(self):
resources = []
schedules = [(self.net_sched, NetworkResource.name), (self.usb_sched, USBResource.name)]
for schedule, resource in schedules:
if schedule:
ratio = float(self.total_steps) / float(self.cpu_steps)
for start, end in schedule:
if ratio >= start and ratio <= end:
resources.append(resource)
break
return resources
class SensorListener(Thread):
LEFT = 0
MID_LEFT = 1
MID_RIGHT = 3
RIGHT = 4
def __init__(self,
left_getter,
mid_left_getter,
mid_right_getter,
right_getter,
timestep=0.1,
threshold=220):
Thread.__init__(self)
self.daemon = True
# Sensors list
self.sensors_list = [
self.LEFT, self.MID_LEFT, self.MID_RIGHT, self.RIGHT
]
# Signals for each sensors
self.__setattr__("signal" + str(self.LEFT), Signal())
self.__setattr__("signal" + str(self.MID_LEFT), Signal())
self.__setattr__("signal" + str(self.MID_RIGHT), Signal())
self.__setattr__("signal" + str(self.RIGHT), Signal())
# Sensors getter
self.__setattr__("getter" + str(self.LEFT), left_getter)
self.__setattr__("getter" + str(self.MID_LEFT), mid_left_getter)
self.__setattr__("getter" + str(self.MID_RIGHT), mid_right_getter)
self.__setattr__("getter" + str(self.RIGHT), right_getter)
# Timestep
self.timestep = timestep
# Stopping event
self.stop = Event()
# Position threshold
self.threshold = threshold
# Atomatically start
self.start()
def bind(self, idx, func):
self.__setattr__("flag" + str(idx), Flag(func))
self.__getattribute__("flag" + str(idx)).bind(
self.__getattribute__("signal" + str(idx)))
def _handle_sensor(self, idx):
try:
a, b = self.__getattribute__("getter" + str(idx))()
except TimeoutError:
pass
if a < self.threshold or b < self.threshold:
self.__getattribute__("signal" + str(idx)).ping()
def run(self):
while not self.stop.is_set():
# Handle each sensors
for sensor in self.sensors_list:
self._handle_sensor(sensor)
sleep(self.timestep)
def test_swipe():
mbd = MockBlueDot()
mbd.mock_client_connected()
def simulate_swipe(pressed_x, pressed_y, moved_x, moved_y, released_x,
released_y):
mbd.mock_blue_dot_pressed(pressed_x, pressed_y)
mbd.mock_blue_dot_moved(moved_x, moved_y)
mbd.mock_blue_dot_released(released_x, released_y)
#wait_for_swipe
delay_function(lambda: simulate_swipe(-1, 0, 0, 0, 1, 0), 0.5)
assert mbd.wait_for_swipe(1)
#when_swiped
event_swiped = Event()
mbd.when_swiped = lambda: event_swiped.set()
assert not event_swiped.is_set()
#simulate swipe left to right
simulate_swipe(-1, 0, 0, 0, 1, 0)
#check event
assert event_swiped.is_set()
#get the swipe
swipe = BlueDotSwipe(mbd.interaction)
assert swipe.right
assert not swipe.left
assert not swipe.up
assert not swipe.down
#right to left
event_swiped.clear()
simulate_swipe(1, 0, 0, 0, -1, 0)
assert event_swiped.is_set()
swipe = BlueDotSwipe(mbd.interaction)
assert not swipe.right
assert swipe.left
assert not swipe.up
assert not swipe.down
#bottom to top
event_swiped.clear()
simulate_swipe(0, -1, 0, 0, 0, 1)
assert event_swiped.is_set()
swipe = BlueDotSwipe(mbd.interaction)
assert not swipe.right
assert not swipe.left
assert swipe.up
assert not swipe.down
#top to bottom
event_swiped.clear()
simulate_swipe(0, 1, 0, 0, 0, -1)
assert event_swiped.is_set()
swipe = BlueDotSwipe(mbd.interaction)
assert not swipe.right
assert not swipe.left
assert not swipe.up
assert swipe.down
# background
event_swiped.clear()
mbd.set_when_swiped(lambda: delay_function(event_swiped.set, 0.2),
background=True)
simulate_swipe(0, 1, 0, 0, 0, -1)
assert not event_swiped.is_set()
assert event_swiped.wait(1)
class RouteManagerBase(ABC):
def __init__(self,
db_wrapper,
coords,
max_radius,
max_coords_within_radius,
path_to_include_geofence,
path_to_exclude_geofence,
routefile,
mode=None,
init=False,
name="unknown",
settings=None):
self.db_wrapper = db_wrapper
self.init = init
self.name = name
self._coords_unstructured = coords
self.geofence_helper = GeofenceHelper(path_to_include_geofence,
path_to_exclude_geofence)
self._routefile = routefile
self._max_radius = max_radius
self._max_coords_within_radius = max_coords_within_radius
self.settings = settings
self.mode = mode
self._is_started = False
# we want to store the workers using the routemanager
self._workers_registered = []
self._workers_registered_mutex = Lock()
self._last_round_prio = False
self._manager_mutex = RLock()
self._round_started_time = None
if coords is not None:
if init:
fenced_coords = coords
else:
fenced_coords = self.geofence_helper.get_geofenced_coordinates(
coords)
self._route = getJsonRoute(fenced_coords, max_radius,
max_coords_within_radius, routefile)
else:
self._route = None
self._current_index_of_route = 0
self._init_mode_rounds = 0
if self.settings is not None:
self.delay_after_timestamp_prio = self.settings.get(
"delay_after_prio_event", None)
self.starve_route = self.settings.get("starve_route", False)
else:
self.delay_after_timestamp_prio = None
self.starve_route = False
# initialize priority queue variables
self._prio_queue = None
self._update_prio_queue_thread = None
self._stop_update_thread = Event()
def __del__(self):
if self._update_prio_queue_thread is not None:
self._stop_update_thread.set()
self._update_prio_queue_thread.join()
def clear_coords(self):
self._manager_mutex.acquire()
self._coords_unstructured = None
self._manager_mutex.release()
def register_worker(self, worker_name):
self._workers_registered_mutex.acquire()
try:
if worker_name in self._workers_registered:
log.info("Worker %s already registered to routemanager %s" %
(str(worker_name), str(self.name)))
return False
else:
log.info("Worker %s registering to routemanager %s" %
(str(worker_name), str(self.name)))
self._workers_registered.append(worker_name)
return True
finally:
self._workers_registered_mutex.release()
def unregister_worker(self, worker_name):
self._workers_registered_mutex.acquire()
try:
if worker_name in self._workers_registered:
log.info("Worker %s unregistering from routemanager %s" %
(str(worker_name), str(self.name)))
self._workers_registered.remove(worker_name)
else:
# TODO: handle differently?
log.info(
"Worker %s failed unregistering from routemanager %s since subscription was previously "
"lifted" % (str(worker_name), str(self.name)))
if len(self._workers_registered) == 0 and self._is_started:
log.info(
"Routemanager %s does not have any subscribing workers anymore, calling stop"
% str(self.name))
self._quit_route()
finally:
self._workers_registered_mutex.release()
def _check_started(self):
return self._is_started
def _start_priority_queue(self):
if (self._update_prio_queue_thread is None
and (self.delay_after_timestamp_prio is not None
or self.mode == "iv_mitm")
and not self.mode == "pokestops"):
self._prio_queue = []
if self.mode not in ["iv_mitm", "pokestops"]:
self.clustering_helper = ClusteringHelper(
self._max_radius, self._max_coords_within_radius,
self._cluster_priority_queue_criteria())
self._update_prio_queue_thread = Thread(
name="prio_queue_update_" + self.name,
target=self._update_priority_queue_loop)
self._update_prio_queue_thread.daemon = False
self._update_prio_queue_thread.start()
# list_coords is a numpy array of arrays!
def add_coords_numpy(self, list_coords):
fenced_coords = self.geofence_helper.get_geofenced_coordinates(
list_coords)
self._manager_mutex.acquire()
if self._coords_unstructured is None:
self._coords_unstructured = fenced_coords
else:
self._coords_unstructured = np.concatenate(
(self._coords_unstructured, fenced_coords))
self._manager_mutex.release()
def add_coords_list(self, list_coords):
to_be_appended = np.zeros(shape=(len(list_coords), 2))
for i in range(len(list_coords)):
to_be_appended[i][0] = list_coords[i][0]
to_be_appended[i][1] = list_coords[i][1]
self.add_coords_numpy(to_be_appended)
@staticmethod
def calculate_new_route(coords,
max_radius,
max_coords_within_radius,
routefile,
delete_old_route,
num_procs=0):
if delete_old_route and os.path.exists(routefile + ".calc"):
log.debug("Deleting routefile...")
os.remove(routefile + ".calc")
new_route = getJsonRoute(coords,
max_radius,
max_coords_within_radius,
num_processes=num_procs,
routefile=routefile)
return new_route
def recalc_route(self,
max_radius,
max_coords_within_radius,
num_procs=1,
delete_old_route=False,
nofile=False):
current_coords = self._coords_unstructured
if nofile:
routefile = None
else:
routefile = self._routefile
new_route = RouteManagerBase.calculate_new_route(
current_coords, max_radius, max_coords_within_radius, routefile,
delete_old_route, num_procs)
self._manager_mutex.acquire()
self._route = new_route
self._current_index_of_route = 0
self._manager_mutex.release()
def _update_priority_queue_loop(self):
if self._priority_queue_update_interval(
) is None or self._priority_queue_update_interval() == 0:
return
while not self._stop_update_thread.is_set():
# retrieve the latest hatches from DB
# newQueue = self._db_wrapper.get_next_raid_hatches(self._delayAfterHatch, self._geofenceHelper)
new_queue = self._retrieve_latest_priority_queue()
self._merge_priority_queue(new_queue)
time.sleep(self._priority_queue_update_interval())
def _merge_priority_queue(self, new_queue):
if new_queue is not None:
self._manager_mutex.acquire()
merged = set(new_queue + self._prio_queue)
merged = list(merged)
merged = self._filter_priority_queue_internal(merged)
heapq.heapify(merged)
self._prio_queue = merged
self._manager_mutex.release()
log.info("New priorityqueue: %s" % merged)
def date_diff_in_seconds(self, dt2, dt1):
timedelta = dt2 - dt1
return timedelta.days * 24 * 3600 + timedelta.seconds
def dhms_from_seconds(self, seconds):
minutes, seconds = divmod(seconds, 60)
hours, minutes = divmod(minutes, 60)
# days, hours = divmod(hours, 24)
return hours, minutes, seconds
def _get_round_finished_string(self):
round_finish_time = datetime.now()
round_completed_in = (
"%d hours, %d minutes, %d seconds" % (self.dhms_from_seconds(
self.date_diff_in_seconds(round_finish_time,
self._round_started_time))))
return round_completed_in
@abstractmethod
def _retrieve_latest_priority_queue(self):
"""
Method that's supposed to return a plain list containing (timestamp, Location) of the next events of interest
:return:
"""
pass
@abstractmethod
def _start_routemanager(self):
"""
Starts priority queue or whatever the implementations require
:return:
"""
pass
@abstractmethod
def _quit_route(self):
"""
Killing the Route Thread
:return:
"""
pass
@abstractmethod
def _get_coords_post_init(self):
"""
Return list of coords to be fetched and used for routecalc
:return:
"""
pass
@abstractmethod
def _check_coords_before_returning(self, lat, lng):
"""
Return list of coords to be fetched and used for routecalc
:return:
"""
pass
@abstractmethod
def _recalc_route_workertype(self):
"""
Return a new route for worker
:return:
"""
pass
@abstractmethod
def _get_coords_after_finish_route(self):
"""
Return list of coords to be fetched after finish a route
:return:
"""
pass
@abstractmethod
def _cluster_priority_queue_criteria(self):
"""
If you do not want to have any filtering, simply return 0, 0, otherwise simply
return timedelta_seconds, distance
:return:
"""
@abstractmethod
def _priority_queue_update_interval(self):
"""
The time to sleep in between consecutive updates of the priority queue
:return:
"""
def _filter_priority_queue_internal(self, latest):
"""
Filter through the internal priority queue and cluster events within the timedelta and distance returned by
_cluster_priority_queue_criteria
:return:
"""
# timedelta_seconds = self._cluster_priority_queue_criteria()
if self.mode == "iv_mitm":
# exclude IV prioQ to also pass encounterIDs since we do not pass additional information through when
# clustering
return latest
delete_seconds_passed = 0
if self.settings is not None:
delete_seconds_passed = self.settings.get(
"remove_from_queue_backlog", 0)
if delete_seconds_passed is not None:
delete_before = time.time() - delete_seconds_passed
else:
delete_before = 0
latest = [
to_keep for to_keep in latest if not to_keep[0] < delete_before
]
# TODO: sort latest by modified flag of event
# merged = self._merge_queue(latest, self._max_radius, 2, timedelta_seconds)
merged = self.clustering_helper.get_clustered(latest)
return merged
def get_next_location(self):
log.debug("get_next_location of %s called" % str(self.name))
if not self._is_started:
log.info("Starting routemanager %s in get_next_location" %
str(self.name))
self._start_routemanager()
next_lat, next_lng = 0, 0
# first check if a location is available, if not, block until we have one...
got_location = False
while not got_location:
log.debug("%s: Checking if a location is available..." %
str(self.name))
self._manager_mutex.acquire()
got_location = (self._prio_queue is not None
and len(self._prio_queue) > 0 or
(self._route is not None and len(self._route) > 0))
self._manager_mutex.release()
if not got_location:
log.debug("%s: No location available yet" % str(self.name))
time.sleep(0.5)
log.debug(
"%s: Location available, acquiring lock and trying to return location"
% str(self.name))
self._manager_mutex.acquire()
# check priority queue for items of priority that are past our time...
# if that is not the case, simply increase the index in route and return the location on route
# determine whether we move to the next location or the prio queue top's item
if (self.delay_after_timestamp_prio is not None
and ((not self._last_round_prio or self.starve_route)
and self._prio_queue and len(self._prio_queue) > 0
and self._prio_queue[0][0] < time.time())):
log.debug("%s: Priority event" % str(self.name))
next_stop = heapq.heappop(self._prio_queue)[1]
next_lat = next_stop.lat
next_lng = next_stop.lng
self._last_round_prio = True
log.info(
"Round of route %s is moving to %s, %s for a priority event" %
(str(self.name), str(next_lat), str(next_lng)))
else:
log.debug("%s: Moving on with route" % str(self.name))
if self._current_index_of_route == 0:
if self._round_started_time is not None:
log.info(
"Round of route %s reached the first spot again. It took %s"
% (str(self.name),
str(self._get_round_finished_string())))
self._round_started_time = datetime.now()
if len(self._route) == 0: return None
log.info("Round of route %s started at %s" %
(str(self.name), str(self._round_started_time)))
# continue as usual
if self._current_index_of_route < len(self._route):
log.info("Moving on with location %s" %
self._route[self._current_index_of_route])
next_lat = self._route[self._current_index_of_route]['lat']
next_lng = self._route[self._current_index_of_route]['lng']
self._current_index_of_route += 1
if self.init and self._current_index_of_route >= len(self._route):
self._init_mode_rounds += 1
if self.init and self._current_index_of_route >= len(self._route) and \
self._init_mode_rounds >= int(self.settings.get("init_mode_rounds", 1)):
# we are done with init, let's calculate a new route
log.warning(
"Init of %s done, it took %s, calculating new route..." %
(str(self.name), self._get_round_finished_string()))
self.clear_coords()
coords = self._get_coords_post_init()
log.debug("Setting %s coords to as new points in route of %s" %
(str(len(coords)), str(self.name)))
self.add_coords_list(coords)
log.debug("Route of %s is being calculated" % str(self.name))
self._recalc_route_workertype()
self.init = False
self.change_init_mapping(self.name)
self._manager_mutex.release()
return self.get_next_location()
elif self._current_index_of_route == len(self._route):
log.info('Reaching last coord of route')
elif self._current_index_of_route > len(self._route):
self._current_index_of_route = 0
coords_after_round = self._get_coords_after_finish_route()
# TODO: check isinstance list?
if coords_after_round is not None:
self.clear_coords()
coords = coords_after_round
self.add_coords_list(coords)
self._recalc_route_workertype()
if len(self._route) == 0: return None
next_lat = self._route[self._current_index_of_route]['lat']
next_lng = self._route[self._current_index_of_route]['lng']
self._manager_mutex.release()
return Location(next_lat, next_lng)
self._manager_mutex.release()
return self.get_next_location()
self._last_round_prio = False
log.info(
"%s done grabbing next coord, releasing lock and returning location: %s, %s"
% (str(self.name), str(next_lat), str(next_lng)))
self._manager_mutex.release()
if self._check_coords_before_returning(next_lat, next_lng):
return Location(next_lat, next_lng)
else:
return self.get_next_location()
def del_from_route(self):
log.debug(
"%s: Location available, acquiring lock and trying to return location"
% str(self.name))
self._manager_mutex.acquire()
log.info('Removing coords from Route')
self._route.pop(int(self._current_index_of_route) - 1)
self._current_index_of_route -= 1
if len(self._route) == 0:
log.info('No more coords are available... Sleeping.')
self._manager_mutex.release()
def change_init_mapping(self, name_area):
with open('configs/mappings.json') as f:
vars = json.load(f)
for var in vars['areas']:
if (var['name']) == name_area:
var['init'] = bool(False)
with open('mappings.json', 'w') as outfile:
json.dump(vars, outfile, indent=4, sort_keys=True)
def get_route_status(self):
if self._route:
return self._current_index_of_route, len(self._route)
return self._current_index_of_route, self._current_index_of_route
class AppEngine(EventSupport):
def __init__(self, tesla_api):
super().__init__(EngineEvents.to_array())
self.commands = SupportedEndpoints
self.events = EngineEvents
self._terminate_poll = Event()
self._poll_terminated = Event()
self._poll_terminated.set()
self._api_active_accesses = ThreadSafeCounter()
self._api_not_updating = Event()
self._api_not_updating.set()
self._tesla_api = tesla_api
self.poll_rate = 3
self._thread_pool = concurrent.futures.ThreadPoolExecutor()
def _thread_api_poller(self):
logger.info(f'Poller thread started')
self.raise_event(self.events.POLL_STARTING)
first_loop = True
while first_loop or not self._terminate_poll.wait(self.poll_rate):
self._api_not_updating.wait()
with self._api_active_accesses:
first_loop = False
logger.info(f'Poll tick started')
new_frames = []
future_results = []
try:
cars = self._tesla_api.send_request(
self._tesla_api.urls.VEHICLE_LIST)
except TeslaApiError as error:
if error.status_code == 401:
self.raise_event(self.events.CREDENTIALS_REQUIRED)
self._terminate_poll.set()
break
else:
raise
with concurrent.futures.ThreadPoolExecutor() as executor:
for car in cars:
if car['state'] == 'online':
future_results.append(
executor.submit(
self._tesla_api.send_request,
self._tesla_api.urls.VEHICLE_DATA,
car['id']))
else:
if car['state'] != 'asleep' and car[
'state'] != 'offline':
logger.critical(car['state'])
new_frames.append(car)
for future in future_results:
try:
new_frames.insert(0, future.result())
except Exception as exception:
logger.error(exception)
if len(new_frames) > 0:
self.raise_event(self.events.NEW_FRAMES_READY,
frames=new_frames)
logger.info(f'Poll tick completed')
self._poll_terminated.set()
logger.info(f'Poller terminated')
self.raise_event(self.events.POLL_STOPPED)
def _thread_api_command(self, command_id, command_name, *args):
logger.info(f'Api command thread started')
try:
self._api_not_updating.wait()
with self._api_active_accesses:
result = self._tesla_api.send_request(command_name, *args)
self.raise_event(self.events.COMMAND_COMPLETED, command_id,
command_name, result, True, '')
except Exception as error:
logger.error(error)
self.raise_event(self.events.COMMAND_COMPLETED, command_id,
command_name, {}, False, str(error))
logger.info(f'Api command thread completed')
def _thread_credentials_load(self, command_id, user_name, password, token):
logger.info(f'Credentials load started')
self._api_not_updating.wait()
self._api_not_updating.clear()
self._api_active_accesses.counter_is_zero.wait()
event_args = []
try:
if token:
# Test the token
self._tesla_api.token = token
self._tesla_api.send_request(self._tesla_api.urls.VEHICLE_LIST)
result = {'access_token': token}
else:
# No token so use the user and password to obtain a new one
result = self._tesla_api.get_token(user_name, password)
self._tesla_api.token = result['access_token']
event_args = [
self.events.CREDENTIALS_RESULT, command_id, result, True, ''
]
except Exception as error:
logger.error(error)
event_args = [
self.events.CREDENTIALS_RESULT, command_id, {}, False,
str(error)
]
self._api_not_updating.set()
self.raise_event(*event_args)
logger.info(f'Credentials load completed')
def poll_start(self):
if self._poll_terminated.is_set():
if not self._tesla_api.token:
logger.debug('poll_start invoked but credentials are required')
self.raise_event(self.events.CREDENTIALS_REQUIRED)
else:
logger.debug('poll_start invoked')
self._poll_terminated.clear()
self._terminate_poll.clear()
self._thread_pool.submit(self._thread_api_poller)
else:
logger.debug('poll_start invoked but terminated flag is not set')
def poll_stop(self):
self._terminate_poll.set()
self.raise_event(self.events.POLL_STOPPING)
def send_car_command(self, command_name, *args):
command_id = uuid.uuid4().hex
self._thread_pool.submit(self._thread_api_command, command_id,
command_name, *args)
return command_id
def load_credentials(self, user_name=None, password=None, token=None):
if not user_name and not token:
raise EngineValidationError('The user name can not be empty')
if not password and not token:
raise EngineValidationError('The password can not be empty')
command_id = uuid.uuid4().hex
self._thread_pool.submit(self._thread_credentials_load, command_id,
user_name, password, token)
return command_id
def get_current_token(self):
return self._tesla_api.token
def sanitize_option_codes(self, codes):
return re.sub(f'[^{VALID_CHARS}]', '', codes.upper())
def translate_option_codes(self, codes):
sanitized_codes = self.sanitize_option_codes(codes)
return translate_codes(sanitized_codes)
def switch_api(self, new_api, is_demo=False):
logger.info(f'Trying to switch api. Demo {is_demo}')
self._api_not_updating.wait()
self._api_not_updating.clear()
self._api_active_accesses.counter_is_zero.wait()
logger.info(f'Switching api now. Demo {is_demo}')
self._tesla_api = new_api
self._api_not_updating.set()
self.raise_event(self.events.API_SWITCHED, is_demo)
def enable_demo_mode(self):
self.raise_event(self.events.REQUEST_DEMO_API)
def enable_real_mode(self):
self.raise_event(self.events.REQUEST_REAL_API)
def reset_api_token(self):
self._api_not_updating.wait()
self._api_not_updating.clear()
self._api_active_accesses.counter_is_zero.wait()
self._tesla_api.token = ''
self._api_not_updating.set()
def get_engine_version(self):
return __version__
class ServerSupervisor(PatternMatchingEventHandler):
def __init__(self,
slave_process_args,
max_seconds_between_restarts,
directories_to_monitor=['.']):
super().__init__(ignore_patterns=[
'.git', '.floo', '*.pyc', '*.pyo', '*/__pycache__/*', '*.db-*',
'*.db'
],
ignore_directories=True)
self.serving_process = SlaveProcess(sys.argv[0], slave_process_args)
self.max_seconds_between_restarts = max_seconds_between_restarts
self.directories_to_monitor = directories_to_monitor
self.directory_observers = []
self.files_changed = Event()
self.request_stop_running = Event()
#override for PatternMatchingEventHandler
def on_any_event(self, event):
self.files_changed.set()
def start_observing_directories(self):
for directory in self.directories_to_monitor:
directory_observer = Observer()
directory_observer.schedule(self, directory, recursive=True)
directory_observer.start()
self.directory_observers.append(directory_observer)
def stop_observing_directories(self):
for directory_observer in self.directory_observers:
directory_observer.stop()
directory_observer.join()
def pause(self):
time.sleep(self.max_seconds_between_restarts)
def stop(self):
self.request_stop_running.set()
def stop_serving(self):
if self.serving_process.is_running():
self.serving_process.terminate()
self.stop_observing_directories()
def run(self):
self.files_changed.clear()
self.start_observing_directories()
self.serving_process.start()
while not self.request_stop_running.is_set():
try:
self.pause()
if self.files_changed.is_set():
logging.getLogger(__name__).debug(
'Changes to filesystem detected, scheduling a restart...'
)
self.files_changed.clear()
if self.serving_process.is_running():
self.serving_process.restart()
else:
self.serving_process.start()
except KeyboardInterrupt:
self.stop()
self.stop_serving()
def adapter(stream, queue: Queue, event: Event):
while not event.is_set():
out = stream.read1()
if out: queue.put(out)
else: break
stream.close()
class WebsocketServer(object):
def __init__(self, args, mitm_mapper, db_wrapper, routemanagers, device_mappings, auths, pogoWindowManager,
configmode=False):
self.__current_users = {}
self.__current_users_mutex = Lock()
self.__stop_server = Event()
self.args = args
self.__listen_address = args.ws_ip
self.__listen_port = int(args.ws_port)
self.__send_queue = queue.Queue()
self.__received = {}
self.__received_mutex = Lock()
self.__requests = {}
self.__requests_mutex = Lock()
self.__db_wrapper = db_wrapper
self.__device_mappings = device_mappings
self.__routemanagers = routemanagers
self.__auths = auths
self.__pogoWindowManager = pogoWindowManager
self.__mitm_mapper = mitm_mapper
self.__next_id = 0
self.__id_mutex = Lock()
self._configmode = configmode
self.__loop = None
def start_server(self):
logger.info("Starting websocket server...")
self.__loop = asyncio.new_event_loop()
# build list of origin IDs
allowed_origins = []
for device in self.__device_mappings.keys():
allowed_origins.append(device)
logger.debug("Device mappings: {}", str(self.__device_mappings))
logger.debug("Allowed origins derived: {}", str(allowed_origins))
asyncio.set_event_loop(self.__loop)
self.__loop.run_until_complete(
websockets.serve(self.handler, self.__listen_address, self.__listen_port, max_size=2 ** 25,
origins=allowed_origins, ping_timeout=10, ping_interval=15))
self.__loop.run_forever()
def stop_server(self):
# TODO: cleanup workers...
self.__stop_server.set()
self.__current_users_mutex.acquire()
for id, worker in self.__current_users.items():
logger.info('Stopping worker {} to apply new mappings.', id)
worker[1].stop_worker()
self.__current_users_mutex.release()
# wait for all workers to be stopped...
while True:
self.__current_users_mutex.acquire()
if len(self.__current_users) == 0:
self.__current_users_mutex.release()
break
else:
self.__current_users_mutex.release()
time.sleep(1)
for routemanager in self.__routemanagers.keys():
area = self.__routemanagers.get(routemanager, None)
if area is None:
continue
area["routemanager"].stop_routemanager()
if self.__loop is not None:
self.__loop.call_soon_threadsafe(self.__loop.stop)
async def handler(self, websocket_client_connection, path):
logger.info("Waiting for connection...")
# wait for a connection...
continue_work = await self.__register(websocket_client_connection)
if not continue_work:
logger.error("Failed registering client, closing connection")
await websocket_client_connection.close()
return
consumer_task = asyncio.ensure_future(
self.__consumer_handler(websocket_client_connection))
producer_task = asyncio.ensure_future(
self.__producer_handler(websocket_client_connection))
done, pending = await asyncio.wait(
[producer_task, consumer_task],
return_when=asyncio.FIRST_COMPLETED,
)
logger.info("consumer or producer of {} stopped, cancelling pending tasks", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
for task in pending:
task.cancel()
logger.info("Awaiting unregister of {}", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
await self.__unregister(websocket_client_connection)
logger.info("All done with {}", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
async def __register(self, websocket_client_connection):
logger.info("Client {} registering", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
if self.__stop_server.is_set():
logger.info(
"MAD is set to shut down, not accepting new connection")
return False
try:
id = str(
websocket_client_connection.request_headers.get_all("Origin")[0])
except IndexError:
logger.warning("Client from {} tried to connect without Origin header", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
return False
if self.__auths:
try:
authBase64 = str(
websocket_client_connection.request_headers.get_all("Authorization")[0])
except IndexError:
logger.warning("Client from {} tried to connect without auth header", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
return False
self.__current_users_mutex.acquire()
try:
logger.debug("Checking if {} is already present", str(id))
user_present = self.__current_users.get(id)
if user_present is not None:
logger.warning("Worker with origin {} is already running, killing the running one and have client reconnect",
str(websocket_client_connection.request_headers.get_all("Origin")[0]))
user_present[1].stop_worker()
return False
elif self.__auths and authBase64 and not check_auth(authBase64, self.args, self.__auths):
logger.warning("Invalid auth details received from {}", str(
websocket_client_connection.request_headers.get_all("Origin")[0]))
return False
if self._configmode:
worker = WorkerConfigmode(self.args, id, self)
logger.debug("Starting worker for {}", str(id))
new_worker_thread = Thread(
name='worker_%s' % id, target=worker.start_worker)
self.__current_users[id] = [
new_worker_thread, worker, websocket_client_connection, 0]
return True
last_known_state = {}
client_mapping = self.__device_mappings[id]
devicesettings = client_mapping["settings"]
logger.info("Setting up routemanagers for {}", str(id))
if client_mapping.get("walker", None) is not None:
if "walker_area_index" not in devicesettings:
devicesettings['walker_area_index'] = 0
devicesettings['finished'] = False
devicesettings['last_action_time'] = None
devicesettings['last_cleanup_time'] = None
walker_index = devicesettings.get('walker_area_index', 0)
if walker_index > 0:
# check status of last area
if not devicesettings.get('finished', False):
logger.info(
'Something wrong with last round - get back to old area')
walker_index -= 1
devicesettings['walker_area_index'] = walker_index
walker_area_array = client_mapping["walker"]
walker_settings = walker_area_array[walker_index]
# preckeck walker setting
while not pre_check_value(walker_settings) and walker_index-1 <= len(walker_area_array):
walker_area_name = walker_area_array[walker_index]['walkerarea']
logger.info(
'{} dont using area {} - Walkervalue out of range', str(id), str(walker_area_name))
if walker_index >= len(walker_area_array) - 1:
logger.error(
'Dont find any working area - check your config')
walker_index = 0
devicesettings['walker_area_index'] = walker_index
walker_settings = walker_area_array[walker_index]
break
walker_index += 1
devicesettings['walker_area_index'] = walker_index
walker_settings = walker_area_array[walker_index]
if devicesettings['walker_area_index'] >= len(walker_area_array):
# check if array is smaller then expected - f.e. on the fly changes in mappings.json
devicesettings['walker_area_index'] = 0
devicesettings['finished'] = False
walker_index = devicesettings.get('walker_area_index', 0)
walker_area_name = walker_area_array[walker_index]['walkerarea']
if walker_area_name not in self.__routemanagers:
raise WrongAreaInWalker()
logger.debug('Devicesettings {}: {}', str(id), devicesettings)
logger.info('{} using walker area {} [{}/{}]', str(id), str(
walker_area_name), str(walker_index+1), str(len(walker_area_array)))
walker_routemanager = \
self.__routemanagers[walker_area_name].get(
"routemanager", None)
devicesettings['walker_area_index'] += 1
devicesettings['finished'] = False
if walker_index >= len(walker_area_array) - 1:
devicesettings['walker_area_index'] = 0
# set global mon_iv
client_mapping['mon_ids_iv'] = \
self.__routemanagers[walker_area_name].get(
"routemanager").settings.get("mon_ids_iv", [])
else:
walker_routemanager = None
if "last_location" not in devicesettings:
devicesettings['last_location'] = Location(0.0, 0.0)
logger.debug("Setting up worker for {}", str(id))
if walker_routemanager is None:
pass
elif walker_routemanager.mode in ["raids_mitm", "mon_mitm", "iv_mitm"]:
worker = WorkerMITM(self.args, id, last_known_state, self, walker_routemanager,
self.__mitm_mapper, devicesettings, db_wrapper=self.__db_wrapper,
pogoWindowManager=self.__pogoWindowManager, walker=walker_settings)
elif walker_routemanager.mode in ["raids_ocr"]:
from worker.WorkerOCR import WorkerOCR
worker = WorkerOCR(self.args, id, last_known_state, self, walker_routemanager,
devicesettings, db_wrapper=self.__db_wrapper,
pogoWindowManager=self.__pogoWindowManager, walker=walker_settings)
elif walker_routemanager.mode in ["pokestops"]:
worker = WorkerQuests(self.args, id, last_known_state, self, walker_routemanager,
self.__mitm_mapper, devicesettings, db_wrapper=self.__db_wrapper,
pogoWindowManager=self.__pogoWindowManager, walker=walker_settings)
elif walker_routemanager.mode in ["idle"]:
worker = WorkerConfigmode(self.args, id, self)
else:
logger.error("Mode not implemented")
sys.exit(1)
logger.debug("Starting worker for {}", str(id))
new_worker_thread = Thread(
name='worker_%s' % id, target=worker.start_worker)
new_worker_thread.daemon = False
self.__current_users[id] = [new_worker_thread,
worker, websocket_client_connection, 0]
new_worker_thread.start()
except WrongAreaInWalker:
logger.error('Unknown Area in Walker settings - check config')
finally:
self.__current_users_mutex.release()
return True
async def __unregister(self, websocket_client_connection):
worker_id = str(
websocket_client_connection.request_headers.get_all("Origin")[0])
self.__current_users_mutex.acquire()
worker = self.__current_users.get(worker_id, None)
if worker is not None:
self.__current_users.pop(worker_id)
self.__current_users_mutex.release()
logger.info("Worker {} unregistered", str(worker_id))
async def __producer_handler(self, websocket_client_connection):
while websocket_client_connection.open:
# logger.debug("Connection still open, trying to send next message")
# retrieve next message from queue to be sent, block if empty
next = None
while next is None and websocket_client_connection.open:
logger.debug("Retrieving next message to send")
next = await self.__retrieve_next_send(websocket_client_connection)
if next is None:
# logger.debug("next is None, stopping connection...")
return
await self.__send_specific(websocket_client_connection, next.id, next.message)
async def __send_specific(self, websocket_client_connection, id, message):
# await websocket_client_connection.send(message)
for key, value in self.__current_users.items():
if key == id and value[2].open:
await value[2].send(message)
async def __retrieve_next_send(self, websocket_client_connection):
found = None
while found is None and websocket_client_connection.open:
try:
found = self.__send_queue.get_nowait()
except Exception as e:
await asyncio.sleep(0.02)
if not websocket_client_connection.open:
logger.warning(
"retrieve_next_send: connection closed, returning None")
return found
async def __consumer_handler(self, websocket_client_connection):
if websocket_client_connection is None:
return
worker_id = str(
websocket_client_connection.request_headers.get_all("Origin")[0])
logger.info("Consumer handler of {} starting", str(worker_id))
while websocket_client_connection.open:
message = None
try:
message = await asyncio.wait_for(websocket_client_connection.recv(), timeout=2.0)
except asyncio.TimeoutError as te:
await asyncio.sleep(0.02)
except websockets.exceptions.ConnectionClosed as cc:
logger.warning(
"Connection to {} was closed, stopping worker", str(worker_id))
self.__current_users_mutex.acquire()
worker = self.__current_users.get(worker_id, None)
self.__current_users_mutex.release()
if worker is not None:
# TODO: do it abruptly in the worker, maybe set a flag to be checked for in send_and_wait to
# TODO: throw an exception
worker[1].stop_worker()
self.clean_up_user(worker_id, None)
return
if message is not None:
await self.__on_message(message)
logger.warning(
"Connection of {} closed in consumer_handler", str(worker_id))
def clean_up_user(self, worker_id, worker_instance):
"""
:param worker_id: The ID/Origin of the worker
:param worker_instance: None if the cleanup is called from within the websocket server
:return:
"""
self.__current_users_mutex.acquire()
if worker_id in self.__current_users.keys() and (worker_instance is None
or self.__current_users[worker_id][1] == worker_instance):
if self.__current_users[worker_id][2].open:
logger.info("Calling close for {}...", str(worker_id))
asyncio.ensure_future(
self.__current_users[worker_id][2].close(), loop=self.__loop)
self.__current_users.pop(worker_id)
logger.info("Info of {} removed in websocket", str(worker_id))
self.__current_users_mutex.release()
async def __on_message(self, message):
id = -1
response = None
if isinstance(message, str):
logger.debug("Receiving message: {}", str(message.strip()))
splitup = message.split(";")
id = int(splitup[0])
response = splitup[1]
else:
logger.debug("Received binary values.")
id = int.from_bytes(message[:4], byteorder='big', signed=False)
response = message[4:]
await self.__set_response(id, response)
if not await self.__set_event(id):
# remove the response again - though that is kinda stupid
self.__pop_response(id)
async def __set_event(self, id):
result = False
self.__requests_mutex.acquire()
if id in self.__requests:
self.__requests[id].set()
result = True
else:
# the request has already been deleted due to a timeout...
logger.error("Request has already been deleted...")
self.__requests_mutex.release()
return result
async def __set_response(self, id, message):
self.__received_mutex.acquire()
self.__received[id] = message
self.__received_mutex.release()
def __pop_response(self, id):
self.__received_mutex.acquire()
message = self.__received.pop(id)
self.__received_mutex.release()
return message
def __get_new_message_id(self):
self.__id_mutex.acquire()
self.__next_id += 1
self.__next_id = int(math.fmod(self.__next_id, 100000))
if self.__next_id == 100000:
self.__next_id = 1
toBeReturned = self.__next_id
self.__id_mutex.release()
return toBeReturned
def __send(self, id, to_be_sent):
next_message = OutgoingMessage(id, to_be_sent)
self.__send_queue.put(next_message)
def send_and_wait(self, id, worker_instance, message, timeout):
logger.debug("{} sending command: {}", str(id), message.strip())
self.__current_users_mutex.acquire()
user_entry = self.__current_users.get(id, None)
self.__current_users_mutex.release()
if user_entry is None or user_entry[1] != worker_instance and worker_instance != 'madmin':
raise WebsocketWorkerRemovedException
message_id = self.__get_new_message_id()
message_event = Event()
message_event.clear()
self.__set_request(message_id, message_event)
to_be_sent = u"%s;%s" % (str(message_id), message)
logger.debug("To be sent: {}", to_be_sent.strip())
self.__send(id, to_be_sent)
# now wait for the response!
result = None
logger.debug("Timeout: {}", str(timeout))
if message_event.wait(timeout):
logger.debug("Received answer in time, popping response")
self.__reset_fail_counter(id)
result = self.__pop_response(message_id)
if isinstance(result, str):
logger.debug("Response to {}: {}",
str(id), str(result.strip()))
else:
logger.debug("Received binary data to {}, starting with {}", str(
id), str(result[:10]))
else:
# timeout reached
logger.warning("Timeout, increasing timeout-counter")
# TODO: why is the user removed here?
new_count = self.__increase_fail_counter(id)
if new_count > 5:
logger.error("5 consecutive timeouts to {}, cleanup", str(id))
# TODO: signal worker to stop and NOT cleanup the websocket by itself!
self.clean_up_user(id, None)
raise WebsocketWorkerTimeoutException
self.__remove_request(message_id)
return result
def __set_request(self, id, event):
self.__requests_mutex.acquire()
self.__requests[id] = event
self.__requests_mutex.release()
def __reset_fail_counter(self, id):
self.__current_users_mutex.acquire()
if id in self.__current_users.keys():
self.__current_users[id][3] = 0
self.__current_users_mutex.release()
def __increase_fail_counter(self, id):
self.__current_users_mutex.acquire()
if id in self.__current_users.keys():
new_count = self.__current_users[id][3] + 1
self.__current_users[id][3] = new_count
else:
new_count = 100
self.__current_users_mutex.release()
return new_count
def __remove_request(self, message_id):
self.__requests_mutex.acquire()
self.__requests.pop(message_id)
self.__requests_mutex.release()
def update_settings(self, routemanagers, device_mappings, auths):
for dev in self.__device_mappings:
if "last_location" in self.__device_mappings[dev]['settings']:
device_mappings[dev]['settings']["last_location"] = \
self.__device_mappings[dev]['settings']["last_location"]
if "walker_area_index" in self.__device_mappings[dev]['settings']:
device_mappings[dev]['settings']["walker_area_index"] = \
self.__device_mappings[dev]['settings']["walker_area_index"]
if "last_mode" in self.__device_mappings[dev]['settings']:
device_mappings[dev]['settings']["last_mode"] = \
self.__device_mappings[dev]['settings']["last_mode"]
self.__current_users_mutex.acquire()
# save reference to old routemanagers to stop them
old_routemanagers = routemanagers
self.__device_mappings = device_mappings
self.__routemanagers = routemanagers
self.__auths = auths
for id, worker in self.__current_users.items():
logger.info('Stopping worker {} to apply new mappings.', id)
worker[1].stop_worker()
self.__current_users_mutex.release()
for routemanager in old_routemanagers.keys():
area = routemanagers.get(routemanager, None)
if area is None:
continue
area["routemanager"].stop_routemanager()
def get_reg_origins(self):
return self.__current_users
def get_origin_communicator(self, origin):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].get_communicator()
return None
def set_geofix_sleeptime_worker(self, origin, sleeptime):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].set_geofix_sleeptime(sleeptime)
return False
class CoversWidget(QWidget): # {{{
chosen = pyqtSignal()
finished = pyqtSignal()
def __init__(self, log, current_cover, parent=None):
QWidget.__init__(self, parent)
self.log = log
self.abort = Event()
self.l = l = QGridLayout()
self.setLayout(l)
self.msg = QLabel()
self.msg.setWordWrap(True)
l.addWidget(self.msg, 0, 0)
self.covers_view = CoversView(current_cover, self)
self.covers_view.chosen.connect(self.chosen)
l.addWidget(self.covers_view, 1, 0)
self.continue_processing = True
def reset_covers(self):
self.covers_view.reset_covers()
def start(self, book, current_cover, title, authors, caches):
self.continue_processing = True
self.abort.clear()
self.book, self.current_cover = book, current_cover
self.title, self.authors = title, authors
self.log('Starting cover download for:', book.title)
self.log('Query:', title, authors, self.book.identifiers)
self.msg.setText(
'<p>' +
_('Downloading covers for <b>%s</b>, please wait...') % book.title)
self.covers_view.start()
self.worker = CoverWorker(self.log, self.abort, self.title,
self.authors, book.identifiers, caches)
self.worker.start()
QTimer.singleShot(50, self.check)
self.covers_view.setFocus(Qt.OtherFocusReason)
def check(self):
if self.worker.is_alive() and not self.abort.is_set():
QTimer.singleShot(50, self.check)
try:
self.process_result(self.worker.rq.get_nowait())
except Empty:
pass
else:
self.process_results()
def process_results(self):
while self.continue_processing:
try:
self.process_result(self.worker.rq.get_nowait())
except Empty:
break
if self.continue_processing:
self.covers_view.clear_failed()
if self.worker.error is not None:
error_dialog(self,
_('Download failed'),
_('Failed to download any covers, click'
' "Show details" for details.'),
det_msg=self.worker.error,
show=True)
num = self.covers_view.model().rowCount()
if num < 2:
txt = _(
'Could not find any covers for <b>%s</b>') % self.book.title
else:
txt = _('Found <b>%(num)d</b> covers of %(title)s. '
'Pick the one you like best.') % dict(num=num - 1,
title=self.title)
self.msg.setText(txt)
self.finished.emit()
def process_result(self, result):
if not self.continue_processing:
return
plugin_name, width, height, fmt, data = result
self.covers_view.model().update_result(plugin_name, width, height,
data)
def cleanup(self):
self.covers_view.delegate.stop_animation()
self.continue_processing = False
def cancel(self):
self.cleanup()
self.abort.set()
def cover_pixmap(self):
idx = None
for i in self.covers_view.selectionModel().selectedIndexes():
if i.isValid():
idx = i
break
if idx is None:
idx = self.covers_view.currentIndex()
return self.covers_view.model().cover_pixmap(idx)
class CraneCuberDaemon(object):
def __init__(self, dev_video, ip, port):
self.shutdown_event = Event()
self.dev_video = dev_video
self.ip = ip
self.port = port
def __str__(self):
return 'CraneCuberDaemon'
def signal_handler(self, signal, frame):
log.info("received SIGINT or SIGTERM")
self.shutdown_event.set()
def main(self):
caught_exception = False
# Log the process ID upon start-up.
pid = os.getpid()
log.info('cranecuberd started with PID %s for /dev/video%d' % (pid, self.dev_video))
tcp_socket = open_tcp_socket(self.ip, self.port)
while True:
# Wrap everything else in try/except so that we can log errors
# and exit cleanly
try:
if self.shutdown_event.is_set():
log.info("Shutdown signal RXed. Breaking out of the loop.")
break
try:
(connection, _) = tcp_socket.accept()
except socket.error as e:
if isinstance(e.args, tuple) and e[0] == 4:
# 4 is 'Interrupted system call', a.k.a. SIGINT.
# The user wants to stop cranecuberd.
log.info("socket.accept() caught signal, starting shutdown")
raise BrokenSocket("socket.accept() caught signal, starting shutdown")
else:
log.info("socket hit error\n%s" % e)
tcp_socket.close()
tcp_socket = open_tcp_socket()
continue
total_data = []
# RX the entire packet
while True:
data = connection.recv(4096)
# If the client is using python2 data will be a str but if they
# are using python3 data will be encoded and must be decoded to
# a str
if not isinstance(data, str):
data = data.decode()
data = data.strip()
log.info("RXed %s" % data)
total_data.append(data)
data = ''.join(total_data)
# Do we have the entire packet?
if data.startswith('<START>') and data.endswith('<END>'):
# Remove the <START> and <END>
data = data[7:-5]
if data.startswith('TAKE_PICTURE'):
side_name = data.strip().split(':')[1]
png_filename = os.path.join(SCRATCHPAD_DIR, 'rubiks-side-%s.png' % side_name)
if side_name == 'F':
for filename in os.listdir(SCRATCHPAD_DIR):
if filename.endswith('.png'):
os.unlink(os.path.join(SCRATCHPAD_DIR, filename))
# Take a pic but throw it away, we do this so the camera adjusts the current lighting conditions
camera = cv2.VideoCapture(self.dev_video)
camera.read()
del(camera)
camera = None
# Now take a pic, keep it and note the brightness, etc
camera = cv2.VideoCapture(self.dev_video)
(retval, img) = camera.read()
brightness = camera.get(cv2.CAP_PROP_BRIGHTNESS)
contrast = camera.get(cv2.CAP_PROP_CONTRAST)
saturation = camera.get(cv2.CAP_PROP_SATURATION)
gain = camera.get(cv2.CAP_PROP_GAIN)
else:
# Set the brightness, etc to be the same as when we took a pic of side F.
# This makes rubiks-color-resolver's job much easier.
camera = cv2.VideoCapture(self.dev_video)
camera.set(cv2.CAP_PROP_BRIGHTNESS, brightness)
camera.set(cv2.CAP_PROP_CONTRAST, contrast)
camera.set(cv2.CAP_PROP_SATURATION, saturation)
camera.set(cv2.CAP_PROP_GAIN, gain)
(retval, img) = camera.read()
# If you do not delete the VideoCapture object opencv2 will sometimes return the
# exact same image when you call read() back-to-back. This is really bad when
# you have flipped the cube to a new side and end up with a pic of the previous
# side.
del(camera)
camera = None
if retval:
# Images for sides U and D need to be rotated 90 degrees
if side_name in ('U', 'D'):
img = rotate_image(img, 90)
# Save the image to disk
cv2.imwrite(png_filename, img)
size = os.path.getsize(png_filename)
if size:
response = 'FINISHED: image %s is %d bytes' % (png_filename, size)
else:
response = 'ERROR: image %s is 0 bytes' % png_filename
else:
response = 'ERROR: image %s camera.read() failed' % png_filename
elif data == 'GET_RGB_COLORS':
cmd = ['rubiks-cube-tracker.py', '--directory', SCRATCHPAD_DIR]
log.info("cmd: %s" % ' '.join(cmd))
response = subprocess.check_output(cmd).strip()
elif data.startswith('GET_CUBE_STATE:'):
cmd = ['rubiks-color-resolver.py', '--json', '--rgb', data[len('GET_CUBE_STATE:'):]]
log.info("cmd: %s" % ' '.join(cmd))
response = subprocess.check_output(cmd).strip()
elif data == 'GET_CUBE_STATE_FROM_PICS':
# Have not tested this
cmd = ['rubiks-cube-tracker.py', '--directory', SCRATCHPAD_DIR]
log.info("cmd: %s" % ' '.join(cmd))
rgb = subprocess.check_output(cmd).strip()
cmd = ['rubiks-color-resolver.py', '--json', '--rgb', rgb]
log.info("cmd: %s" % ' '.join(cmd))
response = subprocess.check_output(cmd).strip()
elif data.startswith('GET_SOLUTION:'):
cmd = "cd ~/rubiks-cube-NxNxN-solver/; ./usr/bin/rubiks-cube-solver.py --state %s" % data.split(':')[1]
log.info("cmd: %s" % cmd)
response = subprocess.check_output(cmd, shell=True).strip()
elif data == 'PING':
response = 'REPLY'
else:
log.warning("RXed %s (not supported)" % data)
# TX our response and close the socket
connection.send(response)
connection.close()
log.info("TXed %s response %s" % (data, response))
# We have the entire msg so break out of the inside 'while True' loop
break
except Exception as e:
log.exception(e)
caught_exception = True
break
log.info('cranecuberd is stopping with PID %s' % pid)
if tcp_socket:
tcp_socket.close()
tcp_socket = None
if caught_exception:
sys.exit(1)
else:
sys.exit(0)
class TrafficGenWorker(object):
"""
Handler which handles all traffic client related operations. Exposes
its APIs via RPCServer. APIs can be accessed by using RPCClient by
providing RPCServers address.
"""
def __init__(self, uid, hb_queue, exchange):
self._hb_queue = hb_queue
self._rule_collections = TrafficRuleCollection()
self._stop_event = Event()
self._run_event = Event()
self._timer = None
self._stop_lock = Lock()
self._exchange = exchange
self._pool = None
self._uid = uid
self._hb_interval = 5
self._metric_cache = MetricsCache()
def initialize(self):
"""
RPCServer calls this method when it starts. This
method starts heartbeat reporter and metrics reporter threads.
"""
HeartBeatSender(self._hb_queue, self._rule_collections,
self._uid).start()
ExchangeReporter(self._metric_cache, self._exchange)
@property
def uid(self):
"""Get the uid of worker"""
return self._uid
@property
def traffic_running(self):
"""Return True if traffic is running otherwise False"""
with self._stop_lock:
return self._run_event.is_set()
def _generate_traffic(self, stop_event, callback):
"""
Generate traffic in infinite loop
:param stop_event: event which control infinite loop
:type stop_event: Event
:param callback: callback to be called after exiting from loop
:type callback: func
"""
for rule in self._rule_collections.round_robin_rule_generator():
with self._stop_lock:
if stop_event.is_set():
break
try:
if rule.protocol == "TCP":
client = TCPClient
elif rule.protocol == "UDP":
client = UDPClient
elif rule.protocol == "HTTP":
client = HTTPClient
else:
print("Invalid protocol")
continue
sender = client(rule.source, rule.destination, rule.port,
self._metric_cache, True, rule.allowed,
rule.request_count)
self._pool.submit(sender.ping)
except Exception as ex:
print(ex)
callback()
def _cleanup_states(self):
"""Get executed as callback when loop exits"""
if self._pool:
self._pool.shutdown()
with self._stop_lock:
self._run_event.clear()
self._stop_event.set()
self._pool = None
def _start_traffic(self):
"""Start traffic thread"""
self._pool = BoundedThreadPoolExecutor(max_workers=10)
thread = Thread(target=self._generate_traffic,
args=(self._stop_event, self._cleanup_states))
thread.daemon = True
thread.start()
def add_clients(self, rules):
"""Add rules to rules collection"""
try:
self._rule_collections.add_rules(rules)
with self._stop_lock:
if not self._run_event.is_set():
self._start_traffic()
self._run_event.set()
except Exception as ex:
print(ex)
def delete_clients(self, rules):
"""Delete rules from rules collection"""
for rule in rules:
self._rule_collections.delete_rule(rule)
def delete_all_clients(self):
""""Delete all rules from collection"""
if not self._stop_event.is_set():
with self._stop_lock:
self._stop_event.set()
self._rule_collections.clear_rules()
def get_rule_count(self):
"""Get the number of rules managed by this worker"""
return self._rule_collections.get_rule_count()
def has_rule(self, rule):
"""Check if this worker is owner of a rule"""
return rule in self._rule_collections
class BridgeProxy:
"""A proxy bridge which can connect to more than one server.
The current implementation has the following limits:
1. All the connections must be alive;
2. It is blocked, which means if one connection is fast and the other
is slow, the overall performance is limited by the slow one.
"""
POLL_TIMEOUT = 100 # timeout of the poller in milliseconds
def __init__(self):
self._context = None
self._client = None
self._frontend = None
self._backend = dict()
self._backend_ready = deque()
self._running = False
self._stopped = Event()
self._stopped.set()
@property
def client(self):
return self._client
def connect(self, endpoints):
"""Connect the backend to one or more endpoints.
:param str/list/tuple endpoints: addresses of endpoints.
"""
if isinstance(endpoints, str):
endpoints = [endpoints]
elif not isinstance(endpoints, (tuple, list)):
raise ValueError("Endpoints must be either a string or "
"a tuple/list of string!")
context = zmq.Context()
for end in endpoints:
backend = context.socket(zmq.DEALER)
backend.connect(end)
self._backend[end] = backend
frontendpoint = "inproc://frontend"
self._frontend = context.socket(zmq.ROUTER)
self._frontend.bind(frontendpoint)
self._client = Client(frontendpoint,
context=context,
timeout=config['BRIDGE_TIMEOUT'])
self._context = context
@run_in_thread()
def start(self):
"""Run the proxy in a thread."""
if self._running:
raise RuntimeError(f"{self.__class__} is already running!")
frontend = self._frontend
poller = zmq.Poller()
poller.register(frontend, zmq.POLLIN)
for address, bk in self._backend.items():
poller.register(bk, zmq.POLLIN)
self._backend_ready.append(address)
self._stopped.clear()
self._running = True
while self._running:
socks = dict(poller.poll(timeout=self.POLL_TIMEOUT))
if socks.get(frontend) == zmq.POLLIN:
message = frontend.recv_multipart()
if len(self._backend_ready) > 0:
address = self._backend_ready.popleft()
self._backend[address].send_multipart(message)
for address, bk in self._backend.items():
if socks.get(bk) == zmq.POLLIN:
message = bk.recv_multipart()
frontend.send_multipart(message)
self._backend_ready.append(address)
# clean up and close all sockets to avoid problems with buffer
poller.unregister(frontend)
for bk in self._backend.values():
poller.unregister(bk)
for bk in self._backend.values():
bk.setsockopt(zmq.LINGER, 0)
self._backend.clear()
self._backend_ready.clear()
self._frontend.setsockopt(zmq.LINGER, 0)
self._frontend = None
self._client = None
self._context.destroy(linger=0)
self._context = None
self._stopped.set()
def stop(self):
"""Stop the proxy running in a thread."""
self._running = False
if not self._stopped.is_set():
self._stopped.wait()
class WolframKernelController(Thread):
""" Control a Wolfram kernel from a Python thread.
A controller can start and stop a Wolfram kernel specified by its path `kernel`. It
can evaluate expression, one at a time.
Most methods from this class return instances of :class:`~concurrent.futures.Future`.
This class is a low level component of the library which is used by local evaluators.
ZMQ sockets are not thread safe, this class ensures encapsulation of them, while enabling
asynchronous operations.
"""
def __init__(
self,
kernel=None,
initfile=None,
consumer=None,
kernel_loglevel=logging.NOTSET,
stdin=PIPE,
stdout=PIPE,
stderr=PIPE,
**kwargs
):
self.id = _thread_counter()
super().__init__(name="wolfram-kernel-%i" % self.id)
self.kernel = kernel or self.default_kernel_path()
if self.kernel:
if not os.isfile(self.kernel):
raise WolframKernelException("Kernel not found at %s." % self.kernel)
if not os.access(self.kernel, os.X_OK):
raise WolframKernelException("Cannot execute kernel %s." % self.kernel)
else:
raise WolframKernelException(
"Cannot locate a kernel automatically. Please provide an explicit kernel path."
)
if initfile is None:
self.initfile = os.path_join(os.dirname(__file__), "initkernel.m")
else:
self.initfile = initfile
if not os.isfile(self.initfile):
raise FileNotFoundError("Kernel initialization file %s not found." % self.initfile)
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
"Initializing kernel %s using script: %s" % (self.kernel, self.initfile)
)
self.tasks_queue = Queue()
self.kernel_socket_in = None
self.kernel_socket_out = None
self.kernel_proc = None
self.consumer = consumer
self.loglevel = kernel_loglevel
self.kernel_logger = None
self.evaluation_count = 0
self._stdin = stdin
self._stdout = stdout
self._stderr = stderr
# some parameters may be passed as kwargs
self.parameters = {}
for k, v in kwargs.items():
try:
self.set_parameter(k, v)
# ignore kwargs unknowns key
except KeyError:
pass
# this is a state: this event is set when the kernel will not serve any more evaluation.
self._state_terminated = False
# lock controlling concurrent access to the state above.
self._state_lock = RLock()
# this is a trigger that will abort most blocking operations.
self.trigger_termination_requested = Event()
def duplicate(self):
""" Build a new object using the same configuration as the current one. """
return self.__class__(
kernel=self.kernel,
initfile=self.initfile,
kernel_loglevel=self.loglevel,
consumer=self.consumer,
stdin=self._stdin,
stdout=self._stdout,
stderr=self._stderr,
**self.parameters
)
_DEFAULT_PARAMETERS = {
"STARTUP_TIMEOUT": 20,
"TERMINATE_TIMEOUT": 3,
"HIDE_SUBPROCESS_WINDOW": True,
}
def get_parameter(self, parameter_name):
"""Return the value of a given session parameter.
Session parameters are:
* ``'STARTUP_TIMEOUT'``: time to wait, in seconds, after the kernel startup is requested. Default is 20 seconds.
* ``'TERMINATE_TIMEOUT'``: time to wait, in seconds, after the ``Quit[]`` command is sent to the kernel. The kernel is killed after this duration. Default is 3 seconds.
"""
try:
return self.parameters.get(
parameter_name, self._DEFAULT_PARAMETERS[parameter_name]
)
except KeyError:
raise KeyError(
"%s is not one of the valid parameters: %s"
% (parameter_name, ", ".join(self._DEFAULT_PARAMETERS.keys()))
)
def set_parameter(self, parameter_name, parameter_value):
"""Set a new value for a given parameter. The new value only applies for this session.
Session parameters are:
* ``'STARTUP_TIMEOUT'``: time to wait, in seconds, after the kernel startup is requested. Default is 20 seconds.
* ``'TERMINATE_TIMEOUT'``: time to wait, in seconds, after the ``Quit[]`` command is sent to the kernel. The kernel is killed after this duration. Default is 3 seconds.
"""
if parameter_name not in self._DEFAULT_PARAMETERS:
raise KeyError(
"%s is not one of the valid parameters: %s"
% (parameter_name, ", ".join(self._DEFAULT_PARAMETERS.keys()))
)
self.parameters[parameter_name] = parameter_value
def default_kernel_path(self):
return find_default_kernel_path()
def _kernel_terminate(self):
self._kernel_stop(gracefully=False)
def _kernel_stop(self, gracefully=True):
"""Stop the kernel process and close sockets.
This function must be called when a given session is no longer useful
to prevent orphan processes and sockets from being generated.
.. note::
Licensing restrictions usually apply to Wolfram kernels and may
prevent new instances from starting if too many kernels are running
simultaneously. Make sure to always terminate sessions to avoid
unexpected start-up errors.
"""
logger.info("Start termination on kernel %s", self)
with self._state_lock:
self._state_terminated = True
self.trigger_termination_requested.set()
if self.kernel_proc is not None:
error = False
if gracefully:
# Graceful stop: first send a Quit command to the kernel.
try:
self.kernel_socket_out.send(b"8:f\x00s\x04Quit", flags=zmq.NOBLOCK)
except:
logger.info("Failed to send Quit[] command to the kernel.")
error = True
if not error:
try:
self.kernel_proc.wait(timeout=self.get_parameter("TERMINATE_TIMEOUT"))
except:
logger.info(
"Kernel process failed to stop after %.02f seconds. Killing it."
% self.get_parameter("TERMINATE_TIMEOUT")
)
error = True
# Kill process if not already terminated.
# Wait for it to cleanly stop if the Quit command was successfully sent,
# otherwise the kernel is likely in a bad state so we kill it immediately.
if self._stdin == PIPE:
try:
self.kernel_proc.stdin.close()
except:
logger.warning("Failed to close kernel process stdin.")
error = True
if self._stdout == PIPE:
try:
self.kernel_proc.stdout.close()
except:
logger.warning("Failed to close kernel process stdout.")
error = True
if self._stderr == PIPE:
try:
self.kernel_proc.stderr.close()
except:
logger.warning("Failed to close kernel process stderr")
error = True
if error or not gracefully:
logger.info("Killing kernel process: %i" % self.kernel_proc.pid)
self.kernel_proc.kill()
self.kernel_proc = None
if self.kernel_socket_out is not None:
try:
self.kernel_socket_out.close()
except Exception as e:
logger.fatal(e)
finally:
self.kernel_socket_out = None
if self.kernel_socket_in is not None:
try:
self.kernel_socket_in.close()
except Exception as e:
logger.fatal(e)
finally:
self.kernel_socket_in = None
if self.kernel_logger is not None:
try:
self.kernel_logger.stopped.set()
self.kernel_logger.join()
except Exception as e:
logger.fatal(e)
finally:
self.kernel_logger = None
assert self.kernel_proc is None
assert self.kernel_socket_in is None
assert self.kernel_socket_out is None
assert self.kernel_logger is None
@property
def started(self):
""" Is the kernel starting or being started. """
with self._state_lock:
return self.is_alive() and not self._state_terminated
@property
def terminated(self):
""" Is the kernel terminated. Terminated kernel no more handle evaluations. """
with self._state_lock:
return self._state_terminated
def is_kernel_alive(self):
""" Return the status of the kernel process. """
try:
# subprocess poll function is thread safe.
return self.kernel_proc is not None and self.kernel_proc.poll() is None
except AttributeError:
# in case kernel_proc was set to None. May not even be possible.
return False
def request_kernel_start(self):
""" Start the thread and the associated kernel. Return a future object indicating the kernel status.
The future object result is True once the kernel is successfully started. Exception raised in the process
and passed to the future object.
Calling this method twice is a no-op."""
with self._state_lock:
future = futures.Future()
if not self.started:
self.enqueue_task(self.START, future, None)
self.start()
else:
future.set_result(True)
return future
def enqueue_task(self, payload, future, callback):
if self.terminated or self.trigger_termination_requested.is_set():
logger.fatal("Cannot enqueue tasks on terminated controller.")
raise RuntimeError("Thread is closing. Cannot queue task")
self.tasks_queue.put((payload, future, callback))
def _safe_kernel_start(self):
""" Start a kernel. If something went wrong, clean-up resources that may have been created. """
try:
self._kernel_start()
except Exception as e:
logger.warning("Failed to start.")
try:
self._kernel_terminate()
finally:
raise e
_KERNEL_OK = b"OK"
_KERNEL_VERSION_NOT_SUPPORTED = 10
def _kernel_start(self):
"""Start a new kernel process and open sockets to communicate with it."""
# Socket to which we push new expressions for evaluation.
if self.kernel_socket_out is None:
self.kernel_socket_out = Socket(zmq_type=zmq.PUSH)
if self.kernel_socket_in is None:
self.kernel_socket_in = Socket(zmq_type=zmq.PULL)
# start the evaluation zmq sockets
self.kernel_socket_out.bind()
self.kernel_socket_in.bind()
if logger.isEnabledFor(logging.INFO):
logger.info("Kernel writes commands to socket: %s", self.kernel_socket_out)
logger.info(
"Kernel receives evaluated expressions from socket: %s", self.kernel_socket_in
)
# start the kernel process
cmd = [self.kernel, "-noprompt", "-initfile", self.initfile]
if self.loglevel != logging.NOTSET:
self.kernel_logger = KernelLogger(
name="wolfram-kernel-logger-%i" % self.id, level=self.loglevel
)
self.kernel_logger.start()
cmd.append("-run")
cmd.append(
'ClientLibrary`Private`SlaveKernelPrivateStart["%s", "%s", "%s", %i];'
% (
self.kernel_socket_out.uri,
self.kernel_socket_in.uri,
self.kernel_logger.socket.uri,
FROM_PY_LOG_LEVEL[self.loglevel],
)
)
else:
cmd.append("-run")
cmd.append(
'ClientLibrary`Private`SlaveKernelPrivateStart["%s", "%s"];'
% (self.kernel_socket_out.uri, self.kernel_socket_in.uri)
)
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Kernel called using command: %s." % " ".join(cmd))
# hide the WolframKernel window.
if six.WINDOWS and self.get_parameter("HIDE_SUBPROCESS_WINDOW"):
startupinfo = STARTUPINFO()
startupinfo.dwFlags |= STARTF_USESHOWWINDOW
else:
startupinfo = None
try:
self.kernel_proc = Popen(
cmd,
stdin=self._stdin,
stdout=self._stdout,
stderr=self._stderr,
startupinfo=startupinfo,
)
if logger.isEnabledFor(logging.INFO):
logger.info("Kernel process started with PID: %s" % self.kernel_proc.pid)
t_start = time.perf_counter()
except Exception as e:
logger.exception(e)
raise WolframKernelException("Failed to start kernel process.")
try:
# First message must be "OK", acknowledging everything is up and running
# on the kernel side.
response = self.kernel_socket_in.recv_abortable(
timeout=self.get_parameter("STARTUP_TIMEOUT"),
abort_event=_StartEvent(self.kernel_proc, self.trigger_termination_requested),
)
if response == self._KERNEL_OK:
if logger.isEnabledFor(logging.INFO):
logger.info(
"Kernel %s is ready. Startup took %.2f seconds."
% (self.pid, time.perf_counter() - t_start)
)
else:
raise WolframKernelException(
"Kernel %s failed to start properly." % self.kernel
)
except (SocketAborted, SocketOperationTimeout) as se:
if self.kernel_proc.returncode == self._KERNEL_VERSION_NOT_SUPPORTED:
raise WolframKernelException(
"Wolfram kernel version is not supported. Please consult library prerequisites."
)
logger.warning("Socket exception: %s", se)
raise WolframKernelException(
"Failed to communicate with kernel: %s." % self.kernel
)
@property
def pid(self):
"""Return the PID of the Wolfram kernel process, if any, or None."""
try:
return self.kernel_proc.pid
except AttributeError:
return None
START = object()
STOP = object()
def stop(self):
future = futures.Future()
with self._state_lock:
if self.terminated:
future.set_result(True)
return future
self.enqueue_task(self.STOP, future, None)
self._state_terminated = True
return future
def terminate(self):
future = futures.Future()
with self._state_lock:
if not self.started:
future.set_result(True)
return future
self.enqueue_task(self.STOP, future, None)
self._state_terminated = True
self.trigger_termination_requested.set()
return future
def join(self, timeout=None):
future = self.stop()
future.result(timeout=timeout)
return super().join(timeout=timeout)
def evaluate_future(self, wxf, future, result_update_callback=None, **kwargs):
self.enqueue_task(wxf, future, result_update_callback)
def _do_evaluate(self, wxf, future, result_update_callback):
start = time.perf_counter()
self.kernel_socket_out.send(wxf)
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Expression sent to kernel in %.06fsec", time.perf_counter() - start)
start = time.perf_counter()
wxf_eval_data = self.kernel_socket_in.recv_abortable(copy=False)
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
"Expression received from kernel after %.06fsec", time.perf_counter() - start
)
self.evaluation_count += 1
result = WolframKernelEvaluationResult(wxf_eval_data.buffer, consumer=self.consumer)
if logger.isEnabledFor(logging.WARNING):
for msg in result.iter_messages():
logger.warning(msg)
if result_update_callback:
result = result_update_callback(result)
future.set_result(result)
def run(self):
future = None
task = None
try:
task = self.tasks_queue.get()
payload, future, result_update_callback = task
# Kernel start requested.
if payload is self.START:
self._safe_kernel_start()
future.set_result(True)
future = None
task = None
self.tasks_queue.task_done()
task = self.tasks_queue.get()
payload, future, result_update_callback = task
elif payload is self.STOP:
future.set_result(True)
future = None
return
# first evaluation. Ensure kernel is started.
else:
self._safe_kernel_start()
while not self.trigger_termination_requested.is_set():
if payload is self.STOP:
# lock controlling concurrent access to state above.
with self._state_lock:
self._state_terminated = True
logger.info(
"Termination requested for kernel controller. Associated kernel PID: %s",
self.pid,
)
task = None
self.tasks_queue.task_done()
break
self._do_evaluate(payload, future, result_update_callback)
future = None
task = None
self.tasks_queue.task_done()
task = self.tasks_queue.get()
payload, future, result_update_callback = task
except (KeyboardInterrupt, RuntimeError, futures.CancelledError) as e:
self.trigger_termination_requested.set()
logger.error("Fatal error in kernel controller: %s", e)
raise e
except Exception as e:
self.trigger_termination_requested.set()
if future and not future.cancelled():
future.set_exception(e)
future = None
else:
raise e
finally:
try:
if task:
self.tasks_queue.task_done()
self._cancel_tasks()
if self.trigger_termination_requested.is_set():
self._kernel_terminate()
else:
self._kernel_stop()
except Exception as e:
if future:
future.set_exception(e)
future = None
finally:
if future:
future.set_result(True)
def _cancel_tasks(self):
while not self.tasks_queue.empty():
task = self.tasks_queue.get()
_, future, _ = task
future.cancel()
def __repr__(self):
if self.started:
return "<%s: pid:%i, kernel sockets: (in:%s, out:%s)>" % (
self.__class__.__name__,
self.kernel_proc.pid,
self.kernel_socket_in.uri,
self.kernel_socket_out.uri,
)
else:
return "<%s: %s>" % (self.__class__.__name__, self.name)
class RunVA(object):
def _test_mqtt_connection(self):
print("testing mqtt connection", flush=True)
mqtt = Client()
while True:
try:
mqtt.connect(mqtthost)
break
except:
print("Waiting for mqtt...", flush=True)
time.sleep(5)
print("mqtt connected", flush=True)
mqtt.disconnect()
def __init__(self, pipeline, version="2"):
super(RunVA, self).__init__()
self._test_mqtt_connection()
self._pipeline = pipeline
self._version = version
self._db = DBIngest(host=dbhost, index="algorithms", office=office)
self._stop = None
def stop(self):
if self._stop:
print("stopping", flush=True)
self._stop.set()
def loop(self,
sensor,
location,
uri,
algorithm,
algorithmName,
options={},
topic="analytics"):
try:
VAServing.start({
'model_dir': '/home/models',
'pipeline_dir': '/home/pipelines',
'max_running_pipelines': 1,
})
try:
source = {
"type": "uri",
"uri": uri,
}
destination = {
"type": "mqtt",
"host": mqtthost,
"clientid": algorithm,
"topic": topic
}
tags = {
"sensor": sensor,
"location": location,
"algorithm": algorithmName,
"office": {
"lat": office[0],
"lon": office[1]
},
}
parameters = {
"inference-interval": every_nth_frame,
"recording_prefix": "/tmp/rec/" + sensor
}
parameters.update(options)
pipeline = VAServing.pipeline(self._pipeline, self._version)
instance_id = pipeline.start(source=source,
destination=destination,
tags=tags,
parameters=parameters)
if instance_id is None:
raise Exception(
"Pipeline {} version {} Failed to Start".format(
self._pipeline, self._version))
self._stop = Event()
while not self._stop.is_set():
status = pipeline.status()
print(status, flush=True)
if status.state.stopped():
print(
"Pipeline {} Version {} Instance {} Ended with {}".
format(self._pipeline, self._version, instance_id,
status.state.name),
flush=True)
break
if status.avg_fps > 0 and status.state is Pipeline.State.RUNNING:
avg_pipeline_latency = status.avg_pipeline_latency
if not avg_pipeline_latency: avg_pipeline_latency = 0
self._db.update(
algorithm, {
"sensor": sensor,
"performance": status.avg_fps,
"latency": avg_pipeline_latency * 1000,
"cpu": psutil.cpu_percent(),
"memory": psutil.virtual_memory().percent,
})
self._stop.wait(3)
self._stop = None
pipeline.stop()
except:
print(traceback.format_exc(), flush=True)
VAServing.stop()
except:
print(traceback.format_exc(), flush=True)
class NemoManager(Thread):
"""
Object to manage NEMO instrument
Saves NEMO data to files and provides high level command interface.
"""
def __init__(self, data_dir='.', port_id=3, reset_gpio_ch=16):
"""NemoManager class constructor"""
self._data_dir = data_dir
self._nemo = nemo.Nemo(log=False,
port_id=port_id,
reset_gpio_ch=reset_gpio_ch)
self._config = util.Configuration(
config_fname=os.path.join(self._data_dir, 'config.json'))
self._shutdown = Event()
self._run = Event()
Path(self._data_dir).mkdir(parents=True, exist_ok=True)
self._config_file = util.RotatingFileManager(
os.path.join(self._data_dir, 'config'),
self._config.config_rotate_period)
self._rate_data_file = util.RotatingFileManager(
os.path.join(self._data_dir, 'rate_data'),
self._config.rate_data_rotate_period)
self._histogram_file = util.RotatingFileManager(
os.path.join(self._data_dir, 'histogram'),
self._config.histogram_rotate_period)
self.set_config(**self._config.get_public_dict())
self._t_last_config = None
self._t_last_data = None
Thread.__init__(self)
self._run.set()
self.start()
@property
def _sec_since_last_config(self):
"""
Elapsed time in seconds since config was last written
Returns None if unknown.
"""
if self._t_last_config is not None:
return (datetime.datetime.now() -
self._t_last_config).total_seconds()
return None
@property
def _sec_since_last_data(self):
"""
Elapsed time in seconds since data (rate & histograms) were last written
Returns None if unknown.
"""
if self._t_last_data is not None:
return (datetime.datetime.now() -
self._t_last_data).total_seconds()
return None
def close(self):
"""Cleanly close the object"""
self._shutdown.set()
def pause(self):
"""
Pause work in the main thread.
Any work in progress will be completed first.
"""
self._run.clear()
def resume(self):
"""
Resume work in the main thread.
"""
self._run.set()
def run(self):
"""
Method representing the thread's activity.
Overrides the Thread class' run() method.
"""
while not self._shutdown.is_set():
if self._run.wait(timeout=10.0):
try:
# if time to write config to file
if (self._sec_since_last_config is None
or (self._sec_since_last_config >
self._config.config_write_period)):
self._config_file.write(
bytes(util.ConfigPacket(self._nemo)))
self._t_last_config = datetime.datetime.now()
logging.info('Wrote config')
# if time to write rate data and histogram to file
if (self._sec_since_last_data is None
or (self._sec_since_last_data >
self._config.data_write_period)):
self._rate_data_file.write(
bytes(util.RateDataPacket(self._nemo)))
self._histogram_file.write(
bytes(util.HistogramPacket(self._nemo)))
self._t_last_data = datetime.datetime.now()
logging.info('Wrote rate data and histogram')
time.sleep(0.25)
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager tread')
time.sleep(1)
def write_register(self, reg_address, values):
"""Direct write of register on NEMO. Allows low-level diagnostics on-orbit."""
try:
self._nemo._write_register(reg_address, values)
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager.write_register')
def read_register(self, reg_address, size):
"""
Direct read of register on NEMO. Allows low-level diagnostics on-orbit.
Results are writen to file config_{datetime}_{reg_address}_{size}
"""
try:
result = self._nemo._read_register(reg_address, size)
dt_str = datetime.datetime.now().strftime('%Y%m%dT%H%M%SZ')
fname = os.path.join(
self._data_dir,
f'read_register_{dt_str}_{reg_address:02X}_{size:02X}')
with open(fname, 'ab+') as file:
file.write(bytes(result))
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager.read_register')
def set_config(self, **kwargs):
"""Set writeable configuration parameters."""
try:
if 'det_enable_uint8' in kwargs:
self._config.set(det_enable_uint8=kwargs['det_enable_uint8'])
self._nemo.det_enable_uint8 = kwargs['det_enable_uint8']
if 'det0_bias_uint8' in kwargs:
self._config.set(det0_bias_uint8=kwargs['det0_bias_uint8'])
self._nemo.det0.bias_uint8 = kwargs['det0_bias_uint8']
if 'det1_bias_uint8' in kwargs:
self._config.set(det1_bias_uint8=kwargs['det1_bias_uint8'])
self._nemo.det1.bias_uint8 = kwargs['det1_bias_uint8']
if 'det0_threshold_uint8' in kwargs:
self._config.set(
det0_threshold_uint8=kwargs['det0_threshold_uint8'])
self._nemo.det0.threshold_uint8 = kwargs[
'det0_threshold_uint8']
if 'det1_threshold_uint8' in kwargs:
self._config.set(
det1_threshold_uint8=kwargs['det1_threshold_uint8'])
self._nemo.det1.threshold_uint8 = kwargs[
'det1_threshold_uint8']
if 'rate_width_min' in kwargs:
self._config.set(rate_width_min=kwargs['rate_width_min'])
self._nemo.rate_width_min = kwargs['rate_width_min']
if 'rate_width_max' in kwargs:
self._config.set(rate_width_max=kwargs['rate_width_max'])
self._nemo.rate_width_max = kwargs['rate_width_max']
if 'bin_width' in kwargs:
self._config.set(bin_width=kwargs['bin_width'])
self._nemo.bin_width = kwargs['bin_width']
if 'bin_0_min_width' in kwargs:
self._config.set(bin_0_min_width=kwargs['bin_0_min_width'])
self._nemo.bin_0_min_width = kwargs['bin_0_min_width']
if 'rate_interval' in kwargs:
self._config.set(rate_interval=kwargs['rate_interval'])
self._nemo.rate_interval = kwargs['rate_interval']
if 'veto_threshold_min' in kwargs:
self._config.set(
veto_threshold_min=kwargs['veto_threshold_min'])
self._nemo.veto_threshold_min = kwargs['veto_threshold_min']
if 'veto_threshold_max' in kwargs:
self._config.set(
veto_threshold_max=kwargs['veto_threshold_max'])
self._nemo.veto_threshold_max = kwargs['veto_threshold_max']
if 'config_write_period' in kwargs:
self._config.set(
config_write_period=kwargs['config_write_period'])
if 'config_rotate_period' in kwargs:
self._config.set(
config_rotate_period=kwargs['config_rotate_period'])
self._config_file.period = kwargs['config_rotate_period']
if 'data_write_period' in kwargs:
self._config.set(data_write_period=kwargs['data_write_period'])
if 'rate_data_rotate_period' in kwargs:
self._config.set(
rate_data_rotate_period=kwargs['rate_data_rotate_period'])
self._rate_data_file.period = kwargs['rate_data_rotate_period']
if 'histogram_rotate_period' in kwargs:
self._config.set(
histogram_rotate_period=kwargs['histogram_rotate_period'])
self._histogram_file.period = kwargs['histogram_rotate_period']
self._config.save()
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager.set_config')
def power_off(self):
"""Turn power off to Nemo (really just holds it in reset"""
try:
self._nemo.hold_in_reset()
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager.power_off')
def power_on(self):
"""Turn power on to Nemo (really just release from reset"""
try:
self._nemo.release_from_reset()
except nemo.I2CTransactionFailure:
logging.error(
'Nemo I2C transaction failure in NemoManager.power_on')
def reboot(self):
"""Reboot and reconfigure Nemo"""
try:
self._nemo.software_reboot()
time.sleep(0.05)
self.set_config(**self._config.get_public_dict())
except nemo.I2CTransactionFailure:
logging.error('Nemo I2C transaction failure in NemoManager.reboot')
def process_rate_data(self, t_start, t_stop, decimation_factor):
"""Process already saved rate data into a lower resolution."""
input_packets = util.RateDataPacket.from_file(os.path.join(
self._data_dir, 'rate_data_*T*Z'),
sc_time_min=t_start,
sc_time_max=t_stop,
sort=True)
fname = f'lores_rate_data_{t_start}_{t_stop}_{decimation_factor}'
with open(os.path.join(self._data_dir, fname), 'wb') as file:
for i in range(0, len(input_packets), decimation_factor):
output_packet = util.LoResRateDataPacket(
input_packets[i:i + decimation_factor])
file.write(bytes(output_packet))
def process_histograms(self, t_start, t_stop, decimation_factor):
"""Process already saved histograms into a lower resolution."""
input_packets = util.HistogramPacket.from_file(os.path.join(
self._data_dir, 'histogram_*T*Z'),
sc_time_min=t_start,
sc_time_max=t_stop,
sort=True)
fname = f'lores_histogram_{t_start}_{t_stop}_{decimation_factor}'
with open(os.path.join(self._data_dir, fname), 'wb') as file:
for i in range(0, len(input_packets), decimation_factor):
output_packet = util.LoResHistogramPacket(
input_packets[i:i + decimation_factor])
file.write(bytes(output_packet))
class OpenBCI_data_manager(Thread):
def __init__(self, queue):
Thread.__init__(self)
### data ###########
self.EEG_buffer = EEG_buffer()
self.all_data_store = []
self.queue = queue
########### SHARED QUEUE ###########
self.filter_bank = filter_bank_class(self.EEG_buffer.LOWCUT,
self.EEG_buffer.HIGHCUT,
self.EEG_buffer.NOTCH,
self.EEG_buffer.ORDER,
self.EEG_buffer.SAMPLE_RATE)
self.filter_bank.set_filters(self.EEG_buffer.LOWCUT,
self.EEG_buffer.HIGHCUT)
self.spectrum = spectrum(self.EEG_buffer.NDIMS,
self.EEG_buffer.SAMPLE_RATE)
# -- flag --
self.exit = Event()
self.exit.set()
###### mutex lock
self.mutexBuffer = Lock()
def run(self):
while self.exit.is_set():
time.sleep(0.001)
while not self.queue.empty():
try:
self.mutexBuffer.acquire()
sample = self.queue.get()
self.mutexBuffer.release()
self.EEG_buffer.append(sample)
self.all_data_store.append(np.asarray(sample).transpose())
except:
print('algo pasa en dmg tronco!!!!!')
def reset_data_store(self):
self.all_data_store = []
def get_allData(self):
return np.asarray(self.all_data_store)
def get_sample(self):
filtered = self.filter_bank.pre_process(self.EEG_buffer.get())
return filtered
def get_short_sample(self, method):
filtered = self.filter_bank.pre_process(self.EEG_buffer.get())
filtered = filtered[:,
int(self.EEG_buffer.pos_ini):int(self.EEG_buffer.
pos_end)]
return filtered
def get_powerSpectrum(self, method):
try:
filtered = self.filter_bank.pre_process(self.EEG_buffer.get())
freqs, spectra = self.spectrum.get_spectrum(filtered)
return [freqs, spectra]
except:
return None
def get_powerSpectrogram(self, method, channel):
try:
filtered = self.filter_bank.pre_process(self.EEG_buffer.get())
spectrogram = self.spectrum.get_spectrogram(filtered[channel, :])
return spectrogram
except:
return None
def kill(self):
self.exit.clear()
def process(self, data):
def shlexjoin():
import shlex
return ' '.join(shlex.quote(cmd) for cmd in commandline)
commandline = [cmd.decode(self.codec) for cmd in self.args.commandline]
self.log_debug(shlexjoin)
posix = 'posix' in sys.builtin_module_names
shell = os.name == 'nt'
process = Popen(commandline,
stdin=PIPE, stdout=PIPE, stderr=PIPE, shell=shell, close_fds=posix)
if self.args.buffer and not self.args.timeout:
out, err = process.communicate(data)
for line in err.splitlines():
self.log_debug(line)
yield out
return
import io
from threading import Thread, Event
from queue import Queue, Empty
from time import process_time, sleep
start = 0
result = None
qerr = Queue()
qout = Queue()
done = Event()
def adapter(stream, queue: Queue, event: Event):
while not event.is_set():
out = stream.read1()
if out: queue.put(out)
else: break
stream.close()
recvout = Thread(target=adapter, args=(process.stdout, qout, done), daemon=True)
recverr = Thread(target=adapter, args=(process.stderr, qerr, done), daemon=True)
recvout.start()
recverr.start()
process.stdin.write(data)
process.stdin.close()
start = process_time()
if self.args.buffer or self.args.timeout:
result = io.BytesIO()
def queue_read(q: Queue):
try: return q.get_nowait()
except Empty: return None
errbuf = io.BytesIO()
while True:
err = queue_read(qerr)
out = queue_read(qout)
if err and self.log_debug():
errbuf.write(err)
errbuf.seek(0)
lines = errbuf.readlines()
errbuf.seek(0)
errbuf.truncate()
if lines:
if not (done.is_set() or lines[~0].endswith(B'\n')):
errbuf.write(lines.pop())
for line in lines:
msg = line.rstrip(B'\n')
if msg: self.log_debug(msg)
if out:
if self.args.buffer or self.args.timeout:
result.write(out)
if not self.args.buffer:
yield out
if done.is_set():
if recverr.is_alive():
self.log_warn('stderr receiver thread zombied')
if recvout.is_alive():
self.log_warn('stdout receiver thread zombied')
break
elif not err and not out and process.poll() is not None:
recverr.join(0.4)
recvout.join(0.4)
done.set()
elif self.args.timeout:
if process_time() - start > self.args.timeout:
self.log_info('terminating process after timeout expired')
done.set()
process.terminate()
for wait in range(4):
if process.poll() is not None: break
sleep(.1)
else:
self.log_warn('process termination may have failed')
recverr.join(0.4)
recvout.join(0.4)
if not len(result.getbuffer()):
result = RuntimeError('timeout reached, process had no output')
else:
result = RefineryPartialResult(
'timeout reached, returning all collected output',
partial=result.getvalue())
if isinstance(result, Exception):
raise result
elif self.args.buffer:
yield result.getvalue()
class WebsocketServer(object):
def __init__(self,
args,
mitm_mapper,
db_wrapper,
mapping_manager,
pogoWindowManager,
data_manager,
configmode=False):
self.__current_users = {}
self.__users_connecting = []
self.__stop_server = Event()
self.args = args
self.__listen_address = args.ws_ip
self.__listen_port = int(args.ws_port)
self.__received = {}
self.__requests = {}
self.__users_mutex: Optional[asyncio.Lock] = None
self.__id_mutex: Optional[asyncio.Lock] = None
self.__send_queue: Optional[asyncio.Queue] = None
self.__received_mutex: Optional[asyncio.Lock] = None
self.__requests_mutex: Optional[asyncio.Lock] = None
self.__db_wrapper = db_wrapper
self.__mapping_manager: MappingManager = mapping_manager
self.__pogoWindowManager = pogoWindowManager
self.__mitm_mapper = mitm_mapper
self.__data_manager = data_manager
self.__next_id = 0
self._configmode = configmode
self.__loop = None
self.__loop_tid = None
self.__loop_mutex = Lock()
self.__worker_shutdown_queue: queue.Queue = queue.Queue()
self.__internal_worker_join_thread: Thread = Thread(
name='worker_join_thread', target=self.__internal_worker_join)
self.__internal_worker_join_thread.daemon = True
self.__internal_worker_join_thread.start()
def _add_task_to_loop(self, coro):
f = functools.partial(self.__loop.create_task, coro)
if current_thread() == self.__loop_tid:
# We can call directly if we're not going between threads.
return f()
else:
# We're in a non-event loop thread so we use a Future
# to get the task from the event loop thread once
# it's ready.
return self.__loop.call_soon_threadsafe(f)
def __internal_worker_join(self):
while not self.__stop_server.is_set():
try:
next_item: Optional[
Thread] = self.__worker_shutdown_queue.get_nowait()
except queue.Empty:
time.sleep(1)
continue
if next_item is not None:
logger.info("Trying to join worker thread")
next_item.join(10)
if next_item.isAlive():
logger.debug(
"Error while joining worker thread - requeue it")
self.__worker_shutdown_queue.put(next_item)
self.__worker_shutdown_queue.task_done()
logger.info("Done joining worker thread")
async def __setup_first_loop(self):
self.__users_mutex = asyncio.Lock()
self.__id_mutex = asyncio.Lock()
self.__send_queue: asyncio.Queue = asyncio.Queue()
self.__received_mutex = asyncio.Lock()
self.__requests_mutex = asyncio.Lock()
def start_server(self):
logger.info("Starting websocket server...")
self.__loop = asyncio.new_event_loop()
logger.debug("Device mappings: {}",
str(self.__mapping_manager.get_all_devicemappings()))
asyncio.set_event_loop(self.__loop)
self._add_task_to_loop(self.__setup_first_loop())
self.__loop.run_until_complete(
websockets.serve(self.handler,
self.__listen_address,
self.__listen_port,
max_size=2**25))
self.__loop_tid = current_thread()
self.__loop.run_forever()
logger.info("Websocketserver stopping...")
async def __internal_stop_server(self):
self.__stop_server.set()
async with self.__users_mutex:
for id, worker in self.__current_users.items():
logger.info('Closing connections to device {}.', id)
await worker[2].close()
if self.__loop is not None:
self.__loop.call_soon_threadsafe(self.__loop.stop)
self.__internal_worker_join_thread.join()
def stop_server(self):
with self.__loop_mutex:
future = asyncio.run_coroutine_threadsafe(
self.__internal_stop_server(), self.__loop)
future.result()
async def handler(self, websocket_client_connection, path):
if self.__stop_server.is_set():
await websocket_client_connection.close()
return
logger.info("Waiting for connection...")
# wait for a connection...
try:
continue_work = await self.__register(websocket_client_connection)
reason = None
if type(continue_work) is tuple:
(continue_work, reason) = continue_work
if not continue_work:
if not reason:
logger.error(
"Failed registering client, closing connection")
else:
logger.info(reason)
await websocket_client_connection.close()
return
except data_manager.dm_exceptions.DataManagerException:
return
consumer_task = asyncio.ensure_future(
self.__consumer_handler(websocket_client_connection))
producer_task = asyncio.ensure_future(
self.__producer_handler(websocket_client_connection))
done, pending = await asyncio.wait(
[producer_task, consumer_task],
return_when=asyncio.FIRST_COMPLETED,
)
logger.debug(
"consumer or producer of {} stopped, cancelling pending tasks",
str(
websocket_client_connection.request_headers.get_all("Origin")
[0]))
for task in pending:
task.cancel()
logger.info(
"Awaiting unregister of {}",
str(
websocket_client_connection.request_headers.get_all("Origin")
[0]))
await self.__unregister(websocket_client_connection)
logger.info(
"All done with {}",
str(
websocket_client_connection.request_headers.get_all("Origin")
[0]))
@logger.catch()
async def __register(self, websocket_client_connection):
try:
origin = str(
websocket_client_connection.request_headers.get_all("Origin")
[0])
except IndexError:
logger.warning(
"Client from {} tried to connect without Origin header",
str(
websocket_client_connection.request_headers.get_all(
"Origin")[0]))
return False
if not self.__data_manager.is_device_active(origin):
return (
False,
'Origin %s is currently paused. Unpause through MADmin to begin working'
% origin)
logger.info("Client {} registering", str(origin))
if self.__mapping_manager is None or origin not in self.__mapping_manager.get_all_devicemappings(
).keys():
logger.warning(
"Register attempt of unknown origin: {}. "
"Have you forgot to hit 'APPLY SETTINGS' in MADmin?".format(
origin))
return False
if origin in self.__users_connecting:
logger.info("Client {} is already connecting".format(origin))
return False
auths = self.__mapping_manager.get_auths()
if auths:
try:
authBase64 = str(
websocket_client_connection.request_headers.get_all(
"Authorization")[0])
except IndexError:
logger.warning(
"Client from {} tried to connect without auth header",
str(
websocket_client_connection.request_headers.get_all(
"Origin")[0]))
return False
async with self.__users_mutex:
logger.debug("Checking if {} is already present", str(origin))
if origin in self.__current_users:
logger.warning(
"Worker with origin {} is already running, killing the running one and have client reconnect",
str(origin))
self.__current_users.get(origin)[1].stop_worker()
## todo: do this better :D
logger.debug(
"Old worker thread is still alive - waiting 20 seconds")
await asyncio.sleep(20)
logger.info("Reconnect ...")
return
self.__users_connecting.append(origin)
# reset pref. error counter if exist
await self.__reset_fail_counter(origin)
try:
if auths and authBase64 and not check_auth(authBase64, self.args,
auths):
logger.warning(
"Invalid auth details received from {}",
str(
websocket_client_connection.request_headers.get_all(
"Origin")[0]))
return False
logger.info("Starting worker {}".format(origin))
if self._configmode:
worker = WorkerConfigmode(
self.args,
origin,
self,
walker=None,
mapping_manager=self.__mapping_manager,
mitm_mapper=self.__mitm_mapper,
db_wrapper=self.__db_wrapper,
routemanager_name=None)
logger.debug("Starting worker for {}", str(origin))
new_worker_thread = Thread(name='worker_%s' % origin,
target=worker.start_worker)
async with self.__users_mutex:
self.__current_users[origin] = [
new_worker_thread, worker, websocket_client_connection,
0
]
return True
last_known_state = {}
client_mapping = self.__mapping_manager.get_devicemappings_of(
origin)
devicesettings = self.__mapping_manager.get_devicesettings_of(
origin)
logger.info("Setting up routemanagers for {}", str(origin))
if client_mapping.get("walker", None) is not None:
if devicesettings is not None and "walker_area_index" not in devicesettings:
logger.debug("Initializing devicesettings")
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', 0)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'finished', False)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'last_action_time', None)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'last_cleanup_time', None)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'job', False)
await asyncio.sleep(
1
) # give the settings a moment... (dirty "workaround" against race condition)
walker_index = devicesettings.get('walker_area_index', 0)
if walker_index > 0:
# check status of last area
if not devicesettings.get('finished', False):
logger.info(
'Something wrong with last round - get back to old area'
)
walker_index -= 1
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', walker_index)
# devicesettings['walker_area_index'] = walker_index
walker_area_array = client_mapping["walker"]
walker_settings = walker_area_array[walker_index]
# preckeck walker setting
while not pre_check_value(
walker_settings) and walker_index - 1 <= len(
walker_area_array):
walker_area_name = walker_area_array[walker_index][
'walkerarea']
logger.info(
'{} not using area {} - Walkervalue out of range',
str(origin),
str(
self.__mapping_manager.routemanager_get_name(
walker_area_name)))
if walker_index >= len(walker_area_array) - 1:
logger.error(
'Could not find any working area at this time - check your mappings for device: {}',
str(origin))
walker_index = 0
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', walker_index)
walker_settings = walker_area_array[walker_index]
await websocket_client_connection.close()
return
walker_index += 1
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', walker_index)
walker_settings = walker_area_array[walker_index]
devicesettings = self.__mapping_manager.get_devicesettings_of(
origin)
logger.debug("Checking walker_area_index length")
if (devicesettings.get("walker_area_index", None) is None
or devicesettings['walker_area_index'] >=
len(walker_area_array)):
# check if array is smaller than expected - f.e. on the fly changes in mappings.json
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', 0)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'finished', False)
walker_index = 0
walker_area_name = walker_area_array[walker_index][
'walkerarea']
if walker_area_name not in self.__mapping_manager.get_all_routemanager_names(
):
await websocket_client_connection.close()
raise WrongAreaInWalker()
logger.debug('Devicesettings {}: {}', str(origin),
devicesettings)
logger.info(
'{} using walker area {} [{}/{}]', str(origin),
str(
self.__mapping_manager.routemanager_get_name(
walker_area_name)), str(walker_index + 1),
str(len(walker_area_array)))
walker_routemanager_mode = self.__mapping_manager.routemanager_get_mode(
walker_area_name)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', walker_index + 1)
self.__mapping_manager.set_devicesetting_value_of(
origin, 'finished', False)
if walker_index >= len(walker_area_array) - 1:
self.__mapping_manager.set_devicesetting_value_of(
origin, 'walker_area_index', 0)
# set global mon_iv
routemanager_settings = self.__mapping_manager.routemanager_get_settings(
walker_area_name)
if routemanager_settings is not None:
client_mapping['mon_ids_iv'] =\
self.__mapping_manager.get_monlist(routemanager_settings.get("mon_ids_iv", None),
walker_area_name)
else:
walker_routemanager_mode = None
if "last_location" not in devicesettings:
devicesettings['last_location'] = Location(0.0, 0.0)
logger.debug("Setting up worker for {}", str(origin))
worker = None
if walker_routemanager_mode is None:
pass
elif walker_routemanager_mode in [
"raids_mitm", "mon_mitm", "iv_mitm"
]:
worker = WorkerMITM(
self.args,
origin,
last_known_state,
self,
routemanager_name=walker_area_name,
mitm_mapper=self.__mitm_mapper,
mapping_manager=self.__mapping_manager,
db_wrapper=self.__db_wrapper,
pogo_window_manager=self.__pogoWindowManager,
walker=walker_settings)
elif walker_routemanager_mode in ["pokestops"]:
worker = WorkerQuests(
self.args,
origin,
last_known_state,
self,
routemanager_name=walker_area_name,
mitm_mapper=self.__mitm_mapper,
mapping_manager=self.__mapping_manager,
db_wrapper=self.__db_wrapper,
pogo_window_manager=self.__pogoWindowManager,
walker=walker_settings)
elif walker_routemanager_mode in ["idle"]:
worker = WorkerConfigmode(
self.args,
origin,
self,
walker=walker_settings,
mapping_manager=self.__mapping_manager,
mitm_mapper=self.__mitm_mapper,
db_wrapper=self.__db_wrapper,
routemanager_name=walker_area_name)
else:
logger.error("Mode not implemented")
sys.exit(1)
if worker is None:
logger.error(
"Invalid walker mode for {}. Closing connection".format(
str(origin)))
await websocket_client_connection.close()
else:
logger.debug("Starting worker for {}", str(origin))
new_worker_thread = Thread(name='worker_%s' % origin,
target=worker.start_worker)
new_worker_thread.daemon = True
async with self.__users_mutex:
self.__current_users[origin] = [
new_worker_thread, worker, websocket_client_connection,
0
]
new_worker_thread.start()
except WrongAreaInWalker:
logger.error('Unknown Area in Walker settings - check config')
await websocket_client_connection.close()
except Exception:
logger.opt(exception=True).error(
"Other unhandled exception during registration of {}.", origin)
await websocket_client_connection.close()
finally:
async with self.__users_mutex:
self.__users_connecting.remove(origin)
await asyncio.sleep(5)
return True
async def __unregister(self, websocket_client_connection):
# worker_thread: Thread = None
async with self.__users_mutex:
worker_id = str(
websocket_client_connection.request_headers.get_all("Origin")
[0])
worker = self.__current_users.get(worker_id, None)
if worker is not None:
worker[1].stop_worker()
self.__current_users.pop(worker_id)
logger.info("Worker {} unregistered", str(worker_id))
self.__worker_shutdown_queue.put(worker[0])
# TODO ? worker_thread.join()
def force_disconnect(self, origin):
worker = self.__current_users.get(origin, None)
if worker is not None:
worker[1]._internal_cleanup()
async def __producer_handler(self, websocket_client_connection):
while websocket_client_connection.open:
# logger.debug("Connection still open, trying to send next message")
# retrieve next message from queue to be sent, block if empty
next = None
while next is None and websocket_client_connection.open:
logger.debug("Fetching next message to send")
next = await self.__retrieve_next_send(
websocket_client_connection)
if next is None:
# logger.debug("next is None, stopping connection...")
return
await self.__send_specific(websocket_client_connection,
next.id, next.message)
async def __send_specific(self, websocket_client_connection, id, message):
# await websocket_client_connection.send(message)
try:
user = None
async with self.__users_mutex:
for key, value in self.__current_users.items():
if key == id and value[2].open:
user = value
if user is not None:
await user[2].send(message)
except Exception as e:
logger.error("Failed sending message in send_specific: {}".format(
str(e)))
async def __retrieve_next_send(self, websocket_client_connection):
found = None
while found is None and websocket_client_connection.open:
try:
found = self.__send_queue.get_nowait()
except Exception as e:
await asyncio.sleep(0.02)
if not websocket_client_connection.open:
logger.debug(
"retrieve_next_send: connection closed, returning None")
return found
async def __consumer_handler(self, websocket_client_connection):
if websocket_client_connection is None:
return
worker_id = str(
websocket_client_connection.request_headers.get_all("Origin")[0])
logger.info("Consumer handler of {} starting", str(worker_id))
while websocket_client_connection.open:
message = None
try:
message = await asyncio.wait_for(
websocket_client_connection.recv(), timeout=4.0)
except asyncio.TimeoutError as te:
await asyncio.sleep(0.02)
except websockets.exceptions.ConnectionClosed as cc:
logger.warning("Connection to {} was closed, stopping worker",
str(worker_id))
async with self.__users_mutex:
worker = self.__current_users.get(worker_id, None)
if worker is not None:
# TODO: do it abruptly in the worker, maybe set a flag to be checked for in send_and_wait to
# TODO: throw an exception
worker[1].stop_worker()
await self.__internal_clean_up_user(worker_id, None)
return
if message is not None:
await self.__on_message(message)
logger.warning("Connection of {} closed in consumer_handler",
str(worker_id))
async def __internal_clean_up_user(self, worker_id, worker_instance):
"""
:param worker_id: The ID/Origin of the worker
:param worker_instance: None if the cleanup is called from within the websocket server
:return:
"""
async with self.__users_mutex:
if worker_id in self.__current_users.keys() and (
worker_instance is None
or self.__current_users[worker_id][1] == worker_instance):
if self.__current_users[worker_id][2].open:
logger.info("Calling close for {}...", str(worker_id))
await self.__current_users[worker_id][2].close()
# self.__current_users.pop(worker_id)
# logger.info("Info of {} removed in websocket", str(worker_id))
def clean_up_user(self, worker_id, worker_instance):
logger.debug2("Cleanup of {} called with ref {}".format(
worker_id, str(worker_instance)))
with self.__loop_mutex:
future = asyncio.run_coroutine_threadsafe(
self.__internal_clean_up_user(worker_id, worker_instance),
self.__loop)
future.result()
async def __on_message(self, message):
id = -1
response = None
if isinstance(message, str):
logger.debug("Receiving message: {}", str(message.strip()))
splitup = message.split(";", 1)
id = int(splitup[0])
response = splitup[1]
else:
logger.debug("Received binary values.")
id = int.from_bytes(message[:4], byteorder='big', signed=False)
response = message[4:]
await self.__set_response(id, response)
if not await self.__set_event(id):
# remove the response again - though that is kinda stupid
await self.__pop_response(id)
async def __set_event(self, id):
result = False
async with self.__requests_mutex:
if id in self.__requests:
self.__requests[id].set()
result = True
else:
# the request has already been deleted due to a timeout...
logger.error("Request has already been deleted...")
return result
async def __set_response(self, id, message):
async with self.__received_mutex:
self.__received[id] = message
async def __pop_response(self, id):
async with self.__received_mutex:
message = self.__received.pop(id)
return message
async def __get_new_message_id(self):
async with self.__id_mutex:
self.__next_id += 1
self.__next_id = int(math.fmod(self.__next_id, 100000))
if self.__next_id == 100000:
self.__next_id = 1
toBeReturned = self.__next_id
return toBeReturned
async def __send(self, id, to_be_sent):
next_message = OutgoingMessage(id, to_be_sent)
await self.__send_queue.put(next_message)
async def __send_and_wait_internal(self,
id,
worker_instance,
message,
timeout,
byte_command: int = None):
async with self.__users_mutex:
user_entry = self.__current_users.get(id, None)
if user_entry is None or user_entry[
1] != worker_instance and worker_instance != 'madmin':
raise WebsocketWorkerRemovedException
message_id = await self.__get_new_message_id()
message_event = asyncio.Event()
message_event.clear()
await self.__set_request(message_id, message_event)
if isinstance(message, str):
to_be_sent: str = u"%s;%s" % (str(message_id), message)
logger.debug("To be sent to {}: {}", id, to_be_sent.strip())
elif byte_command is not None:
to_be_sent: bytes = (int(message_id)).to_bytes(4, byteorder='big')
to_be_sent += (int(byte_command)).to_bytes(4, byteorder='big')
to_be_sent += message
logger.debug("To be sent to {} (message ID: {}): {}", id,
message_id, str(to_be_sent[:10]))
else:
logger.fatal(
"Tried to send invalid message (bytes without byte command or no byte/str passed)"
)
return None
await self.__send(id, to_be_sent)
# now wait for the response!
result = None
logger.debug("Timeout: {}", str(timeout))
event_triggered = None
try:
event_triggered = await asyncio.wait_for(message_event.wait(),
timeout=timeout)
except asyncio.TimeoutError as te:
logger.warning("Timeout, increasing timeout-counter")
# TODO: why is the user removed here?
new_count = await self.__increase_fail_counter(id)
if new_count > 5:
logger.error(
"5 consecutive timeouts to {} or origin is not longer connected, cleanup",
str(id))
await self.__internal_clean_up_user(id, None)
await self.__reset_fail_counter(id)
await self.__remove_request(message_id)
raise WebsocketWorkerTimeoutException
if event_triggered:
logger.debug("Received answer in time, popping response")
await self.__reset_fail_counter(id)
await self.__remove_request(message_id)
result = await self.__pop_response(message_id)
if isinstance(result, str):
logger.debug("Response to {}: {}", str(id),
str(result.strip()))
else:
logger.debug("Received binary data to {}, starting with {}",
str(id), str(result[:10]))
return result
def send_and_wait(self,
id,
worker_instance,
message,
timeout,
byte_command: int = None):
if isinstance(message, bytes):
logger.debug("{} sending binary: {}", str(id), str(message[:10]))
else:
logger.debug("{} sending command: {}", str(id), message.strip())
try:
# future: Handle = self._add_task_to_loop(self.__send_and_wait_internal(id, worker_instance, message,
# timeout))
logger.debug("Appending send_and_wait to {}".format(
str(self.__loop)))
with self.__loop_mutex:
future = asyncio.run_coroutine_threadsafe(
self.__send_and_wait_internal(id,
worker_instance,
message,
timeout,
byte_command=byte_command),
self.__loop)
result = future.result()
except WebsocketWorkerRemovedException:
logger.error(
"Worker {} was removed, propagating exception".format(id))
raise WebsocketWorkerRemovedException
except WebsocketWorkerTimeoutException:
logger.error(
"Sending message failed due to timeout ({})".format(id))
raise WebsocketWorkerTimeoutException
return result
async def __set_request(self, id, event):
async with self.__requests_mutex:
self.__requests[id] = event
async def __reset_fail_counter(self, id):
async with self.__users_mutex:
if id in self.__current_users.keys():
self.__current_users[id][3] = 0
async def __increase_fail_counter(self, id):
async with self.__users_mutex:
if id in self.__current_users.keys():
new_count = self.__current_users[id][3] + 1
self.__current_users[id][3] = new_count
else:
new_count = 100
return new_count
async def __remove_request(self, message_id):
async with self.__requests_mutex:
self.__requests.pop(message_id)
def get_reg_origins(self):
return self.__current_users
def get_origin_communicator(self, origin):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].get_communicator()
return None
def set_geofix_sleeptime_worker(self, origin, sleeptime):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].set_geofix_sleeptime(
sleeptime)
return False
def trigger_worker_check_research(self, origin):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].trigger_check_research()
return False
def set_update_sleeptime_worker(self, origin, sleeptime):
if self.__current_users.get(origin, None) is not None:
return self.__current_users[origin][1].set_geofix_sleeptime(
sleeptime)
return False
def set_job_activated(self, origin):
self.__mapping_manager.set_devicesetting_value_of(origin, 'job', True)
def set_job_deactivated(self, origin):
self.__mapping_manager.set_devicesetting_value_of(origin, 'job', False)
class PHPServer(Thread, RequestManager):
name = "PHPServer"
request_frequency = 0.1
def __init__(self, controller):
super().__init__()
self.cont = controller
self.main_queue = Queue()
self.side_queue = Queue()
self.shutdown_event = Event()
self.serve_event = Event()
self.running_game = Event()
self.server_address = None
self.server_address_messenger = None
self.param = None
self.setup_done = False
def setup(self, param):
self.param = param
self.server_address = self.param["network"]["php_server"]
self.server_address_messenger = self.param["network"]["messenger"]
if not self.setup_done:
self.is_another_server_running()
self.setup_done = True
def is_another_server_running(self):
while True:
response = self.send_request(demand_type="writing",
table="server_is_running",
close_server=0)
self.log("I got the response '{}' from the distant server.".format(
response.text))
if response.text == "I updated is_running variable from server_is_running.":
break
elif response.text == "Another server seems to be running.":
self.log(
"Another server seems to be running! Game could be compromised!",
level=3)
self.cont.queue.put((
"ask_interface", "show_critical_and_ok",
"Another server seems to be running! Game could be compromised!"
))
break
def run(self):
while not self.shutdown_event.is_set():
self.log("Waiting for a message...")
msg = self.main_queue.get()
self.log("I received msg '{}'.".format(msg))
if msg and msg[0] == "serve":
self.serve_event.set()
self.serve()
self.log("I'm dead.")
def serve(self):
while self.serve_event.is_set():
self.treat_sides_requests()
if self.running_game.is_set():
self.treat_game_requests()
Event().wait(self.request_frequency)
def treat_game_requests(self):
response = self.send_request(demand_type="reading", table="request")
if response.text and response.text.split("&")[0] == "request":
requests = [i for i in response.text.split("&")[1:] if i]
if requests:
for request in requests:
self.cont.queue.put(("server_request", request))
self.log("I will treat {} request(s).".format(len(requests)))
self.treat_requests(n_requests=len(requests))
def treat_sides_requests(self):
# check for new msg received
self.receive_messages()
# check for new interactions with sql tables
if not self.side_queue.empty():
msg = self.side_queue.get()
if msg and msg[0] == "send_message":
self.send_message(msg[1], msg[2])
elif msg and msg[0] == "get_waiting_list":
waiting_list = self.get_waiting_list()
self.cont.queue.put(
("server_update_assignment_frame", waiting_list))
elif msg and msg[0] == "erase_sql_tables":
self.ask_for_erasing_tables(tables=msg[1:])
elif msg and msg[0] == "authorize_participants":
self.authorize_participants(*msg[1:])
elif msg and msg[0] == "set_missing_players":
self.set_missing_players(msg[1])
def treat_requests(self, n_requests):
for i in range(n_requests):
self.log("I'm treating the request no {}.".format(i))
should_be_reply, response = self.main_queue.get()
if should_be_reply == "reply":
response = self.send_request(demand_type="writing",
table="response",
gameId=response["game_id"],
response=response["response"])
self.log("Response from distant server is: '{}'.".format(
response.text))
elif should_be_reply == "error":
self.log("I will not send response now (code error is '{}').".
format(response))
else:
raise Exception("Something went wrong...")
def receive_messages(self):
self.log("I send a request for collecting the messages.")
response = self.send_request_messenger(demandType="serverHears",
userName="******",
message="none")
if "reply" in response.text:
args = [i for i in response.text.split("/") if i]
n_messages = int(args[1])
self.log("I received {} new message(s).".format(n_messages))
if n_messages:
for arg in args[2:]:
sep_args = arg.split("<>")
user_name, message = sep_args[0], sep_args[1]
self.cont.queue.put(
("server_new_message", user_name, message))
self.log(
"I send confirmation for message '{}'.".format(arg))
self.send_request_messenger(
demandType="serverReceiptConfirmation",
userName=user_name,
message=message)
def send_message(self, user_name, message):
self.log("I send a message for '{}': '{}'.".format(user_name, message))
response = self.send_request_messenger(demandType="serverSpeaks",
userName=user_name,
message=message)
self.log("I receive: {}".format(response.text))
def set_server_is_not_running_anymore(self):
while True:
self.log("I notify sql tables that server is closed.", level=1)
response = self.send_request(
demand_type="writing",
table="server_is_running",
close_server=1,
)
self.log("I receive: {}".format(response.text))
if response.text == "Updated is_running to 0.":
self.log("Server is not running anymore on SQL tables.",
level=1)
break
def stop_to_serve(self):
self.serve_event.clear()
def end(self):
# when server shutdowns, erase tables and tell
# tables server is not running anymore
if self.setup_done:
tables = ("participants", "waiting_list", "request", "response")
self.ask_for_erasing_tables(tables=tables)
self.set_server_is_not_running_anymore()
self.serve_event.clear()
self.shutdown_event.set()
self.main_queue.put("break")
def get_waiting_list(self):
while True:
self.log(
"I will ask the distant server to the 'waiting_list' table.")
response = self.send_request(demand_type="reading",
table="waiting_list")
if response.text and response.text.split("&")[0] == "waiting_list":
participants = [i for i in response.text.split("&")[1:] if i]
break
return participants
def get_users(self):
while True:
self.log("I will ask the distant server to the 'users' table.")
response = self.send_request(demand_type="reading", table="users")
if response.text and response.text.split("&")[0] == "users":
users = [
i.split("#") for i in response.text.split("&")[1:] if i
]
break
return users
def register_users(self, usernames, passwords):
while True:
self.log(
"I will ask the distant server to write the 'users' table.")
response = self.send_request(demand_type="writing",
table="users",
names=usernames,
passwords=passwords)
self.log("I got the response '{}' from the distant server.".format(
response.text))
if response.text == "I inserted users in 'users' table.":
break
def register_waiting_list(self, usernames):
while True:
self.log(
"I will ask the distant server to write the 'waiting_list' table."
)
response = self.send_request(
demand_type="writing",
table="waiting_list",
names=usernames,
)
self.log("I got the response '{}' from the distant server.".format(
response.text))
if response.text == "I inserted names in 'waiting_list' table.":
break
def authorize_participants(self, participants, roles, game_ids):
while True:
self.log(
"I will ask the distant server to fill the 'participants' table with {}"
.format(participants),
level=1)
response = self.send_request(demand_type="writing",
table="participants",
gameIds=json.dumps(game_ids),
names=json.dumps(participants),
roles=json.dumps(roles))
self.log("I got the response '{}' from the distant server.".format(
response.text),
level=1)
if response.text == "I inserted participants in 'participants' table.":
break
def ask_for_erasing_tables(self, tables):
while True:
self.log("I will ask the distant server to erase tables.", level=1)
response = self.send_request(demand_type="empty_tables",
table_names=json.dumps(tables))
self.log("I got the response '{}' from the distant server.".format(
response.text),
level=1)
if "Tables" in response.text and "have been erased" in response.text:
break
def set_missing_players(self, value):
while True:
self.log(
"I will ask the distant server to update the 'missing_players' variable with {}"
.format(value),
level=1)
response = self.send_request(demand_type="writing",
table="game",
missingPlayers=value)
self.log("I got the response '{}' from the distant server.".format(
response.text),
level=1)
if response.text == "I updated missing players in 'game' table.":
break
class AutoScaler(object):
"""This class implements a simple autoscaling algorithm: if queries queue up for a
configurable duration, a new executor group is started. Likewise, if the number of
concurrently running queries indicated that an executor group can be removed, such
measure is taken.
Users of this class can start an auto scaler by calling start() and must call stop()
before exiting (see main() below for an example).
This class only uses the default admission control pool.
"""
DEFAULT_POOL_NAME = "default-pool"
def __init__(self,
executor_slots,
group_size,
start_batch_size=0,
max_groups=0,
wait_up_s=0,
wait_down_s=0,
coordinator_slots=128):
# Number of queries that can run concurrently on each executor
self.executor_slots = executor_slots
self.coordinator_slots = coordinator_slots
# Number of executors per executor group
self.group_size = group_size
# New executor groups will be started in increments of this size
self.start_batch_size = group_size
if start_batch_size > 0:
self.start_batch_size = start_batch_size
# Maximum number of executor groups. We only have 10 TCP ports free on our
# miniclusters and we need one for the dedicated coordinator.
self.max_groups = 9 / self.group_size
# max_groups can further bound the maximum number of groups we are going to start,
# but we won't start more than possible.
if max_groups > 0 and max_groups < self.max_groups:
self.max_groups = max_groups
# Number of seconds to wait before scaling up/down
self.scale_wait_up_s = 5
if wait_up_s > 0:
self.scale_wait_up_s = wait_up_s
self.scale_wait_down_s = 5
if wait_down_s > 0:
self.scale_wait_down_s = wait_down_s
self.groups = []
self.num_groups = 0
# Stopwatches to track how long the conditions for scaling up/down have been met.
self.scale_up_sw = time.time()
self.scale_down_sw = time.time()
self.loop_thread = None
# Event to signal that the control loop should exit
self.stop_ev = Event()
def get_cluster(self):
return ImpalaCluster.get_e2e_test_cluster()
def get_coordinator(self):
cluster = self.get_cluster()
assert len(cluster.impalads) > 0
return cluster.get_first_impalad()
def get_service(self):
return self.get_coordinator().service
def get_client(self):
return self.get_coordinator().service.create_hs2_client()
def group_name(self, idx):
# By convention, group names must start with their associated resource pool name
# followed by a "-".
return "%s-group-%s" % (self.DEFAULT_POOL_NAME, idx)
def start_base_cluster(self):
"""Starts the base cluster consisting of an exclusive coordinator, catalog, and
statestore. Does not add any executors."""
logging.info(
"Starting base cluster (coordinator, catalog, statestore)")
cluster_args = ["--impalad_args=-executor_groups=coordinator"]
self._start_impala_cluster(cluster_args,
cluster_size=1,
executor_slots=self.coordinator_slots,
expected_num_executors=0,
add_executors=False)
logging.info("Done, number of running executor groups: %s" %
self.num_groups)
def start_group(self):
"""Starts an executor group. The name of the group is automatically determined based
on the current number of total executor groups. Executors in the group will be started
in batches."""
self.num_groups += 1
name = self.group_name(self.num_groups)
desc = "%s:%s" % (name, self.group_size)
logging.info("Starting executor group %s with %s members" %
(name, self.group_size))
cluster_args = ["--impalad_args=-executor_groups=%s" % desc]
batch_size = self.start_batch_size
num_started = 0
num_expected = (self.num_groups - 1) * self.group_size
while (num_started < self.group_size):
to_start = min(batch_size, self.group_size - num_started)
num_expected += to_start
if to_start == 1:
start_msg = "Starting executor %s" % (num_started + 1)
else:
start_msg = "Starting executors %s-%s" % (
num_started + 1, num_started + to_start)
logging.info(start_msg)
self._start_impala_cluster(cluster_args,
cluster_size=to_start,
executor_slots=self.executor_slots,
expected_num_executors=num_expected,
add_executors=True)
num_started += to_start
logging.info("Done, number of running executor groups: %s" %
self.num_groups)
def stop_group(self):
"""Stops the executor group that was added last."""
name = self.group_name(self.num_groups)
group_hosts = self.get_groups()[name]
logging.info("Stopping executor group %s" % name)
for host in group_hosts:
logging.debug("Stopping host %s" % host)
query = ":shutdown('%s');" % host
self.execute(query)
self.wait_for_group_gone(name)
self.num_groups -= 1
logging.info("Done, number of running executor groups: %s" %
self.num_groups)
def wait_for_group_gone(self, group_name, timeout=120):
"""Waits until all executors in group 'group_name' have unregistered themselves from
the coordinator's cluster membership view."""
end = time.time() + timeout
while time.time() < end:
groups = self.get_groups()
if group_name not in groups:
return
time.sleep(0.5)
assert False, "Timeout waiting for group %s to shut down" % group_name
def get_groups(self):
return self.get_service().get_executor_groups()
def execute(self, query):
return self.get_client().execute(query)
def get_num_queued_queries(self):
"""Returns the number of queries currently queued in the default pool on the
coordinator."""
return self.get_service().get_num_queued_queries(
pool_name=self.DEFAULT_POOL_NAME)
def get_num_running_queries(self):
"""Returns the number of queries currently queued in the default pool on the
coordinator."""
return self.get_service().get_num_running_queries(
self.DEFAULT_POOL_NAME)
def loop(self):
"""Controls whether new executor groups need to be started or existing ones need to be
stopped, based on the number of queries that are currently queued and running.
"""
while not self.stop_ev.is_set():
now = time.time()
num_queued = self.get_num_queued_queries()
num_running = self.get_num_running_queries()
capacity = self.executor_slots * self.num_groups
logging.debug("queued: %s, running: %s, capacity: %s" %
(num_queued, num_running, capacity))
if num_queued == 0:
self.scale_up_sw = now
scale_up = self.scale_up_sw < now - self.scale_wait_up_s
if scale_up and self.num_groups < self.max_groups:
self.start_group()
self.scale_up_sw = time.time()
self.scale_down_sw = self.scale_up_sw
continue
surplus = capacity - num_running
if surplus < self.executor_slots:
self.scale_down_sw = now
if self.scale_down_sw < now - self.scale_wait_down_s:
self.stop_group()
self.scale_up_sw = time.time()
self.scale_down_sw = self.scale_up_sw
continue
time.sleep(1)
def start(self):
"""Starts a base cluster with coordinator and statestore and the control loop to start
and stop additional executor groups."""
self.start_base_cluster()
assert self.loop_thread is None
self.loop_thread = Thread(target=self.loop)
self.loop_thread.start()
def stop(self):
"""Stops the AutoScaler and its cluster."""
if self.stop_ev.is_set():
return
self.stop_ev.set()
if self.loop_thread:
self.loop_thread.join()
self.loop_thread = None
self._kill_whole_cluster()
def _start_impala_cluster(self, options, cluster_size, executor_slots,
expected_num_executors, add_executors):
"""Starts an Impala cluster and waits for all impalads to come online.
If 'add_executors' is True, new executors will be added to the cluster and the
existing daemons will not be restarted. In that case 'cluster_size' must specify the
number of nodes that will be added and 'expected_num_executors' must be the total
expected number of executors after the additional ones have started.
If 'add_executors' is false, 'cluster_size' must be 1 and a single exclusive
coordinator will be started (together with catalog and statestore).
"""
assert cluster_size > 0, "cluster_size cannot be 0"
impala_log_dir = os.getenv("LOG_DIR", "/tmp/")
cmd = [
os.path.join(IMPALA_HOME, "bin/start-impala-cluster.py"),
"--cluster_size=%d" % cluster_size,
"--log_dir=%s" % impala_log_dir, "--log_level=1"
]
if add_executors:
cmd.append("--add_executors")
else:
assert expected_num_executors == 0
assert cluster_size == 1
cmd.append("--use_exclusive_coordinators")
impalad_args = [
"-vmodule=admission-controller=3,cluster-membership-mgr=3",
"-max_concurrent_queries=%s" % executor_slots,
"-shutdown_grace_period_s=2"
]
options += ["--impalad_args=%s" % a for a in impalad_args]
logging.debug("Starting cluster with command: %s" %
" ".join(pipes.quote(arg) for arg in cmd + options))
log_debug = logging.getLogger().getEffectiveLevel() == logging.DEBUG
log_file = None
if not log_debug:
log_file = open("/dev/null", "w")
check_call(cmd + options,
close_fds=True,
stdout=log_file,
stderr=log_file)
# The number of statestore subscribers is
# cluster_size (# of impalad) + 1 (for catalogd).
if expected_num_executors > 0:
expected_subscribers = expected_num_executors + 2
expected_backends = expected_num_executors + 1
else:
expected_subscribers = cluster_size + 1
expected_backends = 1
cluster = self.get_cluster()
statestored = cluster.statestored
if statestored is None:
raise Exception("statestored was not found")
logging.debug("Waiting for %s subscribers to come online" %
expected_subscribers)
statestored.service.wait_for_live_subscribers(expected_subscribers,
timeout=60)
for impalad in cluster.impalads:
logging.debug("Waiting for %s executors to come online" %
expected_backends)
impalad.service.wait_for_num_known_live_backends(expected_backends,
timeout=60)
def _kill_whole_cluster(self):
"""Terminates the whole cluster, i.e. all impalads, catalogd, and statestored."""
logging.info("terminating cluster")
check_call([
os.path.join(IMPALA_HOME, "bin/start-impala-cluster.py"),
"--kill_only"
])
class AsyncResult(Future):
"""Class for representing results of non-blocking calls.
Provides the same interface as :py:class:`multiprocessing.pool.AsyncResult`.
"""
msg_ids = None
_targets = None
_tracker = None
_single_result = False
owner = False
def __init__(self, client, children, fname='unknown', targets=None,
owner=False,
):
super(AsyncResult, self).__init__()
if not isinstance(children, list):
children = [children]
self._single_result = True
else:
self._single_result = False
if isinstance(children[0], string_types):
self.msg_ids = children
self._children = []
else:
self._children = children
self.msg_ids = [ f.msg_id for f in children ]
self._client = client
self._fname = fname
self._targets = targets
self.owner = owner
self._ready = False
self._ready_event = Event()
self._output_ready = False
self._output_event = Event()
self._sent_event = Event()
self._success = None
if self._children:
self._metadata = [ f.output.metadata for f in self._children ]
else:
self._metadata = [self._client.metadata[id] for id in self.msg_ids]
self._init_futures()
def _init_futures(self):
"""Build futures for results and output; hook up callbacks"""
if not self._children:
for msg_id in self.msg_ids:
future = self._client._futures.get(msg_id, None)
if not future:
result = self._client.results.get(msg_id, _default)
# result resides in local cache, construct already-resolved Future
if result is not _default:
future = MessageFuture(msg_id)
future.output = Future()
future.output.metadata = self.client.metadata[msg_id]
future.set_result(result)
future.output.set_result(None)
if not future:
raise KeyError("No Future or result for msg_id: %s" % msg_id)
self._children.append(future)
self._result_future = multi_future(self._children)
self._sent_future = multi_future([ f.tracker for f in self._children ])
self._sent_future.add_done_callback(self._handle_sent)
self._output_future = multi_future([self._result_future] + [
f.output for f in self._children
])
# on completion of my constituents, trigger my own resolution
self._result_future.add_done_callback(self._resolve_result)
self._output_future.add_done_callback(self._resolve_output)
self.add_done_callback(self._finalize_result)
def __repr__(self):
if self._ready:
return "<%s: %s:finished>" % (self.__class__.__name__, self._fname)
else:
return "<%s: %s>" % (self.__class__.__name__, self._fname)
def __dir__(self):
keys = dir(self.__class__)
if not _metadata_keys:
from .client import Metadata
_metadata_keys.extend(Metadata().keys())
keys.extend(_metadata_keys)
return keys
def _reconstruct_result(self, res):
"""Reconstruct our result from actual result list (always a list)
Override me in subclasses for turning a list of results
into the expected form.
"""
if self._single_result:
return res[0]
else:
return res
def get(self, timeout=-1):
"""Return the result when it arrives.
If `timeout` is not ``None`` and the result does not arrive within
`timeout` seconds then ``TimeoutError`` is raised. If the
remote call raised an exception then that exception will be reraised
by get() inside a `RemoteError`.
"""
if not self.ready():
self.wait(timeout)
if self._ready:
if self._success:
return self.result()
else:
raise self.exception()
else:
raise error.TimeoutError("Result not ready.")
def _check_ready(self):
if not self.ready():
raise error.TimeoutError("Result not ready.")
def ready(self):
"""Return whether the call has completed."""
if not self._ready:
self.wait(0)
return self._ready
def wait_for_output(self, timeout=-1):
"""Wait for our output to be complete.
AsyncResult.wait only waits for the result,
which may arrive before output is complete.
"""
if self._output_ready:
return True
if timeout and timeout < 0:
timeout = None
return self._output_event.wait(timeout)
def _resolve_output(self, f=None):
"""Callback that fires when outputs are ready"""
if self.owner:
[ self._client.metadata.pop(mid, None) for mid in self.msg_ids ]
self._output_ready = True
self._output_event.set()
def wait(self, timeout=-1):
"""Wait until the result is available or until `timeout` seconds pass.
This method always returns None.
"""
if self._ready:
return True
if timeout and timeout < 0:
timeout = None
self._ready_event.wait(timeout)
self.wait_for_output(0)
return self._ready
def _resolve_result(self, f=None):
try:
if f:
results = f.result()
else:
results = list(map(self._client.results.get, self.msg_ids))
if self._single_result:
r = results[0]
if isinstance(r, Exception):
raise r
else:
results = error.collect_exceptions(results, self._fname)
self._success = True
self.set_result(self._reconstruct_result(results))
except Exception as e:
self._success = False
self.set_exception(e)
def _finalize_result(self,f):
if self.owner:
[ self._client.results.pop(mid, None) for mid in self.msg_ids ]
self._ready = True
self._ready_event.set()
def successful(self):
"""Return whether the call completed without raising an exception.
Will raise ``AssertionError`` if the result is not ready.
"""
assert self.ready()
return self._success
#----------------------------------------------------------------
# Extra methods not in mp.pool.AsyncResult
#----------------------------------------------------------------
def get_dict(self, timeout=-1):
"""Get the results as a dict, keyed by engine_id.
timeout behavior is described in `get()`.
"""
results = self.get(timeout)
if self._single_result:
results = [results]
engine_ids = [ md['engine_id'] for md in self._metadata ]
rdict = {}
for engine_id, result in zip(engine_ids, results):
if engine_id in rdict:
raise ValueError("Cannot build dict, %i jobs ran on engine #%i" % (
engine_ids.count(engine_id), engine_id)
)
else:
rdict[engine_id] = result
return rdict
@property
def r(self):
"""result property wrapper for `get(timeout=-1)`."""
return self.get()
@property
def metadata(self):
"""property for accessing execution metadata."""
if self._single_result:
return self._metadata[0]
else:
return self._metadata
@property
def result_dict(self):
"""result property as a dict."""
return self.get_dict()
def __dict__(self):
return self.get_dict(0)
def abort(self):
"""abort my tasks."""
assert not self.ready(), "Can't abort, I am already done!"
return self._client.abort(self.msg_ids, targets=self._targets, block=True)
def _handle_sent(self, f):
"""Resolve sent Future, build MessageTracker"""
trackers = f.result()
trackers = [t for t in trackers if t is not None]
self._tracker = MessageTracker(*trackers)
self._sent_event.set()
@property
def sent(self):
"""check whether my messages have been sent."""
return self._sent_event.is_set() and self._tracker.done
def wait_for_send(self, timeout=-1):
"""wait for pyzmq send to complete.
This is necessary when sending arrays that you intend to edit in-place.
`timeout` is in seconds, and will raise TimeoutError if it is reached
before the send completes.
"""
if not self._sent_event.is_set():
if timeout and timeout < 0:
# Event doesn't like timeout < 0
timeout = None
# wait for Future to indicate send having been called,
# which means MessageTracker is ready.
tic = time.time()
if not self._sent_event.wait(timeout):
raise error.TimeoutError("Still waiting to be sent")
return False
if timeout:
timeout = max(0, timeout - (time.time() - tic))
try:
if timeout is None:
# MessageTracker doesn't like timeout=None
timeout = -1
return self._tracker.wait(timeout)
except zmq.NotDone:
raise error.TimeoutError("Still waiting to be sent")
#-------------------------------------
# dict-access
#-------------------------------------
def __getitem__(self, key):
"""getitem returns result value(s) if keyed by int/slice, or metadata if key is str.
"""
if isinstance(key, int):
self._check_ready()
return error.collect_exceptions([self.result()[key]], self._fname)[0]
elif isinstance(key, slice):
self._check_ready()
return error.collect_exceptions(self.result()[key], self._fname)
elif isinstance(key, string_types):
# metadata proxy *does not* require that results are done
self.wait(0)
self.wait_for_output(0)
values = [ md[key] for md in self._metadata ]
if self._single_result:
return values[0]
else:
return values
else:
raise TypeError("Invalid key type %r, must be 'int','slice', or 'str'"%type(key))
def __getattr__(self, key):
"""getattr maps to getitem for convenient attr access to metadata."""
try:
return self.__getitem__(key)
except (error.TimeoutError, KeyError):
raise AttributeError("%r object has no attribute %r"%(
self.__class__.__name__, key))
@staticmethod
def _wait_for_child(child, evt, timeout=_FOREVER):
"""Wait for a child to be done"""
evt.clear()
child.add_done_callback(lambda f: evt.set())
evt.wait(timeout)
# asynchronous iterator:
def __iter__(self):
if self._single_result:
raise TypeError("AsyncResults with a single result are not iterable.")
try:
rlist = self.get(0)
except error.TimeoutError:
# wait for each result individually
evt = Event()
for child in self._children:
self._wait_for_child(child, evt=evt)
result = child.result()
error.collect_exceptions([result], self._fname)
yield result
else:
# already done
for r in rlist:
yield r
def __len__(self):
return len(self.msg_ids)
#-------------------------------------
# Sugar methods and attributes
#-------------------------------------
def timedelta(self, start, end, start_key=min, end_key=max):
"""compute the difference between two sets of timestamps
The default behavior is to use the earliest of the first
and the latest of the second list, but this can be changed
by passing a different
Parameters
----------
start : one or more datetime objects (e.g. ar.submitted)
end : one or more datetime objects (e.g. ar.received)
start_key : callable
Function to call on `start` to extract the relevant
entry [default: min]
end_key : callable
Function to call on `end` to extract the relevant
entry [default: max]
Returns
-------
dt : float
The time elapsed (in seconds) between the two selected timestamps.
"""
if not isinstance(start, datetime):
# handle single_result AsyncResults, where ar.stamp is single object,
# not a list
start = start_key(start)
if not isinstance(end, datetime):
# handle single_result AsyncResults, where ar.stamp is single object,
# not a list
end = end_key(end)
return compare_datetimes(end, start).total_seconds()
@property
def progress(self):
"""the number of tasks which have been completed at this point.
Fractional progress would be given by 1.0 * ar.progress / len(ar)
"""
self.wait(0)
return len(self) - len(set(self.msg_ids).intersection(self._client.outstanding))
@property
def elapsed(self):
"""elapsed time since initial submission"""
if self.ready():
return self.wall_time
now = submitted = utcnow()
for msg_id in self.msg_ids:
if msg_id in self._client.metadata:
stamp = self._client.metadata[msg_id]['submitted']
if stamp and stamp < submitted:
submitted = stamp
return compare_datetimes(now, submitted).total_seconds()
@property
@check_ready
def serial_time(self):
"""serial computation time of a parallel calculation
Computed as the sum of (completed-started) of each task
"""
t = 0
for md in self._metadata:
t += compare_datetimes(md['completed'], md['started']).total_seconds()
return t
@property
@check_ready
def wall_time(self):
"""actual computation time of a parallel calculation
Computed as the time between the latest `received` stamp
and the earliest `submitted`.
For similar comparison of other timestamp pairs, check out AsyncResult.timedelta.
"""
return self.timedelta(self.submitted, self.received)
def wait_interactive(self, interval=1., timeout=-1):
"""interactive wait, printing progress at regular intervals"""
if timeout is None:
timeout = -1
N = len(self)
tic = time.time()
while not self.ready() and (timeout < 0 or time.time() - tic <= timeout):
self.wait(interval)
clear_output(wait=True)
print("%4i/%i tasks finished after %4i s" % (self.progress, N, self.elapsed), end="")
sys.stdout.flush()
print("\ndone")
def _republish_displaypub(self, content, eid):
"""republish individual displaypub content dicts"""
ip = get_ipython()
if ip is None:
# displaypub is meaningless outside IPython
return
md = content['metadata'] or {}
md['engine'] = eid
ip.display_pub.publish(data=content['data'], metadata=md)
def _display_stream(self, text, prefix='', file=None):
if not text:
# nothing to display
return
if file is None:
file = sys.stdout
end = '' if text.endswith('\n') else '\n'
multiline = text.count('\n') > int(text.endswith('\n'))
if prefix and multiline and not text.startswith('\n'):
prefix = prefix + '\n'
print("%s%s" % (prefix, text), file=file, end=end)
def _display_single_result(self):
self._display_stream(self.stdout)
self._display_stream(self.stderr, file=sys.stderr)
if get_ipython() is None:
# displaypub is meaningless outside IPython
return
for output in self.outputs:
self._republish_displaypub(output, self.engine_id)
if self.execute_result is not None:
display(self.get())
@check_ready
def display_outputs(self, groupby="type"):
"""republish the outputs of the computation
Parameters
----------
groupby : str [default: type]
if 'type':
Group outputs by type (show all stdout, then all stderr, etc.):
[stdout:1] foo
[stdout:2] foo
[stderr:1] bar
[stderr:2] bar
if 'engine':
Display outputs for each engine before moving on to the next:
[stdout:1] foo
[stderr:1] bar
[stdout:2] foo
[stderr:2] bar
if 'order':
Like 'type', but further collate individual displaypub
outputs. This is meant for cases of each command producing
several plots, and you would like to see all of the first
plots together, then all of the second plots, and so on.
"""
self.wait_for_output()
if self._single_result:
self._display_single_result()
return
stdouts = self.stdout
stderrs = self.stderr
execute_results = self.execute_result
output_lists = self.outputs
results = self.get()
targets = self.engine_id
if groupby == "engine":
for eid,stdout,stderr,outputs,r,execute_result in zip(
targets, stdouts, stderrs, output_lists, results, execute_results
):
self._display_stream(stdout, '[stdout:%i] ' % eid)
self._display_stream(stderr, '[stderr:%i] ' % eid, file=sys.stderr)
if get_ipython() is None:
# displaypub is meaningless outside IPython
continue
if outputs or execute_result is not None:
_raw_text('[output:%i]' % eid)
for output in outputs:
self._republish_displaypub(output, eid)
if execute_result is not None:
display(r)
elif groupby in ('type', 'order'):
# republish stdout:
for eid,stdout in zip(targets, stdouts):
self._display_stream(stdout, '[stdout:%i] ' % eid)
# republish stderr:
for eid,stderr in zip(targets, stderrs):
self._display_stream(stderr, '[stderr:%i] ' % eid, file=sys.stderr)
if get_ipython() is None:
# displaypub is meaningless outside IPython
return
if groupby == 'order':
output_dict = dict((eid, outputs) for eid,outputs in zip(targets, output_lists))
N = max(len(outputs) for outputs in output_lists)
for i in range(N):
for eid in targets:
outputs = output_dict[eid]
if len(outputs) >= N:
_raw_text('[output:%i]' % eid)
self._republish_displaypub(outputs[i], eid)
else:
# republish displaypub output
for eid,outputs in zip(targets, output_lists):
if outputs:
_raw_text('[output:%i]' % eid)
for output in outputs:
self._republish_displaypub(output, eid)
# finally, add execute_result:
for eid,r,execute_result in zip(targets, results, execute_results):
if execute_result is not None:
display(r)
else:
raise ValueError("groupby must be one of 'type', 'engine', 'collate', not %r" % groupby)
class OpenGazeTracker:
def __init__(self, ip='127.0.0.1', port=4242, logfile='default.tsv', \
debug=False):
"""The OpenGazeConnection class communicates to the GazePoint
server through a TCP/IP socket. Incoming samples will be written
to a log at the specified path.
Keyword Arguments
ip - The IP address of the computer that is running the
OpenGaze server. This will usually be the localhost at
127.0.0.1. Type: str. Default = '127.0.0.1'
port - The port number that the OpenGaze server is on; usually
this will be 4242. Type: int. Default = 4242
logfile - The path to the intended log file, including a
file extension ('.tsv'). Type: str. Default =
'default.tsv'
debug - Boolean that determines whether DEBUG mode should be
active (True) or not (False). In DEBUG mode, all sent
and received messages are logged to a file. Type: bool.
Default = False
"""
# DEBUG
self._debug = debug
# Open a new debug file.
if self._debug:
dt = time.strftime("%Y-%m-%d_%H-%M-%S")
self._debuglog = open('debug_%s.txt' % (dt), 'w')
self._debuglog.write("OPENGAZE PYTHON DEBUG LOG %s\n" % (dt))
self._debugcounter = 0
self._debugconsolidatefreq = 1
# CONNECTION
# Save the ip and port numbers.
self.host = ip
self.port = port
# Start a new TCP/IP socket. It is curcial that it has a timeout,
# as timeout exceptions will be handled gracefully, and are in fact
# necessary to prevent the incoming Thread from freezing.
self._debug_print("Connecting to %s (%s)..." % (self.host, self.port))
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.connect((self.host, self.port))
self._sock.settimeout(1.0)
self._debug_print("Successfully connected!")
self._maxrecvsize = 4096
# Create a socket Lock to prevent simultaneous access.
self._socklock = Lock()
# Create an event that should remain set until the connection is
# closed. (This is what keeps the Threads running.)
self._connected = Event()
self._connected.set()
# LOGGING
self._debug_print("Opening new logfile '%s'" % (logfile))
# Open a new log file.
self._logfile = open(logfile, 'w')
# Write the header to the log file.
self._logheader = ['CNT', 'TIME', 'TIME_TICK', \
'FPOGX', 'FPOGY', 'FPOGS', 'FPOGD', 'FPOGID', 'FPOGV', \
'LPOGX', 'LPOGY', 'LPOGV', \
'RPOGX', 'RPOGY', 'RPOGV', \
'BPOGX', 'BPOGY', 'BPOGV', \
'LPCX', 'LPCY', 'LPD', 'LPS', 'LPV', \
'RPCX', 'RPCY', 'RPD', 'RPS', 'RPV', \
'LEYEX', 'LEYEY', 'LEYEZ', 'LPUPILD', 'LPUPILV', \
'REYEX', 'REYEY', 'REYEZ', 'RPUPILD', 'RPUPILV', \
'CX', 'CY', 'CS', \
'USER']
self._n_logvars = len(self._logheader)
self._logfile.write('\t'.join(self._logheader) + '\n')
# The log is consolidated (written to the disk) every N samples.
# This requires an internal counter (because we can't be sure the
# user turned on the 'CNT' sample counter), and a property that
# determines the consolidation frequency. This frequency can also
# be set to None, to never consolidate automatically.
self._logcounter = 0
self._log_consolidation_freq = 60
# Start a Queue for samples that need to be logged.
self._logqueue = Queue()
# Set an event that is set while samples should be logged, and
# unset while they shouldn't.
self._logging = Event()
self._logging.set()
# Set an event that signals is set when the logfile is ready to
# be closed.
self._log_ready_for_closing = Event()
self._log_ready_for_closing.clear()
# Start a Thread that writes queued samples to the log file.
self._logthread = Thread( \
target=self._process_logging,
name='PyGaze_OpenGazeConnection_logging', \
args=[])
# INCOMING
# Start a new dict for the latest incoming messages, and for
# incoming acknowledgements.
self._incoming = {}
self._acknowledgements = {}
# Create a Lock for the incoming message and acknowledgement dicts.
self._inlock = Lock()
self._acklock = Lock()
# Create an empty string for the current unfinished message. This
# is to prevent half a message being parsed when it is cut off
# between two 'self._sock.recv' calls.
self._unfinished = ''
# Start a new Thread that processes the incoming messages.
self._inthread = Thread( \
target=self._process_incoming, \
name='PyGaze_OpenGazeConnection_incoming', \
args=[])
# OUTGOING
# Start a new outgoing Queue (Thread safe, woop!).
self._outqueue = Queue()
# Set an event that is set when all queued outgoing messages have
# been processed.
self._sock_ready_for_closing = Event()
self._sock_ready_for_closing.clear()
# Create a new Thread that processes the outgoing queue.
self._outthread = Thread( \
target=self._process_outgoing, \
name='PyGaze_OpenGazeConnection_outgoing', \
args=[])
# Create a dict that will keep track of at what time which command
# was sent.
self._outlatest = {}
# Create a Lock to prevent simultaneous access to the outlatest
# dict.
self._outlock = Lock()
# RUN THREADS
# Set a signal that will kill all Threads when they receive it.
self._thread_shutdown_signal = 'KILL_ALL_HUMANS'
# Start the threads.
self._debug_print("Starting the logging thread.")
self._logthread.start()
self._debug_print("Starting the incoming thread.")
self._inthread.start()
self._debug_print("Starting the outgoing thread.")
self._outthread.start()
# SET UP LOGGING
# Wait for a bit to allow the Threads to start.
time.sleep(0.5)
# Enable the tracker to send ALL the things.
self.enable_send_counter(True)
self.enable_send_cursor(True)
self.enable_send_eye_left(True)
self.enable_send_eye_right(True)
self.enable_send_pog_best(True)
self.enable_send_pog_fix(True)
self.enable_send_pog_left(True)
self.enable_send_pog_right(True)
self.enable_send_pupil_left(True)
self.enable_send_pupil_right(True)
self.enable_send_time(True)
self.enable_send_time_tick(True)
self.enable_send_user_data(True)
# Reset the user-defined variable.
self.user_data("0")
def calibrate(self, reset=False):
"""Calibrates the eye tracker.
"""
# Reset the calibration to its default points.
if reset:
self.calibrate_reset()
# Show the calibration screen.
self.calibrate_show(True)
# Start the calibration.
self.calibrate_start(True)
# Wait for the calibration result.
result = None
while result == None:
result = self.get_calibration_result()
time.sleep(0.1)
# Hide the calibration window.
self.calibrate_show(False)
return result
def sample(self):
# If there is no current record yet, return None.
if 'REC' not in self._incoming.keys():
return None
elif 'NO_ID' not in self._incoming.keys():
return None
elif 'BPOGX' not in self._incoming['REC']['NO_ID'].keys() or \
'BPOGY' not in self._incoming['REC']['NO_ID'].keys():
return None
# Return the (x,y) coordinate.
return (float(self._incoming['REC']['NO_ID']['BPOGX']), \
float(self._incoming['REC']['NO_ID']['BPOGY']))
def pupil_size(self):
"""Return the current pupil size.
"""
# If there is no current record yet, return None.
if 'REC' not in self._incoming.keys():
return None
elif 'NO_ID' not in self._incoming.keys():
return None
elif 'LPV' not in self._incoming['REC']['NO_ID'].keys() or \
'LPS' not in self._incoming['REC']['NO_ID'].keys() or \
'RPV' not in self._incoming['REC']['NO_ID'].keys() or \
'RPS' not in self._incoming['REC']['NO_ID'].keys():
return None
# Compute the pupil size, and return it if there is valid data.
n = 0
psize = 0
if str(self._incoming['REC']['NO_ID']['LPV']) == '1':
psize += float(self._incoming['REC']['NO_ID']['LPS'])
n += 1
if str(self._incoming['REC']['NO_ID']['RPV']) == '1':
psize += float(self._incoming['REC']['NO_ID']['RPS'])
n += 1
if n == 0:
return None
else:
return psize / float(n)
def log(self, message):
"""Logs a message to the log file. ONLY CALL THIS WHILE RECORDING
DATA!
"""
# Set the user-defined value.
i = copy.copy(self._logcounter)
self.user_data(message)
# Wait until a single sample is logged.
while self._logcounter <= i:
time.sleep(0.0001)
# Reset the user-defined value.
self.user_data("0")
def start_recording(self):
"""Start writing data to the log file.
"""
self.enable_send_data(True)
def stop_recording(self):
"""Pause writing data to the log file.
"""
self.enable_send_data(False)
def _debug_print(self, msg):
if self._debug:
self._debuglog.write('%s: %s\n' % \
(datetime.datetime.now().strftime("%H:%M:%S.%f"), msg))
if self._debugcounter % self._debugconsolidatefreq == 0:
self._debuglog.flush()
os.fsync(self._debuglog.fileno())
self._debugcounter += 1
def _format_msg(self, command, ID, values=None):
# Create the start of the formatted string.
xml = '<%s ID="%s" ' % (command.upper(), ID.upper())
# Add the values for each parameter.
if values:
for par, val in values:
xml += '%s="%s" ' % (par.upper(), val)
# Add the ending.
xml += '/>\r\n'
return xml
def _log_consolidation(self):
# Internal buffer to RAM.
self._logfile.flush()
# RAM to disk.
os.fsync(self._logfile.fileno())
def _log_sample(self, sample):
# Construct an empty line that has the same length as the log's
# header (this was computed in __init__).
line = self._n_logvars * ['']
# Loop through all keys in the dict.
for varname in sample.keys():
# Check if this is a logable variable.
if varname in self._logheader:
# Find the appropriate index in the line
line[self._logheader.index(varname)] = sample[varname]
self._logfile.write('\t'.join(line) + '\n')
def _parse_msg(self, xml):
e = lxml.etree.fromstring(xml)
return (e.tag, e.attrib)
def _process_logging(self):
self._debug_print("Logging Thread started.")
while not self._log_ready_for_closing.is_set():
# Get a new sample from the Queue.
sample = self._logqueue.get()
# Check if this is the shutdown signal.
if sample == self._thread_shutdown_signal:
# Signal that we're done logging all samples.
self._log_ready_for_closing.set()
# Break the while loop.
break
# Log the sample.
self._log_sample(sample)
# Consolidate the log if necessary.
if self._logcounter % self._log_consolidation_freq == 0:
self._log_consolidation()
# Increment the counter.
self._logcounter += 1
self._debug_print("Logging Thread ended.")
return
def _process_incoming(self):
self._debug_print("Incoming Thread started.")
while self._connected.is_set():
# Lock the socket to prevent other Threads from simultaneously
# accessing it.
self._socklock.acquire()
# Get new messages from the OpenGaze Server.
timeout = False
try:
instring = self._sock.recv(self._maxrecvsize)
except socket.timeout:
timeout = True
# Get a received timestamp.
t = time.time()
# Unlock the socket again.
self._socklock.release()
# Skip further processing if no new message came in.
if timeout:
self._debug_print("socket recv timeout")
continue
self._debug_print("Raw instring: %r" % (instring))
# Split the messages (they are separated by '\r\n').
messages = instring.split('\r\n')
# Check if there is currently an unfinished message.
if self._unfinished:
# Combine the currently unfinished message and the
# most recent incoming message.
messages[0] = copy.copy(self._unfinished) + messages[0]
# Reset the unfinished message.
self._unfinished = ''
# Check if the last message was actually complete.
if not messages[-1][-2:] == '/>':
self._unfinished = messages.pop(-1)
# Run through all messages.
for msg in messages:
self._debug_print("Incoming: %r" % (msg))
# Parse the message.
command, msgdict = self._parse_msg(msg)
# Check if the incoming message is an acknowledgement.
# Acknowledgements are also stored in a different dict,
# which is used to monitor whether sent messages are
# properly received.
if command == 'ACK':
self._acklock.acquire()
self._acknowledgements[msg] = copy.copy(t)
self._acklock.release()
# Acquire the Lock for the incoming dict, so that it
# won't be accessed at the same time.
self._inlock.acquire()
# Check if this command is already in the current dict.
if command not in self._incoming.keys():
self._incoming[command] = {}
# Some messages have no ID, for example 'REC' messages.
# We simply assign 'NO_ID' as the ID.
if 'ID' not in msgdict.keys():
msgdict['ID'] = 'NO_ID'
# Check if this ID is already in the current dict.
if msgdict['ID'] not in self._incoming[command].keys():
self._incoming[command][msgdict['ID']] = {}
# Add receiving time stamp, and the values for each
# parameter to the current dict.
self._incoming[command][msgdict['ID']]['t'] = \
copy.copy(t)
for par, val in msgdict.items():
self._incoming[command][msgdict['ID']][par] = \
copy.copy(val)
# Log sample if command=='REC' and when the logging
# event is set.
if command == 'REC' and self._logging.is_set():
self._logqueue.put(copy.deepcopy( \
self._incoming[command][msgdict['ID']]))
# Unlock the incoming dict again.
self._inlock.release()
self._debug_print("Incoming Thread ended.")
return
def _process_outgoing(self):
self._debug_print("Outgoing Thread started.")
while not self._sock_ready_for_closing.is_set():
# Get a new command from the Queue.
msg = self._outqueue.get()
# Check if this is the shutdown signal.
if msg == self._thread_shutdown_signal:
# Signal that we're done processing all the outgoing
# messages.
self._sock_ready_for_closing.set()
# Break the while loop.
break
self._debug_print("Outgoing: %r" % (msg))
# Lock the socket to prevent other Threads from simultaneously
# accessing it.
self._socklock.acquire()
# Send the command to the OpenGaze Server.
t = time.time()
self._sock.send(msg)
# Unlock the socket again.
self._socklock.release()
# Store a timestamp for the latest outgoing message.
self._outlock.acquire()
self._outlatest[msg] = copy.copy(t)
self._outlock.release()
self._debug_print("Outgoing Thread ended.")
return
def _send_message(self, command, ID, values=None, \
wait_for_acknowledgement=True, resend_timeout=1.0, maxwait=5.0):
# Format a message in an XML format that the Open Gaze API needs.
msg = self._format_msg(command, ID, values=values)
# Construct the expected acknowledgement.
ack = msg.replace(command, 'ACK', 1).replace('\r\n', '')
# Run until the message is acknowledged or a timeout occurs (or
# break if we're not supposed to wait for an acknowledgement.)
timeout = False
acknowledged = False
t0 = time.time()
while (not acknowledged) and (not timeout):
# Add the command to the outgoing Queue.
self._debug_print("Outqueue add: %r" % (msg))
self._outqueue.put(msg)
# Wait until an acknowledgement comes in.
if wait_for_acknowledgement:
sent = False
t1 = time.time()
while (time.time() - t1 < resend_timeout) and \
(not acknowledged):
# Check the outgoing queue for the sent message to
# appear.
if not sent:
self._outlock.acquire()
if msg in self._outlatest.keys():
t = copy.copy(self._outlatest[msg])
sent = True
self._debug_print("Outqueue sent: %r" \
% (msg))
self._outlock.release()
time.sleep(0.001)
# Check the incoming queue for the expected
# acknowledgement. (NOTE: This does not check
# whether the values of the incoming acknowlement
# match the sent message. Ideally, they should.)
else:
self._acklock.acquire()
if ack in self._acknowledgements.keys():
if self._acknowledgements[ack] >= t:
acknowledged = True
self._debug_print("Outqueue acknowledged: %r" \
% (msg))
self._acklock.release()
time.sleep(0.001)
# Check if there is a timeout.
if (not acknowledged) and \
(time.time() - t0 > maxwait):
timeout = True
break
# If we're not supposed to wait for an acknowledgement, break
# the while loop.
else:
break
return acknowledged, timeout
def close(self):
"""Closes the connection to the tracker, closes the log files, and
ends the Threads that process the incoming and outgoing messages,
and the logging of samples.
"""
# Reset the user-defined value.
self.user_data('0')
# Unset the self._connected event to stop the incoming Thread.
self._debug_print("Unsetting the connection event")
self._connected.clear()
# Queue the stop signal to stop the outgoing and logging Threads.
self._debug_print("Adding stop signal to outgoing Queue")
self._outqueue.put(self._thread_shutdown_signal)
self._debug_print("Adding stop signal to logging Queue")
self._logqueue.put(self._thread_shutdown_signal)
# Wait for the outgoing Queue to be fully processed.
self._debug_print("Waiting for the socket to close...")
self._sock_ready_for_closing.wait()
# Close the socket connection to the OpenGaze server.
self._debug_print("Closing socket connection...")
self._sock.close()
self._debug_print("Socket connection closed!")
# Wait for the log Queue to be fully processed.
self._debug_print("Waiting for the log to close...")
self._log_ready_for_closing.wait()
# Close the log file.
self._logfile.close()
self._debug_print("Log closed!")
# Join the Threads.
self._debug_print("Waiting for the Threads to join...")
self._outthread.join()
self._debug_print("Outgoing Thread joined!")
self._inthread.join()
self._debug_print("Incoming Thread joined!")
self._logthread.join()
self._debug_print("Logging Thread joined!")
# Close the DEBUG log.
if self._debug:
self._debuglog.write("END OF DEBUG LOG")
self._debuglog.close()
def enable_send_data(self, state):
"""Start (state=True) or stop (state=False) the streaming of data
from the server to the client.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_DATA', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_counter(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the send counter in the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_COUNTER', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_time(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the send time in the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_TIME', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_time_tick(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the send time tick in the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_TIME_TICK', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pog_fix(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the point of gaze as determined by the tracker's fixation filter in
the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_POG_FIX', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pog_left(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the point of gaze of the left eye in the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_POG_LEFT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pog_right(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the point of gaze of the right eye in the data record string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_POG_RIGHT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pog_best(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
the 'best' point of gaze in the data record string. This is based
on the average of the left and right POG if both eyes are available,
or on the value of the one available eye.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_POG_BEST', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pupil_left(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
pupil data on the left eye in the data record string. This data
consists of the following:
LPCX: The horizontal coordinate of the left eye pupil in the camera
image, as a fraction of the camera size.
LPCY: The vertical coordinate of the left eye pupil in the camera
image, as a fraction of the camera size.
LPD: The left eye pupil's diameter in pixels.
LPS: The scale factor of the left eye pupil (unitless). Value
equals 1 at calibration depth, is less than 1 when the user
is closer to the eye tracker and greater than 1 when the user
is further away.
LPV: The valid flag with a value of 1 if the data is valid, and 0
if it is not.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_PUPIL_LEFT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_pupil_right(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
pupil data on the right eye in the data record string. This data
consists of the following:
RPCX: The horizontal coordinate of the right eye pupil in the camera
image, as a fraction of the camera size.
RPCY: The vertical coordinate of the right eye pupil in the camera
image, as a fraction of the camera size.
RPD: The right eye pupil's diameter in pixels.
RPS: The scale factor of the right eye pupil (unitless). Value
equals 1 at calibration depth, is less than 1 when the user
is closer to the eye tracker and greater than 1 when the user
is further away.
RPV: The valid flag with a value of 1 if the data is valid, and 0
if it is not.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_PUPIL_RIGHT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_eye_left(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
3D data on left eye in the data record string. This data consists
of the following:
LEYEX: The horizontal coordinate of the left eye in 3D space with
respect to the camera focal point, in meters.
LEYEY: The vertical coordinate of the left eye in 3D space with
respect to the camera focal point, in meters.
LEYEZ: The depth coordinate of the left eye in 3D space with
respect to the camera focal point, in meters.
LPUPILD: The diameter of the left eye pupil in meters.
LPUPILV: The valid flag with a value of 1 if the data is valid, and
0 if it is not.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_EYE_LEFT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_eye_right(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
3D data on right eye in the data record string. This data consists
of the following:
REYEX: The horizontal coordinate of the right eye in 3D space with
respect to the camera focal point, in meters.
REYEY: The vertical coordinate of the right eye in 3D space with
respect to the camera focal point, in meters.
REYEZ: The depth coordinate of the right eye in 3D space with
respect to the camera focal point, in meters.
RPUPILD: The diameter of the right eye pupil in meters.
RPUPILV: The valid flag with a value of 1 if the data is valid, and
0 if it is not.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_EYE_RIGHT', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_cursor(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
data on the mouse cursor in the data record string. This data
consists of the following:
CX: The horizontal coordinate of the mouse cursor, as a percentage
of the screen resolution.
CY: The vertical coordinate of the mouse cursor, as a percentage
of the screen resolution.
CS: The mouse cursor state, 0 for steady state, 1 for left button
down, 2 for rigght button down.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_CURSOR', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def enable_send_user_data(self, state):
"""Enable (state=True) or disable (state=False) the inclusion of
user-defined variables in the data record string. User-defined
variables can be set with the 'user_data' method.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'ENABLE_SEND_USER_DATA', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_start(self, state):
"""Starts (state=1) or stops (state=0) the calibration procedure.
Make sure to call the 'calibrate_show' function beforehand, or to
implement your own calibration visualisation; otherwise a call to
this function will make the calibration run in the background.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_START', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_show(self, state):
"""Shows (state=1) or hides (state=0) the calibration window on the
tracker's display window. While showing the calibration window, you
can call 'calibrate_start' to run the calibration procedure.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_SHOW', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_timeout(self, value):
"""Set the duration of the calibration point (not including the
animation time) in seconds. The value can be an int or a float.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_TIMEOUT', \
values=[('VALUE', float(value))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_delay(self, value):
"""Set the duration of the calibration animation (before
calibration at a point begins) in seconds. The value can be an int
or a float.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_DELAY', \
values=[('VALUE', float(value))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_result_summary(self):
"""Returns a summary of the calibration results, which consists of
the following values:
AVE_ERROR: Average error over all calibrated points.
VALID_POINTS: Number of successfully calibrated points.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'CALIBRATE_RESULT_SUMMARY', \
values=None, \
wait_for_acknowledgement=True)
# Return the results.
ave_error = None
valid_points = None
if acknowledged:
self._inlock.acquire()
ave_error = copy.copy( \
self._incoming['ACK']['CALIBRATE_RESULT_SUMMARY']['AVE_ERROR'])
valid_points = copy.copy( \
self._incoming['ACK']['CALIBRATE_RESULT_SUMMARY']['VALID_POINTS'])
self._inlock.release()
return ave_error, valid_points
def calibrate_clear(self):
"""Clear the internal list of calibration points.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_CLEAR', \
values=None, \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_reset(self):
"""Reset the internal list of calibration points to the default
values.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_RESET', \
values=None, \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def calibrate_addpoint(self, x, y):
"""Add a calibration point at the passed horizontal (x) and
vertical (y) coordinates. These coordinates should be as a
proportion of the screen resolution, where (0,0) is the top-left,
(0.5,0.5) is the screen centre, and (1,1) is the bottom-right.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'CALIBRATE_ADDPOINT', \
values=[('X', x), ('Y', y)], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def get_calibration_points(self):
"""Returns a list of the current calibration points.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'CALIBRATE_ADDPOINT', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
points = None
if acknowledged:
points = []
self._inlock.acquire()
for i in range(self._incoming['ACK']['CALIBRATE_ADDPOINT']['PTS']):
points.append( \
copy.copy( \
self._incoming['ACK']['CALIBRATE_ADDPOINT']['X%d' % i+1]), \
copy.copy( \
self._incoming['ACK']['CALIBRATE_ADDPOINT']['Y%d' % i+1]) \
)
self._inlock.release()
return points
def clear_calibration_result(self):
"""Clears the internally stored calibration result.
"""
# Clear the calibration results.
self._inlock.acquire()
if 'CAL' in self._incoming.keys():
if 'CALIB_RESULT' in self._incoming['CAL'].keys():
self._incoming['CAL'].pop('CALIB_RESULT')
self._inlock.release()
def get_calibration_result(self):
"""Returns the latest available calibration results as a list of
dicts, each with the following keys:
CALX: Calibration point's horizontal coordinate.
CALY: Calibration point's vertical coordinate
LX: Left eye's recorded horizontal point of gaze.
LY: Left eye's recorded vertical point of gaze.
LV: Left eye's validity status (1=valid, 0=invalid)
RX: Right eye's recorded horizontal point of gaze.
RY: Right eye's recorded vertical point of gaze.
RV: Right eye's validity status (1=valid, 0=invalid)
Returns None if no calibration results are available.
"""
# Parameters of the 'CALIB_RESULT' dict.
params = ['CALX', 'CALY', 'LX', 'LY', 'LV', 'RX', 'RY', 'RV']
# Return the result.
points = None
self._inlock.acquire()
if 'CAL' in self._incoming.keys():
if 'CALIB_RESULT' in self._incoming['CAL'].keys():
# Get the latest calibration results.
cal = copy.deepcopy(self._incoming['CAL']['CALIB_RESULT'])
# Compute the number of fixation points by dividing the
# total number of parameters in the 'CALIB_RESULT' dict
# by 8 (the number of parameters per point). Note that
# the 'CALIB_RESULT' dict also has an 'ID' parameter,
# which we should account for by subtracting 1 from the
# length of the list of keys in the dict.
n_points = (len(cal.keys()) - 1) // len(params)
# Put the results in a different format.
points = []
for i in range(1, n_points + 1):
p = {}
for par in params:
p['%s' % (par, i)] = cal['%s%d' % (par, i)]
points.append(copy.deepcopy(p))
self._inlock.release()
return points
def user_data(self, value):
"""Set the value of the user data field for embedding custom data
into the data stream. The user data value should be a string.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'USER_DATA', \
values=[('VALUE', str(value))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def tracker_display(self, state):
"""Shows (state=1) or hides (state=0) the eye tracker display
window.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'TRACKER_DISPLAY', \
values=[('STATE', int(state))], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def time_tick_frequency(self):
"""Returns the time-tick frequency to convert the TIME_TICK
variable to seconds.
"""
return self.get_time_tick_frequency()
def get_time_tick_frequency(self):
"""Returns the time-tick frequency to convert the TIME_TICK
variable to seconds.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'TIME_TICK_FREQUENCY', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
freq = None
if acknowledged:
self._inlock.acquire()
freq = copy.copy(
self._incoming['ACK']['TIME_TICK_FREQUENCY']['FREQ'])
self._inlock.release()
return freq
def screen_size(self, x, y, w, h):
"""Set the gaze tracking screen position (x,y) and size (w, h). You
can use this to work with multi-monitor systems. All values are in
pixels.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('SET', \
'SCREEN_SIZE', \
values=[('X', x), ('Y', x), ('WIDTH', w), ('HEIGHT', h)], \
wait_for_acknowledgement=True)
# Return a success Boolean.
return acknowledged and (timeout == False)
def get_screen_size(self):
"""Returns the x and y coordinates of the top-left of the screen in
pixels, as well as the screen width and height in pixels. The
result is returned as [x, y, w, h].
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'SCREEN_SIZE', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
x = None
y = None
w = None
h = None
if acknowledged:
self._inlock.acquire()
x = copy.copy(self._incoming['ACK']['SCREEN_SIZE']['X'])
y = copy.copy(self._incoming['ACK']['SCREEN_SIZE']['Y'])
w = copy.copy(self._incoming['ACK']['SCREEN_SIZE']['WIDTH'])
h = copy.copy(self._incoming['ACK']['SCREEN_SIZE']['HEIGHT'])
self._inlock.release()
return [x, y, w, h]
def camera_size(self):
"""Returns the size of the camera sensor in pixels, as [w,h].
"""
return self.get_camera_size()
def get_camera_size(self):
"""Returns the size of the camera sensor in pixels, as [w,h].
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'CAMERA_SIZE', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
w = None
h = None
if acknowledged:
self._inlock.acquire()
w = copy.copy(self._incoming['ACK']['CAMERA_SIZE']['WIDTH'])
h = copy.copy(self._incoming['ACK']['CAMERA_SIZE']['HEIGHT'])
self._inlock.release()
return [w, h]
def product_id(self):
"""Returns the identifier of the connected eye-tracker.
"""
return self.get_product_id()
def get_product_id(self):
"""Returns the identifier of the connected eye-tracker.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'PRODUCT_ID', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
value = None
if acknowledged:
self._inlock.acquire()
value = copy.copy(self._incoming['ACK']['PRODUCT_ID']['VALUE'])
self._inlock.release()
return value
def serial_id(self):
"""Returns the serial number of the connected eye-tracker.
"""
return self.get_serial_id()
def get_serial_id(self):
"""Returns the serial number of the connected eye-tracker.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'SERIAL_ID', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
value = None
if acknowledged:
self._inlock.acquire()
value = copy.copy(self._incoming['ACK']['SERIAL_ID']['VALUE'])
self._inlock.release()
return value
def company_id(self):
"""Returns the identifier of the manufacturer of the connected
eye-tracker.
"""
return self.get_company_id()
def get_company_id(self):
"""Returns the identifier of the manufacturer of the connected
eye-tracker.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'COMPANY_ID', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
value = None
if acknowledged:
self._inlock.acquire()
value = copy.copy(self._incoming['ACK']['COMPANY_ID']['VALUE'])
self._inlock.release()
return value
def api_id(self):
"""Returns the API version number.
"""
return self.get_api_id()
def get_api_id(self):
"""Returns the API version number.
"""
# Send the message (returns after the Server acknowledges receipt).
acknowledged, timeout = self._send_message('GET', \
'API_ID', \
values=None, \
wait_for_acknowledgement=True)
# Return the result.
value = None
if acknowledged:
self._inlock.acquire()
value = copy.copy(self._incoming['ACK']['API_ID']['VALUE'])
self._inlock.release()
return value
class QueueProcessor(Thread):
"""Data Processing Thread
It handles and sorts all API messages relating to
subscription / subscription cancelling, and sorts all data messages into
queues, organized by data type (book, ticker, etc) and pairs
( BTCUSD, ETHBTC, etc).
"""
def __init__(self, data_q, log_level=None, *args, **kwargs):
"""Initialze a QueueProcessor instance.
:param data_q: Queue()
:param log_level: logging level
:param args: Thread *args
:param kwargs: Thread **kwargs
"""
super(QueueProcessor, self).__init__(*args, **kwargs)
self.q = data_q
self._response_handlers = {
'unsubscribed': self._handle_unsubscribed,
'subscribed': self._handle_subscribed,
'conf': self._handle_conf,
'auth': self._handle_auth,
'unauth': self._handle_auth
}
self._data_handlers = {
'ticker': self._handle_ticker,
'book': self._handle_book,
'raw_book': self._handle_raw_book,
'candles': self._handle_candles,
'trades': self._handle_trades
}
# Assigns a channel id to a data handler method.
self._registry = {}
# dict to translate channel ids to channel identifiers and vice versa
self.channel_directory = {}
# dict to register a method with a channel id
self.channel_handlers = {}
# Keeps track of last update to a channel by id.
self.last_update = {}
self.tickers = defaultdict(Queue)
self.books = defaultdict(Queue)
self.raw_books = defaultdict(Queue)
self.trades = defaultdict(Queue)
self.candles = defaultdict(Queue)
self.account = defaultdict(Queue)
# Sentinel Event to kill the thread
self._stopped = Event()
# Internal Logging facilities
self.log = logging.getLogger(self.__module__)
self.log.setLevel(level=logging.INFO if not log_level else log_level)
# Translation dict for Account channels
self.account_channel_names = {
'os': 'Orders',
'ps': 'Positions',
'hos': 'Historical Orders',
'hts': 'Trades',
'fls': 'Loans',
'te': 'Trades',
'tu': 'Trades',
'ws': 'Wallets',
'bu': 'Balance Info',
'miu': 'Margin Info',
'fos': 'Offers',
'fiu': 'Funding Info',
'fcs': 'Credits',
'hfos': 'Historical Offers',
'hfcs': 'Historical Credits',
'hfls': 'Historical Loans',
'htfs': 'Funding Trades',
'n': 'Notifications'
}
def join(self, timeout=None):
"""Set sentinel for run() method and join thread.
:param timeout:
:return:
"""
self._stopped.set()
super(QueueProcessor, self).join(timeout=timeout)
def run(self):
"""Main routine.
:return:
"""
while not self._stopped.is_set():
try:
message = self.q.get(timeout=0.1)
except Empty:
continue
dtype, data, ts = message
if dtype in ('subscribed', 'unsubscribed', 'conf', 'auth',
'unauth'):
try:
self._response_handlers[dtype](dtype, data, ts)
except KeyError:
self.log.error(
"Dtype '%s' does not have a response "
"handler!", dtype)
elif dtype == 'data':
try:
channel_id = data[0]
# Get channel type associated with this data to the
# associated data type (from 'data' to
# 'book', 'ticker' or similar
channel_type, *_ = self.channel_directory[channel_id]
# Run the associated data handler for this channel type.
self._data_handlers[channel_type](channel_type, data, ts)
# Update time stamps.
self.update_timestamps(channel_id, ts)
except KeyError:
self.log.error(
"Channel ID does not have a data handler! %s", message)
else:
self.log.error("Unknown dtype on queue! %s", message)
continue
def _handle_subscribed(
self,
dtype,
data,
ts,
):
"""Handles responses to subscribe() commands
Registers a channel id with the client and assigns a data handler to it.
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_subscribed: %s - %s - %s", dtype, data, ts)
channel_name = data.pop('channel')
channel_id = data.pop('chanId')
config = data
if 'pair' in config:
symbol = config['pair']
elif 'symbol' in config:
symbol = config['symbol']
elif 'key' in config:
symbol = config['key'].split(':')[2][
1:] #layout type:interval:tPair
else:
symbol = None
if 'prec' in config and config['prec'].startswith('R'):
channel_name = 'raw_' + channel_name
self.channel_handlers[channel_id] = self._data_handlers[channel_name]
# Create a channel_name, symbol tuple to identify channels of same type
if 'key' in config:
identifier = (channel_name, symbol, config['key'].split(':')[1])
else:
identifier = (channel_name, symbol)
self.channel_handlers[channel_id] = identifier
self.channel_directory[identifier] = channel_id
self.channel_directory[channel_id] = identifier
self.log.info("Subscription succesful for channel %s", identifier)
def _handle_unsubscribed(self, dtype, data, ts):
"""Handles responses to unsubscribe() commands
Removes a channel id from the client.
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_unsubscribed: %s - %s - %s", dtype, data, ts)
channel_id = data.pop('chanId')
# Unregister the channel from all internal attributes
chan_identifier = self.channel_directory.pop(channel_id)
self.channel_directory.pop(chan_identifier)
self.channel_handlers.pop(channel_id)
self.last_update.pop(channel_id)
self.log.info("Successfully unsubscribed from %s", chan_identifier)
def _handle_auth(self, dtype, data, ts):
"""Handles authentication responses
:param dtype:
:param data:
:param ts:
:return:
"""
# Contains keys status, chanId, userId, caps
if dtype == 'unauth':
raise NotImplementedError
channel_id = data.pop('chanId')
user_id = data.pop('userId')
identifier = ('auth', user_id)
self.channel_handlers[identifier] = channel_id
self.channel_directory[identifier] = channel_id
self.channel_directory[channel_id] = identifier
def _handle_conf(self, dtype, data, ts):
"""Handles configuration messages.
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_conf: %s - %s - %s", dtype, data, ts)
self.log.info("Configuration accepted: %s", dtype)
return
def update_timestamps(self, chan_id, ts):
"""Updates the timestamp for the given channel id.
:param chan_id:
:param ts:
:return:
"""
try:
self.last_update[chan_id] = ts
except KeyError:
self.log.warning(
"Attempted ts update of channel %s, but channel "
"not present anymore.", self.channel_directory[chan_id])
def _handle_account(self, dtype, data, ts):
""" Handles Account related data.
translation table for channel names:
Data Channels
os - Orders
hos - Historical Orders
ps - Positions
hts - Trades (snapshot)
te - Trade Event
tu - Trade Update
ws - Wallets
bu - Balance Info
miu - Margin Info
fiu - Funding Info
fos - Offers
hfos - Historical Offers
fcs - Credits
hfcs - Historical Credits
fls - Loans
hfls - Historical Loans
htfs - Funding Trades
n - Notifications (WIP)
:param dtype:
:param data:
:param ts:
:return:
"""
chan_id, *data = data
channel_identifier = self.account_channel_names[data[0]]
entry = (data, ts)
self.account[channel_identifier].put(entry)
def _handle_ticker(self, dtype, data, ts):
"""Adds received ticker data to self.tickers dict, filed under its channel
id.
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_ticker: %s - %s - %s", dtype, data, ts)
channel_id, *data = data
channel_identifier = self.channel_directory[channel_id]
entry = (data, ts)
self.tickers[channel_identifier].put(entry)
def _handle_book(self, dtype, data, ts):
"""Updates the order book stored in self.books[chan_id]
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_book: %s - %s - %s", dtype, data, ts)
channel_id, *data = data
log.debug("ts: %s\tchan_id: %s\tdata: %s", ts, channel_id, data)
channel_identifier = self.channel_directory[channel_id]
entry = (data, ts)
self.books[channel_identifier].put(entry)
def _handle_raw_book(self, dtype, data, ts):
"""Updates the raw order books stored in self.raw_books[chan_id]
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_raw_book: %s - %s - %s", dtype, data, ts)
channel_id, *data = data
channel_identifier = self.channel_directory[channel_id]
entry = (data, ts)
self.raw_books[channel_identifier].put(entry)
def _handle_trades(self, dtype, data, ts):
"""Files trades in self._trades[chan_id]
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_trades: %s - %s - %s", dtype, data, ts)
channel_id, *data = data
channel_identifier = self.channel_directory[channel_id]
entry = (data, ts)
self.trades[channel_identifier].put(entry)
def _handle_candles(self, dtype, data, ts):
"""Stores OHLC data received via wss in self.candles[chan_id]
:param dtype:
:param data:
:param ts:
:return:
"""
self.log.debug("_handle_candles: %s - %s - %s", dtype, data, ts)
channel_id, *data = data
channel_identifier = self.channel_directory[channel_id]
entry = (data, ts)
self.candles[channel_identifier].put(entry)
class SerialReader(Thread, Observable):
"""
This class is responsible for reading from the serial port and update log lines to its
registered observers. Thread quits if stop() is called. If an exception is raised when
reading from serial due to I/O issues or escape characters received, this thread will
callback to its owner to stop operation. Each log line is timestamped as default.
"""
def __init__(self, serial, callback, do_timestamp=True):
"""
:param serial: A Serial object for communicating with a serial port. It needs to have
a read timeout set. Otherwise, this class object might hang forever if
a end line character is never received.
http://pyserial.readthedocs.io/en/latest/shortintro.html#readline
:type Serial
:param callback: A callback method for calling back to owner when error occurs.
:param do_timestamp: Add a timestamp to each line intercepted from the serial port.
"""
Thread.__init__(self, name=self.__class__.__name__)
Observable.__init__(self)
self.setDaemon(True)
self._stop = Event()
self._do_timestamp = do_timestamp
self._port = serial
self.logger = logging.getLogger(self.__class__.__name__)
self._start_time = None # Is set when first log line arrives from serial port.
self._callback = callback
codecs.register_error('backslashreplace', self.backslash_replace)
@staticmethod
def backslash_replace(error):
"""
An error handler to be called when escape characters are read from the log line queue input.
"""
return u"".join([
u"\\x{:x}".format(ord(error.object[i]))
for i in range(error.start, error.end)
]), error.end
def __repr__(self):
return '{}({!r}, {!r}, {!r}, {!r}, {!r})'.format(
self.__class__.__name__, self.getName(), self.is_alive(),
self._do_timestamp, self._port, self._start_time)
def time_stamp(self, line):
"""
Returns the line with a timestamp suitable for a log file.
:param line: A read line from serial port without timestamp.
:return: timestamp + line
"""
time_delta = datetime.now() - self._start_time
return '(' + ':'.join(str(time_delta).split(':')[1:]) + ') ' + line
def stop(self):
"""
Stop reading from the serial port and commit suicide.
"""
self._stop.set()
self.logger.debug('stop reading from serial port')
if self.is_alive():
self.join()
self.logger.info('reader has terminated')
def run(self):
try:
first_line_received = True
i = 0
self.logger.info('Start reading from serial port.')
while not self._stop.is_set():
# we loop for every line and if no endline is found, then read timeout will occur.
line = self._port.readline().decode(
'ascii', 'backslashreplace').strip()
sleep(0.1) # let in other threads
if first_line_received:
self._start_time = datetime.now()
first_line_received = False
if line:
self.logger.debug('{}: {}'.format(i, line))
if self._do_timestamp:
line = self.time_stamp(line)
self.notify(line) # update listeners
i += 1
except Exception as e: # this may occur if readline() fails handling an escape character
self.logger.error('Error: {}'.format(e))
self._callback('{} has stopped running. error: {}'.format(
self.getName(), e))
self.logger.info('stopped reading from serial port.')
class Rak811Serial(object):
"""Handles serial communication between the RPi and the RAK811 module."""
def __init__(self,
port=PORT,
baudrate=BAUDRATE,
timeout=TIMEOUT,
read_buffer_timeout=READ_BUFFER_TIMEOUT,
keep_untagged=False,
**kwargs):
"""Initialise class.
The serial port is immediately opened and flushed.
All parameters are optional and passed to Serial.
"""
self._read_buffer_timeout = read_buffer_timeout
self._serial = Serial(port=port,
baudrate=baudrate,
timeout=timeout,
**kwargs)
self._serial.reset_input_buffer()
self._keep_untagged = keep_untagged
# Mutex
self._cv_serial = Condition()
self._read_buffer = []
# Read thread
self._read_done = Event()
self._read_thread = Thread(target=self._read_loop, daemon=True)
self._read_thread.start()
self._alive = True
logger.debug('Serial initialized')
def close(self):
"""Release resources."""
if self._alive:
self._read_done.set()
self._read_thread.join()
self._serial.close()
self._alive = False
def _read_loop(self):
"""Read thread.
Continuously read serial. When data is available we want to read all of
it and notify once:
- We need to drain the input after a response. If we notify()
too early the module will miss next command
- We want to catch all events at the same time
"""
while not self._read_done.is_set():
line = self._serial.readline()
if line != b'':
# Not a timeout; process data stream
with self._cv_serial:
while True:
try:
line = line.decode('ascii').rstrip(EOL)
except UnicodeDecodeError:
# Wrong speed or port not configured properly
line = '?'
if match(r'^(OK|ERROR|at+)', line):
logger.debug(f'Received: >{line}<')
self._read_buffer.append(line)
elif self._keep_untagged:
logger.debug(f'Received untagged: >{line}<')
self._read_buffer.append(line)
else:
logger.debug(f'Ignoring untagged: >{line}<')
sleep(0.1)
if self._serial.in_waiting > 0:
line = self._serial.readline()
else:
break
if len(self._read_buffer) > 0:
self._cv_serial.notify()
def receive(self,
single: bool = True,
timeout: int = None) -> Union[str, List[str]]:
"""Receive data from module.
This is a blocking call: it will data or raise Rak811TimeoutError if
nothing is received in time.
Args:
single (optional): Return single line of data when true, otherwise
all available lines are returned. Defaults to True.
timeout (optional): Time to wait for. Defaults to None.
Raises:
Rak811TimeoutError: No data received in time.
Returns:
Single line of data or list of lines.
"""
if timeout is None:
timeout = self._read_buffer_timeout
with self._cv_serial:
while len(self._read_buffer) == 0:
success = self._cv_serial.wait(timeout)
if not success:
raise Rak811TimeoutError('Timeout while waiting for data')
if single:
response = self._read_buffer.pop(0)
else:
response = self._read_buffer
self._read_buffer = []
return response
def send_string(self, string):
"""Send string to the module."""
logger.debug(
f"Sending: >{string.encode('unicode_escape').decode('utf-8')}<")
self._serial.write((bytes)(string, 'utf-8'))
def send_command(self, command):
"""Send AT command to the module."""
self.send_string('at+{0}\r\n'.format(command))
class NitrogenFlower(KerrDevice):
delay = Int(10)
timeout = Int(10)
controller = Any
channel = Int
_ready_signal = None
_timeout_timer = None
_delay_timer = None
_lock = None
_cancel_signal = None
# _cancel = False
flow_button = Event
flow_label = Property(depends_on='_flow_state')
_flow_state = Enum('off', 'purge', 'on')
led = Instance(LED, ())
message = String
def _flow_button_fired(self):
# print self._flowing, 'asdfasdfsa'
if self._flow_state in ('on', 'purge'):
self.stop()
elif self._flow_state == 'off':
self.start()
def _get_flow_label(self):
return FLOW_STATES[self._flow_state]
def stop(self):
if self._delay_timer:
self._delay_timer.cancel()
if self._timeout_timer:
self._timeout_timer.cancel()
self._ready_signal.clear()
self._stop_flow()
self._set_state_off()
def _set_state_off(self):
self._flow_state = 'off'
self.led.state = 0
self.message = ' '
def _set_state_on(self):
self._flow_state = 'on'
self.led.state = 2
t = datetime.now()
d = timedelta(seconds=self.timeout)
t = t + d
st = t.strftime('%H:%M:%S')
self.message = 'Timeout at {}'.format(st)
def _set_state_purge(self):
self._flow_state = 'purge'
self.led.state = 1
self.message = 'Purging for 10s'
def start(self):
if self._flow_state == 'on':
# reset the timeout timer
self._start_timeout_timer()
elif self._flow_state == 'purge':
pass
else:
# start purge
self._start_flow()
self._set_state_purge()
# self._flow_state = 'purge'
# self.led.state = 1
self._start_delay_timer()
# if self._ready_signal is None:
# self._ready_signal=TEvent()
#
# if self._lock is None:
# self._lock=Lock()
#
# if self._cancel_signal is None:
# self._cancel_signal=TEvent()
#
# if not self._ready_signal.is_set():
# self._start_delay_timer()
# self.led.state=1
# do_later(self.trait_set, _flow_state='purge',
# message='Purging for {}s'.format(self.delay)
# )
# else:
# #cancel current timeout timer
# # self._cancel_signal.set()
# # time.sleep(1)
# #reset timer
#
# with self._lock:
# self._timeout_timer=0
# # self._start_timeout_timer()
#
# def stop(self):
# if self._delay_timer:
# self._delay_timer.cancel()
# if self._timeout_timer:
# self._timeout_timer.cancel()
#
# self._cancel_signal.set()
# self._ready_signal.clear()
#
# self._stop_flow()
# self.led.state = 0
# do_later(self.trait_set, _flow_state='off',
# message='',
# # _cancel=True
# )
# # self._flow_state='off'
def _start_flow(self):
self._set_io_state(self.channel, False)
def _stop_flow(self):
self._set_io_state(self.channel, True)
def _start_delay_timer(self):
self._ready_signal = TEvent()
self._ready_signal.clear()
self._delay_timer = Timer(self.delay, self.set_ready, args=(True,))
# self._start_timeout_timer()
self._delay_timer.start()
def _start_timeout_timer(self):
# def _loop(lock, cancel):
# with lock:
# self._timeout_cnt = 0
#
# while self._timeout_cnt < self.timeout and not cancel.is_set():
# v = self.timeout - self._timeout_cnt
# do_later(self.trait_set, message='Timeout after {}s ({}s) '.format(self.timeout, v))
# with lock:
# self._timeout_cnt += 1
#
# time.sleep(1)
#
# if self._timeout_cnt >= self.timeout:
# do_later(self.trait_set, message='Timed out after {}s'.format(self.timeout))
# self.set_ready(False)
#
# t = Thread(target=_loop, args=(self._lock, self._cancel_signal))
# t.start()
if self._timeout_timer:
self._timeout_timer.cancel()
# self._timeout_timer = Timer(10, self.set_ready, args=(False,))
self._timeout_timer = Timer(self.timeout, self.set_ready, args=(False,))
self._timeout_timer.start()
def set_ready(self, onoff):
if onoff:
self._ready_signal.set()
# self.led.state = 2
self._set_state_on()
# do_later(self.trait_set, _flow_state='on',
# message='Timeout at {}'.format(st)
# )
self._start_timeout_timer()
else:
self.stop()
self._ready_signal.clear()
def is_ready(self):
return self._ready_signal.is_set()
def traits_view(self):
v = View(
HGroup(Item('led', editor=LEDEditor(),
show_label=False, style='custom'
),
Item('flow_button',
show_label=False,
editor=ButtonEditor(label_value='flow_label')),
Spring(springy=False, width=20),
CustomLabel('message',
color='maroon', size=14,
springy=True
)
)
)
return v
def main():
"""
main entry point
returns 0 for normal termination (usually SIGTERM)
"""
return_value = 0
log_path = _log_path_template.format(os.environ["NIMBUSIO_LOG_DIR"],
_local_node_name)
initialize_logging(log_path)
log = logging.getLogger("main")
log.info("program starts")
for internal_socket_uri in internal_socket_uri_list:
prepare_ipc_path(internal_socket_uri)
halt_event = Event()
set_signal_handler(halt_event)
database_pool_controller = _launch_database_pool_controller()
io_controller = _launch_io_controller()
zeromq_context = zmq.Context()
rep_socket = _bind_rep_socket(zeromq_context)
db_controller_push_socket = \
_connect_db_controller_push_socket(zeromq_context)
event_push_client = EventPushClient(zeromq_context, "retrieve_source")
event_push_client.info("program-starts", "retrieve source starts")
# we poll the sockets for readability, we assume we can always
# write to the push client sockets
poller = zmq.Poller()
poller.register(rep_socket, zmq.POLLIN | zmq.POLLERR)
last_report_time = 0.0
request_count = 0
try:
while not halt_event.is_set():
poll_subprocess(database_pool_controller)
poll_subprocess(io_controller)
# we've only registered one socket, so we could use an 'if' here,
# but this 'for' works ok and it has the same form as the other
# places where we use poller
for active_socket, event_flags in poller.poll(_poll_timeout):
if event_flags & zmq.POLLERR:
error_message = \
"error flags from zmq {0}".format(active_socket)
log.error(error_message)
raise PollError(error_message)
assert active_socket is rep_socket
_process_one_request(rep_socket, db_controller_push_socket)
request_count += 1
current_time = time.time()
elapsed_time = current_time - last_report_time
if elapsed_time > _reporting_interval:
report_message = "{0:,} requests".format(request_count)
log.info(report_message)
event_push_client.info("request_count",
report_message,
request_count=request_count)
last_report_time = current_time
request_count = 0
except KeyboardInterrupt: # convenience for testing
log.info("keyboard interrupt: terminating normally")
except zmq.ZMQError as zmq_error:
if is_interrupted_system_call(zmq_error) and halt_event.is_set():
log.info("program teminates normally with interrupted system call")
else:
log.exception("zeromq error processing request")
event_push_client.exception(unhandled_exception_topic,
"zeromq_error",
exctype="ZMQError")
return_value = 1
except Exception as instance:
log.exception("error processing request")
event_push_client.exception(unhandled_exception_topic,
str(instance),
exctype=instance.__class__.__name__)
return_value = 1
else:
log.info("program teminates normally")
finally:
terminate_subprocess(database_pool_controller)
terminate_subprocess(io_controller)
rep_socket.close()
db_controller_push_socket.close()
event_push_client.close()
zeromq_context.term()
return return_value
class ConnectionHeartbeat(Thread):
def __init__(self, interval_sec, get_connection_holders):
Thread.__init__(self, name="Connection heartbeat")
self._interval = interval_sec
self._get_connection_holders = get_connection_holders
self._shutdown_event = Event()
self.daemon = True
self.start()
class ShutdownException(Exception):
pass
def run(self):
self._shutdown_event.wait(self._interval)
while not self._shutdown_event.is_set():
start_time = time.time()
futures = []
failed_connections = []
try:
for connections, owner in [
(o.get_connections(), o)
for o in self._get_connection_holders()
]:
for connection in connections:
self._raise_if_stopped()
if not (connection.is_defunct or connection.is_closed):
if connection.is_idle:
try:
futures.append(
HeartbeatFuture(connection, owner))
except Exception:
log.warning(
"Failed sending heartbeat message on connection (%s) to %s",
id(connection),
connection.host,
exc_info=True)
failed_connections.append(
(connection, owner))
else:
connection.reset_idle()
else:
# make sure the owner sees this defunt/closed connection
owner.return_connection(connection)
self._raise_if_stopped()
for f in futures:
self._raise_if_stopped()
connection = f.connection
try:
f.wait(self._interval)
# TODO: move this, along with connection locks in pool, down into Connection
with connection.lock:
connection.in_flight -= 1
connection.reset_idle()
except Exception:
log.warning(
"Heartbeat failed for connection (%s) to %s",
id(connection),
connection.host,
exc_info=True)
failed_connections.append((f.connection, f.owner))
for connection, owner in failed_connections:
self._raise_if_stopped()
connection.defunct(
Exception('Connection heartbeat failure'))
owner.return_connection(connection)
except self.ShutdownException:
pass
except Exception:
log.error("Failed connection heartbeat", exc_info=True)
elapsed = time.time() - start_time
self._shutdown_event.wait(max(self._interval - elapsed, 0.01))
def stop(self):
self._shutdown_event.set()
self.join()
def _raise_if_stopped(self):
if self._shutdown_event.is_set():
raise self.ShutdownException()
class ScreenGear:
"""
ScreenGear is designed exclusively for ultra-fast Screencasting, which means it can grab frames from your monitor in real-time, either by defining an area on the computer screen or full-screen,
at the expense of inconsiderable latency. ScreenGear also seamlessly support frame capturing from multiple monitors as well as supports multiple backends.
ScreenGear API implements a multi-threaded wrapper around pyscreenshot & python-mss python library, and also flexibly supports its internal parameter.
"""
def __init__(self,
monitor=None,
backend="",
colorspace=None,
logging=False,
**options):
"""
This constructor method initializes the object state and attributes of the ScreenGear class.
Parameters:
monitor (int): enables `mss` backend and sets the index of the monitor screen.
backend (str): enables `pyscreenshot` and select suitable backend for extracting frames.
colorspace (str): selects the colorspace of the input stream.
logging (bool): enables/disables logging.
options (dict): provides the flexibility to manually set the dimensions of capture screen area.
"""
# raise error(s) for critical Class imports
import_dependency_safe("mss.mss" if mss is None else "")
import_dependency_safe("pyscreenshot" if pysct is None else "")
# enable logging if specified:
self.__logging = logging if isinstance(logging, bool) else False
# create monitor instance for the user-defined monitor
self.__monitor_instance = None
self.__backend = ""
if monitor is None:
self.__capture_object = pysct
self.__backend = backend.lower().strip()
else:
self.__capture_object = mss()
if backend.strip():
logger.warning(
"Backends are disabled for Monitor Indexing(monitor>=0)!")
try:
self.__monitor_instance = self.__capture_object.monitors[
monitor]
except Exception as e:
logger.exception(str(e))
self.__monitor_instance = None
# assigns special parameter to global variable and clear
# Thread Timeout
self.__thread_timeout = options.pop("THREAD_TIMEOUT", None)
if self.__thread_timeout and isinstance(self.__thread_timeout,
(int, float)):
# set values
self.__thread_timeout = int(self.__thread_timeout)
else:
# defaults to 5mins timeout
self.__thread_timeout = None
# define deque and assign it to global var
self.__queue = queue.Queue(maxsize=96) # max len 96 to check overflow
# log it
if logging:
logger.debug(
"Enabling Threaded Queue Mode by default for ScreenGear!")
if self.__thread_timeout:
logger.debug("Setting Video-Thread Timeout to {}s.".format(
self.__thread_timeout))
# intiate screen dimension handler
screen_dims = {}
# reformat proper mss dict and assign to screen dimension handler
screen_dims = {
k.strip(): v
for k, v in options.items()
if k.strip() in ["top", "left", "width", "height"]
}
# check whether user-defined dimensions are provided
if screen_dims and len(screen_dims) == 4:
key_order = ("top", "left", "width", "height")
screen_dims = OrderedDict((k, screen_dims[k]) for k in key_order)
if logging:
logger.debug(
"Setting Capture-Area dimensions: {}!".format(screen_dims))
else:
screen_dims.clear()
# separately handle colorspace value to int conversion
if colorspace:
self.color_space = capPropId(colorspace.strip())
if logging and not (self.color_space is None):
logger.debug(
"Enabling `{}` colorspace for this video stream!".format(
colorspace.strip()))
else:
self.color_space = None
# intialize mss capture instance
self.__mss_capture_instance = ""
try:
if self.__monitor_instance is None:
if screen_dims:
self.__mss_capture_instance = tuple(screen_dims.values())
# extract global frame from instance
self.frame = np.asanyarray(
self.__capture_object.grab(
bbox=self.__mss_capture_instance,
childprocess=False,
backend=self.__backend,
))
else:
if screen_dims:
self.__mss_capture_instance = {
"top":
self.__monitor_instance["top"] + screen_dims["top"],
"left":
self.__monitor_instance["left"] + screen_dims["left"],
"width":
screen_dims["width"],
"height":
screen_dims["height"],
"mon":
monitor,
}
else:
self.__mss_capture_instance = (
self.
__monitor_instance # otherwise create instance from monitor
)
# extract global frame from instance
self.frame = np.asanyarray(
self.__capture_object.grab(self.__mss_capture_instance))
# initialize and append to queue
self.__queue.put(self.frame)
except Exception as e:
if isinstance(e, ScreenShotError):
# otherwise catch and log errors
if logging:
logger.exception(self.__capture_object.get_error_details())
raise ValueError(
"[ScreenGear:ERROR] :: ScreenShotError caught, Wrong dimensions passed to python-mss, Kindly Refer Docs!"
)
elif isinstance(e, KeyError):
raise ValueError(
"[ScreenGear:ERROR] :: ScreenShotError caught, Invalid backend: `{}`, Kindly Refer Docs!"
.format(backend))
else:
raise SystemError(
"[ScreenGear:ERROR] :: Unable to grab any instance on this system, Are you running headless?"
)
# thread initialization
self.__thread = None
# initialize termination flag
self.__terminate = Event()
def start(self):
"""
Launches the internal *Threaded Frames Extractor* daemon
**Returns:** A reference to the ScreenGear class object.
"""
self.__thread = Thread(target=self.__update,
name="ScreenGear",
args=())
self.__thread.daemon = True
self.__thread.start()
return self
def __update(self):
"""
A **Threaded Frames Extractor**, that keep iterating frames from `mss` API to a internal monitored deque,
until the thread is terminated, or frames runs out.
"""
# intialize frame variable
frame = None
# keep looping infinitely until the thread is terminated
while True:
# if the thread indicator variable is set, stop the thread
if self.__terminate.is_set():
break
try:
if self.__monitor_instance:
frame = np.asanyarray(
self.__capture_object.grab(
self.__mss_capture_instance))
else:
frame = np.asanyarray(
self.__capture_object.grab(
bbox=self.__mss_capture_instance,
childprocess=False,
backend=self.__backend,
))
if not self.__backend or self.__backend == "pil":
frame = frame[:, :, ::-1]
assert not (
frame is None or np.shape(frame) == ()
), "[ScreenGear:ERROR] :: Failed to retreive any valid frames!"
except Exception as e:
if isinstance(e, ScreenShotError):
raise RuntimeError(
self.__capture_object.get_error_details())
else:
logger.exception(str(e))
self.__terminate.set()
continue
if not (self.color_space is None):
# apply colorspace to frames
color_frame = None
try:
if isinstance(self.color_space, int):
color_frame = cv2.cvtColor(frame, self.color_space)
else:
if self.__logging:
logger.warning(
"Global color_space parameter value `{}` is not a valid!"
.format(self.color_space))
self.color_space = None
except Exception as e:
# Catch if any error occurred
self.color_space = None
if self.__logging:
logger.exception(str(e))
logger.warning(
"Input colorspace is not a valid colorspace!")
if not (color_frame is None):
self.frame = color_frame
else:
self.frame = frame
else:
self.frame = frame
# append to queue
self.__queue.put(self.frame)
# finally release mss resources
if self.__monitor_instance:
self.__capture_object.close()
def read(self):
"""
Extracts frames synchronously from monitored deque, while maintaining a fixed-length frame buffer in the memory,
and blocks the thread if the deque is full.
**Returns:** A n-dimensional numpy array.
"""
# check whether or not termination flag is enabled
while not self.__terminate.is_set():
return self.__queue.get(timeout=self.__thread_timeout)
# otherwise return NoneType
return None
def stop(self):
"""
Safely terminates the thread, and release the resources.
"""
if self.__logging:
logger.debug("Terminating ScreenGear Processes.")
# indicate that the thread should be terminate
self.__terminate.set()
# wait until stream resources are released (producer thread might be still grabbing frame)
if self.__thread is not None:
if not (self.__queue is None):
while not self.__queue.empty():
try:
self.__queue.get_nowait()
except queue.Empty:
continue
self.__queue.task_done()
self.__thread.join()
class WorkerBase(AbstractWorker):
def __init__(self,
args,
dev_id,
origin,
last_known_state,
communicator: AbstractCommunicator,
mapping_manager: MappingManager,
area_id: int,
routemanager_name: str,
db_wrapper: DbWrapper,
pogo_window_manager: PogoWindows,
walker=None,
event=None):
AbstractWorker.__init__(self, origin=origin, communicator=communicator)
self._mapping_manager: MappingManager = mapping_manager
self._routemanager_name: str = routemanager_name
self._area_id = area_id
self._dev_id: int = dev_id
self._event = event
self._origin: str = origin
self._applicationArgs = args
self._last_known_state = last_known_state
self._work_mutex = Lock()
self.loop = None
self.loop_started = Event()
self.loop_tid = None
self._async_io_looper_thread = None
self._location_count = 0
self._init: bool = self._mapping_manager.routemanager_get_init(
self._routemanager_name)
self._walker = walker
self._lastScreenshotTaken = 0
self._stop_worker_event = Event()
self._db_wrapper = db_wrapper
self._resocalc = Resocalculator
self._screen_x = 0
self._screen_y = 0
self._geofix_sleeptime = 0
self._pogoWindowManager = pogo_window_manager
self._waittime_without_delays = 0
self._transporttype = 0
self._not_injected_count: int = 0
self._same_screen_count: int = 0
self._last_screen_type: ScreenType = ScreenType.UNDEFINED
self._loginerrorcounter: int = 0
self._wait_again: int = 0
self._mode = self._mapping_manager.routemanager_get_mode(
self._routemanager_name)
self._levelmode = self._mapping_manager.routemanager_get_level(
self._routemanager_name)
self._geofencehelper = self._mapping_manager.routemanager_get_geofence_helper(
self._routemanager_name)
self.current_location = Location(0.0, 0.0)
self.last_location = self.get_devicesettings_value(
"last_location", None)
if self.last_location is None:
self.last_location = Location(0.0, 0.0)
if self.get_devicesettings_value('last_mode', None) is not None and \
self.get_devicesettings_value('last_mode') in ("raids_mitm", "mon_mitm", "iv_mitm"):
# Reset last_location - no useless waiting delays (otherwise stop mode)
self.last_location = Location(0.0, 0.0)
self.set_devicesettings_value(
"last_mode",
self._mapping_manager.routemanager_get_mode(
self._routemanager_name))
self.workerstart = None
self._WordToScreenMatching = WordToScreenMatching(
self._communicator, self._pogoWindowManager, self._origin,
self._resocalc, mapping_manager, self._applicationArgs)
def set_devicesettings_value(self, key: str, value):
self._mapping_manager.set_devicesetting_value_of(
self._origin, key, value)
def get_devicesettings_value(self, key: str, default_value: object = None):
self.logger.debug("Fetching devicemappings")
try:
devicemappings: Optional[
dict] = self._mapping_manager.get_devicemappings_of(
self._origin)
except (EOFError, FileNotFoundError) as e:
self.logger.warning(
"Failed fetching devicemappings with description: {}. Stopping worker",
e)
self._stop_worker_event.set()
return None
if devicemappings is None:
return default_value
return devicemappings.get("settings", {}).get(key, default_value)
def get_screenshot_path(self, fileaddon: bool = False) -> str:
screenshot_ending: str = ".jpg"
addon: str = ""
if self.get_devicesettings_value("screenshot_type", "jpeg") == "png":
screenshot_ending = ".png"
if fileaddon:
addon: str = "_" + str(time.time())
screenshot_filename = "screenshot_{}{}{}".format(
str(self._origin), str(addon), screenshot_ending)
if fileaddon:
self.logger.info("Creating debugscreen: {}", screenshot_filename)
return os.path.join(self._applicationArgs.temp_path,
screenshot_filename)
def check_max_walkers_reached(self):
walkermax = self._walker.get('walkermax', False)
if walkermax is False or (type(walkermax) is str
and len(walkermax) == 0):
return True
reg_workers = self._mapping_manager.routemanager_get_registered_workers(
self._routemanager_name)
if len(reg_workers) > int(walkermax):
return False
return True
@abstractmethod
def _pre_work_loop(self):
"""
Work to be done before the main while true work-loop
Start off asyncio loops etc in here
:return:
"""
pass
@abstractmethod
def _health_check(self):
"""
Health check before a location is grabbed. Internally, a self._start_pogo call is already executed since
that usually includes a topmost check
:return:
"""
pass
@abstractmethod
def _pre_location_update(self):
"""
Override to run stuff like update injections settings in MITM worker
Runs before walk/teleport to the location previously grabbed
:return:
"""
pass
@abstractmethod
def _move_to_location(self):
"""
Location has previously been grabbed, the overriden function will be called.
You may teleport or walk by your choosing
Any post walk/teleport delays/sleeps have to be run in the derived, override method
:return:
"""
pass
@abstractmethod
def _post_move_location_routine(self, timestamp):
"""
Routine called after having moved to a new location. MITM worker e.g. has to wait_for_data
:param timestamp:
:return:
"""
@abstractmethod
def _cleanup(self):
"""
Cleanup any threads you started in derived classes etc
self.stop_worker() and self.loop.stop() will be called afterwards
:return:
"""
@abstractmethod
def _worker_specific_setup_start(self):
"""
Routine preparing the state to scan. E.g. starting specific apps or clearing certain files
Returns:
"""
@abstractmethod
def _worker_specific_setup_stop(self):
"""
Routine destructing the state to scan. E.g. stopping specific apps or clearing certain files
Returns:
"""
def _start_asyncio_loop(self):
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
self.loop_tid = current_thread()
self.loop.call_soon(self.loop_started.set)
self.loop.run_forever()
def _add_task_to_loop(self, coro):
create_task = functools.partial(self.loop.create_task, coro)
if current_thread() == self.loop_tid:
# We can call directly if we're not going between threads.
return create_task()
else:
# We're in a non-event loop thread so we use a Future
# to get the task from the event loop thread once
# it's ready.
return self.loop.call_soon_threadsafe(create_task)
def start_worker(self):
t_main_work = Thread(name=self._origin, target=self._main_work_thread)
t_main_work.daemon = True
t_main_work.start()
# do some other stuff in the main process
while not self._stop_worker_event.is_set():
time.sleep(1)
while t_main_work.is_alive():
time.sleep(1)
t_main_work.join()
self.logger.info("Worker stopped gracefully")
return self._last_known_state
def stop_worker(self):
if self._stop_worker_event.is_set():
self.logger.info(
'Worker stop called, but worker is already stopping...')
else:
self._stop_worker_event.set()
self.logger.info("Worker stop called")
def _internal_pre_work(self):
current_thread().name = self._origin
start_position = self.get_devicesettings_value("startcoords_of_walker",
None)
calc_type = self._mapping_manager.routemanager_get_calc_type(
self._routemanager_name)
if start_position and (self._levelmode and calc_type == "routefree"):
startcoords = self.get_devicesettings_value("startcoords_of_walker").replace(' ', '') \
.replace('_', '').split(',')
if not self._geofencehelper.is_coord_inside_include_geofence(
Location(float(startcoords[0]), float(startcoords[1]))):
self.logger.info(
"Startcoords not in geofence - setting middle of fence as starting position"
)
lat, lng = self._geofencehelper.get_middle_from_fence()
start_position = str(lat) + "," + str(lng)
if start_position is None and \
(self._levelmode and calc_type == "routefree"):
self.logger.info(
"Starting level mode without worker start position")
# setting coords
lat, lng = self._geofencehelper.get_middle_from_fence()
start_position = str(lat) + "," + str(lng)
if start_position is not None:
startcoords = start_position.replace(' ',
'').replace('_',
'').split(',')
if not self._geofencehelper.is_coord_inside_include_geofence(
Location(float(startcoords[0]), float(startcoords[1]))):
self.logger.info(
"Startcoords not in geofence - setting middle of fence as startposition"
)
lat, lng = self._geofencehelper.get_middle_from_fence()
start_position = str(lat) + "," + str(lng)
startcoords = start_position.replace(' ',
'').replace('_',
'').split(',')
self.logger.info('Setting startcoords or walker lat {} / lng {}',
startcoords[0], startcoords[1])
self._communicator.set_location(
Location(startcoords[0], startcoords[1]), 0)
self._mapping_manager.set_worker_startposition(
routemanager_name=self._routemanager_name,
worker_name=self._origin,
lat=float(startcoords[0]),
lon=float(startcoords[1]))
with self._work_mutex:
try:
self._turn_screen_on_and_start_pogo()
self._get_screen_size()
# register worker in routemanager
self.logger.info(
"Try to register in Routemanager {}",
self._mapping_manager.routemanager_get_name(
self._routemanager_name))
self._mapping_manager.register_worker_to_routemanager(
self._routemanager_name, self._origin)
except WebsocketWorkerRemovedException:
self.logger.error("Timeout during init of worker")
# no cleanup required here? TODO: signal websocket server somehow
self._stop_worker_event.set()
return
self._async_io_looper_thread = Thread(name=self._origin,
target=self._start_asyncio_loop)
self._async_io_looper_thread.daemon = True
self._async_io_looper_thread.start()
self.loop_started.wait()
self._pre_work_loop()
def _internal_health_check(self):
# check if pogo is topmost and start if necessary
self.logger.debug4(
"_internal_health_check: Calling _start_pogo routine to check if pogo is topmost"
)
pogo_started = False
with self._work_mutex:
self.logger.debug2("_internal_health_check: worker lock acquired")
self.logger.debug4("Checking if we need to restart pogo")
# Restart pogo every now and then...
restart_pogo_setting = self.get_devicesettings_value(
"restart_pogo", 0)
if restart_pogo_setting > 0:
if self._location_count > restart_pogo_setting:
self.logger.info("scanned {} locations, restarting game",
restart_pogo_setting)
pogo_started = self._restart_pogo()
self._location_count = 0
else:
pogo_started = self._start_pogo()
else:
pogo_started = self._start_pogo()
self.logger.debug4("_internal_health_check: worker lock released")
return pogo_started
def _internal_cleanup(self):
# set the event just to make sure - in case of exceptions for example
self._stop_worker_event.set()
try:
self._mapping_manager.unregister_worker_from_routemanager(
self._routemanager_name, self._origin)
except ConnectionResetError as e:
self.logger.warning(
"Failed unregistering from routemanager, routemanager may have stopped running already."
"Exception: {}", e)
self.logger.info("Internal cleanup of started")
self._cleanup()
self.logger.info("Internal cleanup signaling end to websocketserver")
if self._async_io_looper_thread is not None:
self.logger.info("Stopping worker's asyncio loop")
self.loop.call_soon_threadsafe(self.loop.stop)
self._async_io_looper_thread.join()
self._communicator.cleanup()
self.logger.info("Internal cleanup of finished")
def _main_work_thread(self):
# TODO: signal websocketserver the removal
try:
self._internal_pre_work()
except (InternalStopWorkerException, WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.error(
"Failed initializing worker, connection terminated exceptionally"
)
self._internal_cleanup()
return
if not self.check_max_walkers_reached():
self.logger.warning(
'Max. Walkers in Area {} - closing connections',
self._mapping_manager.routemanager_get_name(
self._routemanager_name))
self.set_devicesettings_value('finished', True)
self._internal_cleanup()
return
while not self._stop_worker_event.is_set():
try:
# TODO: consider getting results of health checks and aborting the entire worker?
walkercheck = self.check_walker()
if not walkercheck:
self.set_devicesettings_value('finished', True)
break
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning("Worker killed by walker settings")
break
try:
# TODO: consider getting results of health checks and aborting the entire worker?
self._internal_health_check()
self._health_check()
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.error(
"Websocket connection to {} lost while running healthchecks, connection terminated "
"exceptionally", self._origin)
break
try:
settings = self._internal_grab_next_location()
if settings is None:
continue
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning(
"Worker of does not support mode that's to be run, connection terminated "
"exceptionally")
break
try:
self.logger.debug('Checking if new location is valid')
valid = self._check_location_is_valid()
if not valid:
break
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning("Worker received invalid coords!")
break
try:
self._pre_location_update()
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning(
"Worker of stopping because of stop signal in pre_location_update, connection "
"terminated exceptionally")
break
try:
self.logger.debug2(
'LastLat: {}, LastLng: {}, CurLat: {}, CurLng: {}',
self.get_devicesettings_value("last_location",
Location(0, 0)).lat,
self.get_devicesettings_value("last_location",
Location(0, 0)).lng,
self.current_location.lat, self.current_location.lng)
time_snapshot, process_location = self._move_to_location()
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning(
"Worker failed moving to new location, stopping worker, connection terminated "
"exceptionally")
break
if process_location:
self._add_task_to_loop(self._update_position_file())
self._location_count += 1
self.logger.debug("Seting new 'scannedlocation' in Database")
self._add_task_to_loop(
self.update_scanned_location(self.current_location.lat,
self.current_location.lng,
time_snapshot))
try:
calculate_waits = settings.get("encounter_all", False)
while self._wait_again > 0:
# We need to wait for data before we're able to do the calculation, otherwise we have wrong
# or missing data
self._post_move_location_routine(time_snapshot)
if calculate_waits:
try:
not_encountered: List[int] = []
latest = self._mitm_mapper.request_latest(
self._origin)
encountered = latest.get(
"ids_encountered", {}).get("values", {})
for cell in latest[106]["values"]["payload"][
"cells"]:
for pokemon in cell["wild_pokemon"]:
encounter_id = pokemon["encounter_id"]
# positive encounter IDs - calculation taken from DbPogoProtoSubmit.mon()
if encounter_id < 0:
encounter_id = encounter_id + 2**64
if encounter_id not in encountered:
monid = pokemon["pokemon_data"][
"id"]
not_encountered.append(monid)
# PD encounters 3 species per GMO
self._wait_again = math.ceil(
len(set(not_encountered)) / 3)
self.logger.debug(
"Found {} unique un-encountered mon IDs: {} - requires {} GMOs to "
"get all encounter data",
len(set(not_encountered)),
set(not_encountered), self._wait_again)
# Do not calculate again on subsequent runs of the while loop
calculate_waits = False
except Exception:
self.logger.warning(
"Exception trying to calculate the number of GMO waits - continue "
"with next location.")
self._wait_again: int = 1
self._wait_again -= 1
if self._wait_again > 0:
self.logger.info(
"Wait for {} more GMOs for more encounter data",
max(self._wait_again, 0))
time_snapshot = time.time()
except (InternalStopWorkerException,
WebsocketWorkerRemovedException,
WebsocketWorkerTimeoutException,
WebsocketWorkerConnectionClosedException):
self.logger.warning(
"Worker failed running post_move_location_routine, stopping worker"
)
break
self.logger.info("Worker finished iteration, continuing work")
self._internal_cleanup()
async def _update_position_file(self):
self.logger.debug2("Updating .position file")
if self.current_location is not None:
with open(
os.path.join(self._applicationArgs.file_path,
self._origin + '.position'), 'w') as outfile:
outfile.write(
str(self.current_location.lat) + ", " +
str(self.current_location.lng))
async def update_scanned_location(self, latitude, longitude, timestamp):
try:
self._db_wrapper.set_scanned_location(str(latitude),
str(longitude))
except Exception as e:
self.logger.error("Failed updating scanned location: {}", e)
return
def check_walker(self):
mode = self._walker['walkertype']
walkereventid = self._walker.get('eventid', None)
if walkereventid is not None and walkereventid != self._event.get_current_event_id(
):
self.logger.warning("Some other Event has started - leaving now")
return False
if mode == "countdown":
self.logger.info("Checking walker mode 'countdown'")
countdown = self._walker['walkervalue']
if not countdown:
self.logger.error(
"No Value for Mode - check your settings! Killing worker")
return False
if self.workerstart is None:
self.workerstart = math.floor(time.time())
else:
if math.floor(
time.time()) >= int(self.workerstart) + int(countdown):
return False
return True
elif mode == "timer":
self.logger.debug("Checking walker mode 'timer'")
exittime = self._walker['walkervalue']
if not exittime or ':' not in exittime:
self.logger.error(
"No or wrong Value for Mode - check your settings! Killing worker"
)
return False
return check_walker_value_type(exittime)
elif mode == "round":
self.logger.debug("Checking walker mode 'round'")
rounds = self._walker['walkervalue']
if len(rounds) == 0:
self.logger.error(
"No Value for Mode - check your settings! Killing worker")
return False
processed_rounds = self._mapping_manager.routemanager_get_rounds(
self._routemanager_name, self._origin)
if int(processed_rounds) >= int(rounds):
return False
return True
elif mode == "period":
self.logger.debug("Checking walker mode 'period'")
period = self._walker['walkervalue']
if len(period) == 0:
self.logger.error(
"No Value for Mode - check your settings! Killing worker")
return False
return check_walker_value_type(period)
elif mode == "coords":
exittime = self._walker['walkervalue']
if len(exittime) > 0:
return check_walker_value_type(exittime)
return True
elif mode == "idle":
self.logger.debug("Checking walker mode 'idle'")
if len(self._walker['walkervalue']) == 0:
self.logger.error(
"Wrong Value for mode - check your settings! Killing worker"
)
return False
sleeptime = self._walker['walkervalue']
self.logger.info('going to sleep')
killpogo = False
if check_walker_value_type(sleeptime):
self._stop_pogo()
killpogo = True
while not self._stop_worker_event.is_set(
) and check_walker_value_type(sleeptime):
time.sleep(1)
self.logger.info('just woke up')
if killpogo:
self._start_pogo()
return False
else:
self.logger.error("Unknown walker mode! Killing worker")
return False
def set_geofix_sleeptime(self, sleeptime: int) -> bool:
self._geofix_sleeptime = sleeptime
return True
def _internal_grab_next_location(self):
# TODO: consider adding runWarningThreadEvent.set()
self._last_known_state["last_location"] = self.last_location
self.logger.debug("Requesting next location from routemanager")
# requesting a location is blocking (iv_mitm will wait for a prioQ item), we really need to clean
# the workers up...
if int(self._geofix_sleeptime) > 0:
self.logger.info(
'Getting a geofix position from MADMin - sleeping for {} seconds',
self._geofix_sleeptime)
time.sleep(int(self._geofix_sleeptime))
self._geofix_sleeptime = 0
self._check_for_mad_job()
self.current_location = self._mapping_manager.routemanager_get_next_location(
self._routemanager_name, self._origin)
self._wait_again: int = 1
return self._mapping_manager.routemanager_get_settings(
self._routemanager_name)
def _check_for_mad_job(self):
if self.get_devicesettings_value("job", False):
self.logger.info("Worker get a job - waiting")
while self.get_devicesettings_value(
"job", False) and not self._stop_worker_event.is_set():
time.sleep(10)
self.logger.info("Worker processed the job and go on ")
def _check_location_is_valid(self):
if self.current_location is None:
# there are no more coords - so worker is finished successfully
self.set_devicesettings_value('finished', True)
return None
elif self.current_location is not None:
self.logger.debug2('Coords are valid')
return True
def _turn_screen_on_and_start_pogo(self):
if not self._communicator.is_screen_on():
self._communicator.start_app("de.grennith.rgc.remotegpscontroller")
self.logger.info("Turning screen on")
self._communicator.turn_screen_on()
time.sleep(
self.get_devicesettings_value("post_turn_screen_on_delay", 2))
# check if pogo is running and start it if necessary
self.logger.info("turnScreenOnAndStartPogo: (Re-)Starting Pogo")
self._start_pogo()
def _ensure_pogo_topmost(self):
self.logger.info('Checking pogo screen...')
while not self._stop_worker_event.is_set():
screen_type: ScreenType = self._WordToScreenMatching.detect_screentype(
)
if screen_type in [ScreenType.POGO, ScreenType.QUEST]:
self._last_screen_type = screen_type
self._loginerrorcounter = 0
self.logger.debug2("Found pogo or questlog to be open")
break
if screen_type != ScreenType.ERROR and self._last_screen_type == screen_type:
self._same_screen_count += 1
self.logger.info("Found {} multiple times in a row ({})",
screen_type, self._same_screen_count)
if self._same_screen_count > 3:
self.logger.warning("Screen is frozen!")
if self._same_screen_count > 4 or not self._restart_pogo():
self.logger.warning(
"Restarting PoGo failed - reboot device")
self._reboot()
break
elif self._last_screen_type != screen_type:
self._same_screen_count = 0
# now handle all screens that may not have been handled by detect_screentype since that only clicks around
# so any clearing data whatsoever happens here (for now)
if screen_type == ScreenType.UNDEFINED:
self.logger.error("Undefined screentype!")
elif screen_type == ScreenType.BLACK:
self.logger.info("Found Black Loading Screen - waiting ...")
time.sleep(20)
elif screen_type == ScreenType.CLOSE:
self.logger.debug(
"screendetection found pogo closed, start it...")
self._start_pogo()
self._loginerrorcounter += 1
elif screen_type == ScreenType.GAMEDATA:
self.logger.info(
'Error getting Gamedata or strange ggl message appears')
self._loginerrorcounter += 1
if self._loginerrorcounter < 2:
self._restart_pogo_safe()
elif screen_type == ScreenType.DISABLED:
# Screendetection is disabled
break
elif screen_type == ScreenType.UPDATE:
self.logger.warning(
'Found update pogo screen - sleeping 5 minutes for another check of the screen'
)
# update pogo - later with new rgc version
time.sleep(300)
elif screen_type in [ScreenType.ERROR, ScreenType.FAILURE]:
self.logger.warning(
'Something wrong with screendetection or pogo failure screen'
)
self._loginerrorcounter += 1
elif screen_type == ScreenType.NOGGL:
self.logger.warning(
'Detected login select screen missing the Google'
' button - likely entered an invalid birthdate previously')
self._loginerrorcounter += 1
elif screen_type == ScreenType.GPS:
self.logger.warning("Detected GPS error - reboot device")
self._reboot()
break
elif screen_type == ScreenType.SN:
self.logger.warning(
'Getting SN Screen - restart PoGo and later PD')
self._restart_pogo_safe()
break
elif screen_type == ScreenType.NOTRESPONDING:
self._reboot()
break
if self._loginerrorcounter > 1:
self.logger.warning(
'Could not login again - (clearing game data + restarting device'
)
self._stop_pogo()
self._communicator.clear_app_cache("com.nianticlabs.pokemongo")
if self.get_devicesettings_value('clear_game_data', False):
self.logger.info('Clearing game data')
self._communicator.reset_app_data(
"com.nianticlabs.pokemongo")
self._loginerrorcounter = 0
self._reboot()
break
self._last_screen_type = screen_type
self.logger.info('Checking pogo screen is finished')
if screen_type in [ScreenType.POGO, ScreenType.QUEST]:
return True
else:
return False
def _restart_pogo_safe(self):
self.logger.info(
"WorkerBase::_restart_pogo_safe restarting pogo the long way")
self._stop_pogo()
time.sleep(1)
if self._applicationArgs.enable_worker_specific_extra_start_stop_handling:
self._worker_specific_setup_stop()
time.sleep(1)
self._communicator.magisk_off()
time.sleep(1)
self._communicator.magisk_on()
time.sleep(1)
self._communicator.start_app("com.nianticlabs.pokemongo")
time.sleep(25)
self._stop_pogo()
time.sleep(1)
if self._applicationArgs.enable_worker_specific_extra_start_stop_handling:
self._worker_specific_setup_start()
time.sleep(1)
return self._start_pogo()
def _switch_user(self):
self.logger.info('Switching User - please wait ...')
self._stop_pogo()
time.sleep(5)
self._communicator.reset_app_data("com.nianticlabs.pokemongo")
self._turn_screen_on_and_start_pogo()
if not self._ensure_pogo_topmost():
self.logger.error('Kill Worker...')
self._stop_worker_event.set()
return False
self.logger.info('Switching finished ...')
return True
def _start_pogo(self) -> bool:
"""
Routine to start pogo.
Return the state as a boolean do indicate a successful start
:return:
"""
pogo_topmost = self._communicator.is_pogo_topmost()
if pogo_topmost:
return True
if not self._communicator.is_screen_on():
self._communicator.start_app("de.grennith.rgc.remotegpscontroller")
self.logger.info("Turning screen on")
self._communicator.turn_screen_on()
time.sleep(
self.get_devicesettings_value("post_turn_screen_on_delay", 7))
cur_time = time.time()
start_result = False
attempts = 0
while not pogo_topmost:
attempts += 1
if attempts > 10:
self.logger.warning("_start_pogo failed 10 times")
return False
start_result = self._communicator.start_app(
"com.nianticlabs.pokemongo")
time.sleep(1)
pogo_topmost = self._communicator.is_pogo_topmost()
if start_result:
self.logger.success("startPogo: Started pogo successfully...")
self._last_known_state["lastPogoRestart"] = cur_time
self._wait_pogo_start_delay()
return start_result
def is_stopping(self) -> bool:
return self._stop_worker_event.is_set()
def _stop_pogo(self):
attempts = 0
stop_result = self._communicator.stop_app("com.nianticlabs.pokemongo")
pogo_topmost = self._communicator.is_pogo_topmost()
while pogo_topmost:
attempts += 1
if attempts > 10:
return False
stop_result = self._communicator.stop_app(
"com.nianticlabs.pokemongo")
time.sleep(1)
pogo_topmost = self._communicator.is_pogo_topmost()
return stop_result
def _reboot(self, mitm_mapper: Optional[MitmMapper] = None):
try:
start_result = self._communicator.reboot()
except WebsocketWorkerRemovedException:
self.logger.error(
"Could not reboot due to client already disconnected")
start_result = False
time.sleep(5)
if mitm_mapper is not None:
mitm_mapper.collect_location_stats(
self._origin, self.current_location, 1, time.time(), 3, 0,
self._mapping_manager.routemanager_get_mode(
self._routemanager_name), 99)
self._db_wrapper.save_last_reboot(self._dev_id)
self._reboot_count = 0
self._restart_count = 0
self.stop_worker()
return start_result
def _restart_pogo(self,
clear_cache=True,
mitm_mapper: Optional[MitmMapper] = None):
successful_stop = self._stop_pogo()
self._db_wrapper.save_last_restart(self._dev_id)
self._restart_count = 0
self.logger.debug("restartPogo: stop game resulted in {}",
str(successful_stop))
if successful_stop:
if clear_cache:
self._communicator.clear_app_cache("com.nianticlabs.pokemongo")
time.sleep(1)
if mitm_mapper is not None:
mitm_mapper.collect_location_stats(
self._origin, self.current_location, 1, time.time(), 4, 0,
self._mapping_manager.routemanager_get_mode(
self._routemanager_name), 99)
return self._start_pogo()
else:
self.logger.warning("Failed restarting PoGo - reboot device")
return self._reboot()
def _get_trash_positions(self, full_screen=False):
self.logger.debug2("_get_trash_positions: Get_trash_position.")
if not self._take_screenshot(
delay_before=self.get_devicesettings_value(
"post_screenshot_delay", 1)):
self.logger.debug(
"_get_trash_positions: Failed getting screenshot")
return None
if os.path.isdir(self.get_screenshot_path()):
self.logger.error(
"_get_trash_positions: screenshot.png is not a file/corrupted")
return None
self.logger.debug2("_get_trash_positions: checking screen")
trashes = self._pogoWindowManager.get_trash_click_positions(
self._origin, self.get_screenshot_path(), full_screen=full_screen)
return trashes
def _take_screenshot(self,
delay_after=0.0,
delay_before=0.0,
errorscreen: bool = False):
self.logger.debug2("Taking screenshot...")
time.sleep(delay_before)
time_since_last_screenshot = time.time() - self._lastScreenshotTaken
self.logger.debug4("Last screenshot taken: {}",
str(self._lastScreenshotTaken))
# TODO: area settings for jpg/png and quality?
screenshot_type: ScreenshotType = ScreenshotType.JPEG
if self.get_devicesettings_value("screenshot_type", "jpeg") == "png":
screenshot_type = ScreenshotType.PNG
screenshot_quality: int = self.get_devicesettings_value(
"screenshot_quality", 80)
take_screenshot = self._communicator.get_screenshot(
self.get_screenshot_path(fileaddon=errorscreen),
screenshot_quality, screenshot_type)
if self._lastScreenshotTaken and time_since_last_screenshot < 0.5:
self.logger.info(
"screenshot taken recently, returning immediately")
return True
elif not take_screenshot:
self.logger.warning("Failed retrieving screenshot")
return False
else:
self.logger.debug("Success retrieving screenshot")
self._lastScreenshotTaken = time.time()
time.sleep(delay_after)
return True
def _check_pogo_main_screen(self, max_attempts, again=False):
self.logger.debug(
"_check_pogo_main_screen: Trying to get to the Mainscreen with {} max attempts...",
max_attempts)
pogo_topmost = self._communicator.is_pogo_topmost()
if not pogo_topmost:
return False
if not self._take_screenshot(
delay_before=self.get_devicesettings_value(
"post_screenshot_delay", 1)):
if again:
self.logger.warning(
"_check_pogo_main_screen: failed getting a screenshot again"
)
return False
attempts = 0
screenshot_path = self.get_screenshot_path()
if os.path.isdir(screenshot_path):
self.logger.error(
"_check_pogo_main_screen: screenshot.png/.jpg is not a file/corrupted"
)
return False
self.logger.debug("_check_pogo_main_screen: checking mainscreen")
while not self._pogoWindowManager.check_pogo_mainscreen(
screenshot_path, self._origin):
self.logger.info("_check_pogo_main_screen: not on Mainscreen...")
if attempts == max_attempts:
# could not reach raidtab in given max_attempts
self.logger.warning(
"_check_pogo_main_screen: Could not get to Mainscreen within {} attempts",
max_attempts)
return False
found = self._pogoWindowManager.check_close_except_nearby_button(
self.get_screenshot_path(),
self._origin,
self._communicator,
close_raid=True)
if found:
self.logger.debug(
"_check_pogo_main_screen: Found (X) button (except nearby)"
)
if not found and self._pogoWindowManager.look_for_button(
self._origin, screenshot_path, 2.20, 3.01,
self._communicator):
self.logger.debug(
"_check_pogo_main_screen: Found button (small)")
found = True
if not found and self._pogoWindowManager.look_for_button(
self._origin, screenshot_path, 1.05, 2.20,
self._communicator):
self.logger.debug(
"_check_pogo_main_screen: Found button (big)")
time.sleep(5)
found = True
self.logger.debug(
"_check_pogo_main_screen: Previous checks found pop ups: {}",
found)
self._take_screenshot(delay_before=self.get_devicesettings_value(
"post_screenshot_delay", 1))
attempts += 1
self.logger.debug("_check_pogo_main_screen: done")
return True
def _check_pogo_button(self):
self.logger.debug("checkPogoButton: Trying to find buttons")
pogo_topmost = self._communicator.is_pogo_topmost()
if not pogo_topmost:
return False
if not self._take_screenshot(
delay_before=self.get_devicesettings_value(
"post_screenshot_delay", 1)):
self.logger.debug("checkPogoButton: Failed getting screenshot")
return False
if os.path.isdir(self.get_screenshot_path()):
self.logger.error(
"checkPogoButton: screenshot.png is not a file/corrupted")
return False
self.logger.debug("checkPogoButton: checking for buttons")
found = self._pogoWindowManager.look_for_button(
self._origin, self.get_screenshot_path(), 2.20, 3.01,
self._communicator)
if found:
time.sleep(1)
self.logger.debug("checkPogoButton: Found button (small)")
if not found and self._pogoWindowManager.look_for_button(
self._origin, self.get_screenshot_path(), 1.05, 2.20,
self._communicator):
self.logger.debug("checkPogoButton: Found button (big)")
found = True
self.logger.debug("checkPogoButton: done")
return found
def _wait_pogo_start_delay(self):
delay_count: int = 0
pogo_start_delay: int = self.get_devicesettings_value(
"post_pogo_start_delay", 60)
self.logger.info('Waiting for pogo start: {} seconds',
pogo_start_delay)
while delay_count <= pogo_start_delay:
if not self._mapping_manager.routemanager_present(self._routemanager_name) \
or self._stop_worker_event.is_set():
self.logger.error("Killed while waiting for pogo start")
raise InternalStopWorkerException
time.sleep(1)
delay_count += 1
def _check_pogo_close(self, takescreen=True):
self.logger.debug("checkPogoClose: Trying to find closeX")
pogo_topmost = self._communicator.is_pogo_topmost()
if not pogo_topmost:
return False
if takescreen:
if not self._take_screenshot(
delay_before=self.get_devicesettings_value(
"post_screenshot_delay", 1)):
self.logger.debug("checkPogoClose: Could not get screenshot")
return False
if os.path.isdir(self.get_screenshot_path()):
self.logger.error(
"checkPogoClose: screenshot.png is not a file/corrupted")
return False
self.logger.debug("checkPogoClose: checking for CloseX")
found = self._pogoWindowManager.check_close_except_nearby_button(
self.get_screenshot_path(), self._origin, self._communicator)
if found:
time.sleep(1)
self.logger.debug(
"checkPogoClose: Found (X) button (except nearby)")
self.logger.debug("checkPogoClose: done")
return True
self.logger.debug("checkPogoClose: done")
return False
def _get_screen_size(self):
if self._stop_worker_event.is_set():
raise WebsocketWorkerRemovedException
screen = self._communicator.get_screensize()
if screen is None:
raise WebsocketWorkerRemovedException
screen = screen.strip().split(' ')
self._screen_x = screen[0]
self._screen_y = screen[1]
x_offset = self.get_devicesettings_value("screenshot_x_offset", 0)
y_offset = self.get_devicesettings_value("screenshot_y_offset", 0)
self.logger.debug(
'Get Screensize: X: {}, Y: {}, X-Offset: {}, Y-Offset: {}',
self._screen_x, self._screen_y, x_offset, y_offset)
self._resocalc.get_x_y_ratio(self, self._screen_x, self._screen_y,
x_offset, y_offset)
def make_graph(path, local_data, compute_data, results, results_opti, opti_names, nb_graph, result_name, show_names, link_vertices, background, save_gif, fps, return_origin):
plt.switch_backend('agg')
# prepare graphs
nb_graph = min(nb_graph, len(results))
cities = [(peak["x"], peak["y"]) for peak in local_data["peak"]]
x, y = zip(*cities) # zip because of tuples : extract couples [x, y]
# take most optimized path for selected graph range
for name, opti in zip(opti_names, results_opti):
convertor = {id_exe: counter for counter, (*_, id_exe) in enumerate(opti) if counter < nb_graph} # id translation
for id_exe, id_convert in convertor.items():
if results[id_exe][1] > opti[id_convert][1]: # lower score = better path
results[id_exe] = opti[id_convert]
del results[id_exe][-1]
results[id_exe].append(name)
for id, line in enumerate(results):
if len(line) < 5:
results[id].append("Generated")
# plot common part
fig = plt.figure(figsize=(6, 3))
axe = fig.gca()
axe.axis('off')
# step1 : optional background (country maps)
if background:
plot_countries = []
# get countries to plot
for city in cities:
res = search(city[0], city[1])
if res and res["ISO3"] not in plot_countries:
plot_countries.append(res["ISO3"])
if len(plot_countries) > MAX_COUNTRIES or len(plot_countries) == 0:
plot_countries = ["WORLD"]
for country in plot_countries: # print background map
country_map = geopandas.read_file(path+MAPS_FOLDER+"\\"+country+".shp")
country_map.plot(ax=axe, color="lightgray", zorder=-1)
# step2 : optional display of link between vertices
if link_vertices:
linked = [(id, peak["link"]) for id, peak in enumerate(compute_data["peak"]) if peak["origin"]]
for origin, dests in linked:
x1, y1 = cities[origin]
for dest in dests:
x2, y2 = cities[dest]
line, = axe.plot([x1, x2], [y1, y2], zorder=0, linestyle="dashed", c="green", alpha=0.3)
line.set_label('Refers')
# step3.1 : display origins and destinations
cities_origin = [(peak["x"], peak["y"]) for peak in local_data["peak"] if peak["origin"]]
cities_dest = [(peak["x"], peak["y"]) for peak in local_data["peak"] if not peak["origin"]]
axe.scatter(*zip(*cities_origin), marker="s", c="red", label="Deposit")
axe.scatter(*zip(*cities_dest), marker="o", c="green", label="Client")
# step3.2 : display travelers origins
nb_trav = len(local_data["traveler"])
if nb_trav < 10:
colors = list(mcolors.TABLEAU_COLORS.values())
else:
colors = mcolors.CSS4_COLORS # 148 colors
colors.pop("lightgray")
colors.pop("lightgrey")
colors.pop("gainsboro")
colors = [color for color in colors.values() if color.count("F") < 4] # 128 left
random.shuffle(colors)
colors = colors[:nb_trav]
travelers = [(trav["x"], trav["y"], trav["name"]) for trav in local_data["traveler"]]
x_trav, y_trav, _ = zip(*travelers)
axe.scatter(x_trav, y_trav, marker="D", c=colors, label="Traveler")
# step4 : optional display of vertices and travelers names
if show_names:
coef = 0.3*(axe.get_xlim()[1]-axe.get_xlim()[0])/15
for i in range(len(cities)):
axe.text(x[i]+coef, y[i]+coef, local_data["peak"][i]["name"],
fontsize='x-small', c="black")
for (x, y, name), color in zip(travelers, colors):
axe.text(x+coef, y+coef, name, fontsize='x-small', c=color)
# step5 : tweak legend order & position
if link_vertices:
handles, labels = axe.get_legend_handles_labels()
handles = handles[1:3]+[handles[0]]+[handles[3]]
labels = labels[1:3]+[labels[0]]+[labels[3]]
axe.legend(handles, labels, ncol=4, loc='lower center', bbox_to_anchor=(0.5, -0.15))
else:
axe.legend(ncol=3, loc='lower center', bbox_to_anchor=(0.5, -0.15))
# step6 : dump base graph and enrich it in each thread
graph_buffer = io.BytesIO()
pickle.dump(fig, graph_buffer)
# step7 : start threads
g_files = []
g_files_lock = Lock()
interrupt_event = Event()
threads = []
for idThread in range(nb_graph):
graph_buffer.seek(0) # crucial
newfig = pickle.load(graph_buffer)
args = (idThread, interrupt_event, newfig, results[idThread], cities,
local_data["traveler"], colors, return_origin,
path, result_name, save_gif, fps, g_files, g_files_lock)
threads.append(Thread(target=plot_path, args=args))
threads[-1].start()
for thread in threads:
thread.join()
if interrupt_event.is_set():
print("Close program")
exit()
return g_files
