Merge pull request #3767 from twz123/inttests-context-aware-logging

Context-aware logging in the inttest wait functions

inttest/common/bootloosesuite.go

@@ -44,6 +44,7 @@ import (
 	"github.com/k0sproject/k0s/internal/pkg/file"
 	apclient "github.com/k0sproject/k0s/pkg/client/clientset"
 	"github.com/k0sproject/k0s/pkg/constant"
+	"github.com/k0sproject/k0s/pkg/k0scontext"
 	"github.com/k0sproject/k0s/pkg/kubernetes/watch"
 	extclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset/typed/apiextensions/v1"
 
@@ -58,6 +59,7 @@ import (
 	"github.com/k0sproject/bootloose/pkg/cluster"
 	"github.com/k0sproject/bootloose/pkg/config"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/multierr"
 	"golang.org/x/sync/errgroup"
@@ -270,11 +272,29 @@ func (s *BootlooseSuite) waitForSSH(ctx context.Context) {
 	s.Require().NoError(g.Wait(), "Failed to ssh into all nodes")
 }
 
-// Context returns this suite's context, which should be passed to all blocking operations.
+// Context returns this suite's context, which should be passed to all blocking
+// operations. It captures the current test as a context value, so that it can
+// be retrieved by helper methods later on.
+//
+// Context should only be called once at the beginning of a test function, and
+// then be passed along to all subsequently called functions in the usual way:
+// as the first function parameter. The test framework itself doesn't have any
+// means of smuggling a context to test functions, hence the suite needs to
+// store it as a field, which is usually considered bad practice. However,
+// relying on the context being passed along implicitly makes the test suite a
+// bit edgy concerning API design and cancellation. This is the main reason why
+// the suite context is being replaced during the suite cleanup: Some functions
+// will obtain a context from the suite again, instead of taking it as their
+// first parameter. That replacement should become obsolete once all functions
+// will take the context as parameter.
 func (s *BootlooseSuite) Context() context.Context {
-	ctx := s.ctx
-	s.Require().NotNil(ctx, "No suite context installed")
-	return ctx
+	ctx, t := s.ctx, s.T()
+	require.NotNil(t, ctx, "No suite context installed")
+	if t == nil {
+		return ctx
+	}
+
+	return k0scontext.WithValue(ctx, t)
 }
 
 // ControllerNode gets the node name of given controller index

inttest/common/util.go

@@ -26,9 +26,11 @@ import (
 	"regexp"
 	"strings"
 	"syscall"
+	"testing"
 	"time"
 
 	"github.com/k0sproject/k0s/pkg/constant"
+	"github.com/k0sproject/k0s/pkg/k0scontext"
 	"github.com/k0sproject/k0s/pkg/kubernetes/watch"
 
 	appsv1 "k8s.io/api/apps/v1"
@@ -46,10 +48,13 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+// LogfFn will be used whenever something needs to be logged.
+type LogfFn func(format string, args ...any)
+
 // Poll tries a condition func until it returns true, an error or the specified
 // context is canceled or expired.
 func Poll(ctx context.Context, condition wait.ConditionWithContextFunc) error {
-	return wait.PollImmediateUntilWithContext(ctx, 100*time.Millisecond, condition)
+	return wait.PollUntilContextCancel(ctx, 100*time.Millisecond, true, condition)
 }
 
 // WaitForKubeRouterReady waits to see all kube-router pods healthy as long as
@@ -106,7 +111,7 @@ func WaitForMetricsReady(ctx context.Context, c *rest.Config) error {
 	watchAPIServices := watch.FromClient[*apiregistrationv1.APIServiceList, apiregistrationv1.APIService]
 	return watchAPIServices(clientset.ApiregistrationV1().APIServices()).
 		WithObjectName("v1beta1.metrics.k8s.io").
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(service *apiregistrationv1.APIService) (bool, error) {
 			for _, c := range service.Status.Conditions {
 				if c.Type == apiregistrationv1.Available {
@@ -125,6 +130,7 @@ func WaitForMetricsReady(ctx context.Context, c *rest.Config) error {
 func WaitForNodeReadyStatus(ctx context.Context, clients kubernetes.Interface, nodeName string, status corev1.ConditionStatus) error {
 	return watch.Nodes(clients.CoreV1().Nodes()).
 		WithObjectName(nodeName).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(node *corev1.Node) (done bool, err error) {
 			for _, cond := range node.Status.Conditions {
 				if cond.Type == corev1.NodeReady {
@@ -145,7 +151,7 @@ func WaitForNodeReadyStatus(ctx context.Context, clients kubernetes.Interface, n
 func WaitForDaemonSet(ctx context.Context, kc *kubernetes.Clientset, name string) error {
 	return watch.DaemonSets(kc.AppsV1().DaemonSets("kube-system")).
 		WithObjectName(name).
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(ds *appsv1.DaemonSet) (bool, error) {
 			return ds.Status.NumberAvailable == ds.Status.DesiredNumberScheduled, nil
 		})
@@ -156,7 +162,7 @@ func WaitForDaemonSet(ctx context.Context, kc *kubernetes.Clientset, name string
 func WaitForDeployment(ctx context.Context, kc *kubernetes.Clientset, name, namespace string) error {
 	return watch.Deployments(kc.AppsV1().Deployments(namespace)).
 		WithObjectName(name).
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(deployment *appsv1.Deployment) (bool, error) {
 			for _, c := range deployment.Status.Conditions {
 				if c.Type == appsv1.DeploymentAvailable {
@@ -177,7 +183,7 @@ func WaitForDeployment(ctx context.Context, kc *kubernetes.Clientset, name, name
 func WaitForStatefulSet(ctx context.Context, kc *kubernetes.Clientset, name, namespace string) error {
 	return watch.StatefulSets(kc.AppsV1().StatefulSets(namespace)).
 		WithObjectName(name).
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(s *appsv1.StatefulSet) (bool, error) {
 			return s.Status.ReadyReplicas == *s.Spec.Replicas, nil
 		})
@@ -203,7 +209,7 @@ func waitForDefaultStorageClass(kc *kubernetes.Clientset) wait.ConditionWithCont
 func WaitForPod(ctx context.Context, kc *kubernetes.Clientset, name, namespace string) error {
 	return watch.Pods(kc.CoreV1().Pods(namespace)).
 		WithObjectName(name).
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(ctx, func(pod *corev1.Pod) (bool, error) {
 			for _, cond := range pod.Status.Conditions {
 				if cond.Type == corev1.PodReady {
@@ -251,7 +257,7 @@ func WaitForLease(ctx context.Context, kc *kubernetes.Clientset, name string, na
 	watchLeases := watch.FromClient[*coordinationv1.LeaseList, coordinationv1.Lease]
 	if err := watchLeases(kc.CoordinationV1().Leases(namespace)).
 		WithObjectName(name).
-		WithErrorCallback(RetryWatchErrors(logrus.Infof)).
+		WithErrorCallback(RetryWatchErrors(logfFrom(ctx))).
 		Until(
 			ctx, func(lease *coordinationv1.Lease) (bool, error) {
 				holderIdentity = *lease.Spec.HolderIdentity
@@ -265,7 +271,7 @@ func WaitForLease(ctx context.Context, kc *kubernetes.Clientset, name string, na
 	return holderIdentity, nil
 }
 
-func RetryWatchErrors(logf func(format string, args ...any)) watch.ErrorCallback {
+func RetryWatchErrors(logf LogfFn) watch.ErrorCallback {
 	return func(err error) (time.Duration, error) {
 		if retryDelay, e := watch.IsRetryable(err); e == nil {
 			logf("Encountered transient watch error, retrying in %s: %v", retryDelay, err)
@@ -327,3 +333,15 @@ func VerifyKubeletMetrics(ctx context.Context, kc *kubernetes.Clientset, node st
 		return false, nil
 	})
 }
+
+// Retrieves the LogfFn stored in context, falling back to use testing.T's Logf
+// if the context has a *testing.T or logrus's Infof as a last resort.
+func logfFrom(ctx context.Context) LogfFn {
+	if logf := k0scontext.Value[LogfFn](ctx); logf != nil {
+		return logf
+	}
+	if t := k0scontext.Value[*testing.T](ctx); t != nil {
+		return t.Logf
+	}
+	return logrus.Infof
+}

pkg/k0scontext/context.go

@@ -14,6 +14,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+// Package k0scontext provides utility functions for working with Go contexts
+// and specific keys for storing and retrieving k0s project-related
+// configuration data.
+//
+// The package also includes functions for setting, retrieving, and checking the
+// presence of values associated with specific types. These functions simplify
+// context-based data management in a type-safe and ergonomic manner. Each
+// distinct type serves as a key for context values, removing the necessity for
+// an extra key constant. The type itself becomes the key.
 package k0scontext
 
 import (
@@ -22,13 +31,18 @@ import (
 	k0sapi "github.com/k0sproject/k0s/pkg/apis/k0s/v1beta1"
 )
 
+// Key represents a string key used for storing and retrieving values in a context.
 type Key string
 
+// Context keys for storing k0s project configuration data in a context.
 const (
 	ContextNodeConfigKey    Key = "k0s_node_config"
 	ContextClusterConfigKey Key = "k0s_cluster_config"
 )
 
+// FromContext retrieves a value from the context associated with the given key
+// and attempts to cast it to the specified type. It returns the value or nil if
+// not found.
 func FromContext[out any](ctx context.Context, key Key) *out {
 	v, ok := ctx.Value(key).(*out)
 	if !ok {
@@ -37,6 +51,7 @@ func FromContext[out any](ctx context.Context, key Key) *out {
 	return v
 }
 
+// GetNodeConfig retrieves the k0s NodeConfig from the context, or nil if not found.
 func GetNodeConfig(ctx context.Context) *k0sapi.ClusterConfig {
 	cfg, ok := ctx.Value(ContextNodeConfigKey).(*k0sapi.ClusterConfig)
 	if !ok {
@@ -45,3 +60,51 @@ func GetNodeConfig(ctx context.Context) *k0sapi.ClusterConfig {
 
 	return cfg
 }
+
+// keyType is used to create unique keys based on the types of the values stored in a context.
+type keyType[T any] struct{}
+
+// bucket is used to wrap values before storing them in a context.
+type bucket[T any] struct{ inner T }
+
+// WithValue adds a value of type T to the context and returns a new context with the added value.
+func WithValue[T any](ctx context.Context, value T) context.Context {
+	var key keyType[T]
+	return context.WithValue(ctx, key, bucket[T]{value})
+}
+
+// HasValue checks if a value of type T is present in the context.
+func HasValue[T any](ctx context.Context) bool {
+	_, ok := value[T](ctx)
+	return ok
+}
+
+// Value retrieves the value of type T from the context.
+// If there's no such value, it returns the zero value of type T.
+func Value[T any](ctx context.Context) T {
+	return ValueOrElse[T](ctx, func() (_ T) { return })
+}
+
+// ValueOr retrieves the value of type T from the context.
+// If there's no such value, it returns the specified fallback value.
+func ValueOr[T any](ctx context.Context, fallback T) T {
+	return ValueOrElse[T](ctx, func() T { return fallback })
+}
+
+// ValueOrElse retrieves the value of type T from the context.
+// If there's no such value, it invokes the fallback function and returns its result.
+func ValueOrElse[T any](ctx context.Context, fallbackFn func() T) T {
+	if val, ok := value[T](ctx); ok {
+		return val.inner
+	}
+
+	return fallbackFn()
+}
+
+// value retrieves a value of type T from the context along with a boolean
+// indicating its presence.
+func value[T any](ctx context.Context) (bucket[T], bool) {
+	var key keyType[T]
+	val, ok := ctx.Value(key).(bucket[T])
+	return val, ok
+}

pkg/k0scontext/context_test.go

@@ -0,0 +1,94 @@
+/*
+Copyright 2023 k0s authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package k0scontext_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/k0sproject/k0s/pkg/k0scontext"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasValue_StructPtrs(t *testing.T) {
+	type Foo struct{}
+	type Bar struct{}
+
+	ctx := context.Background()
+	assert.False(t, k0scontext.HasValue[*Foo](ctx))
+	assert.False(t, k0scontext.HasValue[*Bar](ctx))
+
+	ctxWithFoo := k0scontext.WithValue(ctx, (*Foo)(nil))
+	assert.True(t, k0scontext.HasValue[*Foo](ctxWithFoo))
+	assert.False(t, k0scontext.HasValue[*Bar](ctxWithFoo))
+
+	ctxWithFooAndBar := k0scontext.WithValue(ctxWithFoo, (*Bar)(nil))
+	assert.True(t, k0scontext.HasValue[*Foo](ctxWithFooAndBar))
+	assert.True(t, k0scontext.HasValue[*Bar](ctxWithFooAndBar))
+}
+
+func TestHasValue_Ifaces(t *testing.T) {
+	type Foo interface{}
+	type Bar interface{}
+
+	ctx := context.Background()
+	assert.False(t, k0scontext.HasValue[Foo](ctx))
+	assert.False(t, k0scontext.HasValue[Bar](ctx))
+
+	ctxWithFoo := k0scontext.WithValue(ctx, (Foo)(nil))
+	assert.True(t, k0scontext.HasValue[Foo](ctxWithFoo))
+	assert.False(t, k0scontext.HasValue[Bar](ctxWithFoo))
+
+	ctxWithFooAndBar := k0scontext.WithValue(ctxWithFoo, (Bar)(nil))
+	assert.True(t, k0scontext.HasValue[Foo](ctxWithFooAndBar))
+	assert.True(t, k0scontext.HasValue[Bar](ctxWithFooAndBar))
+}
+
+func TestHasValue_Aliases(t *testing.T) {
+	type Foo string
+	type Bar string
+
+	ctx := context.Background()
+	assert.False(t, k0scontext.HasValue[Foo](ctx))
+	assert.False(t, k0scontext.HasValue[Bar](ctx))
+
+	ctxWithFoo := k0scontext.WithValue(ctx, Foo(""))
+	assert.True(t, k0scontext.HasValue[Foo](ctxWithFoo))
+	assert.False(t, k0scontext.HasValue[Bar](ctxWithFoo))
+
+	ctxWithFooAndBar := k0scontext.WithValue(ctxWithFoo, Bar(""))
+	assert.True(t, k0scontext.HasValue[Foo](ctxWithFooAndBar))
+	assert.True(t, k0scontext.HasValue[Bar](ctxWithFooAndBar))
+}
+
+func TestValue_StructPtrs(t *testing.T) {
+	type Foo struct{}
+	type Bar struct{}
+
+	ctx := context.Background()
+	assert.Zero(t, k0scontext.Value[*Foo](ctx))
+	assert.Zero(t, k0scontext.Value[*Bar](ctx))
+
+	ctxWithFoo := k0scontext.WithValue(ctx, &Foo{})
+	assert.Equal(t, k0scontext.Value[*Foo](ctxWithFoo), &Foo{})
+	assert.Zero(t, k0scontext.Value[*Bar](ctx))
+
+	ctxWithFooAndBar := k0scontext.WithValue(ctxWithFoo, &Bar{})
+	assert.Equal(t, k0scontext.Value[*Foo](ctxWithFooAndBar), &Foo{})
+	assert.Equal(t, k0scontext.Value[*Bar](ctxWithFooAndBar), &Bar{})
+}