diff options
-rw-r--r-- | gear/__init__.py | 36 | ||||
-rw-r--r-- | gear/tests/test_functional.py | 15 |
2 files changed, 42 insertions, 9 deletions
diff --git a/gear/__init__.py b/gear/__init__.py index e338efc..28fef6a 100644 --- a/gear/__init__.py +++ b/gear/__init__.py @@ -2983,7 +2983,8 @@ class Server(BaseClientServer): # the worker disconnected, alert the client try: p = Packet(constants.REQ, constants.WORK_FAIL, job.handle) - job.client_connection.sendPacket(p) + if job.client_connection: + job.client_connection.sendPacket(p) except Exception: self.log.exception("Sending WORK_FAIL to client after " "worker disconnect failed:") @@ -3006,10 +3007,11 @@ class Server(BaseClientServer): def _removeJob(self, job, dequeue=True): # dequeue is tri-state: True, False, or a specific queue - try: - del job.client_connection.related_jobs[job.handle] - except KeyError: - pass + if job.client_connection: + try: + del job.client_connection.related_jobs[job.handle] + except KeyError: + pass if job.worker_connection: try: del job.worker_connection.related_jobs[job.handle] @@ -3303,7 +3305,7 @@ class Server(BaseClientServer): workers += 1 self.statsd.gauge('workers', workers) - def _handleSubmitJob(self, packet, precedence): + def _handleSubmitJob(self, packet, precedence, background=False): name = packet.getArgument(0) unique = packet.getArgument(1) if not unique: @@ -3320,11 +3322,16 @@ class Server(BaseClientServer): self.max_handle += 1 handle = ('H:%s:%s' % (packet.connection.host, self.max_handle)).encode('utf8') - job = ServerJob(handle, name, arguments, packet.connection, unique) + if not background: + conn = packet.connection + else: + conn = None + job = ServerJob(handle, name, arguments, conn, unique) p = Packet(constants.RES, constants.JOB_CREATED, handle) packet.connection.sendPacket(p) self.jobs[handle] = job - packet.connection.related_jobs[handle] = job + if not background: + packet.connection.related_jobs[handle] = job if precedence == PRECEDENCE_HIGH: self.high_queue.append(job) elif precedence == PRECEDENCE_NORMAL: @@ -3343,6 +3350,16 @@ class Server(BaseClientServer): def handleSubmitJobLow(self, packet): return self._handleSubmitJob(packet, PRECEDENCE_LOW) + def handleSubmitJobBg(self, packet): + return self._handleSubmitJob(packet, PRECEDENCE_NORMAL, + background=True) + + def handleSubmitJobHighBg(self, packet): + return self._handleSubmitJob(packet, PRECEDENCE_HIGH, background=True) + + def handleSubmitJobLowBg(self, packet): + return self._handleSubmitJob(packet, PRECEDENCE_LOW, background=True) + def getJobForConnection(self, connection, peek=False): for queue in [self.high_queue, self.normal_queue, self.low_queue]: for job in queue: @@ -3412,7 +3429,8 @@ class Server(BaseClientServer): self.log.info("Received packet %s for unknown job" % (packet,)) return packet.code = constants.RES - job.client_connection.sendPacket(packet) + if job.client_connection: + job.client_connection.sendPacket(packet) if finished: self._removeJob(job, dequeue=False) self._updateStats() diff --git a/gear/tests/test_functional.py b/gear/tests/test_functional.py index e512f2f..09c7d2f 100644 --- a/gear/tests/test_functional.py +++ b/gear/tests/test_functional.py @@ -128,6 +128,21 @@ class TestFunctional(tests.BaseTestCase): self.assertTrue(job.complete) self.assertEqual(job.data, [b'workdata']) + def test_bg_job(self): + self.worker.registerFunction('test') + + job = gear.Job(b'test', b'testdata') + self.client.submitJob(job, background=True) + self.assertNotEqual(job.handle, None) + self.client.shutdown() + del self.client + + workerjob = self.worker.getJob() + self.assertEqual(workerjob.handle, job.handle) + self.assertEqual(workerjob.arguments, b'testdata') + workerjob.sendWorkData(b'workdata') + workerjob.sendWorkComplete() + def test_worker_termination(self): def getJob(): with testtools.ExpectedException(gear.InterruptedError): |