Merge pull request #33 from ajinabraham/toomanyfiles

Configurable workers for multiprocessing

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',

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

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
 

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
 

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 <[email protected]>"]