monitor: add stat for worker queue

Author: criyle

cmd/go-judge/main.go

@@ -508,7 +508,7 @@ func newEnvPool(b pool.EnvBuilder, enableMetrics bool) worker.EnvironmentPool {
 }
 
 func newWorker(conf *config.Config, envPool worker.EnvironmentPool, fs filestore.FileStore) worker.Worker {
-	return worker.New(worker.Config{
+	w := worker.New(worker.Config{
 		FileStore:             fs,
 		EnvironmentPool:       envPool,
 		Parallelism:           conf.Parallelism,
@@ -520,6 +520,10 @@ func newWorker(conf *config.Config, envPool worker.EnvironmentPool, fs filestore
 		OpenFileLimit:         uint64(conf.OpenFileLimit),
 		ExecObserver:          execObserve,
 	})
+	if conf.EnableMetrics {
+		w = newMetricsWorker(w)
+	}
+	return w
 }
 
 func newForceGCWorker(conf *config.Config) {

cmd/go-judge/metrics.go

@@ -16,6 +16,7 @@ const (
 	execSubsystem        = "exec"
 	filestoreSubsystem   = "file"
 	environmentSubsystem = "environment"
+	workerSubsystem      = "worker"
 )
 
 var (
@@ -88,6 +89,16 @@ var (
 		Name:      "current_count",
 		Help:      "Total number of environment currently in use",
 	})
+
+	workerQueue = prometheus.NewDesc(
+		prometheus.BuildFQName(metricsNamespace, workerSubsystem, "queue_count"),
+		"Number of requests waiting in worker queue", nil, nil,
+	)
+
+	workerRunning = prometheus.NewDesc(
+		prometheus.BuildFQName(metricsNamespace, workerSubsystem, "running_count"),
+		"Number of request running by workers", nil, nil,
+	)
 )
 
 func init() {
@@ -205,3 +216,32 @@ func (p *metricsEnvPool) Put(env envexec.Environment) {
 	p.EnvironmentPool.Put(env)
 	envInUse.Dec()
 }
+
+var _ worker.Worker = &metricsWorker{}
+var _ prometheus.Collector = &metricsWorker{}
+
+type metricsWorker struct {
+	worker.Worker
+}
+
+// Collect implements prometheus.Collector.
+func (m *metricsWorker) Collect(ch chan<- prometheus.Metric) {
+	s := m.Stat()
+	ch <- prometheus.MustNewConstMetric(
+		workerQueue, prometheus.GaugeValue, float64(s.Queue),
+	)
+	ch <- prometheus.MustNewConstMetric(
+		workerRunning, prometheus.GaugeValue, float64(s.Running),
+	)
+}
+
+// Describe implements prometheus.Collector.
+func (m *metricsWorker) Describe(ch chan<- *prometheus.Desc) {
+	prometheus.DescribeByCollect(m, ch)
+}
+
+func newMetricsWorker(w worker.Worker) worker.Worker {
+	rt := &metricsWorker{w}
+	prometheus.MustRegister(rt)
+	return rt
+}

env/env_linux.go

@@ -261,15 +261,17 @@ func newSystemdProperty(name string, units any) dbus.Property {
 }
 
 type credGen struct {
-	cur uint32
+	cur atomic.Uint32
 }
 
 func newCredGen(start uint32) *credGen {
-	return &credGen{cur: start}
+	rt := &credGen{}
+	rt.cur.Store(start)
+	return rt
 }
 
 func (c *credGen) Get() syscall.Credential {
-	n := atomic.AddUint32(&c.cur, 1)
+	n := c.cur.Add(1)
 	return syscall.Credential{
 		Uid: n,
 		Gid: n,

envexec/file_util_linux.go

@@ -12,15 +12,15 @@ import (
 
 const memfdName = "input"
 
-var enableMemFd int32
+var enableMemFd atomic.Int32
 
 func readerToFile(reader io.Reader) (*os.File, error) {
-	if atomic.LoadInt32(&enableMemFd) == 0 {
+	if enableMemFd.Load() == 0 {
 		f, err := memfd.DupToMemfd(memfdName, reader)
 		if err == nil {
 			return f, err
 		}
-		atomic.StoreInt32(&enableMemFd, 1)
+		enableMemFd.Store(1)
 	}
 	r, w, err := os.Pipe()
 	if err != nil {

worker/worker.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/criyle/go-judge/envexec"
@@ -39,9 +40,16 @@ type Worker interface {
 	Start()
 	Submit(context.Context, *Request) (<-chan Response, <-chan struct{})
 	Execute(context.Context, *Request) <-chan Response
+	Stat() Stat
 	Shutdown()
 }
 
+// Stat stores the statistic of the Worker
+type Stat struct {
+	Queue   int
+	Running int
+}
+
 // worker defines executor worker
 type worker struct {
 	fs          filestore.FileStore
@@ -62,6 +70,7 @@ type worker struct {
 	wg        sync.WaitGroup
 	workCh    chan workRequest
 	done      chan struct{}
+	running   atomic.Int32
 }
 
 type workRequest struct {
@@ -131,6 +140,13 @@ func (w *worker) Execute(ctx context.Context, req *Request) <-chan Response {
 	return ch
 }
 
+func (w *worker) Stat() Stat {
+	return Stat{
+		Queue:   len(w.workCh),
+		Running: int(w.running.Load()),
+	}
+}
+
 // Shutdown waits all worker to finish
 func (w *worker) Shutdown() {
 	w.stopOnce.Do(func() {
@@ -166,6 +182,9 @@ func (w *worker) loop() {
 }
 
 func (w *worker) workDoCmd(ctx context.Context, req *Request) Response {
+	w.running.Add(1)
+	defer w.running.Add(-1)
+
 	var rt Response
 	if len(req.Cmd) == 1 {
 		rt = w.workDoSingle(ctx, req.Cmd[0])