feat(wait): introduce waitOptions.RunWaitContext() (#136781)

* run wait with context

Signed-off-by: Austin Abro <austinabro321@gmail.com>

* delete other comment

Signed-off-by: Austin Abro <austinabro321@gmail.com>

---------

Signed-off-by: Austin Abro <austinabro321@gmail.com>

staging/src/k8s.io/kubectl/pkg/cmd/delete/delete.go

@@ -17,6 +17,7 @@ limitations under the License.
 package delete
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"strings"
@@ -466,7 +467,7 @@ func (o *DeleteOptions) DeleteResult(r *resource.Result) error {
 		ConditionFn: cmdwait.IsDeleted,
 		IOStreams:   o.IOStreams,
 	}
-	err = waitOptions.RunWait()
+	err = waitOptions.RunWaitContext(context.Background())
 	if errors.IsForbidden(err) || errors.IsMethodNotSupported(err) {
 		// if we're forbidden from waiting, we shouldn't fail.
 		// if the resource doesn't support a verb we need, we shouldn't fail.

staging/src/k8s.io/kubectl/pkg/cmd/logs/logs.go

@@ -182,7 +182,7 @@ func NewCmdLogs(f cmdutil.Factory, streams genericiooptions.IOStreams) *cobra.Co
 		Run: func(cmd *cobra.Command, args []string) {
 			cmdutil.CheckErr(o.Complete(f, cmd, args))
 			cmdutil.CheckErr(o.Validate())
-			cmdutil.CheckErr(o.RunLogs())
+			cmdutil.CheckErr(o.RunLogsContext(cmd.Context()))
 		},
 	}
 	o.AddFlags(cmd)
@@ -356,45 +356,44 @@ func (o LogsOptions) Validate() error {
 	return nil
 }
 
-// RunLogs wraps RunLogsContext with signal handling.
-// When a signal is received, streaming is stopped, then followed by os.Exit(1).
+// Deprecated: Use RunLogsContext instead which allows cancelling.
 func (o LogsOptions) RunLogs() error {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	intr := interrupt.New(nil, cancel)
-	return intr.Run(func() error {
-		return o.RunLogsContext(ctx)
-	})
+	return o.RunLogsContext(context.Background())
 }
 
 // RunLogsContext retrieves a pod log.
 //
-// This function does not handle signals. To interrupt streaming, cancel the context.
+// When a signal is received, streaming is stopped, then followed by os.Exit(1).
 func (o LogsOptions) RunLogsContext(ctx context.Context) error {
-	var requests map[corev1.ObjectReference]rest.ResponseWrapper
-	var err error
-	if o.AllPods {
-		requests, err = o.AllPodLogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
-	} else {
-		requests, err = o.LogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
-	}
-	if err != nil {
-		return err
-	}
-
-	if o.Follow && len(requests) > 1 {
-		if len(requests) > o.MaxFollowConcurrency {
-			return fmt.Errorf(
-				"you are attempting to follow %d log streams, but maximum allowed concurrency is %d, use --max-log-requests to increase the limit",
-				len(requests), o.MaxFollowConcurrency,
-			)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	intr := interrupt.New(nil, cancel)
+	return intr.Run(func() error {
+		var requests map[corev1.ObjectReference]rest.ResponseWrapper
+		var err error
+		if o.AllPods {
+			requests, err = o.AllPodLogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
+		} else {
+			requests, err = o.LogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
+		}
+		if err != nil {
+			return err
 		}
-	}
 
-	if o.Follow && len(requests) > 1 {
-		return o.parallelConsumeRequest(ctx, requests)
-	}
-	return o.sequentialConsumeRequest(ctx, requests)
+		if o.Follow && len(requests) > 1 {
+			if len(requests) > o.MaxFollowConcurrency {
+				return fmt.Errorf(
+					"you are attempting to follow %d log streams, but maximum allowed concurrency is %d, use --max-log-requests to increase the limit",
+					len(requests), o.MaxFollowConcurrency,
+				)
+			}
+		}
+
+		if o.Follow && len(requests) > 1 {
+			return o.parallelConsumeRequest(ctx, requests)
+		}
+		return o.sequentialConsumeRequest(ctx, requests)
+	})
 }
 
 func (o LogsOptions) parallelConsumeRequest(ctx context.Context, requests map[corev1.ObjectReference]rest.ResponseWrapper) error {

staging/src/k8s.io/kubectl/pkg/cmd/logs/logs_test.go

@@ -543,7 +543,7 @@ func TestLog(t *testing.T) {
 			opts.Namespace = "test"
 			opts.Object = testPod()
 			opts.Options = &corev1.PodLogOptions{}
-			err := opts.RunLogs()
+			err := opts.RunLogsContext(t.Context())
 
 			if err == nil && len(test.expectedErr) > 0 {
 				t.Fatalf("expected error %q, got none", test.expectedErr)

staging/src/k8s.io/kubectl/pkg/cmd/portforward/portforward.go

@@ -411,7 +411,6 @@ func (o PortForwardOptions) Validate() error {
 }
 
 // Deprecated: Use RunPortForwardContext instead, which allows canceling.
-// RunPortForward implements all the necessary functionality for port-forward cmd.
 func (o PortForwardOptions) RunPortForward() error {
 	return o.RunPortForwardContext(context.Background())
 }

staging/src/k8s.io/kubectl/pkg/cmd/wait/wait.go

@@ -132,7 +132,7 @@ func NewCmdWait(restClientGetter genericclioptions.RESTClientGetter, streams gen
 		Run: func(cmd *cobra.Command, args []string) {
 			o, err := flags.ToOptions(args)
 			cmdutil.CheckErr(err)
-			cmdutil.CheckErr(o.RunWait())
+			cmdutil.CheckErr(o.RunWaitContext(cmd.Context()))
 		},
 		SuggestFor: []string{"list", "ps"},
 	}
@@ -317,11 +317,15 @@ type WaitOptions struct {
 // ConditionFunc is the interface for providing condition checks
 type ConditionFunc func(ctx context.Context, info *resource.Info, o *WaitOptions) (finalObject runtime.Object, done bool, err error)
 
-// RunWait runs the waiting logic
+// Deprecated: Use RunWaitContext instead, which allows canceling.
 func (o *WaitOptions) RunWait() error {
-	ctx, cancel := watchtools.ContextWithOptionalTimeout(context.Background(), o.Timeout)
-	defer cancel()
+	return o.RunWaitContext(context.Background())
+}
 
+// RunWaitContext runs the waiting logic
+func (o *WaitOptions) RunWaitContext(ctx context.Context) error {
+	ctx, cancel := watchtools.ContextWithOptionalTimeout(ctx, o.Timeout)
+	defer cancel()
 	if strings.ToLower(o.ForCondition) == "create" {
 		// TODO(soltysh): this is not ideal solution, because we're polling every .5s,
 		// and we have to use ResourceFinder, which contains the resource name.

staging/src/k8s.io/kubectl/pkg/cmd/wait/wait_test.go

@@ -599,7 +599,7 @@ func TestWaitForDeletion(t *testing.T) {
 				ConditionFn: IsDeleted,
 				IOStreams:   genericiooptions.NewTestIOStreamsDiscard(),
 			}
-			err := o.RunWait()
+			err := o.RunWaitContext(t.Context())
 			switch {
 			case err == nil && len(test.expectedErr) == 0:
 			case err != nil && len(test.expectedErr) == 0:
@@ -1009,7 +1009,7 @@ func TestWaitForCondition(t *testing.T) {
 				ConditionFn: ConditionalWait{conditionName: "the-condition", conditionStatus: "status-value", errOut: io.Discard}.IsConditionMet,
 				IOStreams:   genericiooptions.NewTestIOStreamsDiscard(),
 			}
-			err := o.RunWait()
+			err := o.RunWaitContext(t.Context())
 			switch {
 			case err == nil && len(test.expectedErr) == 0:
 			case err != nil && len(test.expectedErr) == 0:
@@ -1080,7 +1080,7 @@ func TestWaitForCreate(t *testing.T) {
 				ForCondition: "create",
 				IOStreams:    genericiooptions.NewTestIOStreamsDiscard(),
 			}
-			err := o.RunWait()
+			err := o.RunWaitContext(t.Context())
 			switch {
 			case err == nil && len(test.expectedErr) == 0:
 			case err != nil && len(test.expectedErr) == 0:
@@ -1112,7 +1112,7 @@ func TestWaitForDeletionIgnoreNotFound(t *testing.T) {
 		IOStreams:      genericiooptions.NewTestIOStreamsDiscard(),
 		ForCondition:   "delete",
 	}
-	err := o.RunWait()
+	err := o.RunWaitContext(t.Context())
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -1367,7 +1367,7 @@ func TestWaitForDifferentJSONPathExpression(t *testing.T) {
 				IOStreams: genericiooptions.NewTestIOStreamsDiscard(),
 			}
 
-			err := o.RunWait()
+			err := o.RunWaitContext(t.Context())
 
 			switch {
 			case err == nil && len(test.expectedErr) == 0:
@@ -1630,7 +1630,7 @@ func TestWaitForJSONPathCondition(t *testing.T) {
 				IOStreams: genericiooptions.NewTestIOStreamsDiscard(),
 			}
 
-			err := o.RunWait()
+			err := o.RunWaitContext(t.Context())
 
 			switch {
 			case err == nil && len(test.expectedErr) == 0: