chore: proper Context handling and simplifications

Signed-off-by: Clément Nussbaumer <clement.nussbaumer@postfinance.ch>

internal/kubenurse/handler_test.go

@@ -1,7 +1,6 @@
 package kubenurse
 
 import (
- "context"
  "net/http"
  "net/http/httptest"
  "testing"
@@ -14,7 +13,7 @@ func TestServerHandler(t *testing.T) {
  r := require.New(t)
 
  fakeClient := fake.NewFakeClient()
- kubenurse, err := New(context.Background(), fakeClient)
+ kubenurse, err := New(fakeClient)
 
  r.NoError(err)
  r.NotNil(kubenurse)

internal/kubenurse/server.go

@@ -59,7 +59,7 @@ type Server struct {
 // * KUBENURSE_CHECK_ME_SERVICE
 // * KUBENURSE_CHECK_NEIGHBOURHOOD
 // * KUBENURSE_CHECK_INTERVAL
-func New(ctx context.Context, c client.Client) (*Server, error) { //nolint:funlen // TODO: use a flag parsing library (e.g. ff) to reduce complexity
+func New(c client.Client) (*Server, error) { //nolint:funlen // TODO: use a flag parsing library (e.g. ff) to reduce complexity
  mux := http.NewServeMux()
 
  checkInterval := defaultCheckInterval
@@ -134,7 +134,7 @@ func New(ctx context.Context, c client.Client) (*Server, error) { //nolint:funle
  }
 
  // setup checker
- chk, err := servicecheck.New(ctx, c, promRegistry, server.allowUnschedulable, 1*time.Second, histogramBuckets)
+ chk, err := servicecheck.New(c, promRegistry, server.allowUnschedulable, 1*time.Second, histogramBuckets)
  if err != nil {
  return nil, err
  }
@@ -189,7 +189,7 @@ func New(ctx context.Context, c client.Client) (*Server, error) { //nolint:funle
 }
 
 // Run starts the periodic checker and the http/https server(s) and blocks until Shutdown was called.
-func (s *Server) Run() error {
+func (s *Server) Run(ctx context.Context) error {
  var (
  wg sync.WaitGroup
  errc = make(chan error, 2) // max two errors can happen
@@ -211,7 +211,17 @@ func (s *Server) Run() error {
  go func() {
  defer wg.Done()
 
- s.checker.RunScheduled(s.checkInterval)
+ ticker := time.NewTicker(s.checkInterval)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-ticker.C:
+ s.checker.Run(ctx)
+ case <-ctx.Done():
+ return
+ }
+ }
  }()
 
  wg.Add(1)
@@ -259,7 +269,7 @@ func (s *Server) Run() error {
 }
 
 // Shutdown disables the readiness probe and then gracefully halts the kubenurse http/https server(s).
-func (s *Server) Shutdown(ctx context.Context) error {
+func (s *Server) Shutdown() error {
  s.ready.Store(false)
 
  // wait before actually shutting down the http/s server, as the updated
@@ -268,8 +278,9 @@ func (s *Server) Shutdown(ctx context.Context) error {
  // me_ingress or path errors in other pods
  time.Sleep(s.checker.ShutdownDuration)
 
- // stop the scheduled checker
- s.checker.StopScheduled()
+ // background ctx since, the "root" context is already canceled
+ ctx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer shutdownCancel()
 
  if err := s.http.Shutdown(ctx); err != nil {
  return fmt.Errorf("stop http server: %w", err)

internal/kubenurse/server_test.go

@@ -3,6 +3,7 @@ package kubenurse
 import (
  "context"
  "testing"
+ "time"
 
  "github.com/stretchr/testify/require"
  "sigs.k8s.io/controller-runtime/pkg/client/fake"
@@ -12,24 +13,26 @@ func TestCombined(t *testing.T) {
  r := require.New(t)
 
  fakeClient := fake.NewFakeClient()
- kubenurse, err := New(context.Background(), fakeClient)
+ kubenurse, err := New(fakeClient)
  r.NoError(err)
  r.NotNil(kubenurse)
 
  t.Run("start/stop", func(t *testing.T) {
  r := require.New(t)
  errc := make(chan error, 1)
 
+ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
  go func() {
  // blocks until shutdown is called
- err := kubenurse.Run()
+ err := kubenurse.Run(ctx)
 
  errc <- err
  close(errc)
+ cancel()
  }()
 
  // Shutdown, Run() should stop after function completes
- err := kubenurse.Shutdown(context.Background())
+ err := kubenurse.Shutdown()
  r.NoError(err)
 
  err = <-errc // blocks until kubenurse.Run() finishes and eventually returns an error

internal/servicecheck/servicecheck.go

@@ -25,7 +25,7 @@ const (
 
 // New configures the checker with a httpClient and a cache timeout for check
 // results. Other parameters of the Checker struct need to be configured separately.
-func New(_ context.Context, cl client.Client, promRegistry *prometheus.Registry,
+func New(cl client.Client, promRegistry *prometheus.Registry,
  allowUnschedulable bool, cacheTTL time.Duration, durationHistogramBuckets []float64) (*Checker, error) {
  // setup http transport
  tlsConfig, err := generateTLSConfig(os.Getenv("KUBENURSE_EXTRA_CA"))
@@ -63,13 +63,12 @@ func New(_ context.Context, cl client.Client, promRegistry *prometheus.Registry,
  client: cl,
  httpClient: httpClient,
  cacheTTL: cacheTTL,
- stop: make(chan struct{}),
  }, nil
 }
 
 // Run runs all servicechecks and returns the result togeter with a boolean which indicates success. The cache
 // is respected.
-func (c *Checker) Run() {
+func (c *Checker) Run(ctx context.Context) {
  // Run Checks
  result := sync.Map{}
 
@@ -91,17 +90,17 @@ func (c *Checker) Run() {
 
  wg.Add(4)
 
- go c.measure(&wg, &result, c.APIServerDirect, APIServerDirect)
- go c.measure(&wg, &result, c.APIServerDNS, APIServerDNS)
- go c.measure(&wg, &result, c.MeIngress, meIngress)
- go c.measure(&wg, &result, c.MeService, meService)
+ go c.measure(ctx, &wg, &result, c.APIServerDirect, APIServerDirect)
+ go c.measure(ctx, &wg, &result, c.APIServerDNS, APIServerDNS)
+ go c.measure(ctx, &wg, &result, c.MeIngress, meIngress)
+ go c.measure(ctx, &wg, &result, c.MeService, meService)
 
  if c.SkipCheckNeighbourhood {
  result.Store(NeighbourhoodState, skippedStr)
  return
  }
 
- neighbours, err := c.getNeighbours(context.Background(), c.KubenurseNamespace, c.NeighbourFilter)
+ neighbours, err := c.getNeighbours(ctx, c.KubenurseNamespace, c.NeighbourFilter)
  if err != nil {
  result.Store(NeighbourhoodState, err.Error())
  return
@@ -121,33 +120,12 @@ func (c *Checker) Run() {
  return c.doRequest(ctx, podIPtoURL(neighbour.PodIP, c.UseTLS), true)
  }
 
- go c.measure(&wg, &result, check, "path_"+neighbour.NodeName)
+ go c.measure(ctx, &wg, &result, check, "path_"+neighbour.NodeName)
  }
 
  wg.Wait()
 }
 
-// RunScheduled runs the checks in the specified interval which can be used to keep the metrics up-to-date. This
-// function does not return until StopScheduled is called.
-func (c *Checker) RunScheduled(d time.Duration) {
- ticker := time.NewTicker(d)
- defer ticker.Stop()
-
- for {
- select {
- case <-ticker.C:
- c.Run()
- case <-c.stop:
- return
- }
- }
-}
-
-// StopScheduled is used to stop the scheduled run of checks.
-func (c *Checker) StopScheduled() {
- close(c.stop)
-}
-
 // APIServerDirect checks the /version endpoint of the Kubernetes API Server through the direct link
 func (c *Checker) APIServerDirect(ctx context.Context) string {
  if c.SkipCheckAPIServerDirect {
@@ -189,12 +167,12 @@ func (c *Checker) MeService(ctx context.Context) string {
 }
 
 // measure implements metric collections for the check
-func (c *Checker) measure(wg *sync.WaitGroup, res *sync.Map, check Check, requestType string) {
+func (c *Checker) measure(ctx context.Context, wg *sync.WaitGroup, res *sync.Map, check Check, requestType string) {
  // Add our label (check type) to the context so our http tracer can annotate
  // metrics and errors based with the label
  defer wg.Done()
 
- ctx := context.WithValue(context.Background(), kubenurseTypeKey{}, requestType)
+ ctx = context.WithValue(ctx, kubenurseTypeKey{}, requestType)
  res.Store(requestType, check(ctx))
 }
 

internal/servicecheck/servicecheck_test.go

@@ -36,29 +36,14 @@ func TestCombined(t *testing.T) {
  // fake client, with a dummy neighbour pod
  fakeClient := fake.NewFakeClient(&fakeNeighbourPod)
 
- checker, err := New(context.Background(), fakeClient, prometheus.NewRegistry(), false, 3*time.Second, prometheus.DefBuckets)
+ checker, err := New(fakeClient, prometheus.NewRegistry(), false, 3*time.Second, prometheus.DefBuckets)
  r.NoError(err)
  r.NotNil(checker)
 
  t.Run("run", func(t *testing.T) {
  r := require.New(t)
- checker.Run()
+ checker.Run(context.Background())
 
  r.Equal(okStr, checker.LastCheckResult[NeighbourhoodState])
  })
-
- t.Run("scheduled", func(t *testing.T) {
- stopped := make(chan struct{})
-
- go func() {
- // blocks until StopScheduled()
- checker.RunScheduled(time.Second * 5)
-
- close(stopped)
- }()
-
- checker.StopScheduled()
-
- <-stopped
- })
 }

internal/servicecheck/types.go

@@ -25,7 +25,7 @@ type Checker struct {
  SkipCheckMeIngress bool
  SkipCheckMeService bool
 
- // shutdownDuration defines the time during which kubenurse will wait before stopping
+ // shutdownDuration defines the time during which kubenurse will accept https requests during shutdown
  ShutdownDuration time.Duration
 
  // Kubernetes API
@@ -55,9 +55,6 @@ type Checker struct {
 
  // cacheTTL defines the TTL of how long a cached result is valid
  cacheTTL time.Duration
-
- // stop is used to cancel RunScheduled
- stop chan struct{}
 }
 
 // Check is the signature used by all checks that the checker can execute.

main.go

@@ -8,7 +8,6 @@ import (
  "os"
  "os/signal"
  "syscall"
- "time"
 
  "github.com/postfinance/kubenurse/internal/kubenurse"
  corev1 "k8s.io/api/core/v1"
@@ -62,7 +61,7 @@ func main() {
  return
  }
 
- server, err := kubenurse.New(ctx, c)
+ server, err := kubenurse.New(c)
  if err != nil {
  slog.Error("error in kubenurse.New call", "err", err)
  return
@@ -73,17 +72,13 @@ func main() {
 
  slog.Info("shutting down, received signal to stop")
 
- // background ctx since, the "root" context is already canceled
- shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer shutdownCancel()
-
- if err := server.Shutdown(shutdownCtx); err != nil {
+ if err := server.Shutdown(); err != nil {
  slog.Error("error during graceful shutdown", "err", err)
  }
  }()
 
  // blocks, until the server is stopped by calling Shutdown()
- if err := server.Run(); err != nil {
+ if err := server.Run(ctx); err != nil {
  slog.Error("error while running kubenurse", "err", err)
  }
 }