Use context.WithTimeoutCause and context.WithCancelCause for better readability

Amulyam24 · Amulyam24 · commit 97f8a4f735a8 · 2025-02-06T14:59:24.000+05:30
diff --git a/cmd/clusterctl/client/repository/repository_gitlab.go b/cmd/clusterctl/client/repository/repository_gitlab.go
@@ -147,7 +147,7 @@ func (g *gitLabRepository) GetFile(ctx context.Context, version, path string) ([
 		return content, nil
 	}
 
-	timeoutctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	timeoutctx, cancel := context.WithTimeoutCause(ctx, 30*time.Second, errors.New("http request timeout expired"))
 	defer cancel()
 	request, err := http.NewRequestWithContext(timeoutctx, http.MethodGet, url, http.NoBody)
 	if err != nil {
diff --git a/controllers/clustercache/cluster_accessor_client.go b/controllers/clustercache/cluster_accessor_client.go
@@ -242,7 +242,7 @@ func createCachedClient(ctx, cacheCtx context.Context, clusterAccessorConfig *cl
 
 	// Use a context that is independent of the passed in context, so the cache doesn't get stopped
 	// when the passed in context is canceled.
-	cacheCtx, cacheCtxCancel := context.WithCancel(cacheCtx)
+	cacheCtx, cacheCtxCancel := context.WithCancelCause(cacheCtx)
 
 	// We need to be able to stop the cache's shared informers, so wrap this in a stoppableCache.
 	cache := &stoppableCache{
@@ -275,7 +275,7 @@ func createCachedClient(ctx, cacheCtx context.Context, clusterAccessorConfig *cl
 	go cache.Start(cacheCtx) //nolint:errcheck
 
 	// Wait until the cache is initially synced.
-	cacheSyncCtx, cacheSyncCtxCancel := context.WithTimeout(ctx, clusterAccessorConfig.Cache.InitialSyncTimeout)
+	cacheSyncCtx, cacheSyncCtxCancel := context.WithTimeoutCause(ctx, clusterAccessorConfig.Cache.InitialSyncTimeout, errors.New("initial sync timeout expired"))
 	defer cacheSyncCtxCancel()
 	if !cache.WaitForCacheSync(cacheSyncCtx) {
 		cache.Stop()
@@ -311,13 +311,13 @@ type clientWithTimeout struct {
 var _ client.Client = &clientWithTimeout{}
 
 func (c clientWithTimeout) Get(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
-	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.timeout, errors.New("call timeout expired"))
 	defer cancel()
 	return c.Client.Get(ctx, key, obj, opts...)
 }
 
 func (c clientWithTimeout) List(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
-	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.timeout, errors.New("call timeout expired"))
 	defer cancel()
 	return c.Client.List(ctx, list, opts...)
 }
@@ -328,7 +328,7 @@ type stoppableCache struct {
 
 	lock       sync.Mutex
 	stopped    bool
-	cancelFunc context.CancelFunc
+	cancelFunc context.CancelCauseFunc
 }
 
 // Stop cancels the cache.Cache's context, unless it has already been stopped.
@@ -341,5 +341,5 @@ func (cc *stoppableCache) Stop() {
 	}
 
 	cc.stopped = true
-	cc.cancelFunc()
+	cc.cancelFunc(errors.New("stoppable cache context cancelled"))
 }
diff --git a/controllers/remote/cluster_cache_healthcheck_test.go b/controllers/remote/cluster_cache_healthcheck_test.go
@@ -25,6 +25,7 @@ import (
 	"time"
 
 	. "github.com/onsi/gomega"
+	"github.com/pkg/errors"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -106,7 +107,7 @@ func TestClusterCacheHealthCheck(t *testing.T) {
 
 			testClusterKey = util.ObjectKey(testCluster)
 
-			_, cancel := context.WithCancel(ctx)
+			_, cancel := context.WithCancelCause(ctx)
 			cc = &stoppableCache{cancelFunc: cancel}
 			cct.clusterAccessors[testClusterKey] = &clusterAccessor{cache: cc}
 
@@ -123,7 +124,7 @@ func TestClusterCacheHealthCheck(t *testing.T) {
 			t.Log("Deleting Namespace")
 			g.Expect(env.Delete(ctx, ns)).To(Succeed())
 			t.Log("Stopping the manager")
-			cc.cancelFunc()
+			cc.cancelFunc(errors.New("context cancelled"))
 			mgrCancel()
 		}
 
diff --git a/controllers/remote/cluster_cache_tracker.go b/controllers/remote/cluster_cache_tracker.go
@@ -514,7 +514,7 @@ func (t *ClusterCacheTracker) createCachedClient(ctx context.Context, config *re
 
 	// Use a context that is independent of the passed in context, so the cache doesn't get stopped
 	// when the passed in context is canceled.
-	cacheCtx, cacheCtxCancel := context.WithCancel(context.Background())
+	cacheCtx, cacheCtxCancel := context.WithCancelCause(context.Background())
 
 	// We need to be able to stop the cache's shared informers, so wrap this in a stoppableCache.
 	cache := &stoppableCache{
@@ -547,7 +547,7 @@ func (t *ClusterCacheTracker) createCachedClient(ctx context.Context, config *re
 	go cache.Start(cacheCtx) //nolint:errcheck
 
 	// Wait until the cache is initially synced
-	cacheSyncCtx, cacheSyncCtxCancel := context.WithTimeout(ctx, initialCacheSyncTimeout)
+	cacheSyncCtx, cacheSyncCtxCancel := context.WithTimeoutCause(ctx, initialCacheSyncTimeout, errors.New("initial sync timeout expired"))
 	defer cacheSyncCtxCancel()
 	if !cache.WaitForCacheSync(cacheSyncCtx) {
 		cache.Stop()
@@ -774,13 +774,13 @@ type clientWithTimeout struct {
 var _ client.Client = &clientWithTimeout{}
 
 func (c clientWithTimeout) Get(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
-	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.timeout, errors.New("call timeout expired"))
 	defer cancel()
 	return c.Client.Get(ctx, key, obj, opts...)
 }
 
 func (c clientWithTimeout) List(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
-	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.timeout, errors.New("call timeout expired"))
 	defer cancel()
 	return c.Client.List(ctx, list, opts...)
 }
diff --git a/controllers/remote/stoppable_cache.go b/controllers/remote/stoppable_cache.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"sync"
 
+	"github.com/pkg/errors"
 	"sigs.k8s.io/controller-runtime/pkg/cache"
 )
 
@@ -29,7 +30,7 @@ type stoppableCache struct {
 
 	lock       sync.Mutex
 	stopped    bool
-	cancelFunc context.CancelFunc
+	cancelFunc context.CancelCauseFunc
 }
 
 // Stop cancels the cache.Cache's context, unless it has already been stopped.
@@ -42,5 +43,5 @@ func (cc *stoppableCache) Stop() {
 	}
 
 	cc.stopped = true
-	cc.cancelFunc()
+	cc.cancelFunc(errors.New("stoppable cache context cancelled"))
 }
diff --git a/controlplane/kubeadm/internal/etcd/etcd.go b/controlplane/kubeadm/internal/etcd/etcd.go
@@ -180,7 +180,7 @@ func newEtcdClient(ctx context.Context, etcdClient etcd, callTimeout time.Durati
 		return nil, errors.New("etcd client was not configured with any endpoints")
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, callTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, callTimeout, errors.New("call timeout expired"))
 	defer cancel()
 
 	status, err := etcdClient.Status(ctx, endpoints[0])
@@ -204,7 +204,7 @@ func (c *Client) Close() error {
 
 // Members retrieves a list of etcd members.
 func (c *Client) Members(ctx context.Context) ([]*Member, error) {
-	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.CallTimeout, errors.New("call timeout expired"))
 	defer cancel()
 
 	response, err := c.EtcdClient.MemberList(ctx)
@@ -225,7 +225,7 @@ func (c *Client) Members(ctx context.Context) ([]*Member, error) {
 
 // MoveLeader moves the leader to the provided member ID.
 func (c *Client) MoveLeader(ctx context.Context, newLeaderID uint64) error {
-	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.CallTimeout, errors.New("call timeout expired"))
 	defer cancel()
 
 	_, err := c.EtcdClient.MoveLeader(ctx, newLeaderID)
@@ -234,7 +234,7 @@ func (c *Client) MoveLeader(ctx context.Context, newLeaderID uint64) error {
 
 // RemoveMember removes a given member.
 func (c *Client) RemoveMember(ctx context.Context, id uint64) error {
-	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.CallTimeout, errors.New("call timeout expired"))
 	defer cancel()
 
 	_, err := c.EtcdClient.MemberRemove(ctx, id)
@@ -243,7 +243,7 @@ func (c *Client) RemoveMember(ctx context.Context, id uint64) error {
 
 // Alarms retrieves all alarms on a cluster.
 func (c *Client) Alarms(ctx context.Context) ([]MemberAlarm, error) {
-	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, c.CallTimeout, errors.New("call timeout expired"))
 	defer cancel()
 
 	alarmResponse, err := c.EtcdClient.AlarmList(ctx)
diff --git a/exp/runtime/internal/controllers/warmup.go b/exp/runtime/internal/controllers/warmup.go
@@ -69,7 +69,7 @@ func (r *warmupRunnable) Start(ctx context.Context) error {
 	if r.warmupTimeout == 0 {
 		r.warmupTimeout = defaultWarmupTimeout
 	}
-	ctx, cancel := context.WithTimeout(ctx, r.warmupTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, r.warmupTimeout, errors.New("warmup timeout expired"))
 	defer cancel()
 
 	err := wait.PollUntilContextTimeout(ctx, r.warmupInterval, r.warmupTimeout, true, func(ctx context.Context) (done bool, err error) {
diff --git a/hack/tools/internal/log-push/main.go b/hack/tools/internal/log-push/main.go
@@ -170,7 +170,7 @@ func getLogsFromGCS(ctx context.Context, logPath string, logFileRegex *regexp.Re
 	klog.Infof("Getting logs from gs://%s/%s", bucket, folder)
 
 	// Set timeout.
-	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
+	ctx, cancel := context.WithTimeoutCause(ctx, 2*time.Minute, errors.New("new client timeout expired"))
 	defer cancel()
 
 	// Create GCS client.
diff --git a/hack/tools/internal/tilt-prepare/main.go b/hack/tools/internal/tilt-prepare/main.go
@@ -458,8 +458,8 @@ func runTaskGroup(ctx context.Context, name string, tasks map[string]taskFunctio
 	}(time.Now())
 
 	// Create a context to be used for canceling all the tasks when another fails.
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+	ctx, cancel := context.WithCancelCause(ctx)
+	defer cancel(errors.New("run task group context cancelled"))
 
 	// Make channels to pass fatal errors in WaitGroup
 	errors := make(chan error)
@@ -483,7 +483,7 @@ func runTaskGroup(ctx context.Context, name string, tasks map[string]taskFunctio
 		break
 	case err := <-errors:
 		// cancel all the running tasks
-		cancel()
+		cancel(err)
 		// consumes all the errors from the channel
 		errList := []error{err}
 	Loop:
diff --git a/internal/controllers/machine/drain/drain.go b/internal/controllers/machine/drain/drain.go
@@ -316,7 +316,7 @@ func (d *Helper) EvictPods(ctx context.Context, podDeleteList *PodDeleteList) Ev
 
 	// Trigger evictions for at most 10s. We'll continue on the next reconcile if we hit the timeout.
 	evictionTimeout := 10 * time.Second
-	ctx, cancel := context.WithTimeout(ctx, evictionTimeout)
+	ctx, cancel := context.WithTimeoutCause(ctx, evictionTimeout, errors.New("eviction timeout expired"))
 	defer cancel()
 
 	res := EvictionResult{
diff --git a/internal/runtime/client/client.go b/internal/runtime/client/client.go
@@ -471,7 +471,7 @@ func httpCall(ctx context.Context, request, response runtime.Object, opts *httpC
 		extensionURL.RawQuery = values.Encode()
 
 		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, opts.timeout)
+		ctx, cancel = context.WithTimeoutCause(ctx, opts.timeout, errors.New("http request timeout expired"))
 		defer cancel()
 	}
 
diff --git a/internal/test/envtest/environment.go b/internal/test/envtest/environment.go
@@ -151,7 +151,7 @@ func Run(ctx context.Context, input RunInput) int {
 	// Bootstrapping test environment
 	env := newEnvironment(input.ManagerCacheOptions, input.ManagerUncachedObjs...)
 
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancelCause(ctx)
 	env.cancelManager = cancel
 
 	if input.SetupIndexes != nil {
@@ -227,7 +227,7 @@ type Environment struct {
 	Config *rest.Config
 
 	env           *envtest.Environment
-	cancelManager context.CancelFunc
+	cancelManager context.CancelCauseFunc
 }
 
 // newEnvironment creates a new environment spinning up a local api-server.
@@ -401,7 +401,7 @@ func (e *Environment) start(ctx context.Context) {
 // stop stops the test environment.
 func (e *Environment) stop() error {
 	fmt.Println("Stopping the test environment")
-	e.cancelManager()
+	e.cancelManager(errors.New("test environment context cancelled"))
 	return e.env.Stop()
 }
 

Original file line number	Diff line number	Diff line change
`@@ -147,7 +147,7 @@ func (g *gitLabRepository) GetFile(ctx context.Context, version, path string) ([`
`147`	`147`	`return content, nil`
`148`	`148`	`}`
`149`	`149`
`150`		`- timeoutctx, cancel := context.WithTimeout(ctx, 30*time.Second)`
	`150`	`+ timeoutctx, cancel := context.WithTimeoutCause(ctx, 30*time.Second, errors.New("http request timeout expired"))`
`151`	`151`	`defer cancel()`
`152`	`152`	`request, err := http.NewRequestWithContext(timeoutctx, http.MethodGet, url, http.NoBody)`
`153`	`153`	`if err != nil {`
Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@ import (`
`20`	`20`	`"context"`
`21`	`21`	`"sync"`
`22`	`22`
	`23`	`+ "github.com/pkg/errors"`
`23`	`24`	`"sigs.k8s.io/controller-runtime/pkg/cache"`
`24`	`25`	`)`
`25`	`26`
`@@ -29,7 +30,7 @@ type stoppableCache struct {`
`29`	`30`
`30`	`31`	`lock sync.Mutex`
`31`	`32`	`stopped bool`
`32`		`- cancelFunc context.CancelFunc`
	`33`	`+ cancelFunc context.CancelCauseFunc`
`33`	`34`	`}`
`34`	`35`
`35`	`36`	`// Stop cancels the cache.Cache's context, unless it has already been stopped.`
`@@ -42,5 +43,5 @@ func (cc *stoppableCache) Stop() {`
`42`	`43`	`}`
`43`	`44`
`44`	`45`	`cc.stopped = true`
`45`		`- cc.cancelFunc()`
	`46`	`+ cc.cancelFunc(errors.New("stoppable cache context cancelled"))`
`46`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ func (r *warmupRunnable) Start(ctx context.Context) error {`
`69`	`69`	`if r.warmupTimeout == 0 {`
`70`	`70`	`r.warmupTimeout = defaultWarmupTimeout`
`71`	`71`	`}`
`72`		`- ctx, cancel := context.WithTimeout(ctx, r.warmupTimeout)`
	`72`	`+ ctx, cancel := context.WithTimeoutCause(ctx, r.warmupTimeout, errors.New("warmup timeout expired"))`
`73`	`73`	`defer cancel()`
`74`	`74`
`75`	`75`	`err := wait.PollUntilContextTimeout(ctx, r.warmupInterval, r.warmupTimeout, true, func(ctx context.Context) (done bool, err error) {`