Add dynamic worker management

pkg/agent/agent.go

@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"mrvacommander/pkg/codeql"
@@ -111,8 +112,7 @@ func RunAnalysisJob(job common.AnalyzeJob) (common.AnalyzeResult, error) {
 	}
 
 	// TODO: Upload the archive to storage
-	slog.Info("Results archive size", slog.Int("size", len(resultsArchive)))
-	slog.Info("Analysis job successful.")
+	slog.Debug("Results archive size", slog.Int("size", len(resultsArchive)))
 
 	result = common.AnalyzeResult{
 		RequestId:        job.RequestId,
@@ -125,16 +125,43 @@ func RunAnalysisJob(job common.AnalyzeJob) (common.AnalyzeResult, error) {
 }
 
 // RunWorker runs a worker that processes jobs from queue
-func RunWorker(queue queue.Queue, wg *sync.WaitGroup) {
+func RunWorker(ctx context.Context, stopChan chan struct{}, queue queue.Queue, wg *sync.WaitGroup) {
+	const (
+		WORKER_COUNT_STOP_MESSAGE   = "Worker stopping due to reduction in worker count"
+		WORKER_CONTEXT_STOP_MESSAGE = "Worker stopping due to context cancellation"
+	)
+
 	defer wg.Done()
-	for job := range queue.Jobs() {
-		slog.Info("Running analysis job", slog.Any("job", job))
-		result, err := RunAnalysisJob(job)
-		if err != nil {
-			slog.Error("Failed to run analysis job", slog.Any("error", err))
-			continue
+	slog.Info("Worker started")
+	for {
+		select {
+		case <-stopChan:
+			slog.Info(WORKER_COUNT_STOP_MESSAGE)
+			return
+		case <-ctx.Done():
+			slog.Info(WORKER_CONTEXT_STOP_MESSAGE)
+			return
+		default:
+			select {
+			case job, ok := <-queue.Jobs():
+				if !ok {
+					return
+				}
+				slog.Info("Running analysis job", slog.Any("job", job))
+				result, err := RunAnalysisJob(job)
+				if err != nil {
+					slog.Error("Failed to run analysis job", slog.Any("error", err))
+					continue
+				}
+				slog.Info("Analysis job completed", slog.Any("result", result))
+				queue.Results() <- result
+			case <-stopChan:
+				slog.Info(WORKER_COUNT_STOP_MESSAGE)
+				return
+			case <-ctx.Done():
+				slog.Info(WORKER_CONTEXT_STOP_MESSAGE)
+				return
+			}
 		}
-		slog.Info("Analysis job completed", slog.Any("result", result))
-		queue.Results() <- result
 	}
 }

pkg/codeql/codeql.go

@@ -115,7 +115,7 @@ func RunQuery(database string, nwo string, queryPackPath string, tempDir string)
 		databaseSHA = *dbMetadata.CreationMetadata.SHA
 	}
 
-	cmd := exec.Command(codeql.Path, "database", "run-queries", "--ram=1024", "--additional-packs", queryPackPath, "--", databasePath, queryPackPath)
+	cmd := exec.Command(codeql.Path, "database", "run-queries", "--ram=2048", "--additional-packs", queryPackPath, "--", databasePath, queryPackPath)
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return nil, fmt.Errorf("failed to run queries: %v\nOutput: %s", err, output)
 	}
@@ -133,9 +133,9 @@ func RunQuery(database string, nwo string, queryPackPath string, tempDir string)
 	shouldGenerateSarif := queryPackSupportsSarif(queryPackRunResults)
 
 	if shouldGenerateSarif {
-		slog.Info("Query pack supports SARIF")
+		slog.Debug("Query pack supports SARIF")
 	} else {
-		slog.Info("Query pack does not support SARIF")
+		slog.Debug("Query pack does not support SARIF")
 	}
 
 	var resultCount int
@@ -147,17 +147,17 @@ func RunQuery(database string, nwo string, queryPackPath string, tempDir string)
 			return nil, fmt.Errorf("failed to generate SARIF: %v", err)
 		}
 		resultCount = getSarifResultCount(sarif)
-		slog.Info("Generated SARIF", "resultCount", resultCount)
+		slog.Debug("Generated SARIF", "resultCount", resultCount)
 		sarifFilePath = filepath.Join(resultsDir, "results.sarif")
 		if err := os.WriteFile(sarifFilePath, sarif, 0644); err != nil {
 			return nil, fmt.Errorf("failed to write SARIF file: %v", err)
 		}
 	} else {
 		resultCount = queryPackRunResults.TotalResultsCount
-		slog.Info("Did not generate SARIF", "resultCount", resultCount)
+		slog.Debug("Did not generate SARIF", "resultCount", resultCount)
 	}
 
-	slog.Info("Adjusting BQRS files")
+	slog.Debug("Adjusting BQRS files")
 	bqrsFilePaths, err := adjustBqrsFiles(queryPackRunResults, resultsDir)
 	if err != nil {
 		return nil, fmt.Errorf("failed to adjust BQRS files: %v", err)

pkg/queue/rabbitmq.go

@@ -151,7 +151,7 @@ func (q *RabbitMQQueue) publishResult(queueName string, result interface{}) {
 		return
 	}
 
-	slog.Info("Publishing result", slog.String("result", string(resultBytes)))
+	slog.Debug("Publishing result", slog.String("result", string(resultBytes)))
 	err = q.channel.PublishWithContext(ctx, "", queueName, false, false,
 		amqp.Publishing{
 			ContentType: "application/json",