Repository: knative/pkg
Pull request #1451 (2020-06-27 15:37:04)

Enable use of StatefulSet in sharedmain.
This change allows for (just webhook for now) controllers going through sharedmain to opt into Yanwei's logic by setting several environment variables.

I was able to pull this change in downstream and change the webhook to use a StatefulSet with the following environment:
```
+        # These settings are used for statefulset-based
+        # leader selection.
+        - name: CONTROLLER_ORDINAL
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.name
+        - name: STATEFUL_SERVICE_NAME
+          value: "webhook"
```

Running the above with 10 replicas and 10 buckets worked as intended (keys were evenly distributed across the replicas).

knative-prow-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/knative/pkg/pull/1451#" title="Author self-approved">mattmoor</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=knative%2Fpkg).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/knative/pkg/blob/master/OWNERS)~~ [mattmoor]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
knative-metrics-robot: The following is the coverage report on the affected files.
Say `/test pull-knative-pkg-go-coverage` to re-run this coverage report

File | Old Coverage | New Coverage | Delta
---- |:------------:|:------------:|:-----:
[leaderelection/config.go](https://storage.cloud.google.com/knative-prow/pr-logs/pull/knative_pkg/1451/pull-knative-pkg-go-coverage/1276902500355543040/artifacts/line-cov.html#file0) | 72.0% | 75.8% | 3.8
[leaderelection/context.go](https://storage.cloud.google.com/knative-prow/pr-logs/pull/knative_pkg/1451/pull-knative-pkg-go-coverage/1276902500355543040/artifacts/line-cov.html#file1) | 84.9% | 84.2% | -0.7


Patch:
diff --git a/injection/sharedmain/main.go b/injection/sharedmain/main.go
index fdc910ef0..6c4acf71e 100644
--- a/injection/sharedmain/main.go
+++ b/injection/sharedmain/main.go
@@ -241,7 +241,7 @@ func WebhookMainWithConfig(ctx context.Context, component string, cfg *rest.Conf
 	leConfig := leaderElectionConfig.GetComponentConfig(component)
 	if leConfig.LeaderElect {
 		// Signal that we are executing in a context with leader election.
-		ctx = kle.WithStandardLeaderElectorBuilder(ctx, kubeclient.Get(ctx), leConfig)
+		ctx = kle.WithDynamicLeaderElectorBuilder(ctx, kubeclient.Get(ctx), leConfig)
 	}
 
 	controllers, webhooks := ControllersAndWebhooksFromCtors(ctx, cmw, ctors...)
diff --git a/leaderelection/config.go b/leaderelection/config.go
index 8b7c02058..9002affb7 100644
--- a/leaderelection/config.go
+++ b/leaderelection/config.go
@@ -21,6 +21,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/kelseyhightower/envconfig"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apimachinery/pkg/util/uuid"
@@ -129,9 +130,22 @@ type ComponentConfig struct {
 // StatefulSetConfig represents the required information for a StatefulSet service.
 type StatefulSetConfig struct {
 	StatefulSetName string
-	ServiceName     string
-	Port            string
-	Protocol        string
+	ServiceName     string `envconfig:"STATEFUL_SERVICE_NAME" required:"true"`
+	Port            string `envconfig:"STATEFUL_SERVICE_PORT" default:"80"`
+	Protocol        string `envconfig:"STATEFUL_SERVICE_PROTOCOL" default:"http"`
+}
+
+func NewStatefulSetConfig() (*StatefulSetConfig, error) {
+	var ssc StatefulSetConfig
+	ssn, _, err := ParseControllerOrdinal()
+	if err != nil {
+		return nil, err
+	}
+	ssc.StatefulSetName = ssn
+	if err := envconfig.Process("", &ssc); err != nil {
+		return nil, err
+	}
+	return &ssc, nil
 }
 
 func defaultComponentConfig(name string) ComponentConfig {
diff --git a/leaderelection/context.go b/leaderelection/context.go
index d273a572e..3322b1b87 100644
--- a/leaderelection/context.go
+++ b/leaderelection/context.go
@@ -33,6 +33,16 @@ import (
 	"knative.dev/pkg/system"
 )
 
+// WithDynamicLeaderElectorBuilder sets up the statefulset elector based on environment,
+// falling back on the standard elector.
+func WithDynamicLeaderElectorBuilder(ctx context.Context, kc kubernetes.Interface, cc ComponentConfig) context.Context {
+	ssc, err := NewStatefulSetConfig()
+	if err == nil {
+		return WithStatefulSetLeaderElectorBuilder(ctx, cc, *ssc)
+	}
+	return WithStandardLeaderElectorBuilder(ctx, kc, cc)
+}
+
 // WithStandardLeaderElectorBuilder infuses a context with the ability to build
 // LeaderElectors with the provided component configuration acquiring resource
 // locks via the provided kubernetes client.
diff --git a/leaderelection/ordinal.go b/leaderelection/ordinal.go
index 94a3775a0..28f70f734 100644
--- a/leaderelection/ordinal.go
+++ b/leaderelection/ordinal.go
@@ -27,13 +27,20 @@ import (
 // should be set to the pod name via the downward API.
 const controllerOrdinalEnv = "CONTROLLER_ORDINAL"
 
-// ControllerOrdinal tries to get ordinal from the pod name of a StatefulSet,
-// which is provided from the environment variable CONTROLLER_ORDINAL.
-func ControllerOrdinal() (uint64, error) {
+// ParseControllerOrdinal returns the parts of the StatefulSet's pod name.
+func ParseControllerOrdinal() (string, uint64, error) {
 	v := os.Getenv(controllerOrdinalEnv)
 	if i := strings.LastIndex(v, "-"); i != -1 {
-		return strconv.ParseUint(v[i+1:], 10, 64)
+		ui, err := strconv.ParseUint(v[i+1:], 10, 64)
+		return v[:i], ui, err
 	}
 
-	return 0, fmt.Errorf("ordinal not found in %s=%s", controllerOrdinalEnv, v)
+	return "", 0, fmt.Errorf("ordinal not found in %s=%s", controllerOrdinalEnv, v)
+}
+
+// ControllerOrdinal tries to get ordinal from the pod name of a StatefulSet,
+// which is provided from the environment variable CONTROLLER_ORDINAL.
+func ControllerOrdinal() (uint64, error) {
+	_, ui, err := ParseControllerOrdinal()
+	return ui, err
 }


Test patch:
diff --git a/leaderelection/context_test.go b/leaderelection/context_test.go
index 76e75caae..e124d58ce 100644
--- a/leaderelection/context_test.go
+++ b/leaderelection/context_test.go
@@ -89,7 +89,7 @@ func TestWithBuilder(t *testing.T) {
 		t.Errorf("BuildElector() = %T, wanted an unopposedElector", le)
 	}
 
-	ctx = WithStandardLeaderElectorBuilder(ctx, kc, cc)
+	ctx = WithDynamicLeaderElectorBuilder(ctx, kc, cc)
 	if !HasLeaderElection(ctx) {
 		t.Error("HasLeaderElection() = false, wanted true")
 	}
@@ -164,25 +164,29 @@ func TestWithStatefulSetBuilder(t *testing.T) {
 	}
 	enq := func(reconciler.Bucket, types.NamespacedName) {}
 
-	ctx = WithStatefulSetLeaderElectorBuilder(ctx, cc, StatefulSetConfig{
-		ServiceName:     "autoscaler",
-		StatefulSetName: "as",
-		Protocol:        "ws",
-		Port:            "8080",
-	})
+	if os.Setenv(controllerOrdinalEnv, "as-0") != nil {
+		t.Fatalf("fail to set env var %s=%s", controllerOrdinalEnv, "as-0")
+	}
+	defer os.Unsetenv(controllerOrdinalEnv)
+	if os.Setenv("STATEFUL_SERVICE_NAME", "autoscaler") != nil {
+		t.Fatalf("fail to set env var %s=%s", "STATEFUL_SERVICE_NAME", "autoscaler")
+	}
+	defer os.Unsetenv("STATEFUL_SERVICE_NAME")
+	if os.Setenv("STATEFUL_SERVICE_PORT", "8080") != nil {
+		t.Fatalf("fail to set env var %s=%s", "STATEFUL_SERVICE_PORT", "8080")
+	}
+	defer os.Unsetenv("STATEFUL_SERVICE_PORT")
+	if os.Setenv("STATEFUL_SERVICE_PROTOCOL", "ws") != nil {
+		t.Fatalf("fail to set env var %s=%s", "STATEFUL_SERVICE_PROTOCOL", "ws")
+	}
+	defer os.Unsetenv("STATEFUL_SERVICE_PROTOCOL")
+
+	ctx = WithDynamicLeaderElectorBuilder(ctx, nil, cc)
 	if !HasLeaderElection(ctx) {
 		t.Error("HasLeaderElection() = false, wanted true")
 	}
 
 	le, err := BuildElector(ctx, laf, "name", enq)
-	if err == nil {
-		// controller ordinal env not set
-		t.Error("expected BuildElector() returns error but got none")
-	}
-
-	os.Setenv(controllerOrdinalEnv, "as-0")
-	defer os.Unsetenv(controllerOrdinalEnv)
-	le, err = BuildElector(ctx, laf, "name", enq)
 	if err != nil {
 		t.Fatalf("BuildElector() = %v", err)
 	}
diff --git a/leaderelection/ordinal_test.go b/leaderelection/ordinal_test.go
index ce300a33b..ac1c12a69 100644
--- a/leaderelection/ordinal_test.go
+++ b/leaderelection/ordinal_test.go
@@ -24,10 +24,11 @@ import (
 
 func TestControllerOrdinal(t *testing.T) {
 	testCases := []struct {
-		testname string
-		podName  string
-		want     uint64
-		err      error
+		testname    string
+		podName     string
+		wantName    string
+		wantOrdinal uint64
+		err         error
 	}{{
 		testname: "NotSet",
 		err:      fmt.Errorf("ordinal not found in %s=", controllerOrdinalEnv),
@@ -40,12 +41,15 @@ func TestControllerOrdinal(t *testing.T) {
 		podName:  "as-invalid",
 		err:      fmt.Errorf(`strconv.ParseUint: parsing "invalid": invalid syntax`),
 	}, {
-		testname: "ValidName",
-		podName:  "as-0",
+		testname:    "ValidName",
+		podName:     "as-0",
+		wantName:    "as",
+		wantOrdinal: 0,
 	}, {
-		testname: "ValidName",
-		podName:  "as-1",
-		want:     1,
+		testname:    "ValidName",
+		podName:     "as-1",
+		wantName:    "as",
+		wantOrdinal: 1,
 	}}
 
 	defer os.Unsetenv(controllerOrdinalEnv)
@@ -57,15 +61,26 @@ func TestControllerOrdinal(t *testing.T) {
 				}
 			}
 
-			got, gotErr := ControllerOrdinal()
+			gotOrdinal, gotOrdinalErr := ControllerOrdinal()
 			if tt.err != nil {
-				if gotErr == nil || gotErr.Error() != tt.err.Error() {
-					t.Errorf("got %v, want = %v, ", gotErr, tt.err)
+				if gotOrdinalErr == nil || gotOrdinalErr.Error() != tt.err.Error() {
+					t.Errorf("got %v, want = %v, ", gotOrdinalErr, tt.err)
 				}
-			} else if gotErr != nil {
-				t.Error("ControllerOrdinal() =", gotErr)
-			} else if got != tt.want {
-				t.Errorf("ControllerOrdinal() = %d, want = %d", got, tt.want)
+			} else if gotOrdinalErr != nil {
+				t.Error("ControllerOrdinal() = ", gotOrdinalErr)
+			} else if gotOrdinal != tt.wantOrdinal {
+				t.Errorf("ControllerOrdinal() = %d, want = %d", gotOrdinal, tt.wantOrdinal)
+			}
+
+			gotName, _, gotNameErr := ParseControllerOrdinal()
+			if tt.err != nil {
+				if gotNameErr == nil || gotNameErr.Error() != tt.err.Error() {
+					t.Errorf("got %v, want = %v, ", gotNameErr, tt.err)
+				}
+			} else if gotNameErr != nil {
+				t.Error("ParseControllerOrdinal() = ", gotNameErr)
+			} else if gotName != tt.wantName {
+				t.Errorf("ParseControllerOrdinal() = %s, want = %s", gotName, tt.wantName)
 			}
 		})
 	}