mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2025-05-29 14:05:21 -04:00
144 lines
5.1 KiB
Python
144 lines
5.1 KiB
Python
__package__ = 'archivebox.actors'
|
|
|
|
import os
|
|
import time
|
|
import psutil
|
|
from typing import ClassVar, Generic, TypeVar, Any, cast, Literal
|
|
|
|
from django.db.models import QuerySet
|
|
from multiprocessing import Process, cpu_count
|
|
from threading import Thread, get_native_id
|
|
|
|
# from archivebox.logging_util import TimedProgress
|
|
|
|
ALL_SPAWNED_ACTORS: list[psutil.Process] = []
|
|
|
|
|
|
LaunchKwargs = dict[str, Any]
|
|
|
|
ObjectType = TypeVar('ObjectType')
|
|
|
|
class ActorType(Generic[ObjectType]):
|
|
pid: int
|
|
|
|
MAX_CONCURRENT_ACTORS: ClassVar[int] = min(max(2, int(cpu_count() * 0.7)), 8) # min 2, max 8
|
|
MAX_TICK_TIME: ClassVar[int] = 60
|
|
|
|
def __init__(self, mode: Literal['thread', 'process']='process', **launch_kwargs: LaunchKwargs):
|
|
self.mode = mode
|
|
self.launch_kwargs = launch_kwargs
|
|
|
|
@classmethod
|
|
def get_running_actors(cls) -> list[int]:
|
|
# returns a list of pids of all running actors of this type
|
|
return [
|
|
proc.pid for proc in ALL_SPAWNED_ACTORS
|
|
if proc.is_running() and proc.status() != 'zombie'
|
|
]
|
|
|
|
@classmethod
|
|
def spawn_actor(cls, mode: Literal['thread', 'process']='process', **launch_kwargs: LaunchKwargs) -> int:
|
|
actor = cls(mode=mode, **launch_kwargs)
|
|
# bg_actor_proccess = Process(target=actor.runloop)
|
|
if mode == 'thread':
|
|
bg_actor_thread = Thread(target=actor.runloop)
|
|
bg_actor_thread.start()
|
|
assert bg_actor_thread.native_id is not None
|
|
return bg_actor_thread.native_id
|
|
else:
|
|
bg_actor_process = Process(target=actor.runloop)
|
|
bg_actor_process.start()
|
|
assert bg_actor_process.pid is not None
|
|
ALL_SPAWNED_ACTORS.append(psutil.Process(pid=bg_actor_process.pid))
|
|
return bg_actor_process.pid
|
|
|
|
@classmethod
|
|
def get_queue(cls) -> QuerySet:
|
|
# return ArchiveResult.objects.filter(status='queued', extractor__in=('pdf', 'dom', 'screenshot'))
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def get_next(cls) -> ObjectType | None:
|
|
return cls.get_queue().last()
|
|
|
|
@classmethod
|
|
def get_actors_to_spawn(cls, queue, running_actors) -> list[LaunchKwargs]:
|
|
actors_to_spawn: list[LaunchKwargs] = []
|
|
max_spawnable = cls.MAX_CONCURRENT_ACTORS - len(running_actors)
|
|
queue_length = queue.count()
|
|
|
|
if not queue_length: # queue is empty, spawn 0 actors
|
|
return actors_to_spawn
|
|
elif queue_length > 10: # queue is long, spawn as many as possible
|
|
actors_to_spawn += max_spawnable * [{}]
|
|
elif queue_length > 5: # queue is medium, spawn 1 or 2 actors
|
|
actors_to_spawn += min(2, max_spawnable) * [{}]
|
|
else: # queue is short, spawn 1 actor
|
|
actors_to_spawn += min(1, max_spawnable) * [{}]
|
|
return actors_to_spawn
|
|
|
|
def on_startup(self):
|
|
if self.mode == 'thread':
|
|
self.pid = get_native_id()
|
|
else:
|
|
self.pid = os.getpid()
|
|
print('Actor on_startup()', f'pid={self.pid}')
|
|
# abx.pm.hook.on_actor_startup(self)
|
|
|
|
def on_shutdown(self):
|
|
print('Actor on_shutdown()', f'pid={self.pid}')
|
|
# abx.pm.hook.on_actor_shutdown(self)
|
|
|
|
def runloop(self):
|
|
self.on_startup()
|
|
|
|
rechecks = 30
|
|
|
|
while True:
|
|
obj_to_process: ObjectType | None = None
|
|
try:
|
|
obj_to_process = cast(ObjectType, self.get_next())
|
|
except Exception:
|
|
pass
|
|
|
|
if obj_to_process:
|
|
rechecks = 30
|
|
else:
|
|
if rechecks == 0:
|
|
break # stop looping and exit if queue is empty
|
|
else:
|
|
# print('Actor runloop()', f'pid={self.pid}', 'queue empty, rechecking...')
|
|
rechecks -= 1
|
|
time.sleep(1)
|
|
continue
|
|
|
|
if not self.lock(obj_to_process):
|
|
continue
|
|
|
|
# abx.pm.hook.on_actor_tick_start(self, obj_to_process)
|
|
try:
|
|
# timer = TimedProgress(self.MAX_TICK_TIME, prefix=' ')
|
|
|
|
# run the tick function on the object
|
|
self.tick(obj_to_process)
|
|
except Exception as err:
|
|
# abx.pm.hook.on_actor_tick_exception(self, obj_to_process, err)
|
|
print('ERROR: actor tick failed', err)
|
|
# refresh the db connection
|
|
from django import db
|
|
db.connections.close_all()
|
|
finally:
|
|
# timer.end()
|
|
pass
|
|
# abx.pm.hook.on_actor_tick_end(self, obj_to_process)
|
|
|
|
self.on_shutdown()
|
|
|
|
def tick(self, obj: ObjectType) -> None:
|
|
print('Actor Processing tick()', obj)
|
|
|
|
def lock(self, obj: ObjectType) -> bool:
|
|
print('Actor lock()', obj)
|
|
return True
|
|
|
|
|