gc_collect_manager.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
  1. import gc
  2. import logging
  3. import threading
  4. import time
  5. from typing import Callable, Optional
  6. logger = logging.getLogger(__name__)
  7. class PythonGCThread(threading.Thread):
  8. """A background thread that triggers Python garbage collection.
  9. This thread waits for GC events from CoreWorker and triggers `gc.collect()` when
  10. when requested."""
  11. def __init__(self, *, gc_collect_func: Optional[Callable] = None):
  12. logger.debug("Starting Python GC thread")
  13. super().__init__(name="PythonGCThread", daemon=True)
  14. self._should_exit = False
  15. self._gc_event = threading.Event()
  16. # Sets the gc_collect_func (only for testing), defaults to gc.collect
  17. self._gc_collect_func = gc_collect_func or gc.collect
  18. def trigger_gc(self) -> None:
  19. self._gc_event.set()
  20. def run(self):
  21. while not self._should_exit:
  22. self._gc_event.wait()
  23. self._gc_event.clear()
  24. if self._should_exit:
  25. break
  26. try:
  27. start = time.monotonic()
  28. num_freed = self._gc_collect_func()
  29. if num_freed > 0:
  30. logger.debug(
  31. "gc.collect() freed {} refs in {} seconds".format(
  32. num_freed, time.monotonic() - start
  33. )
  34. )
  35. except Exception as e:
  36. logger.error(f"Error during GC: {e}")
  37. def stop(self):
  38. logger.debug("Stopping Python GC thread")
  39. self._should_exit = True
  40. self._gc_event.set()
  41. self.join()