def LeaseCronJobs(self, cronjob_ids=None, lease_time=None, cursor=None):
"""Leases all available cron jobs."""
now = rdfvalue.RDFDatetime.Now()
now_str = mysql_utils.RDFDatetimeToMysqlString(now)
expiry_str = mysql_utils.RDFDatetimeToMysqlString(now + lease_time)
id_str = utils.ProcessIdString()
query = ("UPDATE cron_jobs "
"SET leased_until=%s, leased_by=%s "
"WHERE (leased_until IS NULL OR leased_until < %s)")
args = [expiry_str, id_str, now_str]
if cronjob_ids:
query += " AND job_id in (%s)" % ", ".join(
["%s"] * len(cronjob_ids))
args += cronjob_ids
updated = cursor.execute(query, args)
if updated == 0:
return []
cursor.execute(
"SELECT job, create_time, disabled, "
"last_run_status, last_run_time, current_run_id, state, "
"leased_until, leased_by "
"FROM cron_jobs WHERE leased_until=%s AND leased_by=%s",
[expiry_str, id_str])
return [self._CronjobFromRow(row) for row in cursor.fetchall()]
def LeaseCronJobs(self, cronjob_ids=None, lease_time=None):
leased_jobs = []
now = rdfvalue.RDFDatetime.Now()
expiration_time = now + lease_time
for job in self.cronjobs.values():
existing_lease = job.leased_until
if existing_lease is None or existing_lease < now:
job.leased_until = expiration_time
job.leased_by = utils.ProcessIdString()
leased_jobs.append(job)
return [copy.copy(job) for job in leased_jobs]
def LeaseCronJobs(self, cronjob_ids=None, lease_time=None):
leased_jobs = []
now = rdfvalue.RDFDatetime.Now()
expiration_time = now + lease_time
for job in self.cronjobs.values():
if cronjob_ids and job.job_id not in cronjob_ids:
continue
existing_lease = self.cronjob_leases.get(job.job_id)
if existing_lease is None or existing_lease[0] < now:
self.cronjob_leases[job.job_id] = (expiration_time,
utils.ProcessIdString())
job = job.Copy()
job.leased_until, job.leased_by = self.cronjob_leases[
job.job_id]
leased_jobs.append(job)
return leased_jobs
def LeaseMessageHandlerRequests(self,
lease_time=None,
limit=1000,
cursor=None):
"""Leases a number of message handler requests up to the indicated limit."""
now = rdfvalue.RDFDatetime.Now()
now_str = mysql_utils.RDFDatetimeToMysqlString(now)
expiry = now + lease_time
expiry_str = mysql_utils.RDFDatetimeToMysqlString(expiry)
query = ("UPDATE message_handler_requests "
"SET leased_until=%s, leased_by=%s "
"WHERE leased_until IS NULL OR leased_until < %s "
"LIMIT %s")
id_str = utils.ProcessIdString()
args = (expiry_str, id_str, now_str, limit)
updated = cursor.execute(query, args)
if updated == 0:
return []
cursor.execute(
"SELECT timestamp, request FROM message_handler_requests "
"WHERE leased_by=%s AND leased_until=%s LIMIT %s",
(id_str, expiry_str, updated))
res = []
for timestamp, request in cursor.fetchall():
req = objects.MessageHandlerRequest.FromSerializedString(request)
req.timestamp = mysql_utils.MysqlToRDFDatetime(timestamp)
req.leased_until = expiry
req.leased_by = id_str
res.append(req)
return res
def LeaseMessageHandlerRequests(self, lease_time=None, limit=1000):
"""Leases a number of message handler requests up to the indicated limit."""
leased_requests = []
now = rdfvalue.RDFDatetime.Now()
zero = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0)
expiration_time = now + lease_time
leases = self.message_handler_leases
for requests in self.message_handler_requests.values():
for r in requests.values():
existing_lease = leases.get(r.handler_name,
{}).get(r.request_id, zero)
if existing_lease < now:
leases.setdefault(r.handler_name,
{})[r.request_id] = expiration_time
r.leased_until = expiration_time
r.leased_by = utils.ProcessIdString()
leased_requests.append(r)
if len(leased_requests) >= limit:
break
return leased_requests
def testMessageHandlerRequestLeasing(self):
requests = [
rdf_objects.MessageHandlerRequest(
client_id="C.1000000000000000",
handler_name="Testhandler",
request_id=i * 100,
request=rdfvalue.RDFInteger(i)) for i in range(10)
]
lease_time = rdfvalue.Duration("5m")
with test_lib.FakeTime(rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)):
self.db.WriteMessageHandlerRequests(requests)
t0 = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100000)
with test_lib.FakeTime(t0):
t0_expiry = t0 + lease_time
leased = self.db.LeaseMessageHandlerRequests(
lease_time=lease_time, limit=5)
self.assertEqual(len(leased), 5)
for request in leased:
self.assertEqual(request.leased_until, t0_expiry)
self.assertEqual(request.leased_by, utils.ProcessIdString())
t1 = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100000 + 100)
with test_lib.FakeTime(t1):
t1_expiry = t1 + lease_time
leased = self.db.LeaseMessageHandlerRequests(
lease_time=lease_time, limit=5)
self.assertEqual(len(leased), 5)
for request in leased:
self.assertEqual(request.leased_until, t1_expiry)
self.assertEqual(request.leased_by, utils.ProcessIdString())
# Nothing left to lease.
leased = self.db.LeaseMessageHandlerRequests(
lease_time=lease_time, limit=2)
self.assertEqual(len(leased), 0)
read = self.db.ReadMessageHandlerRequests()
self.assertEqual(len(read), 10)
for r in read:
self.assertEqual(r.leased_by, utils.ProcessIdString())
self.assertEqual(len([r for r in read if r.leased_until == t0_expiry]), 5)
self.assertEqual(len([r for r in read if r.leased_until == t1_expiry]), 5)
# Half the leases expired.
t2 = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100000 + 350)
with test_lib.FakeTime(t2):
leased = self.db.LeaseMessageHandlerRequests(lease_time=lease_time)
self.assertEqual(len(leased), 5)
# All of them expired.
t3 = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100000 + 10350)
with test_lib.FakeTime(t3):
leased = self.db.LeaseMessageHandlerRequests(lease_time=lease_time)
self.assertEqual(len(leased), 10)