mrva/mrvacommander
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Separate queue and agent logic and refactor

Browse Source
This commit is contained in:
Nicolas Will
2024-06-15 23:12:11 +02:00
parent e107f6cf80
commit 7ea45cb176
5 changed files with 328 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files

304
cmd/agent/main.go
View File

@@ -1,73 +1,41 @@
package main
import (
"io"
"log"
"mrvacommander/pkg/codeql"
"mrvacommander/pkg/common"
"mrvacommander/pkg/agent"
"mrvacommander/pkg/queue"
"mrvacommander/pkg/storage"
"mrvacommander/utils"
"net/http"
"path/filepath"
"runtime"
"context"
"encoding/json"
"flag"
"fmt"
"os"
"os/signal"
"sync"
"strconv"
"syscall"
"time"
"github.com/google/uuid"
amqp "github.com/rabbitmq/amqp091-go"
"golang.org/x/exp/slog"
"flag"
"os"
"runtime"
"sync"
"github.com/elastic/go-sysinfo"
"golang.org/x/exp/slog"
)
func downloadFile(url string, dest string) error {
resp, err := http.Get(url)
if err != nil {
return fmt.Errorf("failed to download file: %w", err)
}
defer resp.Body.Close()
out, err := os.Create(dest)
if err != nil {
return fmt.Errorf("failed to create file: %w", err)
}
defer out.Close()
_, err = io.Copy(out, resp.Body)
if err != nil {
return fmt.Errorf("failed to copy file content: %w", err)
}
return nil
}
func calculateWorkers() int {
const workerMemoryGB = 2
const workerMemoryMB = 2048 // 2 GB
host, err := sysinfo.Host()
if err != nil {
log.Fatalf("failed to get host info: %v", err)
slog.Error("failed to get host info", "error", err)
os.Exit(1)
}
memInfo, err := host.Memory()
if err != nil {
log.Fatalf("failed to get memory info: %v", err)
slog.Error("failed to get memory info", "error", err)
os.Exit(1)
}
// Convert total memory to GB
totalMemoryGB := memInfo.Available / (1024 * 1024 * 1024)
// Get available memory in MB
totalMemoryMB := memInfo.Available / (1024 * 1024)
// Ensure we have at least one worker
workers := int(totalMemoryGB / workerMemoryGB)
workers := int(totalMemoryMB / workerMemoryMB)
if workers < 1 {
workers = 1
}
@@ -81,206 +49,6 @@ func calculateWorkers() int {
return workers
}
type RabbitMQQueue struct {
jobs chan common.AnalyzeJob
results chan common.AnalyzeResult
conn *amqp.Connection
channel *amqp.Channel
}
func InitializeQueue(jobsQueueName, resultsQueueName string) (*RabbitMQQueue, error) {
rabbitMQHost := os.Getenv("MRVA_RABBITMQ_HOST")
rabbitMQPort := os.Getenv("MRVA_RABBITMQ_PORT")
rabbitMQUser := os.Getenv("MRVA_RABBITMQ_USER")
rabbitMQPassword := os.Getenv("MRVA_RABBITMQ_PASSWORD")
if rabbitMQHost == "" || rabbitMQPort == "" || rabbitMQUser == "" || rabbitMQPassword == "" {
return nil, fmt.Errorf("RabbitMQ environment variables not set")
}
rabbitMQURL := fmt.Sprintf("amqp://%s:%s@%s:%s/", rabbitMQUser, rabbitMQPassword, rabbitMQHost, rabbitMQPort)
const (
tryCount = 5
retryDelaySec = 3
)
var conn *amqp.Connection
var err error
for i := 0; i < tryCount; i++ {
slog.Info("Attempting to connect to RabbitMQ", slog.Int("attempt", i+1))
conn, err = amqp.Dial(rabbitMQURL)
if err != nil {
slog.Warn("Failed to connect to RabbitMQ: %w", err)
if i < tryCount-1 {
slog.Info("Retrying in %d seconds", retryDelaySec)
time.Sleep(retryDelaySec * time.Second)
}
}
}
if err != nil {
return nil, fmt.Errorf("failed to connect to RabbitMQ: %w", err)
}
slog.Info("Connected to RabbitMQ")
ch, err := conn.Channel()
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to open a channel: %w", err)
}
_, err = ch.QueueDeclare(jobsQueueName, false, false, false, true, nil)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to declare tasks queue: %w", err)
}
_, err = ch.QueueDeclare(resultsQueueName, false, false, false, true, nil)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to declare results queue: %w", err)
}
err = ch.Qos(1, 0, false)
if err != nil {
conn.Close()
return nil, fmt.Errorf("failed to set QoS: %w", err)
}
return &RabbitMQQueue{
conn: conn,
channel: ch,
jobs: make(chan common.AnalyzeJob),
results: make(chan common.AnalyzeResult),
}, nil
}
func (q *RabbitMQQueue) Jobs() chan common.AnalyzeJob {
return q.jobs
}
func (q *RabbitMQQueue) Results() chan common.AnalyzeResult {
return q.results
}
func (q *RabbitMQQueue) StartAnalyses(analysis_repos *map[common.NameWithOwner]storage.DBLocation, session_id int, session_language string) {
slog.Info("Queueing codeql database analyze jobs")
}
func (q *RabbitMQQueue) Close() {
q.channel.Close()
q.conn.Close()
}
func (q *RabbitMQQueue) ConsumeJobs(queueName string) {
msgs, err := q.channel.Consume(queueName, "", true, false, false, false, nil)
if err != nil {
slog.Error("failed to register a consumer", slog.Any("error", err))
}
for msg := range msgs {
job := common.AnalyzeJob{}
err := json.Unmarshal(msg.Body, &job)
if err != nil {
slog.Error("failed to unmarshal job", slog.Any("error", err))
continue
}
q.jobs <- job
}
close(q.jobs)
}
func (q *RabbitMQQueue) PublishResults(queueName string) {
for result := range q.results {
q.publishResult(queueName, result)
}
}
func (q *RabbitMQQueue) publishResult(queueName string, result interface{}) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
resultBytes, err := json.Marshal(result)
if err != nil {
slog.Error("failed to marshal result", slog.Any("error", err))
return
}
slog.Info("Publishing result", slog.String("result", string(resultBytes)))
err = q.channel.PublishWithContext(ctx, "", queueName, false, false,
amqp.Publishing{
ContentType: "application/json",
Body: resultBytes,
})
if err != nil {
slog.Error("failed to publish result", slog.Any("error", err))
}
}
func RunAnalysisJob(job common.AnalyzeJob) (common.AnalyzeResult, error) {
var result = common.AnalyzeResult{
RequestId: job.RequestId,
ResultCount: 0,
ResultArchiveURL: "",
Status: common.StatusError,
}
// Log job info
slog.Info("Running analysis job", slog.Any("job", job))
// Create a temporary directory
tempDir := filepath.Join(os.TempDir(), uuid.New().String())
if err := os.MkdirAll(tempDir, 0755); err != nil {
return result, fmt.Errorf("failed to create temporary directory: %v", err)
}
defer os.RemoveAll(tempDir)
// Extract the query pack
// TODO: download from the 'job' query pack URL
utils.UntarGz("qp-54674.tgz", filepath.Join(tempDir, "qp-54674"))
// Perform the CodeQL analysis
runResult, err := codeql.RunQuery("google_flatbuffers_db.zip", "cpp", "qp-54674", 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)
}
// TODO: Upload the archive to storage
slog.Info("Results archive size", slog.Int("size", len(resultsArchive)))
slog.Info("Analysis job successful.")
result = common.AnalyzeResult{
RequestId: job.RequestId,
ResultCount: runResult.ResultCount,
ResultArchiveURL: "REPLACE_THIS_WITH_STORED_RESULTS_ARCHIVE",
Status: common.StatusSuccess,
}
return result, nil
}
func RunWorker(queue queue.Queue, wg *sync.WaitGroup) {
defer wg.Done()
for job := range queue.Jobs() {
result, err := RunAnalysisJob(job)
if err != nil {
slog.Error("failed to run analysis job", slog.Any("error", err))
continue
}
queue.Results() <- result
}
}
func main() {
slog.Info("Starting agent")
@@ -297,13 +65,27 @@ func main() {
}
for _, envVar := range requiredEnvVars {
if os.Getenv(envVar) == "" {
log.Fatalf("Fatal: Missing required environment variable %s", envVar)
if _, ok := os.LookupEnv(envVar); !ok {
slog.Error("Missing required environment variable %s", envVar)
os.Exit(1)
}
}
slog.Info("Initializing RabbitMQ connection")
rabbitMQQueue, err := InitializeQueue("tasks", "results")
rmqHost := os.Getenv("MRVA_RABBITMQ_HOST")
rmqPort := os.Getenv("MRVA_RABBITMQ_PORT")
rmqUser := os.Getenv("MRVA_RABBITMQ_USER")
rmqPass := os.Getenv("MRVA_RABBITMQ_PASSWORD")
rmqPortAsInt, err := strconv.Atoi(rmqPort)
if err != nil {
slog.Error("Failed to parse RabbitMQ port", slog.Any("error", err))
os.Exit(1)
}
slog.Info("Initializing RabbitMQ queue")
rabbitMQQueue, err := queue.InitializeRabbitMQQueue(rmqHost, int16(rmqPortAsInt), rmqUser, rmqPass)
if err != nil {
slog.Error("failed to initialize RabbitMQ", slog.Any("error", err))
os.Exit(1)
@@ -318,21 +100,21 @@ func main() {
var wg sync.WaitGroup
for i := 0; i < *workerCount; i++ {
wg.Add(1)
go RunWorker(rabbitMQQueue, &wg)
go agent.RunWorker(rabbitMQQueue, &wg)
}
slog.Info("Starting tasks consumer")
go rabbitMQQueue.ConsumeJobs("tasks")
slog.Info("Starting results publisher")
go rabbitMQQueue.PublishResults("results")
slog.Info("Agent startup complete")
// Gracefully exit on SIGINT/SIGTERM (TODO: add job cleanup)
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
<-sigChan
slog.Info("Shutting down agent")
close(rabbitMQQueue.results)
go func() {
<-sigChan
slog.Info("Shutting down agent")
rabbitMQQueue.Close()
os.Exit(0)
}()
select {}
}