From 5acbeb3d3c9889d623d008cdb8e766c1d198ddfc Mon Sep 17 00:00:00 2001
From: "Ajin.Abraham" <ajin.abraham@chime.com>
Date: Thu, 31 Aug 2023 20:40:54 -0700
Subject: [PATCH 1/2] configurable workers

---
 libsast/__init__.py                     |  2 +-
 libsast/common.py                       | 18 ++++++++++++++++++
 libsast/core_matcher/choice_matcher.py  | 10 ++--------
 libsast/core_matcher/pattern_matcher.py | 10 ++--------
 pyproject.toml                          |  2 +-
 5 files changed, 24 insertions(+), 18 deletions(-)

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..e9b3f1b 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,20 @@ 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."""
+    worker_count = 16
+    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:
+        pass
+    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 <ajin@opensecurity.in>"]

From 9ab0957974e38dac4e47fa1d25e917c258ad07bf Mon Sep 17 00:00:00 2001
From: "Ajin.Abraham" <ajin.abraham@chime.com>
Date: Thu, 31 Aug 2023 20:47:43 -0700
Subject: [PATCH 2/2] qa

---
 libsast/common.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/libsast/common.py b/libsast/common.py
index e9b3f1b..86ba3e1 100644
--- a/libsast/common.py
+++ b/libsast/common.py
@@ -74,7 +74,6 @@ def read_yaml(file_obj, text=False):
 
 def get_worker_count():
     """Get worker count for pool."""
-    worker_count = 16
     try:
         worker_count = os.cpu_count()
         if not worker_count:
@@ -85,5 +84,5 @@ def get_worker_count():
         if os.getenv('LIBSAST_WORKERS'):
             worker_count = int(os.getenv('LIBSAST_WORKERS'))
     except Exception:
-        pass
+        worker_count = 16
     return worker_count