Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import asyncio 

2import os 

3import re 

4import sys 

5import signal 

6import random 

7import yaml 

8import ujson as json 

9import glob 

10import time 

11import shlex 

12import subprocess 

13import psutil 

14# import warnings 

15# warnings.simplefilter('always', DeprecationWarning) 

16 

17from datetime import datetime 

18from logging import debug, info, warn, error 

19from functools import partial 

20 

21# 3rd 

22from daemon import daemon, pidfile 

23from codenamize import codenamize 

24 

25# mine 

26from aiopb.aiopb import Hub 

27from gutools.tools import build_uri, parse_uri, snake_case, \ 

28 update_context, identity, yaml_encode, yaml_decode, wlikelyhood, \ 

29 fileiter, copyfile, dset, load_config, save_config, expandpath, \ 

30 _call, async_call, serializable_container, update_container, \ 

31 walk, rebuild 

32 

33from gutools.system import send_signals, SmartPIDLockFile, get_fingerprint, compare_fingerprint, find_process, test_process, test_alive 

34from gutools.uobjects import UObject, kget, vget 

35from gutools.persistence import DB, DBWorkspace, FSLayout 

36from gutools.unet import TL 

37from gutools.bash_completer import hasheable_cmdline 

38from gutools.uapplication import UApplication # TODO: needed? 

39from gutools.uagents import UAgent 

40from gutools.colors import * 

41 

42async def dummy_main(): 

43 foo = 1 

44 for i in range(20): 

45 print(f"xxx {i}") 

46 await asyncio.sleep(6) 

47 foo = 1 

48 

49 

50class UProcess(UObject): 

51 __args__ = ('name', 'alias', 'uri', 'state', 

52 'pid', 'restart', 'shutdown', 'date') 

53 __kwargs__ = {} 

54 __slots__ = __args__ + tuple(set(__kwargs__).difference(__args__)) 

55 

56class SessionWorkspace(DBWorkspace): 

57 scheme = """ 

58 CREATE TABLE IF NOT EXISTS {} 

59 ( 

60 {}, 

61 date DATETIME DEFAULT (datetime(CURRENT_TIMESTAMP, 'localtime')), 

62 PRIMARY KEY (name), 

63 UNIQUE (alias) 

64 ); 

65 """.format( 

66 snake_case(UProcess.__name__), 

67 ','.join(UProcess.__slots__[:-1]), 

68 ) 

69 

70 

71ext_decoder = { 

72 'yaml': yaml_decode, 

73 'json': json.decode, 

74 'pid': int, 

75 'out': identity, 

76 'err': identity, 

77} 

78 

79ext_encoder = { 

80 'yaml': yaml_encode, 

81 'json': json.encode, 

82 'pid': str, 

83 'out': identity, 

84 'err': identity, 

85} 

86 

87ext_aliases = { 

88 'fp': 'yaml', 

89} 

90 

91for k, v in ext_aliases.items(): 

92 ext_decoder[k] = ext_decoder[v] 

93 ext_encoder[k] = ext_encoder[v] 

94 

95def daemonize(f): 

96 """Decorator for daemonized methods""" 

97 return f 

98 

99ATTACHED_SCHEME = 'attached' 

100class USession(object): 

101 """Create a session workspace to store persistent information related to 

102 a process group. 

103 """ 

104 WORKSPACE = SessionWorkspace # Need to be overrided 

105 

106 # file config 

107 default_config = { 

108 'copy': { 

109 'bin': 'bin/*.*', 

110 'out': '', 

111 'pid': '', 

112 'fp': '', 

113 'err': '', 

114 }, 

115 'templates': { 

116 'scripts': 'bin/*.sh', 

117 }, 

118 'dependences': {}, 

119 } 

120 """Default configuration for a USession. Can be overrided.""" 

121 

122 def __init__(self, path): 

123 self.path = expandpath(path) 

124 self.layout = FSLayout(self.path) 

125 

126 self.db_file = self.layout.get_path('db', 'session', touch=True) 

127 self.db = DB(self.db_file) 

128 self.workspace = self.WORKSPACE(self.db) 

129 self.config = self.layout.get_content('etc', 'session', 

130 default=self.default_config) 

131 

132 def init(self): 

133 """Initialize the file structure, adding missing files from 

134 config['copy'] section. 

135 """ 

136 info(f'Initialize structure in {self.path}') 

137 

138 top = os.path.dirname(os.path.abspath(__file__)) 

139 for name, pattern in self.config.get('copy', {}).items(): 

140 self.layout.get_path('<folder>', name, touch=True) 

141 for path in fileiter(top, wildcard=pattern): 

142 where = os.path.join(self.path, path.split(top)[-1][1:]) 

143 copyfile(path, where, override=False) 

144 

145 def activate(self): 

146 """Activate Session Supervisor and optionally launch attached 

147 processes based on configuration. 

148 """ 

149 scheme = self.__class__.__name__.lstrip('U').lower() 

150 uri = f"{scheme}://supervisor" 

151 self.start_daemon(target=self._supervisor, uri=uri, alias='supervisor') 

152 

153 def deactivate(self): 

154 """Deactivate Session Supervisor and optionally attached processes 

155 based on configuration. 

156 """ 

157 def get_candidates(): 

158 for d, filename, pid in self.layout.iter_file_content('pid'): 

159 if test_alive(pid): 

160 for process in self.workspace.find(UProcess, name=d['name'], shutdown=1): 

161 # assert pid == process.pid 

162 yield d['name'], pid 

163 

164 def done(name): 

165 print(f' {RED}{name}\t\t[done]{RESET}') 

166 

167 # signals = [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM, signal.SIGKILL] 

168 signals = [signal.SIGQUIT, signal.SIGTERM, signal.SIGKILL] 

169 send_signals(get_candidates, signals, done_callback=done) 

170 

171 def start_daemon(self, target, **kw): 

172 """Start a service (target) as a daemon process. 

173 

174 `target` could be a function or an UApplication instance. 

175 """ 

176 context = self._get_daemon_context(**kw) 

177 with _call(daemon.DaemonContext, **context) as bel: 

178 try: 

179 context = self._prepare_start(**context) 

180 # run forever ... 

181 async_call(target, **context) 

182 except Exception as why: 

183 print(f"{YELLOW}ERROR: {RED}{why}{RESET}") 

184 

185 def stop(self, uri=None, name=None, **kw): 

186 """Stop a daemon process.""" 

187 # try to find the process in DB 

188 assert False, "TODO: remove dependeces from DB" 

189 for process in list(self.workspace.find(UProcess, uri=uri, name=name)): 

190 if process.pid: 

191 print(f"{BLUE}Killing '{process.name}' ", end='', flush=True) 

192 for tries in range(100): 

193 if test_process(process.pid, signal.SIGTERM): 

194 time.sleep(0.2) 

195 print(".", end='', flush=True) 

196 else: 

197 process.pid = None 

198 process.state = 'dead' 

199 # self.workspace.update(process) 

200 print(f"{RED}[dead]", flush=True) 

201 break 

202 # give a chance to end and pid has been removed 

203 else: 

204 print(f"{YELLOW}Process '{process.name}' is already dead") 

205 

206 # try to find process in pid file 

207 for filename in fileiter('pid', wildcard=f'{name}.pid'): 

208 test_process(filename, signal.SIGTERM) 

209 

210 print(f"{RESET}") 

211 

212 # alias=None, label='running', shutdown=True): 

213 def attach(self, **kw): 

214 """Attach a process/service to the session. 

215 

216 - `proc`: psutil.Process instance. 

217 - `alias`: human alias for the process. 

218 

219 As optional keywords: 

220 

221 - `restart`: whenever a process must restarted if dies. 

222 - `shutdown`: whenever a process must be shutdown when session is \ 

223 deactivated. 

224 """ 

225 proc = kw.get('proc') 

226 uri = kw.get('uri') 

227 name = kw.get('name') 

228 state = kw.setdefault('state', 'unknown') 

229 alias = kw.get('alias') 

230 

231 # test_alive(self.layout.get_content('etc')) 

232 # if not proc: 

233 # for process in self.workspace.find(**kw): 

234 

235 

236 if proc: 

237 update_context(kw, proc) 

238 fp = get_fingerprint(proc) 

239 kw.setdefault('cmdline', fp['cmdline']) 

240 uri = build_uri(scheme=ATTACHED_SCHEME, 

241 path=hasheable_cmdline(fp['cmdline'])) 

242 kw.setdefault('uri', uri) 

243 

244 if uri: 

245 name = kw.setdefault('name', codenamize(uri)) 

246 

247 if name: 

248 # save process config file 

249 wanted = ['alias', 'cmdline', 'name', ] 

250 conf = dict([(k, v) for (k, v) in kw.items() if k in wanted]) 

251 self.layout.update_content('etc', conf, name) 

252 if alias: 

253 self.layout.set_alias('etc', alias, name) 

254 

255 if proc and name: 

256 # create a fingerprint file for launching the file when needed 

257 # and include a pidfile in pid structure 

258 self.layout.set_content('pid', proc.pid, name) 

259 filename = self.layout.set_content('fp', fp, name, state) 

260 update_context(kw, filename=filename) 

261 

262 # insert process in DB 

263 process = UProcess(**kw) 

264 self.workspace.update(process) 

265 update_context(kw, process=process) 

266 return kw 

267 

268 def status(self): 

269 """**TBD:** Return session processes and status. 

270 """ 

271 raise NotImplemented("implementation pending") 

272 

273 def dettach_process(self, proc: psutil.Process = None, alias=None): 

274 """**TBD:** Dettach a process from the session. 

275 

276 Any of the following arguments can be used: 

277 

278 - `proc`: psutil.Process instance. 

279 - `alias`: human alias of the process. 

280 """ 

281 raise NotImplemented("implementation pending") 

282 

283 def add_dependence(self, parent, child): 

284 """Add a dependence between processes""" 

285 modified = False 

286 for parent in parent.split(','): 

287 for child in child.split(','): 

288 d = self.config.setdefault('dependences', dict()) 

289 # don't use set() to preserve serialization 

290 s = d.setdefault(child, list()) 

291 if parent not in s: 

292 s.append(parent) 

293 modified = True 

294 

295 if modified: 

296 self.layout.set_content('etc', self.config, 'session', ) 

297 

298 def remove_dependence(self, parent, child): 

299 """Remove a dependence between processes""" 

300 modified = False 

301 for parent in parent.split(','): 

302 for child in child.split(','): 

303 d = self.config.get('dependences', dict()) 

304 # don't use set() to preserve serialization 

305 s = d.get(parent, list()) 

306 if child in s: 

307 s.remove(child) 

308 modified = True 

309 

310 if modified: 

311 self.layout.set_content('etc', self.config, 'session', ) 

312 

313 def launch_process(self, match_if_same_cmdline=True, **kw): 

314 """Launch a previous captured process by name. 

315 

316 - `match_if_same_cmdline`: consider the process to be the same if \ 

317 cmdline is the same that stored one. 

318 - `kw`: contains search criteria keywords: 

319 Examples of search keywords are: 

320 

321 - name='adaptable-joint' 

322 - state='operative' 

323 - alias='master.tws' 

324 - restart=1 

325 - shutdown=0 

326 - uri='session://supervisor' 

327 

328 #. try to find a process that match criteria. 

329 #. if score is low but has the same command line (optional) 

330 we consider that is the same process in another non-registered state 

331 

332 #. update DB when process is found 

333 #. launch a new process using the command line from config. 

334 """ 

335 for process, score, proc, state, fp in self.find_processes(**kw): 

336 if proc: 

337 if score >= 0.9 or \ 

338 match_if_same_cmdline and fp['cmdline'] == proc.cmdline(): 

339 # self.layout.set_content('fp', fp, process.name, 'running') 

340 self.workspace.update(process, pid=process.pid, state=state) 

341 continue 

342 

343 # launch a new process if all dependeces are matched 

344 for dep in kget(self.config['dependences'], process.name, process.alias) or []: 

345 parent, state = dep.split(':') 

346 pname, parent = vget(processes, parent, 'name', 'alias') 

347 if not (parent and parent.state == state): 

348 break # any dependence is not matched 

349 else: 

350 # launch the process 

351 conf = self.layout.get_content('etc', process.name) 

352 conf = self._launch_process(**conf) 

353 

354 def _launch_process(self, **data): 

355 if 'cmdline' not in data: 

356 uri = parse_uri(data['uri']) 

357 data['cmdline'] = cmdline = [ 

358 uri['path'], 

359 ] 

360 for k, v in uri['query_'].items(): 

361 cmdline.append(f'{k}={v}') 

362 

363 conf = self._prepare_start(**data) 

364 name = conf['name'] 

365 

366 stdout = open(self.layout.get_path('out', name), 'w') 

367 stderr = open(self.layout.get_path('err', name), 'w') 

368 # conf = self.layout.get_content('etc', name) 

369 p = subprocess.Popen(conf['cmdline'], stdout=stdout, stderr=stderr) 

370 self.layout.set_content('pid', p.pid, name) 

371 

372 process = UProcess(**conf) 

373 self.workspace.update(process, pid=p.pid, state='starting') 

374 

375 # wait until process seems to be 'stable' or timeout 

376 fp0, fp, state = None, None, conf['state'] 

377 for trie in range(15): # timeout=30 secs 

378 proc = psutil.Process(p.pid) 

379 fp0, fp = fp, get_fingerprint(proc) 

380 self.layout.set_content('fp', fp, name, state) 

381 if fp0 == fp: 

382 break 

383 print(f"[{trie}] waiting for process '{process.alias or name}' to be stable...({state})") 

384 time.sleep(5) 

385 state = 'running' 

386 

387 update_context(conf, proc=proc, fp=fp, process=process) 

388 return conf 

389 

390 def _prepare_start(self, uri=None, name=None, alias=None, **conf): 

391 name = name or codenamize(uri) 

392 alias = alias or name 

393 

394 uri_ = parse_uri(uri) 

395 update_context(conf, uri_, uri_['query_']) 

396 realm = conf.setdefault('realm', uri_['path']) 

397 

398 conf.setdefault('cmdline', sys.argv) 

399 update_context(conf, uri=uri, name=name, alias=alias) 

400 self.layout.set_content('etc', conf, name) 

401 self.layout.set_alias('etc', alias, name) 

402 

403 conf['pid'] = pid = os.getpid() 

404 conf['state'] = 'starting' 

405 

406 # TODO: config, tl, hub, db, session, etc 

407 # loop = conf['loop'] = asyncio.get_event_loop() 

408 hub = conf['hub'] = Hub(realm=realm) 

409 

410 # TODO: add any maintenance/domestic procedures 

411 procedures = [dummy_main, ] 

412 app = _call(UApplication, procedures=procedures, 

413 fixtures=conf, **conf) 

414 app.tl = TL(app=app, uid=conf.get('uid', name)) 

415 

416 proc = psutil.Process(pid) 

417 fp = get_fingerprint(proc) 

418 # self.layout.set_content('fp', fp, name, conf['state']) 

419 process = UProcess(**conf) 

420 self.workspace.update(process) 

421 

422 update_context(conf, fp=fp, proc=proc, process=process, 

423 app=app, tl=app.tl, hub=app.hub) 

424 

425 

426 # app.run() 

427 # run(app.tl.start()) 

428 

429 return conf 

430 

431 def launch_service(self, uri, alias=None): 

432 """Launch a service (a method instance) parsing calling arguments\ 

433 from uri. 

434 

435 Example: 

436 

437 ``session.launch_service( 

438 uri="ib://tws:9090/download-historical?uid=historical&account=0&tws=tws.demo", 

439 alias="downloader")`` 

440 """ 

441 uri_ = parse_uri(uri) 

442 func = self 

443 for name in uri_['path'].split('/'): 

444 name = name.replace('-', '_') 

445 func = getattr(func, name, self) 

446 

447 alias = alias or name 

448 if func: 

449 self.start_daemon(target=func, uri=uri, alias=alias) 

450 else: 

451 raise RuntimeError(f"can't find a function linked with f{uri}") 

452 

453 def _supervisor(self): 

454 "Supervisor daemon tasks" 

455 

456 def check_attached_processes(): 

457 self.launch_process(restart=1) 

458 

459 def hide_check_attached_processes(): 

460 processes = dict() 

461 

462 # update fps states 

463 for process in list(self.workspace.find(UProcess, restart=1)): 

464 fps = self._get_available_fps(name=process.name) 

465 score, proc, label, fp = self._best_matching_fps( 

466 fps, process.pid) 

467 if score < 0.90: 

468 process.state, process.pid = None, None 

469 else: 

470 process.state, process.pid = label, proc.pid 

471 self.workspace.update(process) 

472 processes[process.name] = process 

473 

474 for name, process in processes.items(): 

475 if not process.pid: 

476 # check dependeces/pre-requsites 

477 for dep in kget(self.config['dependences'], process.name, process.alias) or []: 

478 parent, state = dep.split(':') 

479 pname, parent = vget(processes, 

480 parent, 'name', 'alias') 

481 if not (parent and parent.state == state): 

482 break 

483 else: 

484 # launch the process/service 

485 self.launch_process(name=process.name) 

486 # # TODO: process and service are the same? 

487 # if process.uri.startswith(ATTACHED_SCHEME): 

488 # self.launch_process(name=process.name) 

489 # else: 

490 # self.launch_service(process.uri) 

491 foo = 1 

492 

493 def clean_dead_pid(): 

494 count = 0 

495 for filename in fileiter('pid', wildcard='*.pid'): 

496 print(f" {YELLOW}Testing {filename:40}", end='') 

497 proc = test_alive(filename) 

498 name = os.path.basename(filename) 

499 name = os.path.splitext(name)[0] 

500 path = self.layout.get_path('etc', name) 

501 if os.path.exists(path) and proc: 

502 count += 1 

503 print(f"{GREEN}[ok] {proc.pid}") 

504 else: 

505 print(f"{RED}[dead]") 

506 os.remove(filename) 

507 print(f" {GREEN}[{count}]{YELLOW} running processes{RESET}") 

508 

509 FUNCTIONS = [clean_dead_pid, check_attached_processes] 

510 functions = [] 

511 

512 # for tries in range(n): 

513 while True: 

514 if not functions: 

515 functions.extend(FUNCTIONS) 

516 # random.shuffle(functions) 

517 if functions: 

518 func = functions.pop() 

519 print(f"{BLUE}- Running: {func.__name__}{RESET}") 

520 func() 

521 time.sleep(5) 

522 

523 print("Bye!") 

524 

525 def _get_daemon_context(self, **kw): 

526 gpath = self.layout.get_path 

527 s = partial(dset, kw) 

528 name = s('name', codenamize(kw['uri'])) 

529 

530 s('pid', gpath('pid', name)) 

531 s('pidfile', SmartPIDLockFile(path=kw['pid'])) 

532 s('files_preserve', list(range(0, 255))) 

533 s('stdin', None) 

534 s('stdout', open, gpath('out', name), 'w') 

535 s('stderr', open, gpath('err', name), 'w') 

536 s('working_directory', self.path) 

537 s('shutdown', 0) 

538 s('restart', 0) 

539 # s('signal_map', {}) 

540 return kw 

541 

542 def _demo_task(self, n): 

543 for i in range(n): 

544 print(f"doing something cool: {i}") 

545 time.sleep(5) 

546 print("-End demo task-") 

547 

548 def _get_available_fps(self, name=None): 

549 fingerprints = dict() 

550 if isinstance(name, str): 

551 regexp = re.compile(name, re.DOTALL) 

552 else: 

553 regexp = None 

554 for info, filename, fp0 in self.layout.iter_file_content('fp'): 

555 name = info['name'] 

556 if regexp and not regexp.match(name): 

557 continue 

558 fps = fingerprints.setdefault(name, dict()) 

559 fps[info['label']] = fp0 

560 

561 return fingerprints 

562 

563 def _best_matching_fp(self, fp0, *pids): 

564 pids = [p for p in pids if p] or psutil.pids() 

565 

566 best_score, best_proc, best_fp = 0, None, None 

567 for pid_ in pids: 

568 proc = psutil.Process(pid_) 

569 fp = get_fingerprint(proc) 

570 s = compare_fingerprint(fp0, fp) 

571 if s > best_score: 

572 best_score, best_proc, best_fp = s, proc, fp 

573 

574 return best_score, best_proc, best_fp 

575 

576 def _best_matching_fps(self, fps, pids=[], name=None, deep_serch=False): 

577 candidates = set([]) 

578 if not isinstance(pids, list): 

579 pids = [pids] 

580 pids = [p for p in pids if p] 

581 if not pids: 

582 deep_serch = True 

583 candidates.update(pids) 

584 

585 if deep_serch: 

586 candidates.update(psutil.pids()) 

587 

588 def do_search(pids): 

589 best_score, best_proc, best_label, best_fp = 0, None, None, None 

590 for pid_ in pids: 

591 # if pid_ < 1024: # not in Linux 

592 # continue # system proceces 

593 try: 

594 proc = psutil.Process(pid_) 

595 fp = get_fingerprint(proc) 

596 except Exception as why: 

597 continue # may be a zomby/protected process 

598 

599 for pname, fps_ in fps.items(): 

600 if name not in (pname, None): 

601 continue 

602 

603 for label, fp0 in fps_.items(): 

604 # if not label and exclude_no_label: 

605 # continue 

606 s = compare_fingerprint(fp0, fp) 

607 if s > best_score: 

608 best_score, best_proc, best_label, best_fp = s, proc, label, fp 

609 if s == 1.0: 

610 return best_score, best_proc, best_label, best_fp 

611 

612 return best_score, best_proc, best_label, best_fp 

613 

614 return do_search(candidates) 

615 

616 

617 def locate_process(self, **data): 

618 for process in self.workspace.find(UProcess, **data): 

619 fps = self._get_available_fps(name=process.name) 

620 score, proc, state, fp = self._best_matching_fps(fps, name=process.name) 

621 if score > 0.9: 

622 return score, proc, state, fp 

623 return [None] * 4 

624 

625 def find_processes(self, join='AND', **query): 

626 """Find running processes that match a query""" 

627 for process in self.workspace.find(UProcess, join=join, **query): 

628 fps = self._get_available_fps(name=process.name) 

629 score, proc, state, fp = self._best_matching_fps( 

630 fps, name=process.name) 

631 yield process, score, proc, state, fp 

632 

633 def _get_preferred_fp(self, name): 

634 desired_labels = ['operative', 'running', 'launch', ] 

635 candidates = dict() 

636 # for d, filename, fp0 in self.layout.iter_file_content('fp', name): 

637 # candidates[d['label']] = d, filename, fp0 

638 

639 for d, filename, fp0 in self.layout.iter_file_content('fp', name): 

640 candidates[d['label']] = d, filename, fp0 

641 

642 if not candidates: 

643 return {}, '', {} 

644 

645 desired_labels.extend(candidates.keys()) 

646 for label in desired_labels: 

647 if label not in candidates: 

648 continue 

649 d, filename, fp0 = candidates[label] 

650 break 

651 

652 return d, filename, fp0 

653 

654 def check_structure(self): 

655 """Check is a USession folder has been Initialized.""" 

656 debug(f'Checking structure in {self.path}') 

657 return os.path.exists(self.db_file) 

658 

659 # def load(self): 

660 # warnings.warn("Load is not longer necessary", DeprecationWarning, stacklevel=2) 

661 # self.config = load_config(self.session_config_file) 

662 # if self.config: 

663 # self.db = DB(self.db_file) 

664 # self.workspace = self.WORKSPACE(self.db) 

665 # else: 

666 # raise RuntimeError("Config not found: maybe need session init here ?") 

667 

668 def update_config(self, value, *keys): 

669 """Helper to save the session configuration.""" 

670 container = self.config 

671 if update_container(container, value, *keys): 

672 self.layout.set_content('etc', container, 'session')