diff --git a/libsast/__init__.py b/libsast/__init__.py index fe84e53..238fc65 100644 --- a/libsast/__init__.py +++ b/libsast/__init__.py @@ -8,7 +8,7 @@ __title__ = 'libsast' __authors__ = 'Ajin Abraham' __copyright__ = 'Copyright 2020 Ajin Abraham, OpenSecurity' -__version__ = '2.0.1' +__version__ = '2.0.2' __version_info__ = tuple(int(i) for i in __version__.split('.')) __all__ = [ 'Scanner', diff --git a/libsast/common.py b/libsast/common.py index c39b573..86ba3e1 100644 --- a/libsast/common.py +++ b/libsast/common.py @@ -1,5 +1,6 @@ # -*- coding: utf_8 -*- """Common Helpers.""" +import os import sys from threading import Thread @@ -69,3 +70,19 @@ def read_yaml(file_obj, text=False): except Exception as gen: raise YamlRuleLoadError( f'Failed to load YAML file: {repr(gen)}') + + +def get_worker_count(): + """Get worker count for pool.""" + try: + worker_count = os.cpu_count() + if not worker_count: + worker_count = 1 + if worker_count != 1 and sys.platform == 'win32': + # Work around https://bugs.python.org/issue26903 + worker_count = min(worker_count, 61) + if os.getenv('LIBSAST_WORKERS'): + worker_count = int(os.getenv('LIBSAST_WORKERS')) + except Exception: + worker_count = 16 + return worker_count diff --git a/libsast/core_matcher/choice_matcher.py b/libsast/core_matcher/choice_matcher.py index 537012f..d650a2c 100644 --- a/libsast/core_matcher/choice_matcher.py +++ b/libsast/core_matcher/choice_matcher.py @@ -1,7 +1,5 @@ # -*- coding: utf_8 -*- """Choice Macher.""" -import os -import sys from pathlib import Path from multiprocessing import ( Pool, @@ -48,15 +46,11 @@ def scan(self, paths: list) -> dict: choice_args.append((scan_paths, rule)) # Multiprocess Pool - worker_count = os.cpu_count() - if sys.platform == 'win32': - # Work around https://bugs.python.org/issue26903 - worker_count = min(worker_count, 61) - with Pool(worker_count) as pool: + with Pool(common.get_worker_count()) as pool: results = pool.starmap( self.choice_matcher, choice_args, - chunksize=10) + chunksize=1) self.add_finding(results) return self.findings diff --git a/libsast/core_matcher/pattern_matcher.py b/libsast/core_matcher/pattern_matcher.py index 5b2fee1..58ffadb 100644 --- a/libsast/core_matcher/pattern_matcher.py +++ b/libsast/core_matcher/pattern_matcher.py @@ -1,7 +1,5 @@ # -*- coding: utf_8 -*- """Pattern Macher.""" -import os -import sys from copy import deepcopy from operator import itemgetter from multiprocessing import ( @@ -50,15 +48,11 @@ def scan(self, paths: list) -> dict: continue files_to_scan.add(sfile) # Multiprocess Pool - worker_count = os.cpu_count() - if sys.platform == 'win32': - # Work around https://bugs.python.org/issue26903 - worker_count = min(worker_count, 61) - with Pool(worker_count) as pool: + with Pool(common.get_worker_count()) as pool: results = pool.map( self.pattern_matcher, files_to_scan, - chunksize=10) + chunksize=1) self.add_finding(results) return self.findings diff --git a/pyproject.toml b/pyproject.toml index 559b893..f13e317 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "libsast" -version = "2.0.1" +version = "2.0.2" description = "A generic SAST library built on top of semgrep and regex" keywords = ["libsast", "SAST", "Python SAST", "SAST API", "Regex SAST", "Pattern Matcher"] authors = ["Ajin Abraham "]