Fix StatefulSet ordinal election leaving buckets unowned

leaderelection/config.go

@@ -177,6 +177,14 @@ type statefulSetConfig struct {
 	ServiceName   string        `envconfig:"STATEFUL_SERVICE_NAME" required:"true"`
 	Port          string        `envconfig:"STATEFUL_SERVICE_PORT" default:"80"`
 	Protocol      string        `envconfig:"STATEFUL_SERVICE_PROTOCOL" default:"http"`
+	ReplicaCount  int           `envconfig:"STATEFUL_REPLICA_COUNT" required:"true"`
+}
+
+// statefulSetConfigured reports whether StatefulSet ordinal leader election
+// has been explicitly selected via STATEFUL_CONTROLLER_ORDINAL.
+func statefulSetConfigured() bool {
+	_, ok := os.LookupEnv("STATEFUL_CONTROLLER_ORDINAL")
+	return ok
 }
 
 // newStatefulSetConfig builds a stateful set LE config.
@@ -185,6 +193,9 @@ func newStatefulSetConfig() (*statefulSetConfig, error) {
 	if err := envconfig.Process("", ssc); err != nil {
 		return nil, err
 	}
+	if ssc.ReplicaCount < 1 {
+		return nil, fmt.Errorf("STATEFUL_REPLICA_COUNT must be >= 1, got %d", ssc.ReplicaCount)
+	}
 	return ssc, nil
 }
 

leaderelection/config_test.go

@@ -34,6 +34,7 @@ const (
 	serviceNameEnv       = "STATEFUL_SERVICE_NAME"
 	servicePortEnv       = "STATEFUL_SERVICE_PORT"
 	serviceProtocolEnv   = "STATEFUL_SERVICE_PROTOCOL"
+	replicaCountEnv      = "STATEFUL_REPLICA_COUNT"
 )
 
 func okConfig() *Config {
@@ -255,43 +256,76 @@ func TestNewStatefulSetConfig(t *testing.T) {
 		service  string
 		port     string
 		protocol string
+		replicas string
 		wantErr  string
 		expected statefulSetConfig
 	}{{
-		name:    "success with default",
-		pod:     "as-42",
-		service: "autoscaler",
+		name:     "success with default",
+		pod:      "as-42",
+		service:  "autoscaler",
+		replicas: "3",
 		expected: statefulSetConfig{
 			StatefulSetID: statefulSetID{
 				ssName:  "as",
 				ordinal: 42,
 			},
-			ServiceName: "autoscaler",
-			Port:        "80",
-			Protocol:    "http",
+			ServiceName:  "autoscaler",
+			Port:         "80",
+			Protocol:     "http",
+			ReplicaCount: 3,
 		},
 	}, {
 		name:     "success with overriding",
 		pod:      "as-42",
 		service:  "autoscaler",
 		port:     "8080",
 		protocol: "ws",
+		replicas: "5",
 		expected: statefulSetConfig{
 			StatefulSetID: statefulSetID{
 				ssName:  "as",
 				ordinal: 42,
 			},
-			ServiceName: "autoscaler",
-			Port:        "8080",
-			Protocol:    "ws",
+			ServiceName:  "autoscaler",
+			Port:         "8080",
+			Protocol:     "ws",
+			ReplicaCount: 5,
 		},
 	}, {
-		name:    "failure with empty envs",
+		name:    "StatefulSet environment completely absent",
 		wantErr: "required key STATEFUL_CONTROLLER_ORDINAL missing value",
 	}, {
-		name:    "failure with invalid name",
+		name:    "malformed ordinal",
 		pod:     "as-abcd",
 		wantErr: `envconfig.Process: assigning STATEFUL_CONTROLLER_ORDINAL to StatefulSetID: converting 'as-abcd' to type leaderelection.statefulSetID. details: strconv.Atoi: parsing "abcd": invalid syntax`,
+	}, {
+		name:     "missing service when ordinal is present",
+		pod:      "as-0",
+		replicas: "1",
+		wantErr:  "required key STATEFUL_SERVICE_NAME missing value",
+	}, {
+		name:    "missing replica when ordinal is present",
+		pod:     "as-0",
+		service: "autoscaler",
+		wantErr: "required key STATEFUL_REPLICA_COUNT missing value",
+	}, {
+		name:     "malformed replica",
+		pod:      "as-0",
+		service:  "autoscaler",
+		replicas: "abc",
+		wantErr:  `envconfig.Process: assigning STATEFUL_REPLICA_COUNT to ReplicaCount: converting 'abc' to type int. details: strconv.ParseInt: parsing "abc": invalid syntax`,
+	}, {
+		name:     "zero replica count",
+		pod:      "as-0",
+		service:  "autoscaler",
+		replicas: "0",
+		wantErr:  "STATEFUL_REPLICA_COUNT must be >= 1, got 0",
+	}, {
+		name:     "negative replica count",
+		pod:      "as-0",
+		service:  "autoscaler",
+		replicas: "-1",
+		wantErr:  "STATEFUL_REPLICA_COUNT must be >= 1, got -1",
 	}}
 
 	for _, tc := range cases {
@@ -308,17 +342,35 @@ func TestNewStatefulSetConfig(t *testing.T) {
 			if tc.protocol != "" {
 				t.Setenv(serviceProtocolEnv, tc.protocol)
 			}
+			if tc.replicas != "" {
+				t.Setenv(replicaCountEnv, tc.replicas)
+			}
 
 			ssc, err := newStatefulSetConfig()
 			if err != nil {
 				if got, want := err.Error(), tc.wantErr; got != want {
 					t.Errorf("Got error: %s. want: %s", got, want)
 				}
 			} else {
+				if tc.wantErr != "" {
+					t.Error("newStatefulSetConfig() = nil, want error")
+				}
 				if got, want := *ssc, tc.expected; !cmp.Equal(got, want, cmp.AllowUnexported(statefulSetID{})) {
 					t.Errorf("Incorrect config: diff(-want,+got):\n%s", cmp.Diff(want, got))
 				}
 			}
 		})
 	}
 }
+
+func TestStatefulSetConfiguredEmptyOrdinal(t *testing.T) {
+	t.Setenv(controllerOrdinalEnv, "")
+
+	if !statefulSetConfigured() {
+		t.Fatal("statefulSetConfigured() = false, want true for empty STATEFUL_CONTROLLER_ORDINAL")
+	}
+
+	if _, err := newStatefulSetConfig(); err == nil {
+		t.Fatal("newStatefulSetConfig() = nil, want error for empty STATEFUL_CONTROLLER_ORDINAL")
+	}
+}

leaderelection/context.go

@@ -35,17 +35,25 @@ import (
 	"knative.dev/pkg/system"
 )
 
-// WithDynamicLeaderElectorBuilder sets up the statefulset elector based on environment,
-// falling back on the standard elector.
+// WithDynamicLeaderElectorBuilder selects StatefulSet ordinal leader election
+// when STATEFUL_CONTROLLER_ORDINAL is set, otherwise standard lease election.
+//
+// StatefulSet ordinal mode is 1:1: ordinal N owns bucket N. cc.Buckets is the
+// hash universe and must equal STATEFUL_REPLICA_COUNT. If StatefulSet mode is
+// explicitly configured but invalid, this helper fails rather than falling
+// back to standard election.
 func WithDynamicLeaderElectorBuilder(ctx context.Context, kc kubernetes.Interface, cc ComponentConfig) context.Context {
 	logger := logging.FromContext(ctx)
+	if !statefulSetConfigured() {
+		logger.Info("Running with Standard leader election")
+		return WithStandardLeaderElectorBuilder(ctx, kc, cc)
+	}
 	b, _, err := NewStatefulSetBucketAndSet(int(cc.Buckets))
-	if err == nil {
-		logger.Info("Running with StatefulSet leader election")
-		return WithStatefulSetElectorBuilder(ctx, cc, b)
+	if err != nil {
+		logger.Fatal("Invalid StatefulSet leader election configuration: ", err)
 	}
-	logger.Info("Running with Standard leader election")
-	return WithStandardLeaderElectorBuilder(ctx, kc, cc)
+	logger.Info("Running with StatefulSet leader election")
+	return WithStatefulSetElectorBuilder(ctx, cc, b)
 }
 
 // WithStandardLeaderElectorBuilder infuses a context with the ability to build
@@ -221,15 +229,28 @@ func (b *statefulSetBuilder) buildElector(ctx context.Context, la reconciler.Lea
 	}, nil
 }
 
-// NewStatefulSetBucketAndSet creates a BucketSet for StatefulSet controller with
-// the given bucket size and the information from environment variables. Then uses
-// the created BucketSet to create a Bucket for this StatefulSet Pod.
+// NewStatefulSetBucketAndSet creates a BucketSet for a StatefulSet controller
+// with the given bucket count and the information from environment variables.
+// The universe has one bucket per configured count, named after StatefulSet
+// pod DNS. The returned Bucket is the one for this process's ordinal
+// (STATEFUL_CONTROLLER_ORDINAL).
+//
+// Assignment is 1:1: ordinal N owns bucket N. STATEFUL_REPLICA_COUNT must
+// equal buckets; extra buckets are not redistributed among fewer pods.
+// An error is returned if the StatefulSet environment is invalid, replica
+// count does not equal buckets, or the local ordinal is out of range
+// [0, buckets).
 func NewStatefulSetBucketAndSet(buckets int) (reconciler.Bucket, *hash.BucketSet, error) {
 	ssc, err := newStatefulSetConfig()
 	if err != nil {
 		return nil, nil, err
 	}
 
+	if ssc.ReplicaCount != buckets {
+		return nil, nil, fmt.Errorf("STATEFUL_REPLICA_COUNT (%d) must equal buckets (%d) for StatefulSet leader election",
+			ssc.ReplicaCount, buckets)
+	}
+
 	if ssc.StatefulSetID.ordinal >= buckets {
 		return nil, nil, fmt.Errorf("ordinal %d is out of range [0, %d)",
 			ssc.StatefulSetID.ordinal, buckets)

leaderelection/context_test.go

@@ -217,6 +217,14 @@ func TestBuilderWithCustomizedLeaseName(t *testing.T) {
 	}
 }
 
+func TestWithDynamicLeaderElectorBuilderWithoutOrdinal(t *testing.T) {
+	cc := ComponentConfig{Component: "the-component", Buckets: 3}
+	ctx := WithDynamicLeaderElectorBuilder(context.Background(), fakekube.NewSimpleClientset(), cc)
+	if _, ok := ctx.Value(builderKey{}).(*standardBuilder); !ok {
+		t.Fatalf("builder = %T, want *standardBuilder", ctx.Value(builderKey{}))
+	}
+}
+
 func TestNewStatefulSetBucketAndSet(t *testing.T) {
 	wantNames := []string{
 		"http://as-0.autoscaler.knative-testing.svc.cluster.local:80",
@@ -226,16 +234,19 @@ func TestNewStatefulSetBucketAndSet(t *testing.T) {
 
 	t.Setenv(controllerOrdinalEnv, "as-2")
 	t.Setenv(serviceNameEnv, "autoscaler")
+	t.Setenv(replicaCountEnv, "2")
 
 	_, _, err := NewStatefulSetBucketAndSet(2)
 	if err == nil {
-		// Ordinal 2 should be range [0, 2)
 		t.Fatal("Expected error from NewStatefulSetBucketAndSet but got nil")
 	}
+	if got, want := err.Error(), "ordinal 2 is out of range [0, 2)"; got != want {
+		t.Errorf("NewStatefulSetBucketAndSet() error = %q, want %q", got, want)
+	}
 
+	t.Setenv(replicaCountEnv, "3")
 	bkt, bs, err := NewStatefulSetBucketAndSet(3)
 	if err != nil {
-		// Ordinal 2 should be range [0, 2)
 		t.Fatal("NewStatefulSetBucketAndSet() = ", err)
 	}
 
@@ -249,6 +260,81 @@ func TestNewStatefulSetBucketAndSet(t *testing.T) {
 	}
 }
 
+func TestNewStatefulSetBucketAndSetReplicaMismatch(t *testing.T) {
+	t.Setenv(serviceNameEnv, "autoscaler")
+
+	t.Run("replicas less than buckets", func(t *testing.T) {
+		t.Setenv(controllerOrdinalEnv, "as-0")
+		t.Setenv(replicaCountEnv, "4")
+		_, _, err := NewStatefulSetBucketAndSet(8)
+		if err == nil {
+			t.Fatal("NewStatefulSetBucketAndSet() = nil, want replica/bucket mismatch error")
+		}
+		want := "STATEFUL_REPLICA_COUNT (4) must equal buckets (8) for StatefulSet leader election"
+		if got := err.Error(); got != want {
+			t.Errorf("error = %q, want %q", got, want)
+		}
+	})
+
+	t.Run("replicas greater than buckets", func(t *testing.T) {
+		t.Setenv(controllerOrdinalEnv, "as-0")
+		t.Setenv(replicaCountEnv, "8")
+		_, _, err := NewStatefulSetBucketAndSet(4)
+		if err == nil {
+			t.Fatal("NewStatefulSetBucketAndSet() = nil, want replica/bucket mismatch error")
+		}
+		want := "STATEFUL_REPLICA_COUNT (8) must equal buckets (4) for StatefulSet leader election"
+		if got := err.Error(); got != want {
+			t.Errorf("error = %q, want %q", got, want)
+		}
+	})
+}
+
+func TestNewStatefulSetBucketAndSetOneToOne(t *testing.T) {
+	const (
+		buckets = 8
+		ordinal = 3
+	)
+	wantNames := make([]string, buckets)
+	for i := range buckets {
+		wantNames[i] = fmt.Sprintf("http://as-%d.autoscaler.knative-testing.svc.cluster.local:80", i)
+	}
+
+	t.Setenv(controllerOrdinalEnv, "as-3")
+	t.Setenv(serviceNameEnv, "autoscaler")
+	t.Setenv(replicaCountEnv, "8")
+
+	bkt, bs, err := NewStatefulSetBucketAndSet(buckets)
+	if err != nil {
+		t.Fatal("NewStatefulSetBucketAndSet() = ", err)
+	}
+
+	if got, want := bkt.Name(), wantNames[ordinal]; got != want {
+		t.Errorf("Bucket.Name() = %s, want = %s", got, want)
+	}
+
+	gotNames := bs.BucketList()
+	if !cmp.Equal(gotNames, wantNames) {
+		t.Errorf("BucketSet.BucketList() = %q, want: %q", gotNames, wantNames)
+	}
+
+	all := bs.Buckets()
+	if len(all) != buckets {
+		t.Fatalf("len(Buckets()) = %d, want %d", len(all), buckets)
+	}
+	if got, want := all[ordinal].Name(), bkt.Name(); got != want {
+		t.Errorf("Buckets()[%d].Name() = %s, want %s", ordinal, got, want)
+	}
+	for i, other := range all {
+		if i == ordinal {
+			continue
+		}
+		if other.Name() == bkt.Name() {
+			t.Errorf("Buckets()[%d].Name() = %s, same as ordinal %d; assignment is 1:1", i, other.Name(), ordinal)
+		}
+	}
+}
+
 func TestWithStatefulSetBuilder(t *testing.T) {
 	cc := ComponentConfig{
 		Component: "the-component",
@@ -268,6 +354,7 @@ func TestWithStatefulSetBuilder(t *testing.T) {
 
 	t.Setenv(controllerOrdinalEnv, "as-2")
 	t.Setenv(serviceNameEnv, "autoscaler")
+	t.Setenv(replicaCountEnv, "3")
 
 	ctx = WithDynamicLeaderElectorBuilder(ctx, nil, cc)
 	if !HasLeaderElection(ctx) {

leaderelection/doc.go

@@ -22,6 +22,12 @@ limitations under the License.
 // management of multiple election strategies (currently, using Kubernetes
 // etcd-based election primitives or StatefulSet indexes and counts).
 //
+// StatefulSet ordinal mode is 1:1: ordinal N owns bucket N, whose identity is
+// the DNS name of pod N. It is selected when STATEFUL_CONTROLLER_ORDINAL is
+// set. STATEFUL_REPLICA_COUNT must equal the configured bucket count; extra
+// buckets are not redistributed. An explicitly configured but invalid
+// StatefulSet topology fails rather than falling back to standard election.
+//
 // For more details, see the original design document:
 // https://docs.google.com/document/d/e/2PACX-1vTh40N-Kk6EPNzYpITiLg8YJk0qZyZv7KgMpcQS72T9Lv_F2PQeGybx4TtH0E1N1aUgLQer7b8u3lDc/pub
 package leaderelection