From db2c0bcea56d02acd1852c7a856a60b2fda6e90e Mon Sep 17 00:00:00 2001
From: marco <marco@crowdsec.net>
Date: Tue, 24 Sep 2024 15:20:14 +0200
Subject: [PATCH] context propagation: SendMetrics, SendUsageMetrics, plugin
 config

---
 pkg/apiserver/apic_metrics.go      | 12 ++++--------
 pkg/apiserver/apic_metrics_test.go | 12 +++++++-----
 pkg/apiserver/apiserver.go         |  4 ++--
 pkg/csplugin/notifier.go           |  4 +---
 pkg/protobufs/plugin_interface.go  |  4 ++--
 5 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/pkg/apiserver/apic_metrics.go b/pkg/apiserver/apic_metrics.go
index 16b2328dbe9..3d9e7b28a79 100644
--- a/pkg/apiserver/apic_metrics.go
+++ b/pkg/apiserver/apic_metrics.go
@@ -251,11 +251,9 @@ func (a *apic) fetchMachineIDs(ctx context.Context) ([]string, error) {
 // Metrics are sent at start, then at the randomized metricsIntervalFirst,
 // then at regular metricsInterval. If a change is detected in the list
 // of machines, the next metrics are sent immediately.
-func (a *apic) SendMetrics(stop chan (bool)) {
+func (a *apic) SendMetrics(ctx context.Context, stop chan (bool)) {
 	defer trace.CatchPanic("lapi/metricsToAPIC")
 
-	ctx := context.TODO()
-
 	// verify the list of machines every <checkInt> interval
 	const checkInt = 20 * time.Second
 
@@ -321,7 +319,7 @@ func (a *apic) SendMetrics(stop chan (bool)) {
 			if metrics != nil {
 				log.Info("capi metrics: sending")
 
-				_, _, err = a.apiClient.Metrics.Add(context.Background(), metrics)
+				_, _, err = a.apiClient.Metrics.Add(ctx, metrics)
 				if err != nil {
 					log.Errorf("capi metrics: failed: %s", err)
 				}
@@ -339,11 +337,9 @@ func (a *apic) SendMetrics(stop chan (bool)) {
 	}
 }
 
-func (a *apic) SendUsageMetrics() {
+func (a *apic) SendUsageMetrics(ctx context.Context) {
 	defer trace.CatchPanic("lapi/usageMetricsToAPIC")
 
-	ctx := context.TODO()
-
 	firstRun := true
 
 	log.Debugf("Start sending usage metrics to CrowdSec Central API (interval: %s once, then %s)", a.usageMetricsIntervalFirst, a.usageMetricsInterval)
@@ -368,7 +364,7 @@ func (a *apic) SendUsageMetrics() {
 				continue
 			}
 
-			_, resp, err := a.apiClient.UsageMetrics.Add(context.Background(), metrics)
+			_, resp, err := a.apiClient.UsageMetrics.Add(ctx, metrics)
 			if err != nil {
 				log.Errorf("unable to send usage metrics: %s", err)
 
diff --git a/pkg/apiserver/apic_metrics_test.go b/pkg/apiserver/apic_metrics_test.go
index 78b16f9c8b7..13a24668f26 100644
--- a/pkg/apiserver/apic_metrics_test.go
+++ b/pkg/apiserver/apic_metrics_test.go
@@ -14,6 +14,8 @@ import (
 )
 
 func TestAPICSendMetrics(t *testing.T) {
+	ctx := context.Background()
+
 	tests := []struct {
 		name            string
 		duration        time.Duration
@@ -34,7 +36,7 @@ func TestAPICSendMetrics(t *testing.T) {
 			metricsInterval: time.Millisecond * 20,
 			expectedCalls:   5,
 			setUp: func(api *apic) {
-				api.dbClient.Ent.Machine.Delete().ExecX(context.Background())
+				api.dbClient.Ent.Machine.Delete().ExecX(ctx)
 				api.dbClient.Ent.Machine.Create().
 					SetMachineId("1234").
 					SetPassword(testPassword.String()).
@@ -42,16 +44,16 @@ func TestAPICSendMetrics(t *testing.T) {
 					SetScenarios("crowdsecurity/test").
 					SetLastPush(time.Time{}).
 					SetUpdatedAt(time.Time{}).
-					ExecX(context.Background())
+					ExecX(ctx)
 
-				api.dbClient.Ent.Bouncer.Delete().ExecX(context.Background())
+				api.dbClient.Ent.Bouncer.Delete().ExecX(ctx)
 				api.dbClient.Ent.Bouncer.Create().
 					SetIPAddress("1.2.3.6").
 					SetName("someBouncer").
 					SetAPIKey("foobar").
 					SetRevoked(false).
 					SetLastPull(time.Time{}).
-					ExecX(context.Background())
+					ExecX(ctx)
 			},
 		},
 	}
@@ -86,7 +88,7 @@ func TestAPICSendMetrics(t *testing.T) {
 
 			httpmock.ZeroCallCounters()
 
-			go api.SendMetrics(stop)
+			go api.SendMetrics(ctx, stop)
 
 			time.Sleep(tc.duration)
 			stop <- true
diff --git a/pkg/apiserver/apiserver.go b/pkg/apiserver/apiserver.go
index 6b5d6803be9..2b2b453348a 100644
--- a/pkg/apiserver/apiserver.go
+++ b/pkg/apiserver/apiserver.go
@@ -357,12 +357,12 @@ func (s *APIServer) initAPIC(ctx context.Context) {
 	}
 
 	s.apic.metricsTomb.Go(func() error {
-		s.apic.SendMetrics(make(chan bool))
+		s.apic.SendMetrics(ctx, make(chan bool))
 		return nil
 	})
 
 	s.apic.metricsTomb.Go(func() error {
-		s.apic.SendUsageMetrics()
+		s.apic.SendUsageMetrics(ctx)
 		return nil
 	})
 }
diff --git a/pkg/csplugin/notifier.go b/pkg/csplugin/notifier.go
index 2b5d57fbcff..ed4a4cc4149 100644
--- a/pkg/csplugin/notifier.go
+++ b/pkg/csplugin/notifier.go
@@ -40,9 +40,7 @@ func (m *GRPCClient) Notify(ctx context.Context, notification *protobufs.Notific
 }
 
 func (m *GRPCClient) Configure(ctx context.Context, config *protobufs.Config) (*protobufs.Empty, error) {
-	_, err := m.client.Configure(
-		context.Background(), config,
-	)
+	_, err := m.client.Configure(ctx, config)
 	return &protobufs.Empty{}, err
 }
 
diff --git a/pkg/protobufs/plugin_interface.go b/pkg/protobufs/plugin_interface.go
index fc89b2fa009..baa76c8941c 100644
--- a/pkg/protobufs/plugin_interface.go
+++ b/pkg/protobufs/plugin_interface.go
@@ -24,12 +24,12 @@ type NotifierPlugin struct {
 type GRPCClient struct{ client NotifierClient }
 
 func (m *GRPCClient) Notify(ctx context.Context, notification *Notification) (*Empty, error) {
-	_, err := m.client.Notify(context.Background(), notification)
+	_, err := m.client.Notify(ctx, notification)
 	return &Empty{}, err
 }
 
 func (m *GRPCClient) Configure(ctx context.Context, config *Config) (*Empty, error) {
-	_, err := m.client.Configure(context.Background(), config)
+	_, err := m.client.Configure(ctx, config)
 	return &Empty{}, err
 }