refactor: extract zbuscontext into its own file and refactor relevant perf tasks to use it consistently

Author: xmonader

cmds/modules/noded/main.go

@@ -167,7 +167,11 @@ func action(cli *cli.Context) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to create a new perfMon")
 	}
-
+	zcl, err := zbus.NewRedisClient(msgBrokerCon)
+	if err != nil {
+		return errors.Wrap(err, "failed to create a zbus client to the msgBroker")
+	}
+	ctx = perf.WithZbusClient(ctx, zcl)
 	healthcheck.RunNTPCheck(ctx)
 	perfMon.AddTask(iperf.NewTask())
 	perfMon.AddTask(cpubench.NewTask())

pkg/perf/cpubench/cpubench_task.go

@@ -59,7 +59,7 @@ func (c *CPUBenchmarkTask) Run(ctx context.Context) (interface{}, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse cpubench output: %w", err)
 	}
-	client := perf.GetZbusClient(ctx)
+	client := perf.MustGetZbusClient(ctx)
 	statistics := stubs.NewStatisticsStub(client)
 
 	workloads, err := statistics.Workloads(ctx)

pkg/perf/healthcheck/healthcheck.go

@@ -63,7 +63,7 @@ func (h *healthcheckTask) Run(ctx context.Context) (interface{}, error) {
 	log.Debug().Msg("starting health check task")
 	errs := make(map[string][]string)
 
-	cl := perf.GetZbusClient(ctx)
+	cl := perf.MustGetZbusClient(ctx)
 	zui := stubs.NewZUIStub(cl)
 
 	var wg sync.WaitGroup

pkg/perf/monitor.go

@@ -78,7 +78,7 @@ func (pm *PerformanceMonitor) runTask(ctx context.Context, task Task) error {
 
 // Run adds the tasks to the cron queue and start the scheduler
 func (pm *PerformanceMonitor) Run(ctx context.Context) error {
-	ctx = withZbusClient(ctx, pm.zbusClient)
+	ctx = WithZbusClient(ctx, pm.zbusClient)
 	for _, task := range pm.tasks {
 		task := task
 		if _, err := pm.scheduler.CronWithSeconds(task.Cron()).Do(func() error {
@@ -105,14 +105,3 @@ func (pm *PerformanceMonitor) Run(ctx context.Context) error {
 	pm.scheduler.StartAsync()
 	return nil
 }
-
-type zbusClient struct{}
-
-func withZbusClient(ctx context.Context, client zbus.Client) context.Context {
-	return context.WithValue(ctx, zbusClient{}, client)
-}
-
-// GetZbusClient gets zbus client from the given context
-func GetZbusClient(ctx context.Context) zbus.Client {
-	return ctx.Value(zbusClient{}).(zbus.Client)
-}

pkg/perf/publicip/publicip_task.go

@@ -78,7 +78,7 @@ func (p *publicIPValidationTask) Run(ctx context.Context) (interface{}, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to get namespace %s: %w", testNamespace, err)
 	}
-	cl := perf.GetZbusClient(ctx)
+	cl := perf.MustGetZbusClient(ctx)
 	substrateGateway := stubs.NewSubstrateGatewayStub(cl)
 	farmID := environment.MustGet().FarmID
 
@@ -191,7 +191,7 @@ func isLeastValidNode(ctx context.Context, farmID uint32, substrateGateway *stub
 	if err != nil {
 		return false, fmt.Errorf("failed to get farm %d nodes: %w", farmID, err)
 	}
-	cl := perf.GetZbusClient(ctx)
+	cl := perf.MustGetZbusClient(ctx)
 	registrar := stubs.NewRegistrarStub(cl)
 	var nodeID uint32
 	err = backoff.Retry(func() error {

pkg/perf/zbusctx.go

@@ -0,0 +1,30 @@
+package perf
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/threefoldtech/zbus"
+)
+
+// WithZbusClient adds a zbus.Client to the provided context, returning a new context.
+// This allows for the retrieval of the client from the context at a later time.
+type zbusClientKey struct{}
+
+func WithZbusClient(ctx context.Context, client zbus.Client) context.Context {
+	return context.WithValue(ctx, zbusClientKey{}, client)
+}
+
+// MustGetZbusClient gets zbus client from the given context
+func MustGetZbusClient(ctx context.Context) zbus.Client {
+	return ctx.Value(zbusClientKey{}).(zbus.Client)
+}
+
+// TryGetZbusClient tries to get zbus client from the given context
+func TryGetZbusClient(ctx context.Context) (zbus.Client, error) {
+	zcl, ok := ctx.Value(zbusClientKey{}).(zbus.Client)
+	if !ok {
+		return zcl, fmt.Errorf("context does not have zbus client")
+	}
+	return zcl, nil
+}