Coverage for grpc_server.py: 100%

1import time

2import uuid

3import grpc

4import redis

5import os

6import signal

7import sys

8import json

9from concurrent import futures

10from threading import Event, Lock, Thread

11from logger import *

12from urllib.parse import urlparse

14import patentflow_pb2 as pf_pb

15import patentflow_pb2_grpc as pf_grpc

17# metrics

18from prometheus_client import start_http_server, disable_created_metrics, Counter, Histogram

19from py_grpc_prometheus.prometheus_server_interceptor import PromServerInterceptor

21disable_created_metrics()

23JOBS_ENQUEUED = Counter(

24 'scraper_server_jobs_enqueued_total',

25 'Total jobs enqueued'

26)

27JOBS_COMPLETED = Counter(

28 'scraper_server_jobs_completed_total',

29 'Total of job misses'

30)

31JOBS_MISSES = Counter(

32 'scraper_server_jobs_misses_total',

33 'Total of job misses'

34)

35JOBS_FAILED = Counter(

36 'scraper_server_jobs_failed_total',

37 'Total jobs failed'

38)

39JOBS_TIMEOUT = Counter(

40 'scraper_server_jobs_timeout_total',

41 'Jobs that timed out waiting for worker'

42)

43JOBS_INVALID = Counter(

44 'scraper_server_jobs_invalid_total',

45 'Jobs rejected due to invalid request'

46)

47JOB_DURATION = Histogram(

48 'scraper_server_job_duration_seconds',

49 'Time from enqueue to completion/failure/timeout',

50 buckets=[5, 15, 30, 60, 120, 300, 600]

51)

52# -----------------------

54redis_conn = None

55shutdown_event = Event()

57class PatentService(pf_grpc.PatentServiceServicer):

58 def __init__(self, redis_conn, logger, shutdown_event=None, timeout=600, sleep_fn=None, jobs_before_cleanup=None):

59 self.redis_conn = redis_conn

60 self.logger = logger

61 self.shutdown_event = shutdown_event or Event()

62 self.timeout = timeout

63 self.sleep_fn = sleep_fn or time.sleep

64 self.jobs_before_cleanup = jobs_before_cleanup

65 self._stale_threshold = timeout + 60

66 self._jobs_since_cleanup = 0

67 self._cleanup_lock = Lock()

69 def _validate_request(self, request) -> tuple[bool, str, str]:

70 if not request.url or not request.url.strip():

71 return False, "INVALID_REQUEST", "URL is required"

73 try:

74 parsed = urlparse(request.url)

76 if not parsed.scheme:

77 raise

79 if not parsed.netloc:

80 raise

82 except Exception:

83 return False, "INVALID_REQUEST", f"Invalid URL format: URL - {request.url}. Check proto file docs"

85 if not request.proxy or not request.proxy.strip():

86 return False, "INVALID_REQUEST", "Proxy is required"

88 try:

89 parsed = urlparse(request.proxy)

91 if not parsed.scheme:

92 raise

94 if not parsed.netloc:

95 raise

97 if parsed.port is None:

98 raise

100 if not (1 <= parsed.port <= 65535):

101 raise

102

103 if bool(parsed.username) != bool(parsed.password):

104 raise

105

106 except Exception:

107 return False, "INVALID_REQUEST", f"Invalid proxy format: Proxy - {request.proxy}. Check proto file docs"

108

109 return True, "", ""

110

111 def GetPatent(self, request, context):

112 job_id = str(uuid.uuid4())

113

114 is_valid, error_code, error_message = self._validate_request(request)

115

116 if not is_valid:

117 JOBS_INVALID.inc()

118 self.logger.warning(f"Invalid request: {error_message}")

119 yield pf_pb.PatentEvent(

120 error=pf_pb.Error(

121 code=error_code,

122 message=error_message

123 )

124 )

125 return

126

127 pubsub = self.redis_conn.pubsub()

128 pubsub.subscribe(f"job_events:{job_id}")

129

130 self.redis_conn.hset(

131 f"job:{job_id}",

132 mapping={

133 "job_id": job_id,

134 "url": request.url,

135 "proxy": request.proxy,

136 "created_at": str(time.time()),

137 "status": "ENQUEUED"

138 }

139 )

140

141 self.redis_conn.lpush("job_queue", job_id)

142 JOBS_ENQUEUED.inc()

143 self.logger.info(f"Enqueued scrape job {job_id} for URL: {request.url}")

144 yield pf_pb.PatentEvent(

145 progress=pf_pb.Progress(

146 stage="ENQUEUED",

147 message="Job enqueued"

148 )

149 )

150

151 last_stage = "ENQUEUED"

152 start_time = time.time()

153

154 try:

155 deadline = start_time + self.timeout

156 while not shutdown_event.is_set():

157 remaining = deadline - time.time()

158 if remaining <= 0:

159 break

160

161 message = pubsub.get_message(timeout=min(5.0, remaining))

162 if message is None or message['type'] != 'message':

163 continue

164

165 job_status = message['data']

166

167 if job_status != last_stage:

168 last_stage = job_status

169

170 if job_status == "FETCHING":

171 yield pf_pb.PatentEvent(

172 progress=pf_pb.Progress(

173 stage="FETCHING",

174 message="Fetching patent page"

175 )

176 )

177

178 elif job_status == "COMPLETED":

179 JOB_DURATION.observe(time.time() - start_time)

180 yield from self._handle_completed(job_id)

181 return

182

183 elif job_status == "FAILED":

184 JOBS_FAILED.inc()

185 JOB_DURATION.observe(time.time() - start_time)

186 yield from self._handle_failed(job_id)

187 return

188

189 except Exception as e:

190 self.logger.exception(f"Critical error processing job {job_id}: {str(e)}")

191

192 if not isinstance(e, redis.ConnectionError):

193 try:

194 redis_conn.delete(f"job:{job_id}")

195 except:

196 pass

197 else:

198 self.logger.critical("Lost Redis connection")

199

200 context.abort(grpc.StatusCode.INTERNAL, str(e))

201

202 finally:

203 pubsub.unsubscribe()

204 pubsub.close()

205 self._maybe_cleanup()

206

207 JOBS_TIMEOUT.inc()

208 JOB_DURATION.observe(time.time() - start_time)

209

210 yield pf_pb.PatentEvent(

211 error=pf_pb.Error(

212 code="TIMEOUT",

213 message="Job processing timed out"

214 )

215 )

216

217 self.redis_conn.delete(f"job:{job_id}")

218

219 def _handle_completed(self, job_id: str):

220 result_json = self.redis_conn.hget(f"job:{job_id}", "result")

221 if result_json:

222 result_data = json.loads(result_json)

223

224 patent_data = result_data['patent_data']

225 patent_info = result_data['patent_info']

226

227 agents = []

228 for agent_data in patent_data.get("agents", []):

229 agents.append(pf_pb.Agent(

230 name=agent_data["name"],

231 address=agent_data["address"],

232 postcode=agent_data["postcode"],

233 country=agent_data["country"]

234 ))

235

236 package = pf_pb.PatentPackage(

237 patent=pf_pb.Patent(

238 name=patent_data["name"],

239 number=patent_data["number"],

240 pub_date=patent_data["pub_date"],

241 priority_date=patent_data["priority_date"],

242 countries=patent_data["countries"],

243 info=dict_to_patent_info(patent_info),

244 ),

245 applicant=pf_pb.Applicant(

246 name=patent_data["applicant"]["name"],

247 address=patent_data["applicant"]["address"],

248 postcode=patent_data["applicant"]["postcode"],

249 country=patent_data["applicant"]["country"]

250 ),

251 agents=agents

252 )

253

254 JOBS_COMPLETED.inc()

255 yield pf_pb.PatentEvent(package=package)

256 self.logger.info(f"Completed scrape job {job_id}")

257

258 self.redis_conn.delete(f"job:{job_id}")

259

260 def _handle_failed(self, job_id: str):

261 error_type = self.redis_conn.hget(f"job:{job_id}", "err_type")

262 error_msg = self.redis_conn.hget(f"job:{job_id}", "err_msg") or "Job failed"

263

264 if error_type is None:

265 raise Exception("No error type detected for failed job")

266

267 self.logger.error(f"Failed job {job_id}: [{error_type}] - {error_msg}")

268

269 if error_type == "CRITICAL":

270 JOBS_FAILED.inc()

271 yield pf_pb.PatentEvent(

272 error=pf_pb.Error(

273 code="JOB_FAILED",

274 message=error_msg

275 )

276 )

277 elif error_type == "RECOVERABLE":

278 JOBS_MISSES.inc()

279 yield pf_pb.PatentEvent(

280 error=pf_pb.Error(

281 code="JOB_MISSED",

282 message=error_msg

283 )

284 )

285 elif error_type == "INTERNAL":

286 JOBS_FAILED.inc()

287 yield pf_pb.PatentEvent(

288 error=pf_pb.Error(

289 code="JOB_FAILED",

290 message="Internal server error"

291 )

292 )

293 else:

294 JOBS_FAILED.inc()

295 raise Exception("Unknown error type detected for failed job")

296

297 self.redis_conn.delete(f"job:{job_id}")

298

299 def _maybe_cleanup(self):

300 if self.jobs_before_cleanup is None:

301 return

302 trigger = False

303 with self._cleanup_lock:

304 self._jobs_since_cleanup += 1

305 if self._jobs_since_cleanup >= self.jobs_before_cleanup:

306 self._jobs_since_cleanup = 0

307 trigger = True

308 if trigger:

309 self.logger.info(f"Reached jobs threshold ({self.jobs_before_cleanup}), intiating stale jobs cleanup")

310 Thread(target=self._cleanup_stale_jobs_bg, daemon=True).start()

311

312 def _cleanup_stale_jobs_bg(self):

313 now = time.time()

314 cursor = 0

315 cleaned = 0

316

317 try:

318 while True:

319 cursor, keys = self.redis_conn.scan(cursor, match="job:*", count=50)

320 for key in keys:

321 created_at = self.redis_conn.hget(key, "created_at")

322 if created_at and (now - float(created_at)) > self._stale_threshold:

323 self.redis_conn.delete(key)

324 cleaned += 1

325 if cursor == 0:

326 break

327 if cleaned:

328 self.logger.info(f"Stale cleanup: removed {cleaned} orphaned job hashes")

329 except Exception:

330 self.logger.exception("Error during background stale job cleanup")

331

332def create_redis_connection(host: str, port: int) -> redis.Redis:

333 pool = redis.ConnectionPool(

334 host=host,

335 port=port,

336 decode_responses=True,

337 socket_connect_timeout=5,

338 socket_timeout=5,

339 max_connections=50

340 )

341 return redis.Redis(connection_pool=pool)

342

343def create_server(patent_service: PatentService, max_workers: int = 10) -> grpc.Server:

344 server = grpc.server(

345 futures.ThreadPoolExecutor(max_workers=max_workers),

346 interceptors=[PromServerInterceptor()]

347 )

348 pf_grpc.add_PatentServiceServicer_to_server(patent_service, server)

349 server.add_insecure_port("[::]:50051")

350

351 return server

352

353def dict_to_patent_info(data: dict) -> pf_pb.PatentInfo:

354 msg = pf_pb.PatentInfo()

355

356 for country, fields in data["np"].items():

357 cd = pf_pb.CountryData()

358 for k, v in fields.items():

359 cd.fields[k] = str(v)

360 msg.np[country].CopyFrom(cd)

361

362 return msg

363

364# def serve():

365# server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))

366# pf_grpc.add_PatentServiceServicer_to_server(PatentService(), server)

367

368# server.add_insecure_port("[::]:50051")

369# server.start()

370# print("Patent Scraper API listening on :50051")

371

372# return server

373

374def cleanup_stale_jobs(redis_conn, logger):

375 redis_conn.delete("job_queue")

376

377 cursor = 0

378 deleted = 0

379 while True:

380 cursor, keys = redis_conn.scan(cursor, match="job:*", count=20)

381 if keys:

382 redis_conn.delete(*keys)

383 deleted += len(keys)

384 if cursor == 0:

385 break

386

387 if deleted > 0:

388 logger.info(f"Cleaned {deleted} stale jobs from previous run")

389

390def main(): # pragma: no cover

391 redis_host = os.getenv('REDIS_HOST', 'localhost')

392 redis_port = int(os.getenv('REDIS_PORT', 6379))

393

394 logger = setup_logging()

395

396 redis_conn = create_redis_connection(redis_host, redis_port)

397

398 try:

399 redis_conn.ping()

400 logger.info(f"Connected to Redis at {redis_host}:{redis_port}")

401 except redis.ConnectionError:

402 logger.exception(f"Failed to connect to Redis at {redis_host}:{redis_port}")

403 sys.exit(1)

404

405 cleanup_stale_jobs(redis_conn, logger)

406

407 shutdown_event = Event()

408 max_jobs_env = os.getenv('JOBS_BEFORE_CLEANUP')

409 max_jobs_before_cleanup = int(max_jobs_env) if max_jobs_env else None

410 patent_service = PatentService(redis_conn, logger, shutdown_event, jobs_before_cleanup=max_jobs_before_cleanup)

411 server = create_server(patent_service)

412

413 start_http_server(50052)

414 logger.info("Prometheus metrics exposed on :50052")

415

416 server.start()

417 logger.info("Patent Scraper API listening on :50051")

418

419 def shutdown(signum, frame):

420 logger.info("Shutdown signal received, draining gRPC server...")

421 patent_service.timeout = 30

422 shutdown_event.set()

423 server.stop(grace=30)

424

425 signal.signal(signal.SIGINT, shutdown)

426 signal.signal(signal.SIGTERM, shutdown)

427

428 try:

429 server.wait_for_termination()

430 finally:

431 if redis_conn:

432 redis_conn.close()

433 logger.info("grpc server stopped cleanly")

434

435if __name__ == "__main__": # pragma: no cover

436 main()