def start_callback():
"""
Starts the collection callback
"""
global getDataCallback
if getDataCallback is None:
getDataCallback = PeriodicCallback(process_data, conf.REFRESH_FREQ)
getDataCallback.start()
elif not getDataCallback.is_running():
getDataCallback.start()
class Timer:
__instance = None
__timer = None
__callback = None
__repeat = True
__args = []
__kwargs = {}
def __init__(self, func, repeat=True, timeout=10, args=None, kwargs=None):
self.__callback = func
self.__repeat = repeat
self.__timeout = timeout
self.__args = args if args else list()
self.__kwargs = kwargs if kwargs else dict()
if self.__repeat:
self.__timer = PeriodicCallback(self.__callback_wrapper,
self.__timeout * 1000)
self.__timer.start()
else:
self.__timer = IOLoop.current().call_later(self.__timeout,
self.__callback_wrapper)
def restart(self):
if self.__timer is not None:
self.__stop()
self.__start()
def start(self):
if self.__timer is None:
self.__start()
def stop(self):
if self.__timer is not None:
self.__stop()
def __start(self):
if self.__repeat and not self.__timer.is_running():
self.__timer.start()
# else:
# print("calling later")
# self.__timer = IOLoop.current().call_later(self.__timeout, self.__callback_wrapper, self)
def __stop(self):
if self.__repeat:
self.__timer.stop()
else:
IOLoop.current().remove_timeout(self.__timer)
def __callback_wrapper(self):
IOLoop.current().add_callback(self.__callback, *self.__args,
**self.__kwargs)
if self.__repeat:
self.__start()
class ProcessWatchdog(object):
def __init__(self, interval=1000):
self.interval = interval
self._process_to_check = {}
self._periodic_callback = PeriodicCallback(self._check_processes, self.interval)
def register_process(self, process, on_dead_callback):
self._process_to_check[process] = on_dead_callback
def start(self):
if not self._periodic_callback.is_running():
self._periodic_callback.start()
def stop(self):
if self._periodic_callback.is_running():
self._periodic_callback.stop()
def _check_processes(self):
for process in self._process_to_check:
if not process.is_running():
self._process_to_check[process](process)
class CallbackHandler():
def __init__(self):
self.callback_fn = None
def start_callback(self):
if self.callback_fn is None:
self.callback_fn = PeriodicCallback(process_data, conf.REFRESH_INTERVAL)
self.callback_fn.start()
elif not self.callback_fn.is_running():
self.callback_fn.start()
def stop_callback(self):
self.callback_fn.stop()
class Timer():
__instance = None
__timer = None
__callback = None
__repeat = True
__args = []
__kwargs = {}
def __init__(self, func, repeat=True, timeout=10, args=None, kwargs=None):
self.__callback = func
self.__repeat = repeat
self.__timeout = timeout
self.__args = args if args else list()
self.__kwargs = kwargs if kwargs else dict()
if self.__repeat:
self.__timer = PeriodicCallback(
partial(self.__callback_wrapper, self), self.__timeout * 1000)
def restart(self):
if self.__timer is not None:
self.__stop()
self.__start()
def start(self):
if self.__timer is None:
self.__start()
def stop(self):
if self.__timer is not None:
self.__stop()
def __start(self):
if self.__repeat and not self.__timer.is_running():
self.__timer.start()
else:
self.__timer = IOLoop.current().call_later(self.__timeout,
self.__callback_wrapper,
self)
def __stop(self):
if self.__repeat:
self.__timer.stop()
else:
IOLoop.current().remove_timeout(self.__timer)
def __callback_wrapper(self):
try:
self.__callback(*self.__args, **self.__kwargs)
except Exception:
pass
if self.__repeat:
self.__start()
class Session():
def __init__(self):
self.cars=dict()
self.timer=PeriodicCallback(self.on_timer,50)
self.users=set()
def publish(self,type,data):
msg={"type":type,"data":data}
for u in self.users:
u.write_message(msg)
def on_timer(self):
for id in self.cars:
car=self.cars[id]
car.step()
self.publish("timer",car.__dict__)
def reset(self):
self.stop()
self.cars.clear()
def start(self):
if not self.timer.is_running():
self.timer.start()
if self.timer.is_running():
print "session started"
else:
print "session failed to started"
self.publish("cars number",len(self.cars))
def stop(self):
self.timer.stop()
if(not self.timer.is_running()):
print "session stopped"
else:
print "session failed to stop"
def __init__(self, *args, **kwargs):
super(EnvironmentKernelSpecManager, self).__init__(*args, **kwargs)
self.log.info("Using EnvironmentKernelSpecManager...")
self._env_data_cache = {}
if self.refresh_interval > 0:
try:
from tornado.ioloop import PeriodicCallback, IOLoop
# Initial loading NOW
IOLoop.current().call_later(0, callback=self._update_env_data, initial=True)
# Later updates
updater = PeriodicCallback(callback=self._update_env_data,
callback_time=1000 * 60 * self.refresh_interval)
updater.start()
if not updater.is_running():
raise Exception()
self._periodic_updater = updater
self.log.info("Started periodic updates of the kernel list (every %s minutes).", self.refresh_interval)
except:
self.log.exception("Error while trying to enable periodic updates of the kernel list.")
else:
self.log.info("Periodical updates the kernel list are DISABLED.")
class PollingHandler(BaseSocketHandler):
"""
This class represents separate websocket connection.
Attributes:
tracker: tornado.ioloop.PeriodicCallback with get_location method as a
callback. Starts when user pushes "track" button. When started, it
runs every 5 seconds to find out and update character's location.
q: tornado.queues.Queue used for running tasks successively.
updating: A flag indicates if router is being updated or not. Required
to avoid race conditions.
"""
def __init__(self, *args, **kwargs):
super(PollingHandler, self).__init__(*args, **kwargs)
# Set Tornado PeriodicCallback with our self.track, we
# will use launch it later on track/untrack commands
self.tracker = PeriodicCallback(self.get_location, 5000)
self.q = Queue(maxsize=5)
self.updating = False
async def get_location(self):
"""
The callback for the `self.tracker`.
Makes an API call, updates router and sends updated data to the
front-end.
"""
# Call API to find out current character location
location = await self.character(self.user_id, '/location/', 'GET')
if location:
# Set `updating` flag to not accept periodic updates
# from front-end, to not overwrite new data
self.updating = True
user = self.user
graph_data = await user['router'].update(
location['solarSystem']['name'])
if graph_data:
message = ['update', graph_data]
logging.warning(graph_data)
await self.safe_write(message)
self.updating = False
else:
message = ['warning', 'Log into game to track your route']
await self.safe_write(message)
async def scheduler(self):
"""
Scheduler for user tasks.
Waits until there is new item in the queue, does task, resolves task.
Tornado queues doc: http://www.tornadoweb.org/en/stable/queues.html
Since we have no guarantee of the order of the incoming messages
(new message from front-end can come before current is done),
we need to ensure all tasks to run successively.
Here comes the asynchronous generator.
"""
logging.info(f"Scheduler started for {self.request.remote_ip}")
# Wait on each iteration until there's actually an item available
async for item in self.q:
logging.debug(f"Started resolving task for {item}...")
user = self.user
try:
if item == 'recover':
# Send saved route
await self.safe_write(['recover', user['router'].recovery])
elif item == 'track':
# Start the PeriodicCallback
if not self.tracker.is_running():
self.tracker.start()
elif item in ['stop', 'reset']:
# Stop the PeriodicCallback
if self.tracker.is_running():
self.tracker.stop()
# Clear all saved data
if item == 'reset':
await user['router'].reset()
elif item[0] == 'backup':
# Do not overwrite user object while it's updating,
# just in case, to avoid race conditions.
if not self.updating:
await user['router'].backup(item[1])
finally:
self.q.task_done()
logging.debug(f'Task "{item}" done.')
async def task(self, item):
"""
Intermediary between `self.on_message` and `self.scheduler`.
Since we cannot do anything asynchronous in the `self.on_message`,
this method can handle any additional non-blocking stuff if we need it.
:argument item: item to pass to the `self.scheduler`.
"""
await self.q.put(item)
#await self.q.join()
def open(self):
"""
Triggers on successful websocket connection.
Ensures user is authorized, spawns `self.scheduler` for user
tasks, adds this websocket object to the connections pool,
spawns the recovery of the saved route.
"""
logging.info(f"Connection received from {self.request.remote_ip}")
if self.user_id:
self.spawn(self.scheduler)
self.vagrants.append(self)
self.spawn(self.task, 'recover')
else:
self.close()
def on_message(self, message):
"""
Triggers on receiving front-end message.
:argument message: front-end message.
Receives user commands and passes them
to the `self.scheduler` via `self.task`.
"""
self.spawn(self.task, json_decode(message))
def on_close(self):
"""
Triggers on closed websocket connection.
Removes this websocket object from the connections pool,
stops `self.tracker` if it is running.
"""
self.vagrants.remove(self)
if self.tracker.is_running():
self.tracker.stop()
logging.info("Connection closed, " + self.request.remote_ip)
class Executor:
"""
Main object communicate to server and control tasks running.
"""
# hearbeat check interval, default: 10 seconds
heartbeat_interval = 10 * 1000
# check task interval, default: 10 seconds
checktask_interval = 10 * 1000
custom_ssl_cert = False
def __init__(self, config=None):
self.ioloop = IOLoop.current()
self.node_id = None
self.status = EXECUTOR_STATUS_OFFLINE
if config is None:
config = AgentConfig()
self.config = config
self.tags = config.get('tags')
self.checktask_callback = None
self.task_slots = TaskSlotContainer(config.getint('slots', 1))
server_base = config.get('server')
if urlparse(server_base).scheme == '':
if config.getint('server_https_port'):
server_https_port = config.getint('server_https_port')
server_base = 'https://%s:%d' % (server_base,
server_https_port)
else:
server_base = 'http://%s:%d' % (server_base,
config.getint('server_port'))
self.service_base = server_base
client_cert = config.get('client_cert') or None
client_key = config.get('client_key') or None
self.keep_job_files = config.getboolean('debug', False)
LOGGER.debug('keep_job_files %s', self.keep_job_files)
httpclient_defaults = {
'request_timeout': config.getfloat('request_timeout', 60)
}
if client_cert:
httpclient_defaults['client_cert'] = client_cert
if client_key:
httpclient_defaults['client_key'] = client_key
if os.path.exists('keys/ca.crt'):
self.custom_ssl_cert = True
httpclient_defaults['ca_certs'] = 'keys/ca.crt'
httpclient_defaults['validate_cert'] = True
LOGGER.debug(httpclient_defaults)
node_key = None
secret_key = None
if os.path.exists('conf/node.conf'):
parser = ConfigParser()
parser.read('conf/node.conf')
node_key = parser.get('agent', 'node_key')
secret_key = parser.get('agent', 'secret_key')
self.httpclient = NodeAsyncHTTPClient(self.service_base,
key=node_key,
secret_key=secret_key,
defaults=httpclient_defaults)
self.runner_factory = RunnerFactory(config)
def start(self):
self.register_node()
# init heartbeat period callback
heartbeat_callback = PeriodicCallback(self.send_heartbeat,
self.heartbeat_interval)
heartbeat_callback.start()
# init checktask period callback
# the new version use HEARTBEAT response to tell client whether there
# is new task on server queue
# so do not start this period method. But if server is an old version
# without that header info
# client sill need to poll GET_TASK.
self.checktask_callback = PeriodicCallback(self.check_task,
self.checktask_interval)
self.ioloop.start()
def check_header_new_task_on_server(self, headers):
try:
new_task_on_server = headers['X-DD-New-Task'] == 'True'
if new_task_on_server:
self.ioloop.call_later(0, self.check_task)
except KeyError:
# if response contains no 'DD-New-Task' header, it might
# be the old version server,
# so use the old GET_TASK pool mode
if not self.checktask_callback.is_running():
self.checktask_callback.start()
@coroutine
def send_heartbeat(self):
if self.status == EXECUTOR_STATUS_OFFLINE:
self.register_node()
return
url = urljoin(self.service_base, '/nodes/%d/heartbeat' % self.node_id)
running_tasks = ','.join([task_executor.task.id for task_executor in
self.task_slots.tasks()])
request = HTTPRequest(url=url, method='POST', body='',
headers={'X-DD-RunningJobs': running_tasks})
try:
res = yield self.httpclient.fetch(request)
if 'X-DD-KillJobs' in res.headers:
LOGGER.info('received kill signal %s',
res.headers['X-DD-KillJobs'])
for job_id in json.loads(res.headers['X-DD-KillJobs']):
task_to_kill = self.task_slots.get_task(job_id)
if task_to_kill:
LOGGER.info('%s', task_to_kill)
task_to_kill.kill()
self.task_slots.remove_task(task_to_kill)
self.check_header_new_task_on_server(res.headers)
except HTTPError as ex:
if ex.code == 400:
logging.warning('Node expired, register now.')
self.status = EXECUTOR_STATUS_OFFLINE
self.register_node()
except URLError as ex:
logging.warning('Cannot connect to server. %s', ex)
except Exception as ex:
logging.warning('Cannot connect to server. %s', ex)
@coroutine
def register_node(self):
if self.custom_ssl_cert and \
self.service_base.startswith('https') and \
not os.path.exists('keys/ca.crt'):
httpclient = AsyncHTTPClient(force_instance=True)
cacertrequest = HTTPRequest(urljoin(self.service_base, 'ca.crt'),
validate_cert=False)
try:
cacertresponse = yield httpclient.fetch(cacertrequest)
if not os.path.exists('keys'):
os.mkdir('keys')
open('keys/ca.crt', 'wb').write(cacertresponse.body)
except HTTPError:
LOGGER.info('Custom ca cert retrieve failed.')
try:
node_id = yield self.httpclient.node_online(self.tags)
self.status = EXECUTOR_STATUS_ONLINE
self.node_id = node_id
LOGGER.info('node %d registered', self.node_id)
except URLError as ex:
logging.warning('Cannot connect to server, %s', ex)
except socket.error as ex:
logging.warning('Cannot connect to server, %s', ex)
def on_new_task_reach(self, task):
if task is not None:
task_executor = self.execute_task(task)
self.task_slots.put_task(task_executor)
@staticmethod
def parse_task_data(response_data):
task = SpiderTask()
task.id = response_data['data']['task']['task_id']
task.spider_id = response_data['data']['task']['spider_id']
task.project_name = response_data['data']['task']['project_name']
task.project_version = response_data['data']['task']['version']
task.spider_name = response_data['data']['task']['spider_name']
if 'extra_requirements' in response_data['data']['task'] and \
response_data['data']['task']['extra_requirements']:
task.extra_requirements = response_data['data']['task']['extra_requirements']
if 'spider_parameters' in response_data['data']['task']:
task.spider_parameters = response_data['data']['task'][
'spider_parameters']
else:
task.spider_parameters = {}
task.settings = response_data['data']
return task
@coroutine
def get_next_task(self):
url = urljoin(self.service_base, '/executing/next_task')
post_data = urlencode({'node_id': self.node_id})
request = HTTPRequest(url=url, method='POST', body=post_data)
try:
response = yield self.httpclient.fetch(request)
response_content = response.body
response_data = json.loads(response_content)
LOGGER.debug(url)
LOGGER.debug(response_content)
if response_data['data'] is not None:
task = Executor.parse_task_data(response_data)
self.on_new_task_reach(task)
except URLError:
LOGGER.warning('Cannot connect to server')
def execute_task(self, task):
egg_downloader = ProjectEggDownloader(service_base=self.service_base,
client=self.httpclient)
executor = TaskExecutor(task, egg_downloader=egg_downloader,
runner_factory=self.runner_factory,
keep_files=self.keep_job_files)
pid = None
future = executor.execute()
self.post_start_task(task, pid)
executor.on_subprocess_start = self.post_start_task
self.ioloop.add_future(future, self.task_finished)
return executor
@coroutine
def post_start_task(self, task, pid):
url = urljoin(self.service_base, '/jobs/%s/start' % task.id)
post_data = urlencode({'pid': pid or ''})
try:
request = HTTPRequest(url=url, method='POST', body=post_data)
yield self.httpclient.fetch(request)
except URLError as ex:
LOGGER.error('Error when post_task_task: %s', ex)
except HTTPError as ex:
LOGGER.error('Error when post_task_task: %s', ex)
def complete_task(self, task_executor, status):
"""
@type task_executor: TaskExecutor
"""
url = urljoin(self.service_base, '/executing/complete')
log_file = open(task_executor.output_file, 'rb')
post_data = {
'task_id': task_executor.task.id,
'status': status,
'log': log_file,
}
items_file = None
if task_executor.items_file and \
os.path.exists(task_executor.items_file):
post_data['items'] = items_file = open(task_executor.items_file,
"rb")
if LOGGER.isEnabledFor(logging.DEBUG):
item_file_size = os.path.getsize(task_executor.items_file)
LOGGER.debug('item file size : %d', item_file_size)
LOGGER.debug(post_data)
datagen, headers = multipart_encode(post_data)
headers['X-DD-Nodeid'] = str(self.node_id)
body_producer = MultipartRequestBodyProducer(datagen)
request = HTTPRequest(url, method='POST', headers=headers,
body_producer=body_producer)
client = self.httpclient
future = client.fetch(request, raise_error=False)
self.ioloop.add_future(future, self.complete_task_done(task_executor,
log_file,
items_file))
LOGGER.info('task %s finished', task_executor.task.id)
def complete_task_done(self, task_executor, log_file, items_file):
def complete_task_done_f(future):
"""
The callback of complete job request, if socket error occurred,
retry commiting complete request in 10 seconds.
If request is completed successfully, remove task from slots.
Always close stream files.
@type future: Future
:return:
"""
response = future.result()
LOGGER.debug(response)
if log_file:
log_file.close()
if items_file:
items_file.close()
if response.error and isinstance(response.error, socket.error):
status = TASK_STATUS_SUCCESS if task_executor.ret_code == 0 \
else TASK_STATUS_FAIL
self.ioloop.call_later(10, self.complete_task, task_executor,
status)
LOGGER.warning('Socket error when completing job, retry in '
'10 seconds.')
return
if response.error:
LOGGER.warning('Error when post task complete request: %s',
response.error)
self.task_slots.remove_task(task_executor)
LOGGER.debug('complete_task_done')
return complete_task_done_f
def task_finished(self, future):
task_executor = future.result()
status = TASK_STATUS_SUCCESS if task_executor.ret_code == 0 \
else TASK_STATUS_FAIL
self.complete_task(task_executor, status)
def check_task(self):
if not self.task_slots.is_full():
self.get_next_task()
class MetadataStorage:
def __init__(self, server):
self.server = server
self.metadata = {}
self.pending_requests = {}
self.events = {}
self.script_response = None
self.busy = False
self.gc_path = os.path.expanduser("~")
self.prune_cb = PeriodicCallback(
self.prune_metadata, METADATA_PRUNE_TIME)
def update_gcode_path(self, path):
if path == self.gc_path:
return
self.metadata = {}
self.gc_path = path
if not self.prune_cb.is_running():
self.prune_cb.start()
def close(self):
self.prune_cb.stop()
def get(self, key, default=None):
if key not in self.metadata:
return default
return dict(self.metadata[key])
def __getitem__(self, key):
return dict(self.metadata[key])
def _handle_script_response(self, result):
try:
proc_resp = json.loads(result.strip())
except Exception:
logging.exception("file_manager: unable to load metadata")
logging.debug(result)
return
proc_log = proc_resp.get('log', [])
for log_msg in proc_log:
logging.info(log_msg)
if 'file' in proc_resp:
self.script_response = proc_resp
def prune_metadata(self):
for fname in list(self.metadata.keys()):
fpath = os.path.join(self.gc_path, fname)
if not os.path.exists(fpath):
del self.metadata[fname]
logging.info(f"Pruned file: {fname}")
continue
def _has_valid_data(self, fname, fsize, modified):
mdata = self.metadata.get(fname, {'size': "", 'modified': 0})
return mdata['size'] == fsize and mdata['modified'] == modified
def remove_file(self, fname):
self.metadata.pop(fname)
def parse_metadata(self, fname, fsize, modified, notify=False):
evt = Event()
if fname in self.pending_requests or \
self._has_valid_data(fname, fsize, modified):
# request already pending or not necessary
evt.set()
return evt
self.pending_requests[fname] = (fsize, modified, notify, evt)
if self.busy:
return evt
self.busy = True
IOLoop.current().spawn_callback(self._process_metadata_update)
return evt
async def _process_metadata_update(self):
while self.pending_requests:
fname, (fsize, modified, notify, evt) = \
self.pending_requests.popitem()
if self._has_valid_data(fname, fsize, modified):
evt.set()
continue
retries = 3
while retries:
try:
await self._run_extract_metadata(fname, notify)
except Exception:
logging.exception("Error running extract_metadata.py")
retries -= 1
else:
break
else:
self.metadata[fname] = {'size': fsize, 'modified': modified}
logging.info(
f"Unable to extract medatadata from file: {fname}")
evt.set()
self.busy = False
async def _run_extract_metadata(self, filename, notify):
# Escape single quotes in the file name so that it may be
# properly loaded
filename = filename.replace("\"", "\\\"")
cmd = " ".join([sys.executable, METADATA_SCRIPT, "-p",
self.gc_path, "-f", f"\"{filename}\""])
shell_command = self.server.lookup_plugin('shell_command')
scmd = shell_command.build_shell_command(
cmd, self._handle_script_response)
self.script_response = None
await scmd.run(timeout=10.)
if self.script_response is None:
raise self.server.error("Unable to extract metadata")
path = self.script_response['file']
metadata = self.script_response['metadata']
if not metadata:
# This indicates an error, do not add metadata for this
raise self.server.error("Unable to extract metadata")
self.metadata[path] = dict(metadata)
metadata['filename'] = path
if notify:
self.server.send_event(
"file_manager:metadata_update", metadata)
class State(tornado.websocket.WebSocketHandler):
def __init__(self, *args, **kwargs):
self.stopped = False
self.t = None
self.car = kwargs.pop('car')
self.rps_ms = kwargs.pop('rps_ms') # Request Per Seconds Milliseconds
self.for_network = kwargs.pop('for_network')
self.timer = [0, 0]
self.timer_index = 0
self.start_time = time()
self.total_requests = 0
self.inf_loop = PeriodicCallback(
self.loop, callback_time=self.rps_ms
) # 50 ms = 15 Request Per Seconds; 65 ms = 15 Request Per Seconds
super(State, self).__init__(*args, **kwargs)
def check_origin(self, origin):
return True
def on_message(self, message):
try:
if message == "send_state":
state = self.car.get_state()
self.write_message(state, True)
logging.debug("Sent send_state")
elif message == "read_state":
if not self.inf_loop.is_running():
self.start_time = time()
self.inf_loop.start()
elif message == "stop_read_state":
self.inf_loop.stop()
except tornado.websocket.WebSocketClosedError:
logging.debug("State WS closed, stopping PeriodicCallback")
self.inf_loop.stop()
def on_close(self):
self.inf_loop.stop()
logging.debug("WebSocket closed")
def loop(self):
self.timer[self.timer_index] += 1
self.total_requests += 1
elapsed = time() - self.start_time
if elapsed > 1:
self.start_time = time()
self.timer_index = 1
self.timer[0] = (self.timer[0] + self.timer[1]) / 2
self.timer[1] = 0
car_rps = self.timer[0]
logging.debug("RPS: " + str(car_rps))
state = {}
state = self.car.get_state_vector(latest=True,
for_network=self.for_network)
state['car_rps'] = car_rps
try:
self.write_message(state, True)
except tornado.websocket.WebSocketClosedError:
logging.debug("State WS closed, stopping PeriodicCallback")
self.inf_loop.stop()
class FutureConnection(object):
def _future_start(self):
self._future_connecting = False
self._future_closed = False
self._future_periodic_connect.start()
self._future_io_loop = self._future_periodic_connect.io_loop
self._future_thread_id = thread.get_ident()
@coroutine
def _future_stream_close(self, shutdown=False):
if shutdown: # positive close issued by user.
self._future_closed = True
self._future_connecting = False
LOG.warn('Closing: threadId:%s!', thread.get_ident())
if self.future_connected():
for callback in self._future_close_callbacks:
if inspect.isgeneratorfunction(callback):
raise SyntaxError(
'wrap generator function with ``@tornado.gen.coroutine`` first :%s '
% callback)
LOG.debug('calling close_callback: %s', callback)
self.future_stream.io_loop.add_callback(callback, self)
def _future_stream_close_callback(conn):
conn.future_stream.close()
conn.future_stream = None
conn._future_packet_bulk = []
conn._future_callback_queue.clear() # registry
if conn.future_stream is None:
if conn._future_closed is False:
conn._future_periodic_connect.start()
return
self.future_stream.io_loop.add_callback(
_future_stream_close_callback, self)
@coroutine
def _future_stream_connect(self,
host=None,
port=None,
af=socket.AF_UNSPEC,
ssl_options=None,
max_buffer_size=None,
source_ip=None,
source_port=None,
timeout=None):
assert self._future_thread_id == thread.get_ident(
), 'this should run in IO thread:%s' % self._future_thread_id
if self.future_connected():
if self._future_periodic_connect.is_running():
self._future_periodic_connect.stop()
raise Return(self)
if self._future_connecting is True:
return
else:
self._future_connecting = True
host = host or self.host
port = port or self.port
LOG.warn('Connecting %s:%s, threadId:%s, id:%s' %
(host, port, thread.get_ident(), id(self)))
try:
self.future_stream = yield self._future_connector.connect(
host,
port,
timeout=timeout,
max_buffer_size=max_buffer_size,
af=af,
ssl_options=ssl_options,
source_ip=source_ip,
source_port=source_port)
# Note: we first register Async Close|Read|Write callback here, since handshake may require Async READ|WRITE
except Exception as e:
self._future_connecting = False
LOG.exception('Error connect to %s:%s, retry.', host, port)
if self.future_stream:
self.future_stream.close()
if not self._future_periodic_connect.is_running():
self._future_periodic_connect.start()
return
self._future_callback_queue.clear() # registry
self._future_packet_bulk = []
self._future_connecting = False
self._future_chunk_size = stream_options(self.future_stream,
blocking=False,
keepalive=True)
self._future_reader_buffer = bytearray(self._future_chunk_size)
self.future_stream.set_close_callback(self._future_stream_close)
self.future_stream.read_into(self._future_reader_buffer,
callback=self._future_socket_read,
partial=True)
LOG.warn('Handshaking: %s:%s, threadId:%s, id:%s', host, port,
thread.get_ident(), id(self))
try:
for callback in self._future_handshake_callbacks:
if not is_coroutine_function(callback):
raise SyntaxError(
'wrap generator function with ``@tornado.gen.coroutine`` first :%s '
% callback)
LOG.debug('calling handshake_callback: %s', callback)
r = yield callback(self)
waiters = tuple(self._future_waiters)
for fut in waiters:
if not fut.done():
fut.set_result(self)
except Exception as e:
LOG.exception('Error Handshaking to %s:%s, retry.', host, port)
self.future_stream.close()
if not self._future_periodic_connect.is_running():
self._future_periodic_connect.start()
return
finally:
self._future_connecting = False
LOG.warn('Connected: %s:%s, threadId:%s, id:%s', host, port,
thread.get_ident(), id(self))
try:
for callback in self._future_initial_callbacks:
if inspect.isgeneratorfunction(callback):
raise SyntaxError(
'wrap generator function with ``@tornado.gen.coroutine`` first :%s '
% callback)
LOG.debug('calling initial_callback: %s', callback)
r = yield callback(self)
except Exception as e:
LOG.exception('Error connect to %s:%s, retry.', host, port)
self.future_stream.close()
if not self._future_periodic_connect.is_running():
self._future_periodic_connect.start()
return
for callback in self._future_close_callbacks:
if inspect.isgeneratorfunction(callback):
raise SyntaxError(
'wrap generator function with ``@tornado.gen.coroutine`` first :%s '
% callback)
LOG.warn('Initialized: %s:%s, threadId:%s, id:%s!', host, port,
thread.get_ident(), id(self))
###########################################################################
def future_initialize(self):
self.host = self.host if hasattr(self, 'host') else '127.0.0.1'
self.port = self.port if hasattr(self, 'port') else '/tmp/redis.sock'
self.future_stream = None
self._future_thread_id = None
self._future_io_loop = None
self._future_connecting = None
self._future_closed = None
self._future_reader_buffer = None
self._future_chunk_size = None
self._future_periodic_connect = PeriodicCallback(
self._future_stream_connect, 1000)
self._future_connector = TCPClient(resolver=CachedOverrideResolver())
self._future_handshake_callbacks = [] # connected, handshaking
self._future_initial_callbacks = [
] # connected, shandshake passed, initializing
self._future_waiters = [
] # connected, shandshake passed, initialization passed, wake up all waiters.
self._future_close_callbacks = []
self._future_packet_bulk = []
self._future_callback_queue = collections.deque() # registry
def future_connected(self):
return self._future_connecting is False and self.future_stream is not None and not self.future_stream.closed(
)
def future_disconnect(self):
if not self.future_connected():
return
self._future_connecting = False
def safe_shutdown_stream_read(stream):
if stream.socket is None:
return
while len(self._future_close_callbacks) > 0:
callback = self._future_close_callbacks.pop(0)
stream.io_loop.add_callback(callback, self)
if stream.reading() is False:
stream.io_loop.update_handler(stream.socket.fileno(),
IOLoop.WRITE | IOLoop.ERROR)
if stream.writing() is False:
stream.io_loop.remove_handler(stream.socket.fileno())
stream.close()
return
stream.io_loop.add_callback(safe_shutdown_stream_read, stream)
self.future_stream.io_loop.add_callback(safe_shutdown_stream_read,
self.future_stream)
def future_connect(self):
connected_future = Future()
if self.future_connected():
if self._future_periodic_connect.is_running():
self._future_periodic_connect.stop()
connected_future.set_result(self)
return connected_future
# most time, Transport should be in connected state.
connected_future.add_done_callback(
lambda cf: self._future_waiters.remove(cf))
for callback in self._future_initial_callbacks:
connected_future.add_done_callback(callback)
if connected_future not in self._future_waiters:
self._future_waiters.append(connected_future)
if not self._future_periodic_connect.is_running():
if self._future_io_loop is None:
self._future_start()
else:
self._future_io_loop.add_callback(
self._future_periodic_connect.start)
return connected_future
###########################################################################
def add_future(self, future, future_callback):
self.future_stream.io_loop.add_future(future, future_callback)
def add_callback(self, callback, *args, **kwargs):
self.future_stream.io_loop.add_callback(callback, *args, **kwargs)
def spawn_callback(self, callback, *args, **kwargs):
self.future_stream.io_loop.spawn_callback(callback, *args, **kwargs)
def add_handshake_callback(self, callback):
self._future_handshake_callbacks.append(callback)
def add_initial_callback(self, callback):
self._future_initial_callbacks.append(callback)
def add_close_callback(self, callback):
self._future_close_callbacks.append(callback)
def clear_connect_callbacks(self):
self._future_close_callbacks = []
def set_future_chunk_size(self, size=None):
if size is None:
size = self._future_chunk_size
if size < 1:
return
cur_size = len(self._future_reader_buffer)
if size > cur_size:
self._future_reader_buffer.extend(bytearray(size - cur_size))
elif size < cur_size:
del self._future_reader_buffer[size:]
###########################################################################
def _future_socket_read(self, num_bytes=0):
if num_bytes < 1:
return
data = memoryview(
self._future_reader_buffer)[:num_bytes] if num_bytes > 0 else None
self.future_packet_read(data)
# this 'read_into() DO MUST call after self.future_packet_read(data)! '
self.future_stream.read_into(self._future_reader_buffer,
callback=self._future_socket_read,
partial=True)
def _future_packet_aggregator(self, packet_num, aggregator_callback,
packet):
self._future_packet_bulk.append(packet)
if len(self._future_packet_bulk) == packet_num:
future_packet_bulk = self._future_packet_bulk
self._future_packet_bulk = []
aggregator_callback(future_packet_bulk)
def _future_socket_send(self,
buffer_or_partial,
packet_callbacks,
send_future=None,
connected_future=None):
if packet_callbacks is None:
packet_callbacks = []
elif callable(packet_callbacks):
packet_callbacks = [packet_callbacks]
if callable(buffer_or_partial):
buffer_or_partial = buffer_or_partial()
if isinstance(buffer_or_partial, (tuple, list)):
buffer_or_partial = BYTES_EMPTY.join(buffer_or_partial)
if send_future is None:
send_future = self.future_stream.write(buffer_or_partial)
else:
self.future_stream.write(buffer_or_partial).add_done_callback(
lambda _: send_future.set_result(_.result()))
self._future_callback_queue.extend(packet_callbacks)
return send_future
def future_packet_send(self, packets, callbacks=None):
packet_send_future = Future()
send_partial = partial(self._future_socket_send, packets, callbacks,
packet_send_future)
if self.future_connected():
return send_partial()
else:
self.future_connect().add_done_callback(send_partial)
return packet_send_future
def future_packet_read(self, data=None):
raise NotImplementedError('Implement this method first !')
class Executor():
heartbeat_interval = 10
checktask_interval = 10
def __init__(self, config=None):
self.ioloop = IOLoop.current()
self.node_id = None
self.status = EXECUTOR_STATUS_OFFLINE
if config is None:
config =AgentConfig()
self.task_slots = TaskSlotContainer(config.getint('slots', 1))
self.config = config
# if server_https_port is configured, prefer to use it.
if config.get('server_https_port'):
self.service_base = 'https://%s:%d'% (config.get('server'), config.getint('server_https_port'))
else:
self.service_base = 'http://%s:%d' % (config.get('server'), config.getint('server_port'))
client_cert = config.get('client_cert') or None
client_key = config.get('client_key') or None
httpclient_defaults = {
'validate_cert': True,
'ca_certs' : 'keys/ca.crt',
'client_cert' : client_cert,
'client_key' : client_key,
'request_timeout': config.getfloat('request_timeout', 60)
}
logger.debug(httpclient_defaults)
self.httpclient = AsyncHTTPClient(defaults=httpclient_defaults)
def start(self):
self.register_node()
#init heartbeat period callback
heartbeat_callback = PeriodicCallback(self.send_heartbeat, self.heartbeat_interval*1000)
heartbeat_callback.start()
#init checktask period callback
# the new version use HEARTBEAT response to tell client whether there is new task on server queue
# so do not start this period method. But if server is an old version without that header info
# client sill need to poll GET_TASK.
self.checktask_callback = PeriodicCallback(self.check_task, self.checktask_interval*1000)
# code for debuging memory leak
# import objgraph
# def check_memory():
# logger.debug('Checking memory.')
# logger.debug(objgraph.by_type('SpiderTask'))
# logger.debug(objgraph.by_type('TaskExecutor'))
# logger.debug(objgraph.by_type('Future'))
# logger.debug(objgraph.by_type('PeriodicCallback'))
# future_objs = objgraph.by_type('Future')
# if future_objs:
# objgraph.show_chain(
# objgraph.find_backref_chain(future_objs[-1], objgraph.is_proper_module),
# filename='chain.png'
# )
# self.check_memory_callback = PeriodicCallback(check_memory, 1*1000)
# self.check_memory_callback.start()
self.ioloop.start()
def check_header_new_task_on_server(self, headers):
try:
new_task_on_server = headers['X-DD-New-Task'] == 'True'
if new_task_on_server:
self.ioloop.call_later(0, self.check_task)
except KeyError:
# if response contains no 'DD-New-Task' header, it might be the old version server,
# so use the old GET_TASK pool mode
if not self.checktask_callback.is_running():
self.checktask_callback.start()
@coroutine
def send_heartbeat(self):
if self.status == EXECUTOR_STATUS_OFFLINE:
self.register_node()
return
url = urlparse.urljoin(self.service_base, '/nodes/%d/heartbeat' % self.node_id)
running_tasks = ','.join([task_executor.task.id for task_executor in self.task_slots.tasks()])
request = HTTPRequest(url=url, method='POST', body='', headers={'X-DD-RunningJobs': running_tasks})
try:
res = yield self.httpclient.fetch(request)
if 'X-DD-KillJobs' in res.headers:
logger.info('received kill signal %s' % res.headers['X-DD-KillJobs'])
for job_id in json.loads(res.headers['X-DD-KillJobs']):
task_to_kill = self.task_slots.get_task(job_id)
if task_to_kill:
logger.info('%s' % task_to_kill)
task_to_kill.kill()
self.check_header_new_task_on_server(res.headers)
except urllib2.HTTPError as e:
if e.code == 400:
logging.warning('Node expired, register now.')
self.status = EXECUTOR_STATUS_OFFLINE
self.register_node()
except urllib2.URLError as e:
logging.warning('Cannot connect to server. %s' % e)
except HTTPError as e:
if e.code == 400:
logging.warning('Node expired, register now.')
self.status = EXECUTOR_STATUS_OFFLINE
self.register_node()
except Exception as e:
logging.warning('Cannot connect to server. %s' % e)
@coroutine
def register_node(self):
if self.service_base.startswith('https') and not os.path.exists('keys/ca.crt') :
httpclient = AsyncHTTPClient(force_instance=True)
cacertrequest = HTTPRequest(urlparse.urljoin(self.service_base, 'ca.crt'), validate_cert=False)
cacertresponse = yield httpclient.fetch(cacertrequest)
if not os.path.exists('keys'):
os.mkdir('keys')
open('keys/ca.crt', 'wb').write(cacertresponse.body)
try:
url = urlparse.urljoin(self.service_base, '/nodes')
request = HTTPRequest(url = url, method='POST', body='')
res = yield self.httpclient.fetch(request)
self.status = EXECUTOR_STATUS_ONLINE
self.node_id = json.loads(res.body)['id']
logger.info('node %d registered' % self.node_id)
except urllib2.URLError as e:
logging.warning('Cannot connect to server, %s' % e )
except socket.error as e:
logging.warning('Cannot connect to server, %s' % e)
def on_new_task_reach(self, task):
if task is not None:
task_executor = self.execute_task(task)
self.task_slots.put_task(task_executor)
@coroutine
def get_next_task(self):
url = urlparse.urljoin(self.service_base, '/executing/next_task')
post_data = urllib.urlencode({'node_id': self.node_id})
request = HTTPRequest(url=url, method='POST', body=post_data)
try:
response = yield self.httpclient.fetch(request)
response_content = response.body
response_data = json.loads(response_content)
logger.debug(url)
logger.debug(response_content)
if response_data['data'] is not None:
task = SpiderTask()
task.id = response_data['data']['task']['task_id']
task.spider_id = response_data['data']['task']['spider_id']
task.project_name = response_data['data']['task']['project_name']
task.project_version = response_data['data']['task']['version']
task.spider_name = response_data['data']['task']['spider_name']
if 'spider_parameters' in response_data['data']['task']:
task.spider_parameters = response_data['data']['task']['spider_parameters']
else:
task.spider_parameters = {}
self.on_new_task_reach(task)
except urllib2.URLError:
logger.warning('Cannot connect to server')
def execute_task(self, task):
executor = TaskExecutor(task, config=self.config)
pid = None
future = executor.execute()
self.post_start_task(task, pid)
executor.on_subprocess_start = self.post_start_task
self.ioloop.add_future(future, self.task_finished)
return executor
@coroutine
def post_start_task(self, task, pid):
url = urlparse.urljoin(self.service_base, '/jobs/%s/start' % task.id)
post_data = urllib.urlencode({'pid':pid or ''})
try:
request = HTTPRequest(url=url, method='POST', body=post_data)
respones = yield self.httpclient.fetch(request)
except urllib2.URLError as e:
logger.error('Error when post_task_task: %s' % e)
except urllib2.HTTPError as e:
logger.error('Error when post_task_task: %s' % e)
def complete_task(self, task_executor, status):
'''
@type task_executor: TaskExecutor
'''
url = urlparse.urljoin(self.service_base, '/executing/complete')
log_file = open(task_executor.output_file, 'rb')
post_data = {
'task_id': task_executor.task.id,
'status': status,
'log': log_file,
}
items_file = None
if task_executor.items_file and os.path.exists(task_executor.items_file):
post_data['items'] = items_file = open(task_executor.items_file, "rb")
logger.debug('item file size : %d' % os.path.getsize(task_executor.items_file))
logger.debug(post_data)
datagen, headers = multipart_encode(post_data)
headers['X-DD-Nodeid'] = str(self.node_id)
request = HTTPRequest(url, method='POST', headers=headers, body_producer=MultipartRequestBodyProducer(datagen))
client = self.httpclient
future = client.fetch(request, raise_error=False)
self.ioloop.add_future(future, self.complete_task_done(task_executor, log_file, items_file))
logger.info('task %s finished' % task_executor.task.id)
def complete_task_done(self, task_executor, log_file, items_file):
def complete_task_done_f(future):
'''
The callback of complete job request, if socket error occurred, retry commiting complete request in 10 seconds.
If request is completed successfully, remove task from slots.
Always close stream files.
@type future: Future
:return:
'''
response = future.result()
logger.debug(response)
if log_file:
log_file.close()
if items_file:
items_file.close()
if response.error and isinstance(response.error, socket.error):
self.ioloop.call_later(10, self.complete_task, task_executor, TASK_STATUS_SUCCESS if task_executor.ret_code == 0 else TASK_STATUS_FAIL)
logger.warning('Socket error when completing job, retry in 10 seconds.')
return
if response.error:
logger.warning('Error when post task complete request: %s' % response.error)
self.task_slots.remove_task(task_executor)
logger.debug('complete_task_done')
return complete_task_done_f
def task_finished(self, future):
task_executor = future.result()
self.complete_task(task_executor, TASK_STATUS_SUCCESS if task_executor.ret_code == 0 else TASK_STATUS_FAIL)
def check_task(self):
if not self.task_slots.is_full():
self.get_next_task()
class UpdateChecker(BotRequestHandler):
def __init__(self, manager, trackers, ioloop=None):
BotRequestHandler.__init__(self)
self.ioloop = ioloop or IOLoop.current()
self.manager = manager
self.trackers = trackers
self.search_renderer = Renderer(
u"<b>{{item.title|e}}</b>\n"
u"Раздел: {{item.category|e}}\n"
u"Размер: {% if item.size | int(-1) == -1 %}{{item.size | e}}{% else %}"
u"{{ item.size | default(0) | int | filesizeformat }}{% endif %}\n"
u"Добавлено: {{ item.added | default(0) | int | todatetime | datetimeformat('%d-%m-%Y') }}\n"
u"Скачать: /download_{{item.id}}\n"
u"\n")
self.torrent_renderer = Renderer(
u"<b>{{ item.name | e }}</b> - {{ \"%0.2f\" | format(item.done*100) }}% done\n\n"
)
self.cache = []
self.update_task = PeriodicCallback(self.do_update, 15 * 60 * 1000)
self.version = telebots.version
pass
@PatternMessageHandler('/show( .*)?', authorized=True)
def cmd_show(self, chat, text):
data = text.split()
search_id = int(data[1])
start = int(data[2])
if 0 <= search_id < len(self.cache):
chat_id = self.cache[search_id]['chat_id']
message_id = self.cache[search_id]['message_id']
results = self.cache[search_id]['results']
self.show_results(search_id, chat_id, message_id, results, start)
return True
@PatternMessageHandler('/status', authorized=True)
def cmd_status(self, message_id, chat):
@gen.coroutine
def execute():
torrents = yield self.manager.get_torrents()
text = [self.torrent_renderer.render(x) for x in torrents]
self.bot.send_message(to=chat['id'],
message="".join(text),
reply_to_message_id=message_id,
parse_mode='HTML')
execute()
return True
@PatternMessageHandler('/update( .*)?', authorized=True)
def cmd_update(self, chat, text):
cmd = text.split()
chat_id = chat['id']
if len(cmd) == 1:
buttons = [
{
'text': 'Now',
'callback_data': '/update now'
},
{
'text': '15m',
'callback_data': '/update 15'
},
{
'text': '30m',
'callback_data': '/update 30'
},
{
'text': '60m',
'callback_data': '/update 60'
},
{
'text': 'Off',
'callback_data': '/update 0'
},
]
markup = {'inline_keyboard': [buttons]}
self.bot.send_message(to=chat_id,
message='Schedule update',
reply_markup=markup)
else:
when = cmd[1]
if when == 'now':
self.do_update(chat_id)
else:
minutes = int(when)
if self.update_task.is_running():
self.update_task.stop()
if minutes > 0:
self.update_task = PeriodicCallback(
self.do_update, minutes * 60 * 1000)
self.update_task.start()
text = 'Schedule updated: Each %d minutes' % minutes
else:
text = 'Schedule updated: off'
self.bot.send_message(to=chat_id, message=text)
pass
return True
def show_results(self,
search_id,
chat_id,
message_id,
results,
start,
page_size=3):
message = 'Sorry, nothing is found'
markup = None
if len(results) > 0:
data = map((lambda i: self.search_renderer.render(results[i])),
range(start - 1, min(start + page_size - 1,
len(results))))
message = u"\n".join(data)
buttons = []
if start > page_size:
buttons.append({
'text':
"Prev %d/%d" % (page_size, start - 1),
'callback_data':
'/show %d %d' % (search_id, start - page_size)
})
if (start + page_size) < len(results):
buttons.append({
'text':
"Next %d/%d" %
(page_size, len(results) - start - page_size + 1),
'callback_data':
'/show %d %d' % (search_id, start + page_size)
})
if len(buttons) > 0:
markup = {'inline_keyboard': [buttons]}
self.bot.edit_message_text(to=chat_id,
message_id=message_id,
text=message,
parse_mode='HTML',
reply_markup=markup)
pass
@PatternMessageHandler('[^/].*', authorized=True)
def do_search(self, text, chat, message_id):
query = text
chat_id = chat['id']
@gen.coroutine
def execute():
# Send reply "search in progress"
msg = yield self.bot.send_message(to=chat_id,
message="Search in progress...",
reply_to_message_id=message_id)
placeholder_message_id = json.loads(
msg.body)['result']['message_id']
self.async_search(None, chat_id, placeholder_message_id, query)
execute()
return True
@PatternMessageHandler('/download_.*', authorized=True)
def do_download(self, chat, text):
query = text
user_id = chat['id']
@gen.coroutine
def execute():
item_id = query.split('_', 1)[1]
logging.info("Search download url for id=%s", item_id)
item = None
for search in self.cache:
for result in search['results']:
if result.id == item_id:
item = result
break
if item is not None:
break
if item is not None:
url = item.link
msg = yield self.bot.send_message(to=user_id,
message='Downloading...')
message_id = json.loads(msg.body)['result']['message_id']
try:
for tracker in self.trackers:
if tracker.check_url(url):
torrent_data = yield tracker.download(url)
tr_info = bencode.bdecode(torrent_data)
torrent_name = tr_info['info']['name']
yield self.manager.add_torrent(torrent_data)
self.bot.edit_message_text(
to=user_id,
message_id=message_id,
text='Torrent "%s" downloaded' % torrent_name)
break
except Exception:
logging.exception("Error while download torrent data")
self.bot.edit_message_text(
to=user_id,
message_id=message_id,
text="Sorry, error occurred!\n%s" %
traceback.format_exc())
else:
logging.warn("Couldn't find download url for id %s", item_id)
execute()
return True
@gen.coroutine
def async_search(self, search_id, chat_id, message_id, query):
if search_id is None:
search_id = len(self.cache)
self.cache.append({
'chat_id': chat_id,
'message_id': message_id,
'query': query,
'results': []
})
logging.info("Start search_id %d for query \"%s\"", search_id, query)
responses = yield [
tracker.do_search(query) for tracker in self.trackers
]
results = []
for response in responses:
results = results + response
logging.info("Found %d torrents", len(results))
self.cache[search_id]['results'] = results
self.show_results(search_id, chat_id, message_id, results, 1)
pass
@gen.coroutine
def do_update(self, reply_chat_id=None):
chat_id = reply_chat_id or self.bot.admins[0]
updated = False
torrents = yield self.manager.get_torrents()
error = False
for torrent in torrents:
url = torrent['url']
logging.info('Checking updates for %s', url)
for tracker in self.trackers:
if tracker.check_url(url):
logging.debug('Download using %s', tracker.__class__)
try:
torrent_data = yield tracker.download(url)
added = yield self.manager.add_torrent(
torrent_data, torrent)
if added:
updated = True
self.bot.send_message(
to=chat_id,
message='Torrent "%s" updated' %
torrent['name'])
except Exception as e:
logging.exception('Error while check updates')
if reply_chat_id is not None:
self.bot.send_message(
to=chat_id, message=traceback.format_exc())
error = True
continue
pass
if not error and not updated and reply_chat_id is not None:
self.bot.send_message(to=chat_id, message='No updates')
pass
def do_notify(self, message):
logging.info("Notify message: \"%s\"", message)
chat_id = self.bot.admins[0]
self.bot.send_message(to=chat_id, message=message)
pass
class MetadataStorage:
def __init__(self, server, gc_path, database):
self.server = server
database.register_local_namespace(METADATA_NAMESPACE)
self.mddb = database.wrap_namespace(METADATA_NAMESPACE,
parse_keys=False)
self.pending_requests = {}
self.events = {}
self.busy = False
self.gc_path = gc_path
self.prune_cb = PeriodicCallback(self.prune_metadata,
METADATA_PRUNE_TIME)
def update_gcode_path(self, path):
if path == self.gc_path:
return
self.mddb.clear()
self.gc_path = path
if not self.prune_cb.is_running():
self.prune_cb.start()
def close(self):
self.prune_cb.stop()
def get(self, key, default=None):
return self.mddb.get(key, default)
def __getitem__(self, key):
return self.mddb[key]
def prune_metadata(self):
for fname in list(self.mddb.keys()):
fpath = os.path.join(self.gc_path, fname)
if not os.path.exists(fpath):
del self.mddb[fname]
logging.info(f"Pruned file: {fname}")
continue
def _has_valid_data(self, fname, fsize, modified):
mdata = self.mddb.get(fname, {'size': "", 'modified': 0})
return mdata['size'] == fsize and mdata['modified'] == modified
def remove_file(self, fname):
try:
del self.mddb[fname]
except Exception:
pass
def parse_metadata(self, fname, fsize, modified, notify=False):
evt = Event()
if fname in self.pending_requests or \
self._has_valid_data(fname, fsize, modified):
# request already pending or not necessary
evt.set()
return evt
self.pending_requests[fname] = (fsize, modified, notify, evt)
if self.busy:
return evt
self.busy = True
IOLoop.current().spawn_callback(self._process_metadata_update)
return evt
async def _process_metadata_update(self):
while self.pending_requests:
fname, (fsize, modified, notify, evt) = \
self.pending_requests.popitem()
if self._has_valid_data(fname, fsize, modified):
evt.set()
continue
retries = 3
while retries:
try:
await self._run_extract_metadata(fname, notify)
except Exception:
logging.exception("Error running extract_metadata.py")
retries -= 1
else:
break
else:
self.mddb[fname] = {'size': fsize, 'modified': modified}
logging.info(
f"Unable to extract medatadata from file: {fname}")
evt.set()
self.busy = False
async def _run_extract_metadata(self, filename, notify):
# Escape single quotes in the file name so that it may be
# properly loaded
filename = filename.replace("\"", "\\\"")
cmd = " ".join([
sys.executable, METADATA_SCRIPT, "-p", self.gc_path, "-f",
f"\"{filename}\""
])
shell_command = self.server.lookup_plugin('shell_command')
scmd = shell_command.build_shell_command(cmd, log_stderr=True)
result = await scmd.run_with_response(timeout=10.)
if result is None:
raise self.server.error(f"Metadata extraction error")
try:
decoded_resp = json.loads(result.strip())
except Exception:
logging.debug(f"Invalid metadata response:\n{result}")
raise
path = decoded_resp['file']
metadata = decoded_resp['metadata']
if not metadata:
# This indicates an error, do not add metadata for this
raise self.server.error("Unable to extract metadata")
self.mddb[path] = dict(metadata)
metadata['filename'] = path
if notify:
self.server.send_event("file_manager:metadata_update", metadata)
class MongoExportPipeline(object):
"""
This pipeline exports tems to MongoDB using async mongo
driver (motor). Interaction with MongoDB doesn't block
the event loop.
On start it creates object in 'jobs' collection and sets
spider.motor_job_id to the ID of this job.
If MONGO_EXPORT_JOBID_KEY option is set, job id is added to
each stored item under the specified key name.
If MONGO_EXPORT_DUMP_PERIOD is non-zero then updated job stats are saved
to Mongo periodically every ``MONGO_EXPORT_DUMP_PERIOD`` seconds
(default is 60).
"""
def __init__(self, crawler):
self.crawler = crawler
settings = self.crawler.settings
if not settings.getbool('MONGO_EXPORT_ENABLED', False):
raise NotConfigured
self.job_id_key = settings.get('MONGO_EXPORT_JOBID_KEY')
self.items_uri = settings.get('MONGO_EXPORT_ITEMS_URI')
self.jobs_uri = settings.get('MONGO_EXPORT_JOBS_URI')
self.items_client, _, _, _, self.items_col = \
motor_from_uri(self.items_uri)
self.jobs_client, _, _, _, self.jobs_col = \
motor_from_uri(self.jobs_uri)
# XXX: spider_closed is used instead of close_spider because
# the latter doesn't provide a closing reason.
crawler.signals.connect(self.spider_closed, signals.spider_closed)
crawler.signals.connect(self.spider_closing, signals.spider_closing)
self.dump_period = settings.getfloat('MONGO_EXPORT_DUMP_PERIOD', 60.0)
self._dump_pc = None
@classmethod
def from_crawler(cls, crawler):
return cls(crawler)
@classmethod
def get_spider_urls(cls, spider):
options = getattr(spider.crawler, 'start_options', None)
if options and "domain" in options:
return [options["domain"]]
else:
return spider.start_urls
@tt_coroutine
def open_spider(self, spider):
try:
yield self.items_col.ensure_index(self.job_id_key)
yield self.jobs_col.ensure_index('id', unique=True)
job = yield self.jobs_col.find_and_modify(
{
'id': spider.crawl_id,
}, {
'id': spider.crawl_id,
'started_at': datetime.datetime.utcnow(),
'status': 'running',
'spider': spider.name,
"urls": self.get_spider_urls(spider),
'options': getattr(spider.crawler, 'start_options', {}),
},
upsert=True,
new=True)
self.job_id = str(job['_id'])
spider.motor_job_id = str(self.job_id)
logger.info("Crawl job generated id: %s",
self.job_id,
extra={'crawler': self.crawler})
if self.dump_period:
self._dump_pc = PeriodicCallback(self.dump_stats,
self.dump_period * 1000)
self._dump_pc.start()
except Exception:
self.job_id = None
logger.error(
"Can't connect to %s. Items won't be stored.",
self.items_uri,
exc_info=True,
extra={'crawler': self.crawler},
)
@tt_coroutine
def spider_closed(self, spider, reason, **kwargs):
self._stop_periodic_tasks()
if self.job_id is None: # what's this?
self.jobs_client.close()
self.items_client.close()
return
# update the job one more time because something might
# have happened while the spider is closing
yield self._update_finished_job(reason)
self.jobs_client.close()
self.items_client.close()
logger.info("Info is saved for a closed job %s",
self.job_id,
extra={'crawler': self.crawler})
@tt_coroutine
def spider_closing(self, spider, reason, **kwargs):
self._stop_periodic_tasks()
if self.job_id is None: # what's this?
return
yield self._update_finished_job(reason)
logger.info("Info is saved for a closing job %s",
self.job_id,
extra={'crawler': self.crawler})
@tt_coroutine
def process_item(self, item, spider):
mongo_item = scrapy_item_to_dict(item)
if self.job_id_key:
mongo_item[self.job_id_key] = self.job_id
try:
yield self.items_col.insert(mongo_item)
self.crawler.stats.inc_value("mongo_export/items_stored_count")
except Exception as e:
self.crawler.stats.inc_value("mongo_export/store_error_count")
self.crawler.stats.inc_value("mongo_export/store_error_count/" +
e.__class__.__name__)
logger.error("Error storing item",
exc_info=True,
extra={'crawler': self.crawler})
raise gen.Return(item)
def _update_finished_job(self, reason):
status = 'finished'
if reason == 'shutdown':
status = 'shutdown'
return self.jobs_col.update({'_id': ObjectId(self.job_id)}, {
'$set': {
'finished_at': datetime.datetime.utcnow(),
'status': status,
'stats': self._get_stats_json(),
'stats_dict': self._get_stats_escaped(),
}
})
def _stop_periodic_tasks(self):
if self._dump_pc is not None and self._dump_pc.is_running():
self._dump_pc.stop()
def _get_stats_json(self):
# json is to fix an issue with dots in key names
return json_encode(self.crawler.stats.get_stats())
def _get_stats_escaped(self):
return replace_dots(copy.deepcopy(self.crawler.stats.get_stats()))
@gen.coroutine
def dump_stats(self):
# json is to fix an issue with dots in key names
stats = self._get_stats_json()
yield self.jobs_col.update({'_id': ObjectId(self.job_id)}, {
'$set': {
'stats': stats,
'stats_dict': self._get_stats_escaped(),
}
})
logger.info("Stats are stored for job %s" % self.job_id,
extra={'crawler': self.crawler})
class HTCondorCluster(htcondor.HTCondorCluster):
job_cls = HTCondorJob
machine_sep = "+"
def __init__(self, *args, maintain=False, **kwargs):
super().__init__(*args, **kwargs)
self.pc_maintain = PeriodicCallback(self._correct_state, callback_time=30000)
self.maintain = maintain
@property
def maintain(self):
return self.pc_maintain.is_running()
@maintain.setter
def maintain(self, val):
if bool(val) != self.maintain:
if val:
self.loop.add_callback(self.pc_maintain.start)
else:
self.pc_maintain.stop()
@classmethod
def worker2machine(cls, worker):
if not isinstance(worker, str):
worker = worker.name
parts = worker.split(cls.machine_sep)
return parts[0] if len(parts) == 2 else None
def _machine_status(self, query="1"):
ret = Counter()
for line in (
check_output(
[
"condor_status",
"-constraint",
'SlotType!="Dynamic"',
"-af",
"Machine",
query,
]
)
.decode()
.splitlines()
):
key, value = line.split()
ret[key] += int(value)
return ret
@property
def machines_known(self):
return self._machine_status("TotalSlotCpus")
@property
def machines_free(self):
return self._machine_status("Cpus")
@property
def machines(self):
ret = Counter(map(self.worker2machine, self.workers.keys()))
ret.pop(None, None)
return ret
@property
def machines_live(self):
ret = Counter(map(self.worker2machine, self.scheduler.workers.values()))
ret.pop(None, None)
return ret
@property
def machines_pending(self):
ml = self.machines_live
return {m for m in self.machines.keys() if m not in ml}
@property
def workers_live(self):
return set(w.name for w in self.scheduler.workers.values())
@property
def workers_pending(self):
return set(self.workers.keys()) - self.workers_live
@property
def workers_machine_pending(self):
mp = self.machines_pending
return set(w for w in self.workers_pending if self.worker2machine(w) in mp)
def status_table(self):
from rich.table import Table, Column
from rich import box
info = (
self.machines,
self.machines_live,
self.machines_free,
self.machines_known,
)
b, l, f, t = [sum(i.values()) for i in info]
def pct(a, b, digits=0):
return "%.*f" % (
max(0, digits),
round(100 * a / (b or 1), digits),
)
tab = Table(
Column("Machine \\ CPUs", "Total"),
Column("live %", pct(l, b), justify="right"),
Column("live", str(l), justify="right"),
Column("booked", str(b), justify="right"),
Column("total", str(t), justify="right"),
Column("free", str(f), justify="right"),
Column("free %", pct(f, t), justify="right"),
show_footer=True,
box=box.SIMPLE,
)
for m in sorted(info[-1]):
b, l, f, t = [i.get(m, 0) for i in info]
tab.add_row(m, pct(l, b), str(l), str(b), str(t), str(f), pct(f, t))
return tab
def scale_machines(
self,
machines,
wait_machines=True,
timeout=np.inf,
poll=1,
stop_pending=False,
):
# remove excess
max_target = make_target(self.machines, machines)
for name in self.worker_spec.keys():
if self.machine_sep not in str(name):
continue
machine, num = name.split(self.machine_sep)
if machine not in max_target:
continue
if max_target[machine] <= num:
del self.worker_spec[name]
# spawn new
name_new = []
for machine, num in make_target(self.machines_known, machines).items():
for i in range(num):
name = self.machine_sep.join((machine, str(i)))
assert name not in self.worker_spec, name
self.worker_spec[name] = self.new_machine_spec(machine)
name_new.append(name)
# sync internal sate
self.sync(self._correct_state)
if not poll:
return
endtime = time() + timeout
while time() < endtime and (
self.machines_pending if wait_machines else self.workers_pending
):
sleep(poll)
if stop_pending:
wp = self.workers_pending
if stop_pending == "machine":
wp = self.workers_machine_pending
for w in wp:
del self.worker_spec[w]
def new_machine_spec(self, machine):
spec = copy.deepcopy(self.new_spec)
req = 'Machine=="%s"' % machine
job_extra = spec.setdefault("options", {}).setdefault("job_extra", {})
if "Requirements" in job_extra:
req = "(%s)&&(%s)" % (req, job_extra["Requirements"])
job_extra["Requirements"] = req
return spec
