CERT-Polska / karton

Distributed malware processing framework based on Python, Redis and S3.
https://karton-core.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
381 stars 45 forks source link

Could be interesting for you: karton-memory-watcher #232

Open rakovskij-stanislav opened 11 months ago

rakovskij-stanislav commented 11 months ago

Hello! I wrote a prehook&posthook addon for karton consumers to automatically exit (to cause restart in docker / systemctl) if process uses too much RAM.

https://github.com/rakovskij-stanislav/karton_memory_watcher

I noticed that some forensic libraries can cause a slow memory leak (in my case ~500 MB for 2-3 month), so the restart is required. To make it softly (for not losing current task) I am making a decicion of restarting in post-hook logic.

Main module

It has 4 configurable rules for restart:

Usage is simple:

from karton.core import Consumer
from karton.memory_watcher import implant_watcher, RestartBehavior, RestartRule

...
class YourConsumer(Consumer):
    pass

if __name__ == "__main__":
    foobar = YourConsumer()
    implant_watcher(
        foobar,
        RestartRule(
            extra_consumed_memory_percent=80,
            # call_before_exit=(close_db_connections, )
        )
    )
    foobar.loop()

I hope this post-hook will help in long-term runs :)