feat(cmd): adding timeout flag (celestiaorg#3580)

Co-authored-by: Hlib Kanunnikov <[email protected]> Co-authored-by: rene <[email protected]> Co-authored-by: Oleg Kovalov <[email protected]>
renaynay · Oct 3, 2024 · 1c62d99 · 1c62d99
1 parent f48223d
commit 1c62d99
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/cmd/rpc.go b/cmd/rpc.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"path/filepath"
+	"time"
 
 	"github.com/spf13/cobra"
 	flag "github.com/spf13/pflag"
@@ -18,6 +19,7 @@ import (
 var (
 	requestURL    string
 	authTokenFlag string
+	timeoutFlag   time.Duration
 )
 
 func RPCFlags() *flag.FlagSet {
@@ -37,6 +39,13 @@ func RPCFlags() *flag.FlagSet {
 		"Authorization token",
 	)
 
+	fset.DurationVar(
+		&timeoutFlag,
+		"timeout",
+		0,
+		"Timeout for RPC requests (e.g. 30s, 1m)",
+	)
+
 	storeFlag := NodeFlags().Lookup(nodeStoreFlag)
 	fset.AddFlag(storeFlag)
 	return fset
@@ -73,6 +82,12 @@ func InitClient(cmd *cobra.Command, _ []string) error {
 		}
 	}
 
+	if timeoutFlag > 0 {
+		// we don't cancel, because we want to keep this context alive outside the InitClient Function
+		ctx, _ := context.WithTimeout(cmd.Context(), timeoutFlag) //nolint:govet
+		cmd.SetContext(ctx)
+	}
+
 	client, err := rpc.NewClient(cmd.Context(), requestURL, authTokenFlag)
 	if err != nil {
 		return err