add timeout for admin requests

deviceinsight · Dec 8, 2020 · 8be68c2 · 8be68c2
1 parent 8309d18
commit 8be68c2
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 0 deletions.
diff --git a/.github/workflows/lint_test.yml b/.github/workflows/lint_test.yml
@@ -52,6 +52,7 @@ jobs:
     env:
       CP_VERSION: ${{ matrix.cp_version }}
       KAFKAVERSION: ${{ matrix.kafka_version }}
+      TIMEOUT: 30
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-go@v2

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Note that kafka >= 2.4.0.0 is required, otherwise the relevant api calls are not available.
 - Added command `alter partition` which currently only enables to manually assign broker replicas to a partition.
   Note that kafka >= 2.4.0.0 is required, otherwise the relevant api calls are not available.
+- Added `timeout` config to control timeout of admin requests.
 
 ## 1.13.3 - 2020-11-11
 

diff --git a/README.md b/README.md
@@ -110,6 +110,9 @@ contexts:
     # optional: kafkaVersion (defaults to 2.0.0)
     kafkaVersion: 1.1.1
 
+    # optional: timeout for admin requests (defaults to 3s)
+    timeout: 10
+
     # optional: avro schema registry
     avro:
       schemaRegistry: localhost:8081

diff --git a/cmd/root.go b/cmd/root.go
@@ -26,6 +26,7 @@ var cfgFile string
 var Verbose bool
 
 var envMapping = map[string]string{
+	"TIMEOUT":             "CONTEXTS_DEFAULT_TIMEOUT",
 	"BROKERS":             "CONTEXTS_DEFAULT_BROKERS",
 	"TLS_ENABLED":         "CONTEXTS_DEFAULT_TLS_ENABLED",
 	"TLS_CA":              "CONTEXTS_DEFAULT_TLS_CA",

diff --git a/cmd/root_test.go b/cmd/root_test.go
@@ -39,6 +39,7 @@ func TestEnvironmentVariableLoadingAliases(t *testing.T) {
 
 	test_util.StartUnitTest(t)
 
+	_ = os.Setenv("TIMEOUT", "30")
 	_ = os.Setenv("BROKERS", "broker1:9092 broker2:9092")
 	_ = os.Setenv("TLS_ENABLED", "true")
 	_ = os.Setenv("TLS_CA", "my-ca")
@@ -65,6 +66,7 @@ func TestEnvironmentVariableLoadingAliases(t *testing.T) {
 
 	test_util.AssertEquals(t, "broker1:9092", viper.GetStringSlice("contexts.default.brokers")[0])
 	test_util.AssertEquals(t, "broker2:9092", viper.GetStringSlice("contexts.default.brokers")[1])
+	test_util.AssertEquals(t, "30", viper.GetString("contexts.default.timeout"))
 	test_util.AssertEquals(t, "true", viper.GetString("contexts.default.tls.enabled"))
 	test_util.AssertEquals(t, "my-ca", viper.GetString("contexts.default.tls.ca"))
 	test_util.AssertEquals(t, "my-cert", viper.GetString("contexts.default.tls.cert"))

diff --git a/operations/common-operation.go b/operations/common-operation.go
@@ -11,6 +11,7 @@ import (
 	"os/user"
 	"regexp"
 	"strings"
+	"time"
 )
 
 var (
@@ -45,6 +46,7 @@ type ClientContext struct {
 	Tls                TlsConfig
 	Sasl               SaslConfig
 	Kubernetes         K8sConfig
+	Timeout            time.Duration
 	ClientID           string
 	KafkaVersion       sarama.KafkaVersion
 	AvroSchemaRegistry string
@@ -75,6 +77,8 @@ func CreateClientContext() (ClientContext, error) {
 	context.Tls.Insecure = viper.GetBool("contexts." + context.Name + ".tls.insecure")
 	context.ClientID = viper.GetString("contexts." + context.Name + ".clientID")
 
+	context.Timeout = viper.GetDuration("contexts." + context.Name + ".timeout")
+
 	if version, err := kafkaVersion(viper.GetString("contexts." + context.Name + ".kafkaVersion")); err == nil {
 		context.KafkaVersion = version
 	} else {
@@ -111,6 +115,11 @@ func CreateClientConfig(context *ClientContext) (*sarama.Config, error) {
 	config.Version = context.KafkaVersion
 	config.ClientID = getClientID(context)
 
+	if context.Timeout > 0 {
+		output.Debugf("using timeout: %d", context.Timeout)
+		config.Admin.Timeout = context.Timeout
+	}
+
 	if context.Tls.Enabled {
 		output.Debugf("TLS is enabled.")
 		config.Net.TLS.Enable = true