def request_handler(self, i, context, uri):
sock = context.socket(zmq.REP)
sock.connect(uri)
while True:
try:
obj = util.recv_bson(sock)
log.debug('Request %s in %s', obj, i)
command = obj.get('$command')
log.debug('Req %s in %s', command, i)
if command == 'echo':
util.send_bson(sock, dict(message=obj))
elif command == 'terminate':
self._terminate = True
elif command in ('mapreduce', 'xmapreduce', 'basic'):
job = Job.from_request(self, obj)
self.job_manager.alloc_job(job)
log.debug('Start job %s', job.id)
if obj.get('wait'):
job.run()
self.job_manager.dealloc_job(job.id)
else:
gevent.spawn(util.ending_request(self.mongodb), job.run)
util.send_bson(sock, dict(job_id=job.id, status=job.status))
continue
elif command == 'status':
try:
jid = obj['job_id']
status = self.job_manager.job_status(jid)
util.send_bson(sock, dict(job_id=jid, status=status))
if status['state'] in ('complete', 'error'):
self.job_manager.dealloc_job(jid)
except:
log.exception('Error getting status for job')
util.send_bson(sock, dict(
job_id=obj['job_id'], status=dict(
state='UNKNOWN')))
continue
elif command == 'config':
util.send_bson(sock, dict(
self.options.zmr,
src_port=self._src_port,
sink_port=self._sink_port))
except:
log.exception('Error in request handler')
util.send_bson(sock, dict(error=traceback.format_exc()))
continue
finally:
self.mongodb.end_request()
def run(self, requests):
for req_num in xrange(requests):
connection = ZM.doc_session.db.connection
header = util.recv_bson(self._src)
log.debug('RECV %s', header.keys())
more = self._src.getsockopt(zmq.RCVMORE)
parts = []
while more:
if header['compress']:
parts.append(zlib.decompress(self._src.recv()))
else:
parts.append(self._src.recv())
more = self._src.getsockopt(zmq.RCVMORE)
with util.ending_request(connection):
self.handle_message(header, parts)
r = resource.getrusage(resource.RUSAGE_SELF)
if r.ru_maxrss > self._suicide_level:
log.error('Suicide (painlessly) worker at RSS %s / %s',
r.ru_maxrss, self._suicide_level)
break
log.info('Closing worker after %d requests', req_num)
self._src.close()
self._sink.close()
def _g_sink_handler(self):
while True:
header = util.recv_bson(self._sink_sock)
log.debug('SINK RECV %s', header.keys())
# Retrieve the payload
more = self._sink_sock.getsockopt(zmq.RCVMORE)
parts = []
while more:
parts.append(self._sink_sock.recv())
more = self._sink_sock.getsockopt(zmq.RCVMORE)
# Grab headers for routing
job_id = header.pop('job_id', None)
listener_id = header.pop('listener_id', None)
# Lookup the job
j = self.jobs.get(job_id, None)
if j is None:
log.error('Discard message for unknown job_id: %s', job_id)
continue
# Route the message
listener = j.listeners.get(listener_id)
if listener is None:
log.error('Discard header for unknown listener %s', listener_id)
continue
listener.put((header, parts))