package agent import ( "context" "fmt" "log/slog" "os" "path/filepath" "sync" "github.com/hohn/mrvacommander/pkg/artifactstore" "github.com/hohn/mrvacommander/pkg/codeql" "github.com/hohn/mrvacommander/pkg/common" "github.com/hohn/mrvacommander/pkg/qldbstore" "github.com/hohn/mrvacommander/pkg/queue" "github.com/hohn/mrvacommander/utils" "github.com/google/uuid" ) /* type RunnerSingle struct { queue queue.Queue } func NewAgentSingle(numWorkers int, v *Visibles) *RunnerSingle { r := RunnerSingle{queue: v.Queue} for id := 1; id <= numWorkers; id++ { go r.worker(id) } return &r } func (r *RunnerSingle) worker(wid int) { var job common.AnalyzeJob for { job = <-r.queue.Jobs() result, err := RunAnalysisJob(job) if err != nil { slog.Error("Failed to run analysis job", slog.Any("error", err)) continue } r.queue.Results() <- result } } */ const ( workerMemoryMB = 2048 // 2 GB ) func StartAndMonitorWorkers(ctx context.Context, artifacts artifactstore.Store, databases qldbstore.Store, queue queue.Queue, desiredWorkerCount int, wg *sync.WaitGroup) { var workerCount int if desiredWorkerCount > 0 { workerCount = desiredWorkerCount slog.Info("Starting fixed number of workers", slog.Int("count", workerCount)) } else { workerCount = 1 slog.Info("Starting preset number of workers", slog.Int("count", workerCount)) } stopChans := make([]chan struct{}, workerCount) for i := 0; i < workerCount; i++ { stopChan := make(chan struct{}) stopChans[i] = stopChan wg.Add(1) go RunWorker(ctx, artifacts, databases, queue, stopChan, wg) } // Wait for context cancellation <-ctx.Done() for _, stopChan := range stopChans { close(stopChan) } } // RunAnalysisJob runs a CodeQL analysis job (AnalyzeJob) returning an AnalyzeResult func RunAnalysisJob( job queue.AnalyzeJob, artifacts artifactstore.Store, dbs qldbstore.Store) (queue.AnalyzeResult, error) { var result = queue.AnalyzeResult{ Spec: job.Spec, ResultCount: 0, ResultLocation: artifactstore.ArtifactLocation{}, Status: common.StatusFailed, } // Create a temporary directory tempDir := filepath.Join(os.TempDir(), uuid.New().String()) if err := os.MkdirAll(tempDir, 0600); err != nil { return result, fmt.Errorf("failed to create temporary directory: %v", err) } defer os.RemoveAll(tempDir) // Download the query pack as a byte slice queryPackData, err := artifacts.GetQueryPack(job.QueryPackLocation) if err != nil { return result, fmt.Errorf("failed to download query pack: %w", err) } // Write the query pack data to the filesystem queryPackArchivePath := filepath.Join(tempDir, "query-pack.tar.gz") if err := os.WriteFile(queryPackArchivePath, queryPackData, 0600); err != nil { return result, fmt.Errorf("failed to write query pack archive to disk: %w", err) } // Make a directory and extract the query pack queryPackPath := filepath.Join(tempDir, "pack") if err := os.Mkdir(queryPackPath, 0600); err != nil { return result, fmt.Errorf("failed to create query pack directory: %w", err) } if err := utils.UntarGz(queryPackArchivePath, queryPackPath); err != nil { return result, fmt.Errorf("failed to extract query pack: %w", err) } databaseData, err := dbs.GetDatabase(job.Spec.NameWithOwner) if err != nil { slog.Error("Failed to get database", slog.String("owner", job.Spec.Owner), slog.String("repo", job.Spec.Repo), slog.Int("session_id", job.Spec.SessionID), slog.String("operation", "GetDatabase"), slog.Any("error", err), ) return result, fmt.Errorf("failed to get database for %s/%s: %w", job.Spec.Owner, job.Spec.Repo, err) } // Write the CodeQL database data to the filesystem databasePath := filepath.Join(tempDir, "database.zip") if err := os.WriteFile(databasePath, databaseData, 0600); err != nil { return result, fmt.Errorf("failed to write CodeQL database to disk: %w", err) } // Perform the CodeQL analysis runResult, err := codeql.RunQuery(databasePath, job.QueryLanguage, queryPackPath, tempDir) if err != nil { return result, fmt.Errorf("failed to run analysis: %w", err) } // Generate a ZIP archive containing SARIF and BQRS files resultsArchive, err := codeql.GenerateResultsZipArchive(runResult) if err != nil { return result, fmt.Errorf("failed to generate results archive: %w", err) } // Upload the archive to storage slog.Debug("Results archive size", slog.Int("size", len(resultsArchive))) resultsLocation, err := artifacts.SaveResult(job.Spec, resultsArchive) if err != nil { return result, fmt.Errorf("failed to save results archive: %w", err) } result = queue.AnalyzeResult{ Spec: job.Spec, ResultCount: runResult.ResultCount, ResultLocation: resultsLocation, Status: common.StatusSucceeded, SourceLocationPrefix: runResult.SourceLocationPrefix, DatabaseSHA: runResult.DatabaseSHA, } return result, nil } // RunWorker runs a worker that processes jobs from queue func RunWorker(ctx context.Context, artifacts artifactstore.Store, databases qldbstore.Store, queue queue.Queue, stopChan chan struct{}, 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() 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, artifacts, databases) 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 } } } }