Drop dynamic worker count; set default to 1
This commit is contained in:
committed by
=Michael Hohn
parent
75e57dc0a8
commit
bde8ac2db7
@@ -16,7 +16,7 @@ import (
|
|||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
slog.Info("Starting agent")
|
slog.Info("Starting agent")
|
||||||
workerCount := flag.Int("workers", 0, "number of workers")
|
workerCount := flag.Int("workers", 1, "number of workers")
|
||||||
logLevel := flag.String("loglevel", "info", "Set log level: debug, info, warn, error")
|
logLevel := flag.String("loglevel", "info", "Set log level: debug, info, warn, error")
|
||||||
flag.Parse()
|
flag.Parse()
|
||||||
|
|
||||||
|
|||||||
@@ -6,9 +6,7 @@ import (
|
|||||||
"log/slog"
|
"log/slog"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"runtime"
|
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
|
||||||
|
|
||||||
"github.com/hohn/mrvacommander/pkg/artifactstore"
|
"github.com/hohn/mrvacommander/pkg/artifactstore"
|
||||||
"github.com/hohn/mrvacommander/pkg/codeql"
|
"github.com/hohn/mrvacommander/pkg/codeql"
|
||||||
@@ -17,7 +15,6 @@ import (
|
|||||||
"github.com/hohn/mrvacommander/pkg/queue"
|
"github.com/hohn/mrvacommander/pkg/queue"
|
||||||
"github.com/hohn/mrvacommander/utils"
|
"github.com/hohn/mrvacommander/utils"
|
||||||
|
|
||||||
"github.com/elastic/go-sysinfo"
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -52,40 +49,8 @@ func (r *RunnerSingle) worker(wid int) {
|
|||||||
|
|
||||||
const (
|
const (
|
||||||
workerMemoryMB = 2048 // 2 GB
|
workerMemoryMB = 2048 // 2 GB
|
||||||
monitorIntervalSec = 10 // Monitor every 10 seconds
|
|
||||||
)
|
)
|
||||||
|
|
||||||
func calculateWorkers() int {
|
|
||||||
host, err := sysinfo.Host()
|
|
||||||
if err != nil {
|
|
||||||
slog.Error("failed to get host info", "error", err)
|
|
||||||
os.Exit(1)
|
|
||||||
}
|
|
||||||
|
|
||||||
memInfo, err := host.Memory()
|
|
||||||
if err != nil {
|
|
||||||
slog.Error("failed to get memory info", "error", err)
|
|
||||||
os.Exit(1)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Get available memory in MB
|
|
||||||
totalMemoryMB := memInfo.Available / (1024 * 1024)
|
|
||||||
|
|
||||||
// Ensure we have at least one worker
|
|
||||||
workers := int(totalMemoryMB / workerMemoryMB)
|
|
||||||
if workers < 1 {
|
|
||||||
workers = 1
|
|
||||||
}
|
|
||||||
|
|
||||||
// Limit the number of workers to the number of CPUs
|
|
||||||
cpuCount := runtime.NumCPU()
|
|
||||||
if workers > cpuCount {
|
|
||||||
workers = max(cpuCount, 1)
|
|
||||||
}
|
|
||||||
|
|
||||||
return workers
|
|
||||||
}
|
|
||||||
|
|
||||||
func StartAndMonitorWorkers(ctx context.Context,
|
func StartAndMonitorWorkers(ctx context.Context,
|
||||||
artifacts artifactstore.Store,
|
artifacts artifactstore.Store,
|
||||||
databases qldbstore.Store,
|
databases qldbstore.Store,
|
||||||
@@ -93,58 +58,30 @@ func StartAndMonitorWorkers(ctx context.Context,
|
|||||||
desiredWorkerCount int,
|
desiredWorkerCount int,
|
||||||
wg *sync.WaitGroup) {
|
wg *sync.WaitGroup) {
|
||||||
|
|
||||||
currentWorkerCount := 0
|
var workerCount int
|
||||||
stopChans := make([]chan struct{}, 0)
|
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))
|
||||||
|
}
|
||||||
|
|
||||||
if desiredWorkerCount != 0 {
|
stopChans := make([]chan struct{}, workerCount)
|
||||||
slog.Info("Starting workers", slog.Int("count", desiredWorkerCount))
|
|
||||||
for i := 0; i < desiredWorkerCount; i++ {
|
for i := 0; i < workerCount; i++ {
|
||||||
stopChan := make(chan struct{})
|
stopChan := make(chan struct{})
|
||||||
stopChans = append(stopChans, stopChan)
|
stopChans[i] = stopChan
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go RunWorker(ctx, artifacts, databases, queue, stopChan, wg)
|
go RunWorker(ctx, artifacts, databases, queue, stopChan, wg)
|
||||||
}
|
}
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
slog.Info("Worker count not specified, managing based on available memory and CPU")
|
// Wait for context cancellation
|
||||||
|
<-ctx.Done()
|
||||||
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case <-ctx.Done():
|
|
||||||
// signal all workers to stop
|
|
||||||
for _, stopChan := range stopChans {
|
for _, stopChan := range stopChans {
|
||||||
close(stopChan)
|
close(stopChan)
|
||||||
}
|
}
|
||||||
return
|
|
||||||
default:
|
|
||||||
newWorkerCount := calculateWorkers()
|
|
||||||
|
|
||||||
if newWorkerCount != currentWorkerCount {
|
|
||||||
slog.Info(
|
|
||||||
"Modifying worker count",
|
|
||||||
slog.Int("current", currentWorkerCount),
|
|
||||||
slog.Int("new", newWorkerCount))
|
|
||||||
}
|
|
||||||
|
|
||||||
if newWorkerCount > currentWorkerCount {
|
|
||||||
for i := currentWorkerCount; i < newWorkerCount; i++ {
|
|
||||||
stopChan := make(chan struct{})
|
|
||||||
stopChans = append(stopChans, stopChan)
|
|
||||||
wg.Add(1)
|
|
||||||
go RunWorker(ctx, artifacts, databases, queue, stopChan, wg)
|
|
||||||
}
|
|
||||||
} else if newWorkerCount < currentWorkerCount {
|
|
||||||
for i := newWorkerCount; i < currentWorkerCount; i++ {
|
|
||||||
close(stopChans[i])
|
|
||||||
}
|
|
||||||
stopChans = stopChans[:newWorkerCount]
|
|
||||||
}
|
|
||||||
currentWorkerCount = newWorkerCount
|
|
||||||
|
|
||||||
time.Sleep(monitorIntervalSec * time.Second)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// RunAnalysisJob runs a CodeQL analysis job (AnalyzeJob) returning an AnalyzeResult
|
// RunAnalysisJob runs a CodeQL analysis job (AnalyzeJob) returning an AnalyzeResult
|
||||||
|
|||||||
Reference in New Issue
Block a user