Coverage for dispatch/job.py: 97%

1"""Job lifecycle: enqueue / claim / renew / reclaim (spec §7.2, §7.3, §8, §9). 

2 

3This is the Python skin over the Lua state transitions in ``scripts.py``. Every 

4transition is atomic in Redis; this layer only picks the ``now`` clock (Python 

5``time.time()``, passed into every script via ARGV — never a server-side TIME), 

6walks ``jobs:leased`` for the time-gated orphan scan, and shapes payloads. 

7 

8Slice 2 of the delivery plan (§15.2) — dark: nothing calls these yet. The 

9``complete``/``abort``/JE-landing transitions are slice 3; presigned ``code_url`` 

10signing on the payload is slice 4. Both are called out at their seams below. 

11""" 

12 

13import time 

14from typing import Dict, Optional 

15 

16from ulid import ULID 

17 

18from mongo.utils import RedisCache 

19from . import params 

20from . import redis_keys 

21from . import scripts 

22 

23JOB_ID_PREFIX = 'jb_' 

24 

25# One shared RedisCache instance, mirroring runner.py: under fakeredis each 

26# RedisCache owns an isolated dataset, so enqueue/claim/renew must share one to 

27# stay coherent; in production they all share the pooled real Redis anyway. 

28_cache: Optional[RedisCache] = None 

29 

30 

31def _redis(): 

32 global _cache 

33 if _cache is None: 

34 _cache = RedisCache() 

35 return _cache.client 

36 

37 

38def _now() -> float: 

39 return time.time() 

40 

41 

42def _decode(value) -> Optional[str]: 

43 if value is None: 

44 return None 

45 return value.decode() if isinstance(value, bytes) else value 

46 

47 

48def enqueue_job( 

49 submission_id: str, 

50 problem_id: str, 

51 language: int, 

52 code_minio_path: str, 

53 checker: str, 

54 tasks_meta_json: str, 

55) -> str: 

56 """Enqueue a fresh judging job for ``submission_id`` (spec §8, §9). 

57 

58 Writes the job hash, points ``submission:<sid>:current_job`` at it (last 

59 enqueue wins — that is exactly how a rejudge supersedes an in-flight job; the 

60 stale one is destroyed lazily on the dispatch side by ``claim_pending``, INV4, 

61 never proactively here), then LPUSHes onto the FIFO pending queue. 

62 """ 

63 now = _now() 

64 job_id = JOB_ID_PREFIX + str(ULID()) 

65 

66 client = _redis() 

67 job_key = redis_keys.job(job_id) 

68 

69 pipe = client.pipeline() 

70 pipe.hset( 

71 job_key, 

72 mapping={ 

73 'submission_id': submission_id, 

74 'problem_id': problem_id, 

75 'language': language, 

76 'code_minio_path': code_minio_path, 

77 'checker': checker, 

78 'tasks_meta_json': tasks_meta_json, 

79 'leased_by': '', 

80 'lease_deadline': '0', 

81 'state': 'pending', 

82 'attempts': 0, 

83 'created_at': repr(now), 

84 'last_error': '', 

85 }, 

86 ) 

87 pipe.set(redis_keys.submission_current_job(submission_id), job_id) 

88 pipe.lpush(redis_keys.JOBS_PENDING, job_id) 

89 pipe.execute() 

90 

91 return job_id 

92 

93 

94def claim_next_job(runner_id: str) -> Optional[Dict]: 

95 """Hand ``runner_id`` its next job, or None when there is nothing to do. 

96 

97 Order matches spec §7.3: first a time-gated orphan scan (at most one runner 

98 per ``ORPHAN_SCAN_INTERVAL_SEC`` window), whose first successful reclaim goes 

99 straight to this caller; otherwise fall through to the pending queue. 

100 """ 

101 now = _now() 

102 client = _redis() 

103 

104 reclaimed = _orphan_scan(client, runner_id, now) 

105 if reclaimed is not None: 

106 return _payload(client, reclaimed) 

107 

108 job_id = scripts.load(client).claim_pending( 

109 keys=[redis_keys.JOBS_PENDING, redis_keys.JOBS_LEASED], 

110 args=[now, runner_id, params.LEASE_TTL_SEC], 

111 ) 

112 if job_id is None: 

113 return None 

114 return _payload(client, _decode(job_id)) 

115 

116 

117def renew_lease(runner_id: str, job_id: str) -> bool: 

118 """Extend the lease on ``job_id`` iff ``runner_id`` still owns it (spec §7.2). 

119 

120 Thin wrapper over ``renew_lease`` Lua; heartbeat wiring is slice 4. 

121 """ 

122 now = _now() 

123 result = scripts.load(_redis()).renew_lease( 

124 keys=[redis_keys.JOBS_LEASED, 

125 redis_keys.job(job_id)], 

126 args=[job_id, runner_id, now, params.LEASE_TTL_SEC], 

127 ) 

128 return result == 1 

129 

130 

131def _orphan_scan(client, runner_id: str, now: float) -> Optional[str]: 

132 """Time-gated sweep of ``jobs:leased``; return an id reclaimed for the caller. 

133 

134 The ``SET NX EX`` gate on ``dispatch:last_recovery`` lets at most one scan run 

135 per window (spec §7.3 step 1). Within a winning scan, each expired candidate 

136 is offered to ``reclaim_expired``: the first ``1`` (reclaimed to this caller) 

137 is returned immediately; ``-1`` (attempts exhausted) candidates are skipped 

138 and left in place — the JE landing sweep that removes them is slice 3. 

139 """ 

140 acquired = client.set( 

141 redis_keys.DISPATCH_LAST_RECOVERY, 

142 '1', 

143 nx=True, 

144 ex=params.ORPHAN_SCAN_INTERVAL_SEC, 

145 ) 

146 if not acquired: 

147 return None 

148 

149 reclaim = scripts.load(client).reclaim_expired 

150 for member in client.smembers(redis_keys.JOBS_LEASED): 

151 job_id = _decode(member) 

152 result = reclaim( 

153 keys=[redis_keys.JOBS_LEASED, 

154 redis_keys.job(job_id)], 

155 args=[ 

156 job_id, 

157 runner_id, 

158 now, 

159 params.LEASE_TTL_SEC, 

160 params.MAX_ATTEMPTS, 

161 ], 

162 ) 

163 if result == 1: 

164 return job_id 

165 return None 

166 

167 

168def _payload(client, job_id: str) -> Optional[Dict]: 

169 """Shape a job hash into a claim payload (spec §7.3), or None if it vanished. 

170 

171 Returns ``job_id`` plus the raw hash fields the HTTP layer will need. Typed 

172 coercion and the presigned ``code_url`` are slice 4's concern, not this one's. 

173 """ 

174 raw = client.hgetall(redis_keys.job(job_id)) 

175 if not raw: 

176 return None 

177 fields: Dict[str, Optional[str]] = { 

178 _decode(k): _decode(v) 

179 for k, v in raw.items() 

180 } 

181 fields['job_id'] = job_id 

182 return fields