From 78ccf1e6b3d9e44604faf283b65faa3a5e0825f9 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Mon, 14 Oct 2024 22:44:07 +0200
Subject: pkg/rpcserver: take stats as a dependency

It will enable collecting statistics for several simultaneous RPCServer
objects.
---
 pkg/rpcserver/local.go          |  1 +
 pkg/rpcserver/rpcserver.go      | 44 ++++++++++++++++++++++++++++++-----------
 pkg/rpcserver/rpcserver_test.go |  4 ++--
 3 files changed, 35 insertions(+), 14 deletions(-)

(limited to 'pkg/rpcserver')

diff --git a/pkg/rpcserver/local.go b/pkg/rpcserver/local.go
index 09cd1868d..5faa8334b 100644
--- a/pkg/rpcserver/local.go
+++ b/pkg/rpcserver/local.go
@@ -43,6 +43,7 @@ func RunLocal(cfg *LocalConfig) error {
 	cfg.FilterSignal = true
 	cfg.RPC = ":0"
 	cfg.PrintMachineCheck = log.V(1)
+	cfg.Stats = NewStats()
 	ctx := &local{
 		cfg:       cfg,
 		setupDone: make(chan bool),
diff --git a/pkg/rpcserver/rpcserver.go b/pkg/rpcserver/rpcserver.go
index 53181dd2b..76f1ddc68 100644
--- a/pkg/rpcserver/rpcserver.go
+++ b/pkg/rpcserver/rpcserver.go
@@ -31,6 +31,8 @@ import (
 
 type Config struct {
 	vminfo.Config
+	Stats
+
 	VMArch string
 	VMType string
 	RPC    string
@@ -85,15 +87,33 @@ type server struct {
 	canonicalModules *cover.Canonicalizer
 	coverFilter      []uint64
 
-	mu             sync.Mutex
-	runners        map[int]*Runner
-	execSource     *queue.Distributor
-	triagedCorpus  atomic.Bool
-	statVMRestarts *stat.Val
+	mu            sync.Mutex
+	runners       map[int]*Runner
+	execSource    *queue.Distributor
+	triagedCorpus atomic.Bool
+
+	Stats
 	*runnerStats
 }
 
-func New(cfg *mgrconfig.Config, mgr Manager, debug bool) (Server, error) {
+type Stats struct {
+	StatExecs      *stat.Val
+	StatNumFuzzing *stat.Val
+	StatVMRestarts *stat.Val
+}
+
+func NewStats() Stats {
+	return Stats{
+		StatExecs: stat.New("exec total", "Total test program executions",
+			stat.Console, stat.Rate{}, stat.Prometheus("syz_exec_total")),
+		StatNumFuzzing: stat.New("fuzzing VMs",
+			"Number of VMs that are currently fuzzing", stat.Graph("fuzzing VMs")),
+		StatVMRestarts: stat.New("vm restarts", "Total number of VM starts",
+			stat.Rate{}, stat.NoGraph),
+	}
+}
+
+func New(cfg *mgrconfig.Config, mgr Manager, stats Stats, debug bool) (Server, error) {
 	var pcBase uint64
 	if cfg.KernelObj != "" {
 		var err error
@@ -121,6 +141,7 @@ func New(cfg *mgrconfig.Config, mgr Manager, debug bool) (Server, error) {
 			Sandbox:    sandbox,
 			SandboxArg: cfg.SandboxArg,
 		},
+		Stats:  stats,
 		VMArch: cfg.TargetVMArch,
 		RPC:    cfg.RPC,
 		VMLess: cfg.VMLess,
@@ -153,8 +174,7 @@ func newImpl(ctx context.Context, cfg *Config, mgr Manager) *server {
 		baseSource: baseSource,
 		execSource: queue.Distribute(queue.Retry(baseSource)),
 
-		statVMRestarts: stat.New("vm restarts", "Total number of VM starts",
-			stat.Rate{}, stat.NoGraph),
+		Stats: cfg.Stats,
 		runnerStats: &runnerStats{
 			statExecRetries: stat.New("exec retries",
 				"Number of times a test program was restarted because the first run failed",
@@ -162,7 +182,7 @@ func newImpl(ctx context.Context, cfg *Config, mgr Manager) *server {
 			statExecutorRestarts: stat.New("executor restarts",
 				"Number of times executor process was restarted", stat.Rate{}, stat.Graph("executor")),
 			statExecBufferTooSmall: queue.StatExecBufferTooSmall,
-			statExecs:              queue.StatExecs,
+			statExecs:              cfg.Stats.StatExecs,
 			statNoExecRequests:     queue.StatNoExecRequests,
 			statNoExecDuration:     queue.StatNoExecDuration,
 		},
@@ -205,7 +225,7 @@ func (serv *server) handleConn(conn *flatrpc.Conn) {
 	} else if err := checkRevisions(connectReq, serv.cfg.Target); err != nil {
 		log.Fatal(err)
 	}
-	serv.statVMRestarts.Add(1)
+	serv.StatVMRestarts.Add(1)
 
 	serv.mu.Lock()
 	runner := serv.runners[id]
@@ -312,8 +332,8 @@ func (serv *server) connectionLoop(runner *Runner) error {
 		}
 	}
 
-	queue.StatNumFuzzing.Add(1)
-	defer queue.StatNumFuzzing.Add(-1)
+	serv.StatNumFuzzing.Add(1)
+	defer serv.StatNumFuzzing.Add(-1)
 
 	return runner.ConnectionLoop()
 }
diff --git a/pkg/rpcserver/rpcserver_test.go b/pkg/rpcserver/rpcserver_test.go
index 3252ddd4f..760ae5e13 100644
--- a/pkg/rpcserver/rpcserver_test.go
+++ b/pkg/rpcserver/rpcserver_test.go
@@ -80,7 +80,7 @@ func TestNew(t *testing.T) {
 			cfg.Target, err = prog.GetTarget(cfg.TargetOS, cfg.TargetArch)
 			assert.NoError(t, err)
 
-			serv, err := New(cfg, nil, tt.debug)
+			serv, err := New(cfg, nil, NewStats(), tt.debug)
 			if tt.expectedErr != nil {
 				assert.Equal(t, tt.expectedErr, err)
 			} else if tt.expectsErr {
@@ -195,7 +195,7 @@ func TestHandleConn(t *testing.T) {
 			cfg.Target.Revision = tt.req.SyzRevision
 			assert.NoError(t, err)
 
-			s, err := New(cfg, managerMock, debug)
+			s, err := New(cfg, managerMock, NewStats(), debug)
 			assert.NoError(t, err)
 			serv := s.(*server)
 
-- 
cgit mrf-deployment