CM-45588 - Make batching more configurable and friendly in logs (#284)

Author: MarshalX

Dockerfile

@@ -6,7 +6,7 @@ FROM base AS builder
 ENV POETRY_VERSION=1.8.3
 
 # deps are required to build cffi
-RUN apk add --no-cache --virtual .build-deps gcc=14.2.0-r4 libffi-dev=3.4.6-r0 musl-dev=1.2.5-r9 &&  \
+RUN apk add --no-cache --virtual .build-deps gcc=14.2.0-r4 libffi-dev=3.4.7-r0 musl-dev=1.2.5-r9 &&  \
     pip install --no-cache-dir "poetry==$POETRY_VERSION" "poetry-dynamic-versioning[plugin]" &&  \
     apk del .build-deps gcc libffi-dev musl-dev
 

cycode/cli/commands/scan/code_scanner.py

@@ -171,7 +171,7 @@ def _scan_batch_thread_func(batch: List[Document]) -> Tuple[str, CliError, Local
         should_use_sync_flow = _should_use_sync_flow(command_scan_type, scan_type, sync_option, scan_parameters)
 
         try:
-            logger.debug('Preparing local files, %s', {'batch_size': len(batch)})
+            logger.debug('Preparing local files, %s', {'batch_files_count': len(batch)})
             zipped_documents = zip_documents(scan_type, batch)
             zip_file_size = zipped_documents.size
             scan_result = perform_scan(

cycode/cli/consts.py

@@ -145,7 +145,11 @@
 # scan in batches
 DEFAULT_SCAN_BATCH_MAX_SIZE_IN_BYTES = 9 * 1024 * 1024
 SCAN_BATCH_MAX_SIZE_IN_BYTES = {SAST_SCAN_TYPE: 50 * 1024 * 1024}
+SCAN_BATCH_MAX_SIZE_IN_BYTES_ENV_VAR_NAME = 'SCAN_BATCH_MAX_SIZE_IN_BYTES'
+
 DEFAULT_SCAN_BATCH_MAX_FILES_COUNT = 1000
+SCAN_BATCH_MAX_FILES_COUNT_ENV_VAR_NAME = 'SCAN_BATCH_MAX_FILES_COUNT'
+
 # if we increase this values, the server doesn't allow connecting (ConnectionError)
 SCAN_BATCH_MAX_PARALLEL_SCANS = 5
 SCAN_BATCH_SCANS_PER_CPU = 1

cycode/cli/utils/scan_batch.py

@@ -5,25 +5,83 @@
 from cycode.cli import consts
 from cycode.cli.models import Document
 from cycode.cli.utils.progress_bar import ScanProgressBarSection
+from cycode.cyclient import logger
 
 if TYPE_CHECKING:
     from cycode.cli.models import CliError, LocalScanResult
     from cycode.cli.utils.progress_bar import BaseProgressBar
 
 
+def _get_max_batch_size(scan_type: str) -> int:
+    logger.debug(
+        'You can customize the batch size by setting the environment variable "%s"',
+        consts.SCAN_BATCH_MAX_SIZE_IN_BYTES_ENV_VAR_NAME,
+    )
+
+    custom_size = os.environ.get(consts.SCAN_BATCH_MAX_SIZE_IN_BYTES_ENV_VAR_NAME)
+    if custom_size:
+        logger.debug('Custom batch size is set, %s', {'custom_size': custom_size})
+        return int(custom_size)
+
+    return consts.SCAN_BATCH_MAX_SIZE_IN_BYTES.get(scan_type, consts.DEFAULT_SCAN_BATCH_MAX_SIZE_IN_BYTES)
+
+
+def _get_max_batch_files_count(_: str) -> int:
+    logger.debug(
+        'You can customize the batch files count by setting the environment variable "%s"',
+        consts.SCAN_BATCH_MAX_FILES_COUNT_ENV_VAR_NAME,
+    )
+
+    custom_files_count = os.environ.get(consts.SCAN_BATCH_MAX_FILES_COUNT_ENV_VAR_NAME)
+    if custom_files_count:
+        logger.debug('Custom batch files count is set, %s', {'custom_files_count': custom_files_count})
+        return int(custom_files_count)
+
+    return consts.DEFAULT_SCAN_BATCH_MAX_FILES_COUNT
+
+
 def split_documents_into_batches(
+    scan_type: str,
     documents: List[Document],
-    max_size: int = consts.DEFAULT_SCAN_BATCH_MAX_SIZE_IN_BYTES,
-    max_files_count: int = consts.DEFAULT_SCAN_BATCH_MAX_FILES_COUNT,
 ) -> List[List[Document]]:
+    max_size = _get_max_batch_size(scan_type)
+    max_files_count = _get_max_batch_files_count(scan_type)
+
+    logger.debug(
+        'Splitting documents into batches, %s',
+        {'document_count': len(documents), 'max_batch_size': max_size, 'max_files_count': max_files_count},
+    )
+
     batches = []
 
     current_size = 0
     current_batch = []
     for document in documents:
         document_size = len(document.content.encode('UTF-8'))
 
-        if (current_size + document_size > max_size) or (len(current_batch) >= max_files_count):
+        exceeds_max_size = current_size + document_size > max_size
+        if exceeds_max_size:
+            logger.debug(
+                'Going to create new batch because current batch size exceeds the limit, %s',
+                {
+                    'batch_index': len(batches),
+                    'current_batch_size': current_size + document_size,
+                    'max_batch_size': max_size,
+                },
+            )
+
+        exceeds_max_files_count = len(current_batch) >= max_files_count
+        if exceeds_max_files_count:
+            logger.debug(
+                'Going to create new batch because current batch files count exceeds the limit, %s',
+                {
+                    'batch_index': len(batches),
+                    'current_batch_files_count': len(current_batch),
+                    'max_batch_files_count': max_files_count,
+                },
+            )
+
+        if exceeds_max_size or exceeds_max_files_count:
             batches.append(current_batch)
 
             current_batch = [document]
@@ -35,6 +93,8 @@ def split_documents_into_batches(
     if current_batch:
         batches.append(current_batch)
 
+    logger.debug('Documents were split into batches %s', {'batches_count': len(batches)})
+
     return batches
 
 
@@ -49,9 +109,8 @@ def run_parallel_batched_scan(
     documents: List[Document],
     progress_bar: 'BaseProgressBar',
 ) -> Tuple[Dict[str, 'CliError'], List['LocalScanResult']]:
-    max_size = consts.SCAN_BATCH_MAX_SIZE_IN_BYTES.get(scan_type, consts.DEFAULT_SCAN_BATCH_MAX_SIZE_IN_BYTES)
-
-    batches = [documents] if scan_type == consts.SCA_SCAN_TYPE else split_documents_into_batches(documents, max_size)
+    # batching is disabled for SCA; requested by Mor
+    batches = [documents] if scan_type == consts.SCA_SCAN_TYPE else split_documents_into_batches(scan_type, documents)
 
     progress_bar.set_section_length(ScanProgressBarSection.SCAN, len(batches))  # * 3
     # TODO(MarshalX): we should multiply the count of batches in SCAN section because each batch has 3 steps:
@@ -61,9 +120,13 @@ def run_parallel_batched_scan(
     # it's not possible yet because not all scan types moved to polling mechanism
     # the progress bar could be significant improved (be more dynamic) in the future
 
+    threads_count = _get_threads_count()
     local_scan_results: List['LocalScanResult'] = []
     cli_errors: Dict[str, 'CliError'] = {}
-    with ThreadPool(processes=_get_threads_count()) as pool:
+
+    logger.debug('Running parallel batched scan, %s', {'threads_count': threads_count, 'batches_count': len(batches)})
+
+    with ThreadPool(processes=threads_count) as pool:
         for scan_id, err, result in pool.imap(scan_function, batches):
             if result:
                 local_scan_results.append(result)