class Collector(CompileStateListener):
"""
Collect all state updates, optionally hang the processing of listeners
"""
def __init__(self):
self.seen = []
self.preseen = []
self.lock = Semaphore(1)
def reset(self):
self.seen = []
self.preseen = []
async def compile_done(self, compile: data.Compile):
self.preseen.append(compile)
print("Got compile done for ", compile.remote_id)
async with self.lock:
self.seen.append(compile)
async def hang(self):
await self.lock.acquire()
def release(self):
self.lock.release()
def verify(self, envs: uuid.UUID):
assert sorted([x.remote_id for x in self.seen]) == sorted(envs)
self.reset()
async def worker(xs: List[int], pred: Callable[[int], bool],
res: List[List[int]], semaphore: asyncio.Semaphore) -> None:
for x in xs:
await semaphore.acquire()
if pred(x):
res[0].append(x)
semaphore.release()
await asyncio.sleep(0.000000001)
async def get_some_dress(semaphore: Semaphore):
await semaphore.acquire(
) # занимаем примерочную, счетчик свободных примерочных уменьшится на 1
start = time.time()
await main_task()
time_of_work = time.time() - start
print("ВРЕМЯ РАБОТЫ:", time_of_work)
await asyncio.sleep(1 - time_of_work)
semaphore.release(
) # освобождаем примерочную, счетчик свободных примерочных увеличится на 1
async def analyze(cli_args: argparse.Namespace, domain: str,
recursion_level: int, results_queue: asyncio.Queue,
sem: asyncio.Semaphore, input_domains_queue: asyncio.Queue):
tasks = []
try:
#
# Getting info from AWS
#
t1 = asyncio.create_task(
get_s3(cli_args, domain, recursion_level, input_domains_queue,
results_queue))
tasks.append(t1)
except Exception as e:
print(e)
try:
#
# Get web links?
#
if not cli_args.no_links:
t2 = asyncio.create_task(
get_links(cli_args, domain, recursion_level,
input_domains_queue, results_queue))
tasks.append(t2)
except Exception as e:
print(e)
try:
#
# Get cnames
#
# if cli_args.dns:
if not cli_args.no_dnsdiscover:
t3 = asyncio.create_task(
get_dns_info(cli_args, domain, recursion_level,
input_domains_queue))
tasks.append(t3)
except Exception as e:
print(e)
try:
await asyncio.gather(*tasks)
finally:
sem.release()
input_domains_queue.task_done()
class Barreira:
def __init__(self, n):
self.n = n # Número de threads aguardadas
self.contador = 0 # Número de threads até o momento
self.barreira_1 = Semaphore(0) # Proteção de entrada da barreira
self.barreira_2 = Semaphore(1) # Proteção de saída da barreira
self.mutex = Lock() # Protege acesso à variável contador
'''
Aguarda entrada na barreira.
'''
async def acquire_1(self):
async with self.mutex:
self.contador += 1
if self.contador == self.n:
await self.barreira_2.acquire()
self.barreira_1.release()
await self.barreira_1.acquire()
self.barreira_1.release()
'''
Aguarda saída da barreira.
'''
async def acquire_2(self):
async with self.mutex:
self.contador -= 1
if self.contador == 0:
await self.barreira_1.acquire()
self.barreira_2.release()
await self.barreira_2.acquire()
self.barreira_2.release()
class FifoQueue:
def __init__(self):
self._queue = deque()
self._semaphore = Semaphore(0)
def __len__(self):
return len(self._queue)
async def push(self, request):
self._queue.append(request)
self._semaphore.release()
async def pop(self):
await self._semaphore.acquire()
return self._queue.popleft()
async def perform_request(conn: Optional[aiohttp.TCPConnector], i: int,
sem: asyncio.Semaphore, url: str):
t0 = time.time()
try:
async with aiohttp.request('get', url, connector=conn) as resp:
# Read wjhole response without buffering.
chunk = True
while chunk:
chunk = await resp.content.read(1024 * 1024)
delta = time.time() - t0
LOG.info("[%d] HTTP %d after %s seconds", i, resp.status, delta)
except aiohttp.client_exceptions.ClientConnectorError as e:
LOG.info("[%d] %s", i, e)
finally:
sem.release()
class PriorityQueue:
def __init__(self):
self._queue = []
self._semaphore = Semaphore(0)
def __len__(self):
return len(self._queue)
async def push(self, request):
heappush(self._queue, _PriorityQueueItem(request))
self._semaphore.release()
async def pop(self):
await self._semaphore.acquire()
item = heappop(self._queue)
return item.request
async def frequency_limiter(sem: Semaphore) -> None:
"""
Function that must be used only inside router module
to wait some time before answer to the client and let the background task to start
:param sem: semaphore to store iot or user data
"""
try:
await wait_for(sem.acquire(), 1)
sem.release()
except TimeoutError:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
headers={"Retry-After": str(randrange(1, 29) * random() + 1)},
)
class HangRunner(object):
"""
compile runner mock, hang until released
"""
def __init__(self):
self.lock = Semaphore(0)
self.started = False
self.done = False
self.version = None
async def run(self, force_update: Optional[bool] = False):
self.started = True
await self.lock.acquire()
self.done = True
return True, None
def release(self):
self.lock.release()
class IPQuery(object):
def __init__(self, concurrent_limit=0):
self._futs = set()
self.smaphore = Semaphore(
concurrent_limit) if concurrent_limit > 0 else None
async def query_ip(self, ip, batch_count=3):
if isinstance(ip, str):
return await self._query_one(ip)
else:
rtn = dict()
query_iterable = _IPQueryIterable(self, ip, batch_count)
try:
while True:
finished_ip, ip_info = await query_iterable._once()
rtn[finished_ip] = ip_info
except _FinishException:
pass
return rtn
def iter(self, ip, batch_count=3):
if isinstance(ip, str):
ip = (ip, )
return _IPQueryIterable(self, ip, batch_count)
async def _query_one(self, ip):
if self.smaphore:
await self.smaphore.acquire()
try:
async with aiohttp.ClientSession() as session:
async with session.get(
"http://ip.taobao.com/service/getIpInfo.php",
params={"ip": ip}) as resp:
rtn_data = json.loads(await resp.text())
if rtn_data["code"] == 0:
return ip, rtn_data['data']
except asyncio.CancelledError:
raise
except Exception as e:
return ip, e
finally:
if self.smaphore:
self.smaphore.release()
class PriorityQueue:
def __init__(self):
self.queue = []
self.items = Semaphore(value=0)
async def push(self, data, priority=0):
self.queue.append((priority, json.dumps(data)))
self.queue.sort()
self.items.release()
async def pop(self, timeout: int = 1) -> Any:
try:
await wait_for(self.items.acquire(), timeout)
return json.loads(self.queue.pop(-1)[1])
except TimeoutError:
return None
async def pop_ready(self) -> Any:
if self.items.locked():
return None
await self.items.acquire()
return json.loads(self.queue.pop(-1)[1])
async def score(self, data):
data = json.dumps(data)
for priority, item in self.queue:
if data == item:
return priority
return None
async def rank(self, data):
data = json.dumps(data)
for index, (_, item) in enumerate(self.queue):
if data == item:
return len(self.queue) - index - 1
return None
async def clear(self):
self.queue = []
self.items = Semaphore(value=0)
async def length(self):
return len(self.queue)
async def _fetch_and_parse(
self,
sem: Semaphore,
pool: Executor,
sess: ClientSession,
url: str,
parse_fn: Optional[Callable[[str], T]] = None) -> Optional[str]:
try:
async with sess.get(url) as resp:
content = await resp.text()
# Run `parse_fn` in subprocess from process-pool for parallelism.
if parse_fn is not None:
content = await asyncio.get_event_loop().run_in_executor(
pool, parse_fn, content)
except Exception:
content = None
sem.release()
return content
async def hello(x: int, sem: asyncio.Semaphore):
r = random.randint(1, 5)
async for key in redis.iscan(match='something*'):
print('Matched:', key)
print("Waiting {}".format(x))
await asyncio.sleep(r)
print("Finish: " + str(x))
print("Waiting {}".format(x))
await asyncio.sleep(r)
print("Finish: " + str(x))
print("Waiting {}".format(x))
await asyncio.sleep(r)
print("Finish: " + str(x))
f = asyncio.Future()
sem.release()
class MockTransportSender(TransportSenderBase):
def __init__(self):
super().__init__()
self.send_called = Semaphore(0)
async def send(self, buffer: List[int], offset: int, count: int) -> int:
# Assert
if count == 48: # Header
print("Validating Header...")
header = HeaderSerializer.deserialize(buffer, offset, count)
assert header.type == "A"
assert header.payload_length == 3
assert header.end
else: # Payload
print("Validating Payload...")
assert count == 3
self.send_called.release()
return count
def close(self):
pass
class SmartPool:
"""Pool which utilizes alive connections."""
def __init__(self, connector, pool_size, connection_cls):
self.pool_size = pool_size
self.pool = set()
self.sem = Semaphore(pool_size)
for _ in range(pool_size):
self.pool.add(connection_cls(connector))
async def acquire(self, urlparsed: ParseResult = None):
"""Acquire connection."""
await self.sem.acquire()
if urlparsed:
key = f'{urlparsed.hostname}-{urlparsed.port}'
for item in self.pool:
if item.key == key:
self.pool.remove(item)
return item
return self.pool.pop()
def release(self, conn) -> None:
"""Release connection."""
self.pool.add(conn)
self.sem.release()
def free_conns(self) -> int:
return len(self.pool)
def is_all_free(self):
"""Indicates if all pool is free."""
return self.pool_size == self.sem._value
async def cleanup(self) -> None:
"""Get all conn and close them, this method let this pool unusable."""
for count in range(self.pool_size):
conn = await self.acquire()
conn.close()
def _read(self, r_content: _ReadContent, over_semaphore: asyncio.Semaphore):
# 读操作（阻塞）
r_content.content = self._f.read()
# 让父协程从等待队列中唤醒
over_semaphore.release()
class AsyncioSubscriptionManager(SubscriptionManager):
def __init__(self, pubnub_instance):
subscription_manager = self
self._message_worker = None
self._message_queue = Queue()
self._subscription_lock = Semaphore(1)
self._subscribe_loop_task = None
self._heartbeat_periodic_callback = None
self._reconnection_manager = AsyncioReconnectionManager(pubnub_instance)
super(AsyncioSubscriptionManager, self).__init__(pubnub_instance)
self._start_worker()
class AsyncioReconnectionCallback(ReconnectionCallback):
def on_reconnect(self):
subscription_manager.reconnect()
pn_status = PNStatus()
pn_status.category = PNStatusCategory.PNReconnectedCategory
pn_status.error = False
subscription_manager._subscription_status_announced = True
subscription_manager._listener_manager.announce_status(pn_status)
self._reconnection_listener = AsyncioReconnectionCallback()
self._reconnection_manager.set_reconnection_listener(self._reconnection_listener)
def _set_consumer_event(self):
if not self._message_worker.cancelled():
self._message_worker.cancel()
def _message_queue_put(self, message):
self._message_queue.put_nowait(message)
def _start_worker(self):
consumer = AsyncioSubscribeMessageWorker(self._pubnub,
self._listener_manager,
self._message_queue, None)
self._message_worker = asyncio.ensure_future(consumer.run(),
loop=self._pubnub.event_loop)
def reconnect(self):
# TODO: method is synchronized in Java
self._should_stop = False
self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop())
self._register_heartbeat_timer()
def disconnect(self):
# TODO: method is synchronized in Java
self._should_stop = True
self._stop_heartbeat_timer()
self._stop_subscribe_loop()
def stop(self):
super(AsyncioSubscriptionManager, self).stop()
self._reconnection_manager.stop_polling()
if self._subscribe_loop_task is not None and not self._subscribe_loop_task.cancelled():
self._subscribe_loop_task.cancel()
@asyncio.coroutine
def _start_subscribe_loop(self):
self._stop_subscribe_loop()
yield from self._subscription_lock.acquire()
combined_channels = self._subscription_state.prepare_channel_list(True)
combined_groups = self._subscription_state.prepare_channel_group_list(True)
if len(combined_channels) == 0 and len(combined_groups) == 0:
self._subscription_lock.release()
return
self._subscribe_request_task = asyncio.ensure_future(Subscribe(self._pubnub)
.channels(combined_channels)
.channel_groups(combined_groups)
.timetoken(self._timetoken).region(self._region)
.filter_expression(self._pubnub.config.filter_expression)
.future())
e = yield from self._subscribe_request_task
if self._subscribe_request_task.cancelled():
self._subscription_lock.release()
return
if e.is_error():
if e.status is not None and e.status.category == PNStatusCategory.PNCancelledCategory:
self._subscription_lock.release()
return
if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory:
self._pubnub.event_loop.call_soon(self._start_subscribe_loop)
self._subscription_lock.release()
return
logger.error("Exception in subscribe loop: %s" % str(e))
if e.status is not None and e.status.category == PNStatusCategory.PNAccessDeniedCategory:
e.status.operation = PNOperationType.PNUnsubscribeOperation
# TODO: raise error
self._listener_manager.announce_status(e.status)
self._reconnection_manager.start_polling()
self._subscription_lock.release()
self.disconnect()
return
else:
self._handle_endpoint_call(e.result, e.status)
self._subscription_lock.release()
self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop())
self._subscription_lock.release()
def _stop_subscribe_loop(self):
if self._subscribe_request_task is not None and not self._subscribe_request_task.cancelled():
self._subscribe_request_task.cancel()
def _stop_heartbeat_timer(self):
if self._heartbeat_periodic_callback is not None:
self._heartbeat_periodic_callback.stop()
def _register_heartbeat_timer(self):
super(AsyncioSubscriptionManager, self)._register_heartbeat_timer()
self._heartbeat_periodic_callback = AsyncioPeriodicCallback(
self._perform_heartbeat_loop,
self._pubnub.config.heartbeat_interval * 1000,
self._pubnub.event_loop)
if not self._should_stop:
self._heartbeat_periodic_callback.start()
@asyncio.coroutine
def _perform_heartbeat_loop(self):
if self._heartbeat_call is not None:
# TODO: cancel call
pass
cancellation_event = Event()
state_payload = self._subscription_state.state_payload()
presence_channels = self._subscription_state.prepare_channel_list(False)
presence_groups = self._subscription_state.prepare_channel_group_list(False)
if len(presence_channels) == 0 and len(presence_groups) == 0:
return
try:
heartbeat_call = (Heartbeat(self._pubnub)
.channels(presence_channels)
.channel_groups(presence_groups)
.state(state_payload)
.cancellation_event(cancellation_event)
.future())
envelope = yield from heartbeat_call
heartbeat_verbosity = self._pubnub.config.heartbeat_notification_options
if envelope.status.is_error:
if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL or \
heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL:
self._listener_manager.announce_stateus(envelope.status)
else:
if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL:
self._listener_manager.announce_stateus(envelope.status)
except PubNubAsyncioException as e:
pass
# TODO: check correctness
# if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory:
# self._start_subscribe_loop()
# else:
# self._listener_manager.announce_status(e.status)
finally:
cancellation_event.set()
def _send_leave(self, unsubscribe_operation):
asyncio.ensure_future(self._send_leave_helper(unsubscribe_operation))
@asyncio.coroutine
def _send_leave_helper(self, unsubscribe_operation):
envelope = yield from Leave(self._pubnub) \
.channels(unsubscribe_operation.channels) \
.channel_groups(unsubscribe_operation.channel_groups).future()
self._listener_manager.announce_status(envelope.status)
class Queue:
def __init__(self, n_slots):
self._n_slots = n_slots
self._wait_tx = OrderedDict()
self._wait_tx_sem = Semaphore(0)
self._wait_rx = OrderedDict()
@property
def n_slots(self):
return self._n_slots
@property
def is_full(self):
return len(self._wait_tx) + len(self._wait_rx) >= self._n_slots
@property
def qsize(self):
return self._n_slots
def add(self, request):
instance_id = request.instanceId
analysis_id = request.analysisId
key = (instance_id, analysis_id)
existing = self._wait_tx.get(key)
if existing is not None:
ex_request, ex_stream = existing
log.debug('%s %s', 'cancelling', req_str(ex_request))
ex_stream.cancel()
else:
existing = self._wait_rx.get(key)
if existing is not None:
ex_request, ex_stream = existing
log.debug('%s %s', 'cancelling', req_str(ex_request))
ex_stream.cancel()
if self.is_full:
raise QueueFull
stream = Stream()
log.debug('%s %s', 'queueing', req_str(request))
self._wait_tx[key] = (request, stream)
if self._wait_tx_sem.locked():
self._wait_tx_sem.release()
stream.add_complete_listener(self._stream_complete)
return stream
def get(self, key):
value = self._wait_tx.get(key)
if value is not None:
return value
value = self._wait_rx.get(key)
if value is not None:
return value
return None
def _stream_complete(self):
for key, value in self._wait_rx.items():
request, stream = value
if stream.is_complete:
del self._wait_rx[key]
break
log.debug('%s %s', 'removing', req_str(request))
def stream(self):
# # this isn't compatible with python 3.5:
# while True:
# await self._wait_tx_sem.acquire()
# while len(self._wait_tx) > 0:
# key, value = self._wait_tx.popitem()
# self._wait_rx[key] = value
# request, stream = value
# log.debug('%s %s', 'yielding', req_str(request))
# yield value
# # so we had to do this:
class AsyncGenerator:
def __init__(self, parent):
self._parent = parent
def __aiter__(self):
return self
async def __anext__(self):
if len(self._parent._wait_tx) == 0:
await self._parent._wait_tx_sem.acquire()
key, value = self._parent._wait_tx.popitem()
self._parent._wait_rx[key] = value
request, stream = value
log.debug('%s %s', 'yielding', req_str(request))
return value
return AsyncGenerator(self)
def __contains__(self, value):
return value in self._wait_tx or value in self._wait_rx
class Worker:
def __init__(self, loop=None):
logger.info('Worker initialising...')
loop = loop or asyncio.get_event_loop()
self.loop = loop
logger.debug('Connecting to db: "%s"', DB_DSN)
self._pool = loop.run_until_complete(create_pool(dsn=DB_DSN, loop=loop, minsize=2, maxsize=10))
self.wkh2p_sema = Semaphore(value=MAX_WORKER_THREADS, loop=loop)
self.worker_sema = Semaphore(value=MAX_WORKER_JOBS, loop=loop)
self.redis = None
self.exc_info = None
def run_forever(self):
self.loop.run_until_complete(self.work_loop())
async def work_loop(self):
# TODO deal with SIGTERM gracefully
logger.debug('Connecting to redis on: "%s"', REDIS_HOST)
self.redis = await aioredis.create_redis((REDIS_HOST, 6379), loop=self.loop)
logger.info('Worker started...')
try:
while True:
await self.worker_sema.acquire()
queue, data = await self.redis.blpop(*QUEUES)
self.loop.create_task(self.work_handler(queue, data))
finally:
self.redis.close()
async def work_handler(self, queue, raw_data):
try:
await self.work(queue, raw_data)
except:
logger.error('error processing job: %s', sys.exc_info()[1])
self.exc_info = sys.exc_info()
raise
async def work(self, queue_raw, raw_data):
"""
Do job, data shape:
{
'job_id': UUID of job,
THEN
'content': JSON object to use in template
OR
'html': HTML to generate pdf for
}
:param queue_raw: queue name bytes
:param raw_data: json bytes
:return:
"""
queue = queue_raw.decode()
data = raw_data.decode()
logger.debug('starting job from queue "%s" with data "%s"', queue, data)
data = json.loads(data)
job_id = data['job_id']
org_code, env_id = await self.get_basic_info(job_id)
logger.info('starting job - %s for %s', job_id, org_code)
await self.job_in_progress(job_id)
content = data.get('content')
if content:
raise NotImplementedError()
# TODO generate html
else:
html = data['html']
await self.wkh2p_sema.acquire()
pdf_file = await self.loop.run_in_executor(None, generate_pdf, html)
self.wkh2p_sema.release()
logger.info('pdf generated - %s for %s', job_id, org_code)
# the temporary file is not automatically deleted, so we need to make sure we do it here
try:
file_size = os.path.getsize(pdf_file)
await store_file(job_id, org_code, pdf_file)
finally:
os.remove(pdf_file)
await self.job_finished(job_id, html, file_size)
logger.info('finishing job - %s for %s', job_id, org_code)
self.worker_sema.release()
async def job_in_progress(self, job_id):
ctx = [JobStatus.STATUS_IN_PROGRESS, job_id]
await self.execute('UPDATE jobs_job SET status=%s, timestamp_started=current_timestamp WHERE id=%s;', ctx)
async def job_finished(self, job_id, html, file_size):
ctx = [JobStatus.STATUS_COMPLETE, html, file_size, job_id]
await self.execute('UPDATE jobs_job SET status=%s, timestamp_complete=current_timestamp, '
'html=%s, file_size=%s WHERE id=%s;', ctx)
async def get_basic_info(self, job_id):
cur = await self.execute(
'SELECT orgs_organisation.code, resources_env.id FROM orgs_organisation '
'INNER JOIN resources_env ON orgs_organisation.id = resources_env.org_id '
'INNER JOIN jobs_job ON resources_env.id = jobs_job.env_id WHERE '
'jobs_job.id = %s', [job_id])
org, env_id = await cur.fetchone()
return org, env_id
async def get_env_info(self, job_id, env_id):
cur = await self.execute(
'SELECT r_main.ref AS main_ref, r_main.file AS main_file, '
'r_base.ref as base_ref, r_base.file as base_file, '
'r_base.ref as header_ref, r_header.file as header_file, '
'r_base.ref as footer_ref, r_footer.file as footer_file '
'FROM resources_env '
'JOIN resources_file AS r_main ON r_main.id = resources_env.main_template_id '
'JOIN resources_file AS r_base ON r_base.id = resources_env.base_template_id '
'JOIN resources_file AS r_header ON r_header.id = resources_env.header_template_id '
'JOIN resources_file AS r_footer ON r_footer.id = resources_env.footer_template_id '
'WHERE resources_env.id = %s', [env_id], dict_cursor=True)
data = dict(await cur.fetchone())
print(data)
# FIXME, work stopped here
@asyncio.coroutine
def execute(self, *args, **kwargs):
cursor_factory = kwargs.pop('dict_cursor', None) and DictCursor
with (yield from self._pool) as conn:
cur = yield from conn.cursor(cursor_factory=cursor_factory)
yield from cur.execute(*args, **kwargs)
return cur
class Overseer:
def __init__(self, manager):
self.log = get_logger('overseer')
self.workers = []
self.manager = manager
self.things_count = deque(maxlen=9)
self.paused = False
self.coroutines_count = 0
self.skipped = 0
self.visits = 0
self.coroutine_semaphore = Semaphore(conf.COROUTINES_LIMIT, loop=LOOP)
self.redundant = 0
self.running = True
self.all_seen = False
self.idle_seconds = 0
self.log.info('Overseer initialized')
self.pokemon_found = ''
def start(self, status_bar):
self.captcha_queue = self.manager.captcha_queue()
Worker.captcha_queue = self.manager.captcha_queue()
self.extra_queue = self.manager.extra_queue()
Worker.extra_queue = self.manager.extra_queue()
if conf.MAP_WORKERS:
Worker.worker_dict = self.manager.worker_dict()
for username, account in ACCOUNTS.items():
account['username'] = username
if account.get('banned'):
continue
if account.get('captcha'):
self.captcha_queue.put(account)
else:
self.extra_queue.put(account)
self.workers = tuple(Worker(worker_no=x) for x in range(conf.GRID[0] * conf.GRID[1]))
db_proc.start()
LOOP.call_later(10, self.update_count)
LOOP.call_later(max(conf.SWAP_OLDEST, conf.MINIMUM_RUNTIME), self.swap_oldest)
LOOP.call_soon(self.update_stats)
if status_bar:
LOOP.call_soon(self.print_status)
def update_count(self):
self.things_count.append(str(db_proc.count))
self.pokemon_found = (
'Pokemon found count (10s interval):\n'
+ ' '.join(self.things_count)
+ '\n')
LOOP.call_later(10, self.update_count)
def swap_oldest(self, interval=conf.SWAP_OLDEST, minimum=conf.MINIMUM_RUNTIME):
if not self.paused and not self.extra_queue.empty():
oldest, minutes = self.longest_running()
if minutes > minimum:
LOOP.create_task(oldest.lock_and_swap(minutes))
LOOP.call_later(interval, self.swap_oldest)
def print_status(self, refresh=conf.REFRESH_RATE):
try:
self._print_status()
except CancelledError:
return
except Exception as e:
self.log.exception('{} occurred while printing status.', e.__class__.__name__)
self.print_handle = LOOP.call_later(refresh, self.print_status)
async def exit_progress(self):
while self.coroutines_count > 2:
try:
self.update_coroutines_count(simple=False)
pending = len(db_proc)
# Spaces at the end are important, as they clear previously printed
# output - \r doesn't clean whole line
print(
'{} coroutines active, {} DB items pending '.format(
self.coroutines_count, pending),
end='\r'
)
await sleep(.5)
except CancelledError:
return
except Exception as e:
self.log.exception('A wild {} appeared in exit_progress!', e.__class__.__name__)
def update_stats(self, refresh=conf.STAT_REFRESH, med=median, count=conf.GRID[0] * conf.GRID[1]):
visits = []
seen_per_worker = []
after_spawns = []
speeds = []
for w in self.workers:
after_spawns.append(w.after_spawn)
seen_per_worker.append(w.total_seen)
visits.append(w.visits)
speeds.append(w.speed)
self.stats = (
'Seen per worker: min {}, max {}, med {:.0f}\n'
'Visits per worker: min {}, max {}, med {:.0f}\n'
'Visit delay: min {:.1f}, max {:.1f}, med {:.1f}\n'
'Speed: min {:.1f}, max {:.1f}, med {:.1f}\n'
'Extra accounts: {}, CAPTCHAs needed: {}\n'
).format(
min(seen_per_worker), max(seen_per_worker), med(seen_per_worker),
min(visits), max(visits), med(visits),
min(after_spawns), max(after_spawns), med(after_spawns),
min(speeds), max(speeds), med(speeds),
self.extra_queue.qsize(), self.captcha_queue.qsize()
)
self.sighting_cache_size = len(SIGHTING_CACHE.store)
self.mystery_cache_size = len(MYSTERY_CACHE.store)
self.update_coroutines_count()
self.counts = (
'Known spawns: {}, unknown: {}, more: {}\n'
'{} workers, {} coroutines\n'
'sightings cache: {}, mystery cache: {}, DB queue: {}\n'
).format(
len(spawns), len(spawns.unknown), spawns.cells_count,
count, self.coroutines_count,
len(SIGHTING_CACHE), len(MYSTERY_CACHE), len(db_proc)
)
LOOP.call_later(refresh, self.update_stats)
def get_dots_and_messages(self):
"""Returns status dots and status messages for workers
Dots meaning:
. = visited more than a minute ago
, = visited less than a minute ago, no pokemon seen
0 = visited less than a minute ago, no pokemon or forts seen
: = visited less than a minute ago, pokemon seen
! = currently visiting
| = cleaning bag
$ = spinning a PokéStop
* = sending a notification
~ = encountering a Pokémon
I = initial, haven't done anything yet
» = waiting to log in (limited by SIMULTANEOUS_LOGINS)
° = waiting to start app simulation (limited by SIMULTANEOUS_SIMULATION)
∞ = bootstrapping
L = logging in
A = simulating app startup
T = completing the tutorial
X = something bad happened
C = CAPTCHA
Other letters: various errors and procedures
"""
dots = []
messages = []
row = []
for i, worker in enumerate(self.workers):
if i > 0 and i % conf.GRID[1] == 0:
dots.append(row)
row = []
if worker.error_code in BAD_STATUSES:
row.append('X')
messages.append(worker.status.ljust(20))
elif worker.error_code:
row.append(worker.error_code[0])
else:
row.append('.')
if row:
dots.append(row)
return dots, messages
def update_coroutines_count(self, simple=True, loop=LOOP):
try:
tasks = Task.all_tasks(loop)
self.coroutines_count = len(tasks) if simple else sum(not t.done() for t in tasks)
except RuntimeError:
# Set changed size during iteration
self.coroutines_count = '-1'
def _print_status(self, _ansi=ANSI, _start=datetime.now(), _notify=conf.NOTIFY):
running_for = datetime.now() - _start
seconds_since_start = running_for.seconds - self.idle_seconds or 0.1
hours_since_start = seconds_since_start / 3600
output = [
'{}Monocle running for {}'.format(_ansi, running_for),
self.counts,
self.stats,
self.pokemon_found,
('Visits: {}, per second: {:.2f}\n'
'Skipped: {}, unnecessary: {}').format(
self.visits, self.visits / seconds_since_start,
self.skipped, self.redundant)
]
try:
seen = Worker.g['seen']
captchas = Worker.g['captchas']
output.append('Seen per visit: {v:.2f}, per minute: {m:.0f}'.format(
v=seen / self.visits, m=seen / (seconds_since_start / 60)))
if captchas:
captchas_per_request = captchas / (self.visits / 1000)
captchas_per_hour = captchas / hours_since_start
output.append(
'CAPTCHAs per 1K visits: {r:.1f}, per hour: {h:.1f}, total: {t:d}'.format(
r=captchas_per_request, h=captchas_per_hour, t=captchas))
except ZeroDivisionError:
pass
try:
hash_status = HashServer.status
output.append('Hashes: {}/{}, refresh in {:.0f}'.format(
hash_status['remaining'],
hash_status['maximum'],
hash_status['period'] - time()
))
except (KeyError, TypeError):
pass
if _notify:
sent = Worker.notifier.sent
output.append('Notifications sent: {}, per hour {:.1f}'.format(
sent, sent / hours_since_start))
output.append('')
if not self.all_seen:
no_sightings = ', '.join(str(w.worker_no)
for w in self.workers
if w.total_seen == 0)
if no_sightings:
output += ['Workers without sightings so far:', no_sightings, '']
else:
self.all_seen = True
dots, messages = self.get_dots_and_messages()
output += [' '.join(row) for row in dots]
previous = 0
for i in range(4, len(messages) + 4, 4):
output.append('\t'.join(messages[previous:i]))
previous = i
if self.paused:
output.append('\nCAPTCHAs are needed to proceed.')
if not _ansi:
system('cls')
print('\n'.join(output))
def longest_running(self):
workers = (x for x in self.workers if x.start_time)
worker = next(workers)
earliest = worker.start_time
for w in workers:
if w.start_time < earliest:
worker = w
earliest = w.start_time
minutes = ((time() * 1000) - earliest) / 60000
return worker, minutes
def get_start_point(self):
smallest_diff = float('inf')
now = time() % 3600
closest = None
for spawn_id, spawn_time in spawns.known.values():
time_diff = now - spawn_time
if 0 < time_diff < smallest_diff:
smallest_diff = time_diff
closest = spawn_id
if smallest_diff < 3:
break
return closest
async def update_spawns(self, initial=False):
while True:
try:
await run_threaded(spawns.update)
LOOP.create_task(run_threaded(spawns.pickle))
except OperationalError as e:
self.log.exception('Operational error while trying to update spawns.')
if initial:
raise OperationalError('Could not update spawns, ensure your DB is set up.') from e
await sleep(15, loop=LOOP)
except CancelledError:
raise
except Exception as e:
self.log.exception('A wild {} appeared while updating spawns!', e.__class__.__name__)
await sleep(15, loop=LOOP)
else:
break
async def launch(self, bootstrap, pickle):
exceptions = 0
self.next_mystery_reload = 0
if not pickle or not spawns.unpickle():
await self.update_spawns(initial=True)
if not spawns or bootstrap:
try:
await self.bootstrap()
await self.update_spawns()
except CancelledError:
return
update_spawns = False
self.mysteries = spawns.mystery_gen()
while True:
try:
await self._launch(update_spawns)
update_spawns = True
except CancelledError:
return
except Exception:
exceptions += 1
if exceptions > 25:
self.log.exception('Over 25 errors occured in launcher loop, exiting.')
return False
else:
self.log.exception('Error occured in launcher loop.')
update_spawns = False
async def _launch(self, update_spawns):
if update_spawns:
await self.update_spawns()
LOOP.create_task(run_threaded(dump_pickle, 'accounts', ACCOUNTS))
spawns_iter = iter(spawns.items())
else:
start_point = self.get_start_point()
if start_point and not spawns.after_last():
spawns_iter = dropwhile(
lambda s: s[1][0] != start_point, spawns.items())
else:
spawns_iter = iter(spawns.items())
current_hour = get_current_hour()
if spawns.after_last():
current_hour += 3600
captcha_limit = conf.MAX_CAPTCHAS
skip_spawn = conf.SKIP_SPAWN
for point, (spawn_id, spawn_seconds) in spawns_iter:
try:
if self.captcha_queue.qsize() > captcha_limit:
self.paused = True
self.idle_seconds += await run_threaded(self.captcha_queue.full_wait, conf.MAX_CAPTCHAS)
self.paused = False
except (EOFError, BrokenPipeError, FileNotFoundError):
pass
spawn_time = spawn_seconds + current_hour
# negative = hasn't happened yet
# positive = already happened
time_diff = time() - spawn_time
while time_diff < 0.5:
try:
mystery_point = next(self.mysteries)
await self.coroutine_semaphore.acquire()
LOOP.create_task(self.try_point(mystery_point))
except StopIteration:
if self.next_mystery_reload < monotonic():
self.mysteries = spawns.mystery_gen()
self.next_mystery_reload = monotonic() + conf.RESCAN_UNKNOWN
else:
await sleep(min(spawn_time - time() + .5, self.next_mystery_reload - monotonic()), loop=LOOP)
time_diff = time() - spawn_time
if time_diff > 5 and spawn_id in SIGHTING_CACHE.store:
self.redundant += 1
continue
elif time_diff > skip_spawn:
self.skipped += 1
continue
await self.coroutine_semaphore.acquire()
LOOP.create_task(self.try_point(point, spawn_time, spawn_id))
async def try_again(self, point):
async with self.coroutine_semaphore:
worker = await self.best_worker(point, False)
async with worker.busy:
if await worker.visit(point):
self.visits += 1
async def bootstrap(self):
try:
self.log.warning('Starting bootstrap phase 1.')
await self.bootstrap_one()
except CancelledError:
raise
except Exception:
self.log.exception('An exception occurred during bootstrap phase 1.')
try:
self.log.warning('Starting bootstrap phase 2.')
await self.bootstrap_two()
except CancelledError:
raise
except Exception:
self.log.exception('An exception occurred during bootstrap phase 2.')
self.log.warning('Starting bootstrap phase 3.')
unknowns = list(spawns.unknown)
shuffle(unknowns)
tasks = (self.try_again(point) for point in unknowns)
await gather(*tasks, loop=LOOP)
self.log.warning('Finished bootstrapping.')
async def bootstrap_one(self):
async def visit_release(worker, num, *args):
async with self.coroutine_semaphore:
async with worker.busy:
point = get_start_coords(num, *args)
self.log.warning('start_coords: {}', point)
self.visits += await worker.bootstrap_visit(point)
if bounds.multi:
areas = [poly.polygon.area for poly in bounds.polygons]
area_sum = sum(areas)
percentages = [area / area_sum for area in areas]
tasks = []
for i, workers in enumerate(percentage_split(
self.workers, percentages)):
grid = best_factors(len(workers))
tasks.extend(visit_release(w, n, grid, bounds.polygons[i])
for n, w in enumerate(workers))
else:
tasks = (visit_release(w, n) for n, w in enumerate(self.workers))
await gather(*tasks, loop=LOOP)
async def bootstrap_two(self):
async def bootstrap_try(point):
async with self.coroutine_semaphore:
randomized = randomize_point(point, randomization)
LOOP.call_later(1790, LOOP.create_task, self.try_again(randomized))
worker = await self.best_worker(point, False)
async with worker.busy:
self.visits += await worker.bootstrap_visit(point)
# randomize to within ~140m of the nearest neighbor on the second visit
randomization = conf.BOOTSTRAP_RADIUS / 155555 - 0.00045
tasks = (bootstrap_try(x) for x in get_bootstrap_points(bounds))
await gather(*tasks, loop=LOOP)
async def try_point(self, point, spawn_time=None, spawn_id=None):
try:
point = randomize_point(point)
skip_time = monotonic() + (conf.GIVE_UP_KNOWN if spawn_time else conf.GIVE_UP_UNKNOWN)
worker = await self.best_worker(point, skip_time)
if not worker:
if spawn_time:
self.skipped += 1
return
async with worker.busy:
if spawn_time:
worker.after_spawn = time() - spawn_time
if await worker.visit(point, spawn_id):
self.visits += 1
except CancelledError:
raise
except Exception:
self.log.exception('An exception occurred in try_point')
finally:
self.coroutine_semaphore.release()
async def best_worker(self, point, skip_time):
good_enough = conf.GOOD_ENOUGH
while self.running:
gen = (w for w in self.workers if not w.busy.locked())
try:
worker = next(gen)
lowest_speed = worker.travel_speed(point)
except StopIteration:
lowest_speed = float('inf')
for w in gen:
speed = w.travel_speed(point)
if speed < lowest_speed:
lowest_speed = speed
worker = w
if speed < good_enough:
break
if lowest_speed < conf.SPEED_LIMIT:
worker.speed = lowest_speed
return worker
if skip_time and monotonic() > skip_time:
return None
await sleep(conf.SEARCH_SLEEP, loop=LOOP)
def refresh_dict(self):
while not self.extra_queue.empty():
account = self.extra_queue.get()
username = account['username']
ACCOUNTS[username] = account
class scheduler(threading.Thread):
def __init__(self, worker_q, ws, map_jobs, n_reducer=30, url=""):
'''
worker_q: we put availabe workers into this priority queue.
ws: a websocket connection object.
map_jobs: a list storing all map tasks
(more precisely: which part of the input file)
n_reducer: Number of reducer.
'''
super(scheduler, self).__init__()
self.url_list = []
self.url = url
self.mutex = Semaphore()
self.map_jobs = map_jobs
self.n_reducer = n_reducer
self.stoprequest = threading.Event()
self.worker_q = worker_q
self.ws = ws
self.dead_worker = set()
self.mapCount = 0
self.reduceCount = n_reducer
# a dict used to track status of all map/reduce jobs,
# entities in map_jobs are keys,
# a list of worker(if not finished) or None object(if finished)
# is the corresponding value
self.map_status = {}
self.reduce_status = {}
self.tid_map = {}
for job in map_jobs:
self.map_status[job] = []
for i in range(n_reducer):
self.reduce_status[i] = []
def removeWorker(self, uid):
self.dead_worker.add(uid)
def jobFinished(self, tid, url, type):
'''
Nth slice of the map job is finished,
mark them as done in map_status
'''
if type == "m":
if self.mapCount == 0:
return
self.mutex.acquire()
n = self.tid_map[tid]
if self.map_status[n] == None:
self.mutex.release()
return
for worker in self.map_status[n]:
self.worker_q.put(worker)
self.url_list.append(url)
self.map_status[n] = None
self.mapCount -= 1
self.map_jobs.remove(n)
self.mutex.release()
return
if type == "r":
self.mutex.acquire()
n = self.tid_map[tid]
if self.reduce_status[n] == None:
self.mutex.release()
return
for worker in self.reduce_status[n]:
self.worker_q.put(worker)
self.reduce_status[n] = None
self.reduceCount -= 1
self.mutex.release()
return
def schedule_reduce(self):
print("start reduce")
counter = 0
while (True):
new_worker = self.worker_q.get(True)
if new_worker[1] in self.dead_worker:
self.dead_worker.remove(new_worker[1])
continue
self.mutex.acquire()
if self.reduceCount == 0:
return
counter %= self.n_reducer
while (self.reduce_status[counter] == None):
counter += 1
self.reduce_status[counter].append(new_worker)
self.mutex.release()
tid = (int(time.time() * 1000) + new_worker[0])
self.tid_map[tid] = counter
data = {
'type': 'r',
'uid': new_worker[1],
'tid': tid,
'slice': counter,
'url': self.url_list
}
self.ws.send(msg_generator(1, "", "task", data))
counter += 1
def schedule_map(self):
'''
Map tasks are scheduled in this function,
it will return iff all map tasks are finished.
'''
self.mapCount = len(self.map_jobs)
counter = 0
while (True):
new_worker = self.worker_q.get(True)
if new_worker[1] in self.dead_worker:
self.dead_worker.remove(new_worker[1])
continue
self.mutex.acquire()
if self.mapCount == 0:
self.worker_q.put(new_worker)
return
counter %= self.mapCount
job = self.map_jobs[counter]
counter += 1
self.map_status[job].append(new_worker)
self.mutex.release()
tid = (int(time.time() * 1000) + new_worker[0])
#tid = counter
self.tid_map[tid] = job
data = {
'type': 'm',
'uid': new_worker[1],
'tid': tid,
'slice': job,
'url': [self.url]
}
self.ws.send(msg_generator(1, "", "task", data))
def run(self):
while not self.stoprequest.isSet():
self.schedule_map()
print("map done!")
self.schedule_reduce()
print("reduce done!")
class TaskGroup:
'''A class representing a group of executing tasks. tasks is an optional set of existing
tasks to put into the group. New tasks can later be added using the spawn() method
below.
wait specifies the policy used for waiting for tasks by the join() method. If wait is
all then wait for all tasks to complete. If wait is any then wait for any task to
complete and then cancel tasks that are still running. If wait is object then wait
for the first task to return a non-None result and cancel tasks that are still
runnning. None means wait for no tasks and cancel all still running.
When join() is called, if any of the tasks in the group raises an exception or is
cancelled then all tasks in the group, including daemon tasks, are cancelled. If the
join() operation itself is cancelled then all running tasks in the group are also
cancelled. Once join() returns all tasks have completed and new tasks may not be
added. Tasks can be added while join() is waiting.
A TaskGroup is often used as a context manager, which calls the join() method on
context-exit. Each TaskGroup is an independent entity. Task groups do not form a
hierarchy or any kind of relationship to other previously created task groups or
tasks. Moreover, Tasks created by the top level spawn() function are not placed into
any task group. To create a task in a group, it should be created using
TaskGroup.spawn() or explicitly added using TaskGroup.add_task().
A task group has the following public attributes:
completed: initially None, and set by join() to the first task in the group that
finished. Tasks removed from the group by calls to next_done() (and if wait is object
tasks returning None) do not count.
joined: true if the task group join() operation has completed
daemons: a set of all running daemonic tasks in the group.
tasks: a set of all non-daemonic tasks in the group.
'''
def __init__(self, tasks=(), *, wait=all):
if wait not in (any, all, object, None):
raise ValueError('invalid wait argument')
# Tasks that have not yet finished
self._pending = set()
# All non-daemonic tasks tracked by the group
self.tasks = set()
# All running deamonic tasks in the group
self.daemons = set()
# Non-daemonic tasks that have completed
self._done = deque()
self._wait = wait
self.joined = False
self._semaphore = Semaphore(0)
self.completed = None
for task in tasks:
self._add_task(task)
def _on_done(self, task):
task._task_group = None
if getattr(task, '_daemon', False):
self.daemons.discard(task)
else:
self._pending.discard(task)
self._done.append(task)
self._semaphore.release()
def _add_task(self, task):
'''Add an already existing task to the task group.'''
if hasattr(task, '_task_group'):
raise RuntimeError('task is already part of a group')
if self.joined:
raise RuntimeError('task group terminated')
task._task_group = self
daemon = getattr(task, '_daemon', False)
if not daemon:
self.tasks.add(task)
if task.done():
self._on_done(task)
elif daemon:
self.daemons.add(task)
else:
self._pending.add(task)
task.add_done_callback(self._on_done)
@property
def result(self):
''' The result of the first completed task. Should only be called after join()
has returned.'''
if not self.joined:
raise RuntimeError('task group not yet terminated')
if not self.completed:
raise RuntimeError('no task successfully completed')
return self.completed.result()
@property
def exception(self):
''' The exception of the first completed task. Should only be called after join()
has returned.'''
if not self.joined:
raise RuntimeError('task group not yet terminated')
return safe_exception(self.completed) if self.completed else None
@property
def results(self):
'''A list of all results collected by join() in no particular order.
If a task raised an exception or was cancelled then that exception will be raised.
'''
if not self.joined:
raise RuntimeError('task group not yet terminated')
return [task.result() for task in self.tasks]
@property
def exceptions(self):
'''A list of all exceptions collected by join() in no particular order.'''
if not self.joined:
raise RuntimeError('task group not yet terminated')
return [safe_exception(task) for task in self.tasks]
async def spawn(self, coro, *args, daemon=False):
'''Create a new task and put it in the group. Returns a Task instance.
Daemonic tasks are both ignored and cancelled by join().
'''
task = await spawn(coro, *args, daemon=daemon)
self._add_task(task)
return task
async def add_task(self, task):
'''Add an already existing task to the task group.'''
self._add_task(task)
async def next_done(self):
'''Return the next completed task and remove it from the group. Return None if no more
tasks remain. A TaskGroup may also be used as an asynchronous iterator.
'''
if self._done or self._pending:
await self._semaphore.acquire()
if self._done:
return self._done.popleft()
return None
async def next_result(self):
'''Return the result of the next completed task and remove it from the group. If the task
failed with an exception, that exception is raised. A RuntimeError exception is
raised if no tasks remain.
'''
task = await self.next_done()
if not task:
raise NoRemainingTasksError('no tasks remain')
return task.result()
async def join(self):
'''Wait for tasks in the group to terminate according to the wait policy for the group.
'''
try:
# Wait for no-one; all tasks are cancelled
if self._wait is None:
return
while True:
task = await self.next_done()
if task is None:
return
# Set self.completed if not yet set; unless wait is object and
if self.completed is None:
if not (self._wait is object and not safe_exception(task)
and task.result() is None):
self.completed = task
if (safe_exception(task) or self._wait is any
or (self._wait is object and self.completed)):
return
finally:
self.joined = True
# Cancel everything but don't block
await self._cancel_tasks(self._pending.union(self.daemons), False)
# Ensure the event loop has processed the cancellations when we return
# This is mainly for no-surprises and cleanliness, including in the testsuite
await sleep(0)
async def _cancel_tasks(self, tasks, blocking):
'''Cancel the passed set of tasks. Wait for them to complete if blocking.'''
for task in tasks:
task.cancel()
if blocking and tasks:
def pop_task(task):
unfinished.remove(task)
if not unfinished:
all_done.set()
unfinished = set(tasks)
all_done = Event()
for task in tasks:
task.add_done_callback(pop_task)
await all_done.wait()
async def cancel_remaining(self):
'''Cancel all remaining non-daemonic tasks and wait for them to complete.
If any task blocks cancellation this routine will not return.
'''
await self._cancel_tasks(self._pending, True)
def __aiter__(self):
return self
async def __anext__(self):
task = await self.next_done()
if task:
return task
raise StopAsyncIteration
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc_value, traceback):
if exc_type:
await self.cancel_remaining()
await self.join()
async def test_multi_upgrade_lockout(postgresql_pool, get_columns_in_db_table,
hard_clean_db):
async with postgresql_pool.acquire() as postgresql_client:
async with postgresql_pool.acquire() as postgresql_client2:
# schedule 3 updates, hang on second, unblock one, verify, unblock other, verify
corev: Set[int] = await get_core_versions(postgresql_client)
db_schema = schema.DBSchema("test_multi_upgrade_lockout",
inmanta.db.versions, postgresql_client)
db_schema2 = schema.DBSchema("test_multi_upgrade_lockout",
inmanta.db.versions,
postgresql_client2)
await db_schema.ensure_self_update()
lock = Semaphore(0)
async def update_function_a(connection):
# Fix syntax issue
await connection.execute(
"CREATE TABLE public.taba(id integer primary key, val varchar NOT NULL);"
)
async def update_function_b(connection):
# Syntax error should trigger database rollback
await lock.acquire()
await connection.execute(
"CREATE TABLE public.tabb(id integer primary key, val varchar NOT NULL);"
)
async def update_function_c(connection):
# Fix syntax issue
await connection.execute(
"CREATE TABLE public.tabc(id integer primary key, val varchar NOT NULL);"
)
current_db_versions: Set[
int] = await db_schema.get_installed_versions()
assert len(current_db_versions) == 0
update_function_map = make_versions(1, update_function_a,
update_function_b,
update_function_c)
r1 = asyncio.ensure_future(
db_schema._update_db_schema(update_function_map))
r2 = asyncio.ensure_future(
db_schema2._update_db_schema(update_function_map))
both = asyncio.as_completed([r1, r2]).__iter__()
await asyncio.sleep(0.1)
first = next(both)
lock.release()
await first
# second one doesn't even hit the lock, as it never sees schema version 0
# lock.release()
second = next(both)
await second
# Assert done
assert (await db_schema.get_installed_versions()) == {1, 2, 3}
assert (await postgresql_client.fetchval(
"SELECT table_name FROM information_schema.tables "
"WHERE table_schema='public' AND table_name='taba'"
)) is not None
assert (await postgresql_client.fetchval(
"SELECT table_name FROM information_schema.tables "
"WHERE table_schema='public' AND table_name='tabb'"
)) is not None
assert (await postgresql_client.fetchval(
"SELECT table_name FROM information_schema.tables "
"WHERE table_schema='public' AND table_name='tabc'"
)) is not None
await assert_core_untouched(postgresql_client, corev)
class AsyncioSubscriptionManager(SubscriptionManager):
def __init__(self, pubnub_instance):
subscription_manager = self
self._message_worker = None
self._message_queue = Queue()
self._subscription_lock = Semaphore(1)
self._subscribe_loop_task = None
self._heartbeat_periodic_callback = None
self._reconnection_manager = AsyncioReconnectionManager(pubnub_instance)
super(AsyncioSubscriptionManager, self).__init__(pubnub_instance)
self._start_worker()
class AsyncioReconnectionCallback(ReconnectionCallback):
def on_reconnect(self):
subscription_manager.reconnect()
pn_status = PNStatus()
pn_status.category = PNStatusCategory.PNReconnectedCategory
pn_status.error = False
subscription_manager._subscription_status_announced = True
subscription_manager._listener_manager.announce_status(pn_status)
self._reconnection_listener = AsyncioReconnectionCallback()
self._reconnection_manager.set_reconnection_listener(self._reconnection_listener)
def _set_consumer_event(self):
if not self._message_worker.cancelled():
self._message_worker.cancel()
def _message_queue_put(self, message):
self._message_queue.put_nowait(message)
def _start_worker(self):
consumer = AsyncioSubscribeMessageWorker(self._pubnub,
self._listener_manager,
self._message_queue, None)
self._message_worker = asyncio.ensure_future(consumer.run(),
loop=self._pubnub.event_loop)
def reconnect(self):
# TODO: method is synchronized in Java
self._should_stop = False
self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop())
self._register_heartbeat_timer()
def disconnect(self):
# TODO: method is synchronized in Java
self._should_stop = True
self._stop_heartbeat_timer()
self._stop_subscribe_loop()
def stop(self):
super(AsyncioSubscriptionManager, self).stop()
self._reconnection_manager.stop_polling()
if self._subscribe_loop_task is not None and not self._subscribe_loop_task.cancelled():
self._subscribe_loop_task.cancel()
@asyncio.coroutine
def _start_subscribe_loop(self):
self._stop_subscribe_loop()
yield from self._subscription_lock.acquire()
combined_channels = self._subscription_state.prepare_channel_list(True)
combined_groups = self._subscription_state.prepare_channel_group_list(True)
if len(combined_channels) == 0 and len(combined_groups) == 0:
self._subscription_lock.release()
return
self._subscribe_request_task = asyncio.ensure_future(Subscribe(self._pubnub)
.channels(combined_channels)
.channel_groups(combined_groups)
.timetoken(self._timetoken).region(self._region)
.filter_expression(self._pubnub.config.filter_expression)
.future())
e = yield from self._subscribe_request_task
if self._subscribe_request_task.cancelled():
self._subscription_lock.release()
return
if e.is_error():
if e.status is not None and e.status.category == PNStatusCategory.PNCancelledCategory:
self._subscription_lock.release()
return
if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory:
self._pubnub.event_loop.call_soon(self._start_subscribe_loop)
self._subscription_lock.release()
return
logger.error("Exception in subscribe loop: %s" % str(e))
if e.status is not None and e.status.category == PNStatusCategory.PNAccessDeniedCategory:
e.status.operation = PNOperationType.PNUnsubscribeOperation
# TODO: raise error
self._listener_manager.announce_status(e.status)
self._reconnection_manager.start_polling()
self._subscription_lock.release()
self.disconnect()
return
else:
self._handle_endpoint_call(e.result, e.status)
self._subscription_lock.release()
self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop())
self._subscription_lock.release()
def _stop_subscribe_loop(self):
if self._subscribe_request_task is not None and not self._subscribe_request_task.cancelled():
self._subscribe_request_task.cancel()
def _stop_heartbeat_timer(self):
if self._heartbeat_periodic_callback is not None:
self._heartbeat_periodic_callback.stop()
def _register_heartbeat_timer(self):
super(AsyncioSubscriptionManager, self)._register_heartbeat_timer()
self._heartbeat_periodic_callback = AsyncioPeriodicCallback(
self._perform_heartbeat_loop,
self._pubnub.config.heartbeat_interval * 1000,
self._pubnub.event_loop)
if not self._should_stop:
self._heartbeat_periodic_callback.start()
@asyncio.coroutine
def _perform_heartbeat_loop(self):
if self._heartbeat_call is not None:
# TODO: cancel call
pass
cancellation_event = Event()
state_payload = self._subscription_state.state_payload()
presence_channels = self._subscription_state.prepare_channel_list(False)
presence_groups = self._subscription_state.prepare_channel_group_list(False)
if len(presence_channels) == 0 and len(presence_groups) == 0:
return
try:
heartbeat_call = (Heartbeat(self._pubnub)
.channels(presence_channels)
.channel_groups(presence_groups)
.state(state_payload)
.cancellation_event(cancellation_event)
.future())
envelope = yield from heartbeat_call
heartbeat_verbosity = self._pubnub.config.heartbeat_notification_options
if envelope.status.is_error:
if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL or \
heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL:
self._listener_manager.announce_status(envelope.status)
else:
if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL:
self._listener_manager.announce_status(envelope.status)
except PubNubAsyncioException as e:
pass
# TODO: check correctness
# if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory:
# self._start_subscribe_loop()
# else:
# self._listener_manager.announce_status(e.status)
finally:
cancellation_event.set()
def _send_leave(self, unsubscribe_operation):
asyncio.ensure_future(self._send_leave_helper(unsubscribe_operation))
@asyncio.coroutine
def _send_leave_helper(self, unsubscribe_operation):
envelope = yield from Leave(self._pubnub) \
.channels(unsubscribe_operation.channels) \
.channel_groups(unsubscribe_operation.channel_groups).future()
self._listener_manager.announce_status(envelope.status)
class WorkerRaider(Worker):
workers = []
gyms = {}
gym_scans = 0
skipped = 0
visits = 0
hash_burn = 0
workers_needed = 0
job_queue = PriorityQueue()
last_semaphore_value = workers_needed
coroutine_semaphore = Semaphore(len(workers), loop=LOOP)
def __init__(self,
worker_no,
overseer,
captcha_queue,
account_queue,
worker_dict,
account_dict,
start_coords=None):
super().__init__(worker_no,
overseer,
captcha_queue,
account_queue,
worker_dict,
account_dict,
start_coords=start_coords)
self.scan_delayed = 0
def needs_sleep(self):
return False
def get_start_coords(self):
return bounds.center
def required_extra_accounts(self):
return super().required_extra_accounts() + self.workers_needed
@classmethod
def preload(self):
log.info("Preloading forts")
with session_scope() as session:
forts = session.query(Fort) \
.options(joinedload(Fort.sightings)) \
.filter(Fort.lat.between(bounds.south, bounds.north),
Fort.lon.between(bounds.west, bounds.east))
try:
for fort in forts:
if (fort.lat, fort.lon) not in bounds:
continue
obj = {
'id': fort.id,
'external_id': fort.external_id,
'lat': fort.lat,
'lon': fort.lon,
'name': fort.name,
'url': fort.url,
'last_modified': 0,
'updated': 0,
}
if len(fort.sightings) > 0:
sighting = fort.sightings[0]
obj['last_modified'] = sighting.last_modified or 0
obj['updated'] = sighting.updated or 0
self.add_gym(obj)
except Exception as e:
log.error("ERROR: {}", e)
log.info("Loaded {} forts", self.job_queue.qsize())
@classmethod
def add_job(self, gym):
updated = gym.get('updated', gym.get('last_modified', 0)) or 0
self.job_queue.put_nowait((updated, random(), gym))
@classmethod
def add_gym(self, gym):
if gym['external_id'] in self.gyms:
return
self.gyms[gym['external_id']] = {'miss': 0}
self.workers_needed = int(ceil(conf.RAIDERS_PER_GYM * len(self.gyms)))
if len(self.workers) < self.workers_needed:
try:
self.workers.append(
WorkerRaider(worker_no=len(self.workers),
overseer=self.overseer,
captcha_queue=self.overseer.captcha_queue,
account_queue=self.overseer.extra_queue,
worker_dict=self.overseer.worker_dict,
account_dict=self.overseer.account_dict))
except Exception as e:
log.error("WorkerRaider initialization error: {}", e)
traceback.print_exc()
self.add_job(gym)
@classmethod
def obliterate_gym(self, gym):
external_id = gym['external_id']
if external_id not in self.gyms:
return
with session_scope() as session:
fort_id = get_fort_internal_id(session, external_id)
if not fort_id:
return
session.query(GymDefender).filter(
GymDefender.fort_id == fort_id).delete()
session.query(Raid).filter(Raid.fort_id == fort_id).delete()
session.query(FortSighting).filter(
FortSighting.fort_id == fort_id).delete()
session.query(Fort).filter(Fort.id == fort_id).delete()
del self.gyms[external_id]
FORT_CACHE.remove_gym(external_id)
log.warning("Fort {} obliterated.", external_id)
@classmethod
async def launch(self, overseer):
self.overseer = overseer
self.preload()
try:
await sleep(5)
log.info("Couroutine launched.")
log.info("WorkerRaider count: ({}/{})", len(self.workers),
self.workers_needed)
while True:
try:
while self.last_semaphore_value > 0 and self.last_semaphore_value == len(
self.workers) and not self.job_queue.empty():
priority_job = self.job_queue.get()
updated = priority_job[0]
job = priority_job[2]
log.debug("Job: {}", job)
if (time() - updated) < 30:
await sleep(1)
self.add_job(job)
continue
await self.coroutine_semaphore.acquire()
LOOP.create_task(self.try_point(job))
except CancelledError:
raise
except Exception as e:
log.warning("A wild error appeared in launcher loop: {}",
e)
worker_count = len(self.workers)
if self.last_semaphore_value != worker_count:
self.last_semaphore_value = worker_count
self.coroutine_semaphore = Semaphore(worker_count,
loop=LOOP)
log.info("Semaphore updated with value {}", worker_count)
await sleep(1)
except CancelledError:
log.info("Coroutine cancelled.")
except Exception as e:
log.warning("A wild error appeared in launcher: {}", e)
@classmethod
async def try_point(self, job):
try:
point = (job['lat'], job['lon'])
fort_external_id = job['external_id']
updated = job.get('updated', job.get('last_modified', 0)) or 0
point = randomize_point(point,
amount=0.00003) # jitter around 3 meters
skip_time = monotonic() + (conf.SEARCH_SLEEP)
worker = await self.best_worker(point, job, updated, skip_time)
if not worker:
return
async with worker.busy:
visit_result = await worker.visit(point, gym=job)
if visit_result == -1:
self.hash_burn += 1
await sleep(1.0, loop=LOOP)
point = randomize_point(
point, amount=0.00001) # jitter around 3 meters
visit_result = await worker.visit(point, gym=job)
if visit_result:
if visit_result == -1:
miss = self.gyms[fort_external_id]['miss']
miss += 1
self.gyms[fort_external_id]['miss'] = miss
raise GymNotFoundError(
"Gym {} disappeared. Total misses: {}".format(
fort_external_id, miss))
else:
if worker and worker.account and 'gym_nothing_seen' in worker.account:
del worker.account['gym_nothing_seen']
self.gyms[fort_external_id]['miss'] = 0
now = int(time())
worker.scan_delayed = now - updated
job['updated'] = now
self.visits += 1
else:
if worker and worker.account:
username = worker.username
account_miss = worker.account.get(
'gym_nothing_seen', 0)
account_miss += 1
worker.account['gym_nothing_seen'] = account_miss
else:
username = None
account_miss = 1
raise NothingSeenAtGymSpotError(
"Nothing seen while scanning {} by {} for {} times.".
format(fort_external_id, username, account_miss))
except CancelledError:
raise
except (GymNotFoundError, NothingSeenAtGymSpotError) as e:
self.skipped += 1
if worker:
worker.log.error('Gym visit error: {}', e)
if isinstance(e, GymNotFoundError):
miss = self.gyms[fort_external_id]['miss']
if miss >= 10:
self.obliterate_gym(job)
if isinstance(e, NothingSeenAtGymSpotError):
if worker and worker.account:
account_miss = worker.account.get('gym_nothing_seen', 0)
await sleep(account_miss * 5, loop=LOOP)
except Exception as e:
self.skipped += 1
log.exception('An exception occurred in try_point: {}', e)
finally:
if fort_external_id in self.gyms:
if 'updated' in job:
job['updated'] += 5
self.add_job(job)
self.coroutine_semaphore.release()
@classmethod
async def best_worker(self, point, job, updated, skip_time):
while self.overseer.running:
gen = (w for w in self.workers if not w.busy.locked())
worker = None
try:
worker = next(gen)
lowest_speed = worker.travel_speed(point)
except StopIteration:
lowest_speed = float('inf')
for w in gen:
speed = w.travel_speed(point)
if speed < lowest_speed:
lowest_speed = speed
worker = w
time_diff = max(int(time() - updated), 0)
speed_factor = (1.0 + (time_diff / 10))
speed_limit = (conf.SPEED_LIMIT * speed_factor)
if worker and lowest_speed < speed_limit:
worker.speed = lowest_speed
return worker
if skip_time and monotonic() > skip_time:
return None
await sleep(conf.SEARCH_SLEEP, loop=LOOP)
class BoltPool:
""" A pool of connections to a single address.
:param opener: a function to which an address can be passed that
returns an open and ready Bolt connection
:param address: the remote address for which this pool operates
:param max_size: the maximum permitted number of simultaneous
connections that may be owned by this pool, both in-use and
free
:param max_age: the maximum permitted age, in seconds, for
connections to be retained in this pool
"""
def __init__(self, opener, address, max_size=1, max_age=None):
self._opener = opener
self._address = address
self._max_size = max_size
self._max_age = max_age
self._in_use_list = deque()
self._free_list = deque()
self._slots = Semaphore(self._max_size)
def __contains__(self, cx):
return cx in self._in_use_list or cx in self._free_list
def __len__(self):
return self.size
@property
def address(self):
""" The remote address for which this pool operates.
"""
return self._address
@property
def max_size(self):
""" The maximum permitted number of simultaneous connections
that may be owned by this pool, both in-use and free.
"""
return self._max_size
@property
def max_age(self):
""" The maximum permitted age, in seconds, for connections to
be retained in this pool.
"""
return self._max_age
@property
def in_use(self):
""" The number of connections in this pool that are currently
in use.
"""
return len(self._in_use_list)
@property
def free(self):
""" The number of free connections available in this connection
pool.
"""
return len(self._free_list)
@property
def size(self):
""" The number of connections (both in-use and free) currently
owned by this connection pool.
"""
return self.in_use + self.free
async def acquire(self):
""" Acquire a connection from the pool.
In the simplest case, this will return an existing open
connection, if one is free. If not, and the pool is not full,
a new connection will be created. If the pool is full and no
free connections are available, this will block until a
connection is released, or until the acquire call is cancelled.
"""
cx = None
while cx is None or cx.broken or cx.closed:
try:
# Plan A: select a free connection from the pool
cx = self._free_list.popleft()
except IndexError:
if self.size < self.max_size:
# Plan B: if the pool isn't full, open a new connection
cx = await self._opener(self.address)
else:
# Plan C: wait for an in-use connection to become available
await self._slots.acquire()
else:
expired = self.max_age is not None and cx.age > self.max_age
if expired:
await cx.close()
else:
await cx.reset(force=True)
self._in_use_list.append(cx)
return cx
async def release(self, cx):
""" Release a Bolt connection back into the pool.
:param cx: the connection to release
:raise ValueError: if the connection is not currently in use,
or if it does not belong to this pool
"""
if cx in self._in_use_list:
self._in_use_list.remove(cx)
await cx.reset()
self._free_list.append(cx)
self._slots.release()
elif cx in self._free_list:
raise ValueError("Connection is not in use")
else:
raise ValueError("Connection does not belong to this pool")
async def prune(self):
""" Close all free connections.
"""
await self.__close(self._free_list)
async def close(self):
""" Close all connections.
This does not permanently disable the connection pool, merely
ensures all open connections are shut down, including those in
use. It is perfectly acceptable to re-acquire connections after
pool closure, which will have the implicit affect of reopening
the pool.
"""
await self.prune()
await self.__close(self._in_use_list)
@classmethod
async def __close(cls, connections):
""" Close all connections in the given list.
"""
closers = deque()
while True:
try:
cx = connections.popleft()
except IndexError:
break
else:
closers.append(cx.close())
await wait(closers)
class WordleWaiter(Waiter.create([GroupMessage])):
""" wordle Waiter """
def __init__(self,
wordle_instance: Wordle,
group: Union[Group, int],
member: Optional[Union[Member, int]] = None):
self.wordle = wordle_instance
self.group = group if isinstance(group, int) else group.id
self.member = (member if isinstance(member, int) else
member.id) if member else None
self.member_list = set()
self.member_list_mutex = Semaphore(1)
async def detected_event(self, app: Ariadne, group: Group, member: Member,
message: MessageChain):
word = message.asDisplay().strip()
message_source = message.getFirst(Source)
if self.group == group.id and (self.member == member.id
or not self.member):
if message.asDisplay().strip() in ("/wordle -giveup",
"/wordle -g"):
dic = group_word_dic[group.id]
word_data = word_list[dic][len(
self.wordle.word)][self.wordle.word]
explain = '\n'.join(
[f"【{key}】：{word_data[key]}" for key in word_data])
await app.sendGroupMessage(
group,
MessageChain([
Image(data_bytes=self.wordle.get_board_bytes()),
Plain("很遗憾，没有人猜出来呢"
f"单词：{self.wordle.word}\n{explain}")
]),
quote=message_source)
await self.member_list_mutex.acquire()
for member in self.member_list:
await update_member_statistic(group, member,
StatisticType.lose)
await update_member_statistic(group, member,
StatisticType.game)
self.member_list_mutex.release()
await mutex.acquire()
group_running[group.id] = False
mutex.release()
return True
if message.asDisplay().strip() == "/wordle -hint":
await update_member_statistic(group, member,
StatisticType.hint)
hint = self.wordle.get_hint()
if not hint:
await app.sendGroupMessage(
group,
MessageChain("你还没有猜对过一个字母哦~再猜猜吧~"),
quote=message.getFirst(Source))
else:
await app.sendGroupMessage(
group,
MessageChain(
[Image(data_bytes=self.wordle.draw_hint())]),
quote=message.getFirst(Source))
return False
if len(word) == self.wordle.length and word.encode(
'utf-8').isalpha():
await self.member_list_mutex.acquire()
self.member_list.add(member.id)
self.member_list_mutex.release()
await self.wordle.draw_mutex.acquire()
print("required")
result = self.wordle.guess(word)
print(result)
print("released")
self.wordle.draw_mutex.release()
if not result:
return True
if result[0]:
await update_member_statistic(
group, member, StatisticType.correct
if result[1] else StatisticType.wrong)
await self.member_list_mutex.acquire()
for member in self.member_list:
await update_member_statistic(
group, member, StatisticType.win
if result[1] else StatisticType.lose)
await update_member_statistic(group, member,
StatisticType.game)
self.member_list_mutex.release()
dic = group_word_dic[group.id]
word_data = word_list[dic][len(
self.wordle.word)][self.wordle.word]
explain = '\n'.join(
[f"【{key}】：{word_data[key]}" for key in word_data])
await app.sendGroupMessage(
group,
MessageChain([
Image(data_bytes=self.wordle.get_board_bytes()),
Plain(
f"\n{'恭喜你猜出了单词！' if result[1] else '很遗憾，没有人猜出来呢'}\n"
f"【单词】：{self.wordle.word}\n{explain}")
]),
quote=message_source)
await mutex.acquire()
group_running[group.id] = False
mutex.release()
return True
elif not result[2]:
await app.sendGroupMessage(
group,
MessageChain(f"你确定 {word} 是一个合法的单词吗？"),
quote=message_source)
elif result[3]:
await app.sendGroupMessage(group,
MessageChain("你已经猜过这个单词了呢"),
quote=message_source)
else:
await update_member_statistic(group, member,
StatisticType.wrong)
await app.sendGroupMessage(
group,
MessageChain(
[Image(data_bytes=self.wordle.get_board_bytes())]),
quote=message_source)
return False
class PayloadStream:
def __init__(self, assembler: PayloadStreamAssembler):
self._assembler = assembler
self._buffer_queue: List[List[int]] = []
self._lock = Lock()
self._data_available = Semaphore(0)
self._producer_length = 0 # total length
self._consumer_position = 0 # read position
self._active: List[int] = []
self._active_offset = 0
self._end = False
def __len__(self):
return self._producer_length
def give_buffer(self, buffer: List[int]):
self._buffer_queue.append(buffer)
self._producer_length += len(buffer)
self._data_available.release()
def done_producing(self):
self.give_buffer([])
def write(self, buffer: List[int], offset: int, count: int):
buffer_copy = buffer[offset:offset + count]
self.give_buffer(buffer_copy)
async def read(self, buffer: List[int], offset: int, count: int):
if self._end:
return 0
if not self._active:
await self._data_available.acquire()
async with self._lock:
self._active = self._buffer_queue.pop(0)
available_count = min(len(self._active) - self._active_offset, count)
for index in range(available_count):
buffer[offset + index] = self._active[self._active_offset]
self._active_offset += 1
self._consumer_position += available_count
if self._active_offset >= len(self._active):
self._active = []
self._active_offset = 0
if (self._assembler
and self._consumer_position >= self._assembler.content_length):
self._end = True
return available_count
async def read_until_end(self):
result = [None] * self._assembler.content_length
current_size = 0
while not self._end:
count = await self.read(result, current_size,
self._assembler.content_length)
current_size += count
return result
def _write(self, content, over_semaphore: asyncio.Semaphore):
# 写操作（阻塞）
self._f.write(content)
# 让父协程从等待队列中唤醒
over_semaphore.release()
class Overseer:
def __init__(self, manager):
self.log = get_logger('overseer')
self.workers = []
self.manager = manager
self.things_count = deque(maxlen=9)
self.paused = False
self.coroutines_count = 0
self.skipped = 0
self.visits = 0
self.coroutine_semaphore = Semaphore(conf.COROUTINES_LIMIT, loop=LOOP)
self.redundant = 0
self.running = True
self.all_seen = False
self.idle_seconds = 0
self.log.info('Overseer initialized')
self.status_log_at = 0
self.pokemon_found = ''
self.login_semaphore = Semaphore(conf.SIMULTANEOUS_LOGINS, loop=LOOP)
self.sim_semaphore = Semaphore(conf.SIMULTANEOUS_SIMULATION, loop=LOOP)
def start(self, status_bar):
self.captcha_queue = self.manager.captcha_queue()
self.extra_queue = self.manager.extra_queue()
if conf.MAP_WORKERS:
self.worker_dict = self.manager.worker_dict()
else:
self.worker_dict = None
self.account_dict = get_accounts()
self.add_accounts_to_queue(self.account_dict, self.captcha_queue,
self.extra_queue)
for x in range(conf.GRID[0] * conf.GRID[1]):
try:
self.workers.append(
Worker(worker_no=x,
overseer=self,
captcha_queue=self.captcha_queue,
account_queue=self.extra_queue,
worker_dict=self.worker_dict,
account_dict=self.account_dict))
except Exception as e:
self.log.error("Worker initialization error: {}", e)
traceback.print_exc()
self.log.info("Worker count: ({}/{})", len(self.workers),
conf.GRID[0] * conf.GRID[1])
db_proc.Weather = Weather
db_proc.start()
LOOP.call_later(10, self.update_count)
LOOP.call_later(max(conf.SWAP_OLDEST, conf.MINIMUM_RUNTIME * 60),
self.swap_oldest)
LOOP.call_soon(self.update_stats)
if status_bar:
LOOP.call_soon(self.print_status)
def add_accounts_to_queue(self, account_dict, captcha_queue,
account_queue):
for username, account in account_dict.items():
account['username'] = username
if account.get('banned') or account.get('warn') or account.get(
'sbanned'):
continue
if account.get('captcha'):
captcha_queue.put(account)
else:
account_queue.put(account)
def update_count(self):
self.things_count.append(str(db_proc.count))
self.pokemon_found = ('Pokemon found count (10s interval):\n' +
' '.join(self.things_count) + '\n')
LOOP.call_later(10, self.update_count)
def swap_oldest(self,
interval=conf.SWAP_OLDEST,
minimum=conf.MINIMUM_RUNTIME):
if (not self.paused and conf.EXTRA_ACCOUNT_PERCENT > 0.0
and (Account.estimated_extra_accounts() > 0
or not self.extra_queue.empty())):
try:
oldest, minutes = self.longest_running()
if minutes > minimum:
LOOP.create_task(oldest.lock_and_swap(minutes))
except StopIteration as e:
pass
LOOP.call_later(interval, self.swap_oldest)
def print_status(self, refresh=conf.REFRESH_RATE):
try:
self._print_status()
except CancelledError:
return
except Exception as e:
self.log.exception('{} occurred while printing status.',
e.__class__.__name__)
self.print_handle = LOOP.call_later(refresh, self.print_status)
async def exit_progress(self):
while self.coroutines_count > 2:
try:
self.update_coroutines_count(simple=False)
pending = len(db_proc)
# Spaces at the end are important, as they clear previously printed
# output - \r doesn't clean whole line
print('{} coroutines active, {} DB items pending '.format(
self.coroutines_count, pending),
end='\r')
await sleep(.5)
except CancelledError:
return
except Exception as e:
self.log.exception('A wild {} appeared in exit_progress!',
e.__class__.__name__)
def update_stats(self,
refresh=conf.STAT_REFRESH,
med=median,
count=conf.GRID[0] * conf.GRID[1]):
visits = []
seen_per_worker = []
after_spawns = []
speeds = []
for w in self.workers:
after_spawns.append(w.after_spawn)
seen_per_worker.append(w.total_seen)
visits.append(w.visits)
speeds.append(w.speed)
try:
account_stats = Account.stats()
account_reasons = ', '.join(
['%s: %s' % (k, v) for k, v in account_stats[1].items()])
account_refresh = datetime.fromtimestamp(
account_stats[0]).strftime('%Y-%m-%d %H:%M:%S')
account_clean = account_stats[2].get('clean')
account_test = account_stats[2].get('test')
account30_clean = account_stats[2].get('clean30')
account30_test = account_stats[2].get('test30')
except Exception as e:
self.log.error("Unexpected error in overseer.update_stats: {}", e)
account_reasons = None
account_refresh = None
account_clean = None
account_test = None
account30_clean = None
account30_test = None
stats_template = (
'Seen per worker: min {}, max {}, med {:.0f}\n'
'Visits per worker: min {}, max {}, med {:.0f}\n'
'Visit delay: min {:.1f}, max {:.1f}, med {:.1f}\n'
'Speed: min {:.1f}, max {:.1f}, med {:.1f}\n'
'Worker30: {}, jobs: {}, caches: {}, encounters: {}, visits: {}, skips: {}, late: {}, hash wastes: {}\n'
'Extra accounts: {}, CAPTCHAs needed: {}\n'
'Accounts (this instance) {} (refreshed: {})\n'
'Accounts (DB-wide) fresh/clean: {}, hibernated: {}, (Lv.30) fresh/clean: {}, hibernated: {}\n'
)
try:
self.stats = stats_template.format(
min(seen_per_worker),
max(seen_per_worker),
med(seen_per_worker),
min(visits),
max(visits),
med(visits),
min(after_spawns),
max(after_spawns),
med(after_spawns),
min(speeds),
max(speeds),
med(speeds),
len(Worker30.workers),
Worker30.job_queue.qsize(),
len(ENCOUNTER_CACHE),
Worker30.encounters,
Worker30.visits,
Worker30.skipped,
Worker30.lates,
Worker30.hash_burn,
self.extra_queue.qsize(),
self.captcha_queue.qsize(),
account_reasons,
account_refresh,
account_clean,
account_test,
account30_clean,
account30_test,
)
except Exception as e:
self.stats = stats_template.format(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
None, None, 0, 0, 0, 0)
if Worker.has_raiders:
smallest = nsmallest(1, WorkerRaider.job_queue.queue)
if len(smallest) > 0:
oldest_gym_raided = int(
time() - smallest[0][0]) if len(smallest) > 0 else 0
else:
oldest_gym_raided = None
self.stats += 'Raider workers: {}, gyms: {}, queue: {}, oldest: {}s\n'.format(
len(WorkerRaider.workers), len(WorkerRaider.gyms),
WorkerRaider.job_queue.qsize(), oldest_gym_raided)
self.update_coroutines_count()
counts_template = (
'Known spawns: {}, unknown: {}, more: {}\n'
'workers: {}, coroutines: {}\n'
'sightings cache: {}, mystery cache: {}, DB queue: {}\n')
try:
self.counts = counts_template.format(len(spawns),
len(spawns.unknown),
spawns.cells_count, count,
self.coroutines_count,
len(SIGHTING_CACHE),
len(MYSTERY_CACHE),
len(db_proc))
except Exception as e:
self.counts = counts_template.format(0, 0, 0, 0, 0, 0, 0, 0)
if self.status_log_at < time() - 15.0:
self.status_log_at = time()
visits = """Visits: {}, Skipped: {}, Unnecessary: {}""".format(
self.visits, self.skipped, self.redundant)
for line in self.stats.split('\n'):
if line.strip():
self.log.info(line)
for line in self.counts.split('\n'):
if line.strip():
self.log.info(line)
for line in visits.split('\n'):
if line.strip():
self.log.info(line)
self.log.info(
"BorderCache: {}, MorePointTestCache: {}, Worker30SemLock: {}",
Worker.in_bounds.cache_info(), len(spawns.have_point_cache),
Worker30.coroutine_semaphore.locked())
LOOP.call_later(refresh, self.update_stats)
def get_dots_and_messages(self):
"""Returns status dots and status messages for workers
Dots meaning:
. = visited more than a minute ago
, = visited less than a minute ago, no pokemon seen
0 = visited less than a minute ago, no pokemon or forts seen
: = visited less than a minute ago, pokemon seen
! = currently visiting
| = cleaning bag
$ = spinning a PokéStop
* = sending a notification
~ = encountering a Pokémon
I = initial, haven't done anything yet
» = waiting to log in (limited by SIMULTANEOUS_LOGINS)
° = waiting to start app simulation (limited by SIMULTANEOUS_SIMULATION)
∞ = bootstrapping
L = logging in
A = simulating app startup
T = completing the tutorial
X = something bad happened
C = CAPTCHA
G = scanning a Gym for details
Other letters: various errors and procedures
"""
dots = []
messages = []
row = []
for i, worker in enumerate(self.workers):
if i > 0 and i % conf.GRID[1] == 0:
dots.append(row)
row = []
if worker.error_code in BAD_STATUSES:
row.append('X')
messages.append(worker.status.ljust(20))
elif worker.error_code:
row.append(worker.error_code[0])
else:
row.append('.')
if row:
dots.append(row)
return dots, messages
def update_coroutines_count(self, simple=True, loop=LOOP):
try:
tasks = Task.all_tasks(loop)
self.coroutines_count = len(tasks) if simple else sum(
not t.done() for t in tasks)
except RuntimeError:
# Set changed size during iteration
self.coroutines_count = '-1'
def _print_status(self,
_ansi=ANSI,
_start=datetime.now(),
_notify=conf.NOTIFY or conf.NOTIFY_RAIDS):
running_for = datetime.now() - _start
seconds_since_start = running_for.seconds - self.idle_seconds or 0.1
hours_since_start = seconds_since_start / 3600
try:
percent_skip = (self.skipped * 100) // (self.visits + self.skipped)
except (ZeroDivisionError):
percent_skip = '?'
try:
output = [
'{}Monocle/Alternate ({}) running for {}'.format(
_ansi, conf.INSTANCE_ID, running_for), self.counts,
self.stats, self.pokemon_found,
('Visits: {}, per second: {:.2f}\n'
'Skipped: {} ({}%), unnecessary: {}').format(
self.visits, self.visits / seconds_since_start,
self.skipped, percent_skip, self.redundant)
]
except (AttributeError):
output = []
try:
seen = Worker.g['seen']
captchas = Worker.g['captchas']
output.append(
'Seen per visit: {v:.2f}, per minute: {m:.0f}'.format(
v=seen / self.visits, m=seen / (seconds_since_start / 60)))
if captchas:
captchas_per_request = captchas / (self.visits / 1000)
captchas_per_hour = captchas / hours_since_start
output.append(
'CAPTCHAs per 1K visits: {r:.1f}, per hour: {h:.1f}, total: {t:d}'
.format(r=captchas_per_request,
h=captchas_per_hour,
t=captchas))
except ZeroDivisionError:
pass
try:
hash_status = HashServer.status
output.append('Hashes: {}/{}, refresh in {:.0f}'.format(
hash_status['remaining'], hash_status['maximum'],
hash_status['period'] - time()))
except (KeyError, TypeError):
pass
if _notify:
sent = Worker.notifier.sent
output.append('Notifications sent: {}, per hour {:.1f}'.format(
sent, sent / hours_since_start))
output.append('')
if not self.all_seen:
no_sightings = ', '.join(
str(w.worker_no) for w in self.workers if w.total_seen == 0)
if no_sightings:
output += [
'Workers without sightings so far:', no_sightings, ''
]
else:
self.all_seen = True
dots, messages = self.get_dots_and_messages()
output += [' '.join(row) for row in dots]
previous = 0
for i in range(4, len(messages) + 4, 4):
output.append('\t'.join(messages[previous:i]))
previous = i
if self.paused:
output.append('\nCAPTCHAs are needed to proceed.')
if not _ansi:
system('cls')
print('\n'.join(output))
def longest_running(self):
workers = (x for x in self.workers if x.start_time)
worker = next(workers)
earliest = worker.start_time
for w in workers:
if w.start_time < earliest:
worker = w
earliest = w.start_time
minutes = ((time() * 1000) - earliest) / 60000
return worker, minutes
def get_start_point(self):
smallest_diff = float('inf')
now = time() % 3600
closest = None
for spawn_id, spawn_time in spawns.known.values():
time_diff = now - spawn_time
if 0 < time_diff < smallest_diff:
smallest_diff = time_diff
closest = spawn_id
if smallest_diff < 3:
break
return closest
async def update_spawns(self, initial=False):
while True:
try:
await run_threaded(spawns.update)
LOOP.create_task(run_threaded(spawns.pickle))
except OperationalError as e:
self.log.exception(
'Operational error while trying to update spawns.')
if initial:
raise OperationalError(
'Could not update spawns, ensure your DB is set up.'
) from e
await sleep(15, loop=LOOP)
except CancelledError:
raise
except Exception as e:
self.log.exception('A wild {} appeared while updating spawns!',
e.__class__.__name__)
await sleep(15, loop=LOOP)
else:
break
async def launch(self, bootstrap, pickle):
exceptions = 0
self.next_mystery_reload = 0
#if not pickle or not spawns.unpickle():
await self.update_spawns(initial=True)
FORT_CACHE.preload()
FORT_CACHE.pickle()
SIGHTING_CACHE.preload()
ENCOUNTER_CACHE.preload()
RAID_CACHE.preload()
self.Worker30 = Worker30
self.ENCOUNTER_CACHE = ENCOUNTER_CACHE
self.worker30 = LOOP.create_task(Worker30.launch(overseer=self))
self.WorkerRaider = WorkerRaider
self.worker_raider = LOOP.create_task(
WorkerRaider.launch(overseer=self))
if not spawns or bootstrap:
try:
await self.bootstrap()
await self.update_spawns()
except CancelledError:
return
update_spawns = False
self.mysteries = spawns.mystery_gen()
while True:
try:
await self._launch(update_spawns)
update_spawns = True
except CancelledError:
return
except Exception:
exceptions += 1
if exceptions > 25:
self.log.exception(
'Over 25 errors occured in launcher loop, exiting.')
return False
else:
self.log.exception('Error occured in launcher loop.')
update_spawns = False
async def _launch(self, update_spawns):
if update_spawns:
await self.update_spawns()
ACCOUNTS = get_accounts()
LOOP.create_task(run_threaded(dump_pickle, 'accounts', ACCOUNTS))
spawns_iter = iter(spawns.items())
else:
start_point = self.get_start_point()
if start_point and not spawns.after_last():
spawns_iter = dropwhile(lambda s: s[1][0] != start_point,
spawns.items())
else:
spawns_iter = iter(spawns.items())
current_hour = get_current_hour()
if spawns.after_last():
current_hour += 3600
captcha_limit = conf.MAX_CAPTCHAS
skip_spawn = conf.SKIP_SPAWN
for point, (spawn_id, spawn_seconds) in spawns_iter:
try:
if self.captcha_queue.qsize() > captcha_limit:
self.paused = True
self.idle_seconds += await run_threaded(
self.captcha_queue.full_wait, conf.MAX_CAPTCHAS)
self.paused = False
except (EOFError, BrokenPipeError, FileNotFoundError):
pass
spawn_time = spawn_seconds + current_hour
spawns.spawn_timestamps[spawn_id] = spawn_time
# negative = hasn't happened yet
# positive = already happened
time_diff = time() - spawn_time
while time_diff < 0.5:
try:
mystery_point = next(self.mysteries)
await self.coroutine_semaphore.acquire()
LOOP.create_task(self.try_point(mystery_point))
except StopIteration:
if self.next_mystery_reload < monotonic():
self.mysteries = spawns.mystery_gen()
self.next_mystery_reload = monotonic(
) + conf.RESCAN_UNKNOWN
else:
await sleep(
min(spawn_time - time() + .5,
self.next_mystery_reload - monotonic()),
loop=LOOP)
time_diff = time() - spawn_time
if time_diff > 5 and spawn_id in SIGHTING_CACHE.spawn_ids:
self.redundant += 1
continue
elif time_diff > skip_spawn:
self.skipped += 1
continue
await self.coroutine_semaphore.acquire()
LOOP.create_task(self.try_point(point, spawn_time, spawn_id))
async def try_again(self, point):
async with self.coroutine_semaphore:
worker = await self.best_worker(point, False)
async with worker.busy:
if await worker.visit(point):
self.visits += 1
async def bootstrap(self):
notifier = Notifier()
try:
self.log.warning('Starting bootstrap phase 1.')
LOOP.create_task(
notifier.scan_log_webhook('Bootstrap Status Change',
'Starting bootstrap phase 1.',
'65300'))
await self.bootstrap_one()
except CancelledError:
raise
except Exception:
self.log.exception(
'An exception occurred during bootstrap phase 1.')
LOOP.create_task(
notifier.scan_log_webhook(
'Bootstrap Status Change',
'An exception occurred during bootstrap phase 1.',
'16060940'))
try:
self.log.warning('Starting bootstrap phase 2.')
LOOP.create_task(
notifier.scan_log_webhook('Bootstrap Status Change',
'Starting bootstrap phase 2.',
'65300'))
await self.bootstrap_two()
except CancelledError:
raise
except Exception:
self.log.exception(
'An exception occurred during bootstrap phase 2.')
LOOP.create_task(
notifier.scan_log_webhook(
'Bootstrap Status Change',
'An exception occurred during bootstrap phase 2.',
'16060940'))
self.log.warning('Starting bootstrap phase 3.')
LOOP.create_task(
notifier.scan_log_webhook('Bootstrap Status Change',
'Starting bootstrap phase 3.', '65300'))
unknowns = list(spawns.unknown)
shuffle(unknowns)
tasks = (self.try_again(point) for point in unknowns)
await gather(*tasks, loop=LOOP)
self.log.warning('Finished bootstrapping. Updating gym parks')
LOOP.create_task(
notifier.scan_log_webhook('Bootstrap Status Change',
'Finished bootstrapping.', '65300'))
with Parks() as parks:
parks_thread = Thread(target=parks.reset_parks)
parks_thread.start()
self.bootstrapping = False
async def bootstrap_one(self):
async def visit_release(worker, num, *args):
async with self.coroutine_semaphore:
async with worker.busy:
point = get_start_coords(num, *args)
self.log.warning('start_coords: {}', point)
self.visits += await worker.bootstrap_visit(point)
if bounds.multi:
areas = [poly.polygon.area for poly in bounds.polygons]
area_sum = sum(areas)
percentages = [area / area_sum for area in areas]
tasks = []
for i, workers in enumerate(
percentage_split(self.workers, percentages)):
grid = best_factors(len(workers))
tasks.extend(
visit_release(w, n, grid, bounds.polygons[i])
for n, w in enumerate(workers))
else:
tasks = (visit_release(w, n) for n, w in enumerate(self.workers))
await gather(*tasks, loop=LOOP)
async def bootstrap_two(self):
async def bootstrap_try(point):
async with self.coroutine_semaphore:
randomized = randomize_point(point, randomization)
LOOP.call_later(1790, LOOP.create_task,
self.try_again(randomized))
worker = await self.best_worker(point, False)
async with worker.busy:
self.visits += await worker.bootstrap_visit(point)
# randomize to within ~140m of the nearest neighbor on the second visit
randomization = conf.BOOTSTRAP_RADIUS / 155555 - 0.00045
tasks = (bootstrap_try(x) for x in get_bootstrap_points(bounds))
await gather(*tasks, loop=LOOP)
async def try_point(self, point, spawn_time=None, spawn_id=None):
try:
point = randomize_point(point)
skip_time = monotonic() + (conf.GIVE_UP_KNOWN
if spawn_time else conf.GIVE_UP_UNKNOWN)
worker = await self.best_worker(point, skip_time)
if not worker:
if spawn_time:
self.skipped += 1
return
async with worker.busy:
if spawn_time:
worker.after_spawn = time() - spawn_time
if await worker.visit(point, spawn_id):
self.visits += 1
except CancelledError:
raise
except Exception:
self.log.exception('An exception occurred in try_point')
finally:
self.coroutine_semaphore.release()
async def best_worker(self, point, skip_time):
good_enough = conf.GOOD_ENOUGH
while self.running:
gen = (w for w in self.workers if not w.busy.locked())
try:
worker = next(gen)
lowest_speed = worker.travel_speed(point)
except StopIteration:
lowest_speed = float('inf')
for w in gen:
speed = w.travel_speed(point)
if speed < lowest_speed:
lowest_speed = speed
worker = w
if speed < good_enough:
break
if lowest_speed < conf.SPEED_LIMIT:
worker.speed = lowest_speed
return worker
if skip_time and monotonic() > skip_time:
return None
await sleep(conf.SEARCH_SLEEP, loop=LOOP)
def refresh_dict(self):
ACCOUNTS = get_accounts()
while not self.extra_queue.empty():
account = self.extra_queue.get()
username = account['username']
ACCOUNTS[username] = account
class DNSCollector:
"""
Retrieve and parse DNS records asynchronously.
**Parameters**
``concurrent_limit``
Set limit on number of concurrent DNS requests to avoid hitting system limits
"""
# Configure the DNS resolver to be asynchronous and use specific nameservers
resolver = asyncresolver.Resolver()
resolver.lifetime = 1
resolver.nameservers = ["8.8.8.8", "8.8.4.4", "1.1.1.1"]
def __init__(self, concurrent_limit=50):
# Limit used for Semaphore to avoid hitting system limits on open requests
self.semaphore = Semaphore(value=concurrent_limit)
async def _query(self, domain: str,
record_type: str) -> Union[Answer, NXDOMAIN, NoAnswer]:
"""
Execute a DNS query for the target domain and record type.
**Parameters**
``domain``
Domain to be used for DNS record collection
``record_type``
DNS record type to collect
"""
try:
# Wait to acquire the semaphore to avoid too many concurrent DNS requests
await self.semaphore.acquire()
answer = await self.resolver.resolve(domain, record_type)
except Exception as e:
answer = e
# Release semaphore to allow next request
self.semaphore.release()
return answer
async def _parse_answer(self, dns_record: Answer) -> list:
"""
Parse the provided instance of ``dns.resolver.Answer``.
**Parameters**
``dns_record``
Instance of ``dns.resolve.Answer``
"""
record_list = []
for rdata in dns_record.response.answer:
for item in rdata.items:
record_list.append(item.to_text())
return record_list
async def _fetch_record(self, domain: str, record_type: str) -> dict:
"""
Fetch a DNS record for the given domain and record type.
**Parameters**
``domain``
Domain to be used for DNS record collection
``record_type``
DNS record type to collect (A, NS, SOA, TXT, MX, CNAME, DMARC)
"""
logger.debug("Fetching %s records for %s", record_type, domain)
# Prepare the results dictionary
result = {}
result[domain] = {}
result[domain]["domain"] = domain
# Handle DMARC as a special record type
if record_type.lower() == "dmarc":
record_type = "A"
query_type = "dmarc_record"
query_domain = "_dmarc." + domain
else:
query_type = record_type.lower() + "_record"
query_domain = domain
# Execute query and await completion
response = await self._query(query_domain, record_type)
# Only parse result if it's an ``Answer``
if isinstance(response, Answer):
record = await self._parse_answer(response)
result[domain][query_type] = record
else:
# Return the type of exception (e.g., NXDOMAIN)
result[domain][query_type] = type(response).__name__
return result
async def _prepare_async_dns(self, domains: list,
record_types: list) -> list:
"""
Prepare asynchronous DNS queries for a list of domain names.
**Parameters**
``domains``
Queryset of :model:`shepherd.Domain` entries
``record_types``
List of record types represented as strings (e.g., ["A", "TXT"])
"""
tasks = []
# For each domain, create a task for each DNS record of interest
for domain in domains:
for record_type in record_types:
tasks.append(
self._fetch_record(domain=domain.name,
record_type=record_type))
# Gather all tasks for execution
all_tasks = await asyncio.gather(*tasks)
return all_tasks
def run_async_dns(self, domains: list, record_types: list) -> dict:
"""
Execute asynchronous DNS queries for a list of domain names.
**Parameters**
``domains``
List of domain names
``record_types``
List of record types represented as strings (e.g., ["A", "TXT"])
"""
# Setup an event loop
event_loop = asyncio.get_event_loop()
# Use an event loop (instead of ``asyncio.run()``) to easily get list of results
results = event_loop.run_until_complete(
self._prepare_async_dns(domains=domains,
record_types=record_types))
# Result is a list of dicts – seven for each domain name
combined = {}
# Combine all dicts with the same domain name
for res in results:
for key, value in res.items():
if key in combined:
combined[key].update(value)
else:
combined[key] = {}
combined[key].update(value)
return combined
async def release(semaphore: Semaphore):
print('Releasing as a one off!')
semaphore.release()
print('Released as a one off!')
class Jobs:
__storage_sem = None
__path_user_files = None
__logger = None
__naas_folder = ".naas"
__json_name = "jobs.json"
def __init__(self, logger, clean=False, init_data=[]):
self.__path_user_files = os.environ.get("JUPYTER_SERVER_ROOT", "/home/ftp")
self.__path_naas_files = os.path.join(
self.__path_user_files, self.__naas_folder
)
self.__json_secrets_path = os.path.join(
self.__path_naas_files, self.__json_name
)
# self.__storage_sem = BoundedSemaphore(1)
self.__storage_sem = Semaphore(1)
self.__logger = logger
if not os.path.exists(self.__path_naas_files):
try:
print("Init Naas folder Jobs")
os.makedirs(self.__path_naas_files)
except OSError as exc: # Guard against race condition
print("__path_naas_files", self.__path_naas_files)
if exc.errno != errno.EEXIST:
raise
except Exception as e:
print("Exception", e)
if not os.path.exists(self.__json_secrets_path) or clean:
uid = str(uuid.uuid4())
try:
print("Init Job Storage", self.__json_secrets_path)
self.__save(uid, init_data)
except Exception as e:
print("Exception", e)
self.__logger.error(
{
"id": uid,
"type": "init_job_storage",
"status": "error",
"error": str(e),
}
)
def __save(self, uid, data):
try:
with open(self.__json_secrets_path, "w+") as f:
f.write(
json.dumps(data, sort_keys=True, indent=4).replace("NaN", "null")
)
f.close()
except Exception as err:
print("__save Exception", err)
self.__logger.error(
{
"id": str(uid),
"type": "set_job_storage",
"status": "exception",
"filepath": self.__json_secrets_path,
"error": str(err),
}
)
def find_by_value(self, uid, value, target_type):
data = self.list(uid)
if len(data) > 0:
job_list = pd.DataFrame(data)
cur_jobs = job_list[
(job_list.type == target_type) & (job_list.value == value)
]
cur_job = cur_jobs.to_dict("records")
if len(cur_job) == 1:
return cur_job[0]
return None
def find_by_path(self, uid, filepath, target_type):
data = self.list(uid)
if len(data) > 0:
job_list = pd.DataFrame(data)
cur_jobs = job_list[
(job_list.type == target_type) & (job_list.path == filepath)
]
cur_job = cur_jobs.to_dict("records")
if len(cur_job) == 1:
return cur_job[0]
return None
def is_running(self, uid, notebook_filepath, target_type):
cur_job = self.find_by_path(uid, notebook_filepath, target_type)
if cur_job:
status = cur_job.get("status", None)
if status and status == t_start:
return True
return False
def list(self, uid):
data = []
try:
with open(self.__json_secrets_path, "r") as f:
data = json.load(f)
f.close()
except Exception as err:
print("cannot open")
self.__logger.error(
{
"id": uid,
"type": "get_job_storage",
"status": "exception",
"filepath": self.__json_secrets_path,
"error": str(err),
}
)
data = []
return data
async def update(self, uid, path, target_type, value, params, status, runTime=0):
await self.__storage_sem.acquire()
data = None
res = t_error
try:
if len(self.list(uid)) != 0:
df = pd.DataFrame(self.list(uid))
else:
df = pd.DataFrame(
columns=[
"id",
"type",
"value",
"path",
"status",
"params",
"lastUpdate",
"lastRun",
"totalRun",
]
)
res = status
cur_elem = df[(df.type == target_type) & (df.path == path)]
now = datetime.datetime.now()
dt_string = now.strftime("%Y-%m-%d %H:%M:%S")
if len(cur_elem) == 1:
if status == t_delete:
self.__logger.info(
{
"id": uid,
"type": target_type,
"value": value,
"status": t_delete,
"path": path,
"params": params,
}
)
df = df.drop(cur_elem.index)
else:
self.__logger.info(
{
"id": uid,
"type": target_type,
"value": value,
"status": t_update,
"path": path,
"params": params,
}
)
index = cur_elem.index[0]
df.at[index, "id"] = uid
df.at[index, "status"] = status
df.at[index, "value"] = value
df.at[index, "params"] = params
df.at[index, "lastUpdate"] = dt_string
if runTime > 0 and status != t_add:
df.at[index, "lastRun"] = runTime
df.at[index, "totalRun"] = runTime + (
df.at[index, "totalRun"] if df.at[index, "totalRun"] else 0
)
elif status == t_add:
df.at[index, "lastRun"] = 0
df.at[index, "totalRun"] = 0
res = t_update
elif status == t_add and len(cur_elem) == 0:
self.__logger.info(
{
"id": uid,
"type": target_type,
"value": value,
"status": t_update,
"path": path,
"params": params,
}
)
new_row = [
{
"id": uid,
"type": target_type,
"value": value,
"status": t_add,
"path": path,
"params": params,
"lastRun": runTime,
"totalRun": runTime,
"lastUpdate": dt_string,
}
]
df_new = pd.DataFrame(new_row)
df = pd.concat([df, df_new], axis=0)
else:
res = t_skip
data = df.to_dict("records")
if res != t_skip:
self.__save(uid, data)
except Exception as e:
print("cannot update", e)
self.__logger.error(
{
"id": uid,
"type": target_type,
"value": value,
"status": t_error,
"path": path,
"params": params,
"error": str(e),
}
)
self.__storage_sem.release()
return {"status": res, "data": data}
