Repository: knative-sandbox/eventing-kafka
Pull request #427 (2021-03-03 01:53:43)

[master] Upgrade to latest dependencies
Cron. -knative-prow-robot

/cc knative-sandbox/channel-wg-leads
/assign knative-sandbox/channel-wg-leads

Produced via:
```shell
./hack/update-deps.sh --upgrade --release v0.22
./hack/update-codegen.sh
```

Details:
```
bumping knative.dev/pkg 9039dc1...2bc944b:
  > 2bc944b Fix race: Make informed watcher start wait for Add event 🏎️ (# 2036)
bumping knative.dev/eventing f220668...0f2612e:
  > 0f2612e add test for trigger having a wrong UID (# 5004)
  > 4282f45 fix broken in-mem ns mode (# 4906)
  > 9089b8e adding code coverage to adapters v2 (# 5000)
  > eb12f8a update to breaking change for feature.T in reconciler-test (# 5003)
  > c5f92fa upgrade to latest dependencies (# 5001)
```

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

This pull-request has been approved by:
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **slinkydeveloper** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @slinkydeveloper` in a comment when ready.

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

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

- **[OWNERS](https://github.com/knative-sandbox/eventing-kafka/blob/master/OWNERS)**

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":["slinkydeveloper"]} -->
codecov[bot]: # [Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=h1) Report
> Merging [#427](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=desc) (5c0afc6) into [master](https://codecov.io/gh/knative-sandbox/eventing-kafka/commit/51bc90b19d0a5f30fe7445a96c1583193d61d1b5?el=desc) (51bc90b) will **increase** coverage by `0.34%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427/graphs/tree.svg?width=650&height=150&src=pr&token=W765CT93QC)](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #427      +/-   ##
==========================================
+ Coverage   72.73%   73.07%   +0.34%     
==========================================
  Files         129      129              
  Lines        4995     4995              
==========================================
+ Hits         3633     3650      +17     
+ Misses       1127     1109      -18     
- Partials      235      236       +1     
```


| [Impacted Files](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [...el/distributed/dispatcher/dispatcher/dispatcher.go](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427/diff?src=pr&el=tree#diff-cGtnL2NoYW5uZWwvZGlzdHJpYnV0ZWQvZGlzcGF0Y2hlci9kaXNwYXRjaGVyL2Rpc3BhdGNoZXIuZ28=) | `76.80% <0.00%> (+5.59%)` | :arrow_up: |
| [pkg/common/scheduler/statefulset/autoscaler.go](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427/diff?src=pr&el=tree#diff-cGtnL2NvbW1vbi9zY2hlZHVsZXIvc3RhdGVmdWxzZXQvYXV0b3NjYWxlci5nbw==) | `63.82% <0.00%> (+21.27%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=footer). Last update [51bc90b...5c0afc6](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/427?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/go.mod b/go.mod
index 5c53514d4..1bf8bfa32 100644
--- a/go.mod
+++ b/go.mod
@@ -33,7 +33,7 @@ require (
 	k8s.io/apimachinery v0.19.7
 	k8s.io/client-go v0.19.7
 	k8s.io/utils v0.0.0-20200729134348-d5654de09c73
-	knative.dev/eventing v0.21.1-0.20210302004649-f22066824960
+	knative.dev/eventing v0.21.1-0.20210302231422-0f2612ef8aba
 	knative.dev/hack v0.0.0-20210203173706-8368e1f6eacf
-	knative.dev/pkg v0.0.0-20210226182947-9039dc189ced
+	knative.dev/pkg v0.0.0-20210302195815-2bc944bb68d0
 )
diff --git a/go.sum b/go.sum
index f9f7aee0a..d192c33d8 100644
--- a/go.sum
+++ b/go.sum
@@ -1244,15 +1244,15 @@ k8s.io/kube-openapi v0.0.0-20200805222855-6aeccd4b50c6 h1:+WnxoVtG8TMiudHBSEtrVL
 k8s.io/kube-openapi v0.0.0-20200805222855-6aeccd4b50c6/go.mod h1:UuqjUnNftUyPE5H64/qeyjQoUZhGpeFDVdxjTeEVN2o=
 k8s.io/utils v0.0.0-20200729134348-d5654de09c73 h1:uJmqzgNWG7XyClnU/mLPBWwfKKF1K8Hf8whTseBgJcg=
 k8s.io/utils v0.0.0-20200729134348-d5654de09c73/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
-knative.dev/eventing v0.21.1-0.20210302004649-f22066824960 h1:E5lX5F6PNv2GwBcrJyrq+uHls2SnKLX4TLH1BeLq6j4=
-knative.dev/eventing v0.21.1-0.20210302004649-f22066824960/go.mod h1:KtN2FEhPzxyqXKxpJufdq0lzL3Iu+5qjFqBal8Lg2tQ=
+knative.dev/eventing v0.21.1-0.20210302231422-0f2612ef8aba h1:eelYU6hdCV/J7VxKqLXhgqB5ASNpmNx770NARDJiGn4=
+knative.dev/eventing v0.21.1-0.20210302231422-0f2612ef8aba/go.mod h1:xx42mnYqB+iEMYny6/M3vbe58BZAo+uLrWbOQF5Qrm0=
 knative.dev/hack v0.0.0-20210203173706-8368e1f6eacf h1:u4cY4jr2LYvhoz/1HBWEPsMiLkm0HMdDTfmmw1RE8zE=
 knative.dev/hack v0.0.0-20210203173706-8368e1f6eacf/go.mod h1:PHt8x8yX5Z9pPquBEfIj0X66f8iWkWfR0S/sarACJrI=
-knative.dev/pkg v0.0.0-20210215165523-84c98f3c3e7a/go.mod h1:TJSdebQOWX5N2bszohOYVi0H1QtXbtlYLuMghAFBMhY=
-knative.dev/pkg v0.0.0-20210225221047-07b5ddfaea60/go.mod h1:TJSdebQOWX5N2bszohOYVi0H1QtXbtlYLuMghAFBMhY=
 knative.dev/pkg v0.0.0-20210226182947-9039dc189ced h1:OCtTbXsJhiR35qCDxFj9Xc8sU3i87ePbc7YlHCt0al4=
 knative.dev/pkg v0.0.0-20210226182947-9039dc189ced/go.mod h1:TJSdebQOWX5N2bszohOYVi0H1QtXbtlYLuMghAFBMhY=
-knative.dev/reconciler-test v0.0.0-20210223181346-209f32c8cdc4/go.mod h1:RP/K5xJylB72Go6eAsXYEsQHp4zCCNMNjmsqhvq7wko=
+knative.dev/pkg v0.0.0-20210302195815-2bc944bb68d0 h1:1GaHRzKydEHOpuTGI28wTOLA2upINB4LX2Zo6+K0j+g=
+knative.dev/pkg v0.0.0-20210302195815-2bc944bb68d0/go.mod h1:TJSdebQOWX5N2bszohOYVi0H1QtXbtlYLuMghAFBMhY=
+knative.dev/reconciler-test v0.0.0-20210302173618-5f01af2a17ac/go.mod h1:70MF1WzRT29YtTBujNmB58UUlRGLlfYElOTNmO7MxdA=
 pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U=
 pgregory.net/rapid v0.3.3/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
diff --git a/pkg/common/scheduler/scheduler.go b/pkg/common/scheduler/scheduler.go
index eb47174d9..3bd40cf2a 100644
--- a/pkg/common/scheduler/scheduler.go
+++ b/pkg/common/scheduler/scheduler.go
@@ -25,8 +25,7 @@ import (
 )
 
 var (
-	ErrNoReplicas        = errors.New("scheduling failed (no pod replicas available)")
-	ErrNotEnoughReplicas = errors.New("scheduling partially failed (not enough pod replicas)")
+	ErrNotEnoughReplicas = errors.New("scheduling failed (not enough pod replicas)")
 )
 
 // VPodLister is the function signature for returning a list of VPods
diff --git a/pkg/common/scheduler/statefulset/autoscaler.go b/pkg/common/scheduler/statefulset/autoscaler.go
index 47332745f..67ddc2304 100644
--- a/pkg/common/scheduler/statefulset/autoscaler.go
+++ b/pkg/common/scheduler/statefulset/autoscaler.go
@@ -23,108 +23,132 @@ import (
 
 	"go.uber.org/zap"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
+	clientappsv1 "k8s.io/client-go/kubernetes/typed/apps/v1"
+
+	kubeclient "knative.dev/pkg/client/injection/kube/client"
+	"knative.dev/pkg/logging"
 )
 
-func (s *StatefulSetScheduler) autoscale(ctx context.Context) {
-	// Regularly compute the average number of replicas per pods.
-	// When the average goes above a certain ratio, scale up
-	// otherwise scale down.
+type Autoscaler interface {
+	// Start runs the autoscaler until cancelled.
+	Start(ctx context.Context)
 
-	for {
-		func() {
-			s.lock.Lock()
-			defer s.lock.Unlock()
-
-			// Refresh state as the scheduler does not receive deletion notifications.
-			snapshot, err := s.snapshot(s.logger)
-			if err != nil {
-				s.logger.Info("error while refreshing scheduler state (will retry)", zap.Error(err))
-				return
-			}
-
-			s.logger.Infow("checking adapter capacity",
-				zap.Int32("free", s.freeCapacity(snapshot)),
-				zap.Int32("used", s.usedCapacity(snapshot)),
-				zap.Int32("pending", s.pendingVReplicas()),
-				zap.Int32("replicas", s.replicas),
-				zap.Int32("last ordinal", snapshot.lastOrdinal))
-
-			var ratio float64
-			if s.replicas == 0 {
-				if s.pendingVReplicas() > 0 {
-					ratio = 1.0
-				} else {
-					ratio = 0.5
-				}
-			} else {
-				ratio = s.avgUsedCapacityPerPod(snapshot) / float64(s.capacity)
-			}
-
-			if ratio > 0.7 || ratio < 0.3 {
-				// Scale up when capacity is above 80% (TODO: configurable)
-				// Scale down when capacity is below 30% (TODO: configurable)
-				s.logger.Infow("autoscaling statefulset", zap.Int("avg used capacity per pod", int(ratio*100.0)))
-
-				scale, err := s.statefulSetClient.GetScale(ctx, s.statefulSetName, metav1.GetOptions{})
-				if err != nil {
-					// skip a beat
-					s.logger.Infow("failed to get scale subresource", zap.Error(err))
-					return
-				}
-
-				// Desired ratio is 0.5 (TODO: configurable)
-				new := int32(math.Ceil(float64(s.usedCapacity(snapshot)+s.pendingVReplicas()) / (float64(s.capacity) * 0.5)))
-
-				// Make sure not to scale down past the last pod with placed vpods
-				if new < snapshot.lastOrdinal+1 {
-					new = snapshot.lastOrdinal + 1
-				}
-
-				if new != scale.Spec.Replicas {
-					scale.Spec.Replicas = new
-					s.logger.Infow("updating adapter replicas", zap.Int32("replicas", scale.Spec.Replicas))
-
-					_, err = s.statefulSetClient.UpdateScale(ctx, s.statefulSetName, scale, metav1.UpdateOptions{})
-					if err != nil {
-						s.logger.Errorw("updating scale subresource failed", zap.Error(err))
-					}
-				}
-			}
-		}()
+	// Autoscale is used to immediately trigger the autoscaler with the hint
+	// that pending number of vreplicas coudn't be scheduled.
+	Autoscale(pending int32)
+}
+
+type autoscaler struct {
+	statefulSetClient clientappsv1.StatefulSetInterface
+	statefulSetName   string
+	logger            *zap.SugaredLogger
+	stateAccessor     stateAccessor
+	trigger           chan int32
+
+	// capacity is the total number of virtual replicas available per pod.
+	capacity int32
+
+	// refreshPeriod is how often to try attempting scaling down the statefulset
+	refreshPeriod time.Duration
+}
 
+func NewAutoscaler(ctx context.Context, namespace, name string, stateAccessor stateAccessor, refreshPeriod time.Duration) Autoscaler {
+	return &autoscaler{
+		logger:            logging.FromContext(ctx),
+		statefulSetClient: kubeclient.Get(ctx).AppsV1().StatefulSets(namespace),
+		statefulSetName:   name,
+		stateAccessor:     stateAccessor,
+		trigger:           make(chan int32, 1),
+		capacity:          int32(10),
+		refreshPeriod:     refreshPeriod,
+	}
+}
+
+func (a *autoscaler) Start(ctx context.Context) {
+	attemptScaleDown := false
+	pending := int32(0)
+	for {
 		select {
 		case <-ctx.Done():
 			return
-		case <-time.After(10 * time.Second): // TODO: configurable refresh period
+		case <-time.After(a.refreshPeriod):
+			attemptScaleDown = true
+		case pending = <-a.trigger:
+			attemptScaleDown = false
 		}
-	}
-}
 
-// freeCapacity returns the free capacity across all pods
-func (s *StatefulSetScheduler) freeCapacity(snapshot *Snapshot) int32 {
-	t := int32(0)
-	for i := int32(0); i < s.replicas; i++ {
-		t += snapshot.free[podNameFromOrdinal(s.statefulSetName, i)]
+		// Retry a few times, just so that we don't have to wait for the next beat when
+		// a transient error occurs
+		wait.Poll(500*time.Millisecond, 5*time.Second, func() (bool, error) {
+			err := a.doautoscale(ctx, attemptScaleDown, pending)
+			return err == nil, nil
+		})
+
+		pending = int32(0)
 	}
-	return t
 }
 
-// usedCapacity returns the used capacity across all pods
-func (s *StatefulSetScheduler) usedCapacity(snapshot *Snapshot) int32 {
-	return s.capacity*s.replicas - s.freeCapacity(snapshot)
+func (a *autoscaler) Autoscale(pending int32) {
+	a.trigger <- pending
 }
 
-// pendingReplicas returns the total number of virtual replicas
-// that haven't been scheduled yet
-func (s *StatefulSetScheduler) pendingVReplicas() int32 {
-	t := int32(0)
-	for _, v := range s.pending {
-		t += v
+func (a *autoscaler) doautoscale(ctx context.Context, attemptScaleDown bool, pending int32) error {
+	state, err := a.stateAccessor.State()
+	if err != nil {
+		a.logger.Info("error while refreshing scheduler state (will retry)", zap.Error(err))
+		return err
+	}
+
+	scale, err := a.statefulSetClient.GetScale(ctx, a.statefulSetName, metav1.GetOptions{})
+	if err != nil {
+		// skip a beat
+		a.logger.Infow("failed to get scale subresource", zap.Error(err))
+		return err
+	}
+
+	a.logger.Infow("checking adapter capacity",
+		zap.Int32("pending", pending),
+		zap.Int32("replicas", scale.Spec.Replicas),
+		zap.Int32("last ordinal", state.lastOrdinal))
+
+	newreplicas := state.lastOrdinal + 1 // Ideal number
+
+	// Take into account pending replicas
+	if pending > 0 {
+		// The number of replicas may be lower than the last ordinal, for instance
+		// when the statefulset is manually scaled down. In that case, replicas above
+		// scale.Spec.Replicas have not been considered when scheduling vreplicas.
+		// Adjust accordindly
+		pending -= state.freeCapacity()
+
+		// Still need more?
+		if pending > 0 {
+			// Make sure to allocate enough pods for holding all pending replicas.
+			newreplicas += int32(math.Ceil(float64(pending) / float64(a.capacity)))
+		}
+	}
+
+	// Make sure to never scale down past the last ordinal
+	if newreplicas <= state.lastOrdinal {
+		newreplicas = state.lastOrdinal + 1
+	}
+
+	// Only scale down if permitted
+	if !attemptScaleDown && newreplicas < scale.Spec.Replicas {
+		newreplicas = scale.Spec.Replicas
+	}
+
+	if newreplicas != scale.Spec.Replicas {
+		scale.Spec.Replicas = newreplicas
+		a.logger.Infow("updating adapter replicas", zap.Int32("replicas", scale.Spec.Replicas))
+
+		_, err = a.statefulSetClient.UpdateScale(ctx, a.statefulSetName, scale, metav1.UpdateOptions{})
+		if err != nil {
+			a.logger.Errorw("updating scale subresource failed", zap.Error(err))
+			return err
+		}
 	}
-	return t
-}
 
-// avgUsedCapacityPerPod returns the average used capacity per replica
-func (s *StatefulSetScheduler) avgUsedCapacityPerPod(snapshot *Snapshot) float64 {
-	return float64(s.usedCapacity(snapshot)) / float64(s.replicas)
+	return nil
 }
diff --git a/pkg/common/scheduler/statefulset/scheduler.go b/pkg/common/scheduler/statefulset/scheduler.go
index 8796bad45..a05a22757 100644
--- a/pkg/common/scheduler/statefulset/scheduler.go
+++ b/pkg/common/scheduler/statefulset/scheduler.go
@@ -19,14 +19,15 @@ package statefulset
 import (
 	"context"
 	"sync"
+	"time"
 
 	"go.uber.org/zap"
 	appsv1 "k8s.io/api/apps/v1"
+	"k8s.io/apimachinery/pkg/types"
 	clientappsv1 "k8s.io/client-go/kubernetes/typed/apps/v1"
 	"k8s.io/client-go/tools/cache"
 	"k8s.io/utils/integer"
 
-	"k8s.io/apimachinery/pkg/types"
 	kubeclient "knative.dev/pkg/client/injection/kube/client"
 	statefulsetinformer "knative.dev/pkg/client/injection/kube/informers/apps/v1/statefulset"
 	"knative.dev/pkg/controller"
@@ -36,38 +37,46 @@ import (
 	"knative.dev/eventing-kafka/pkg/common/scheduler"
 )
 
-// StatefulSetScheduler is a scheduler placing VPod into statefulset-managed pods
+// NewScheduler creates a new scheduler with pod autoscaling enabled.
+func NewScheduler(ctx context.Context, namespace, name string, lister scheduler.VPodLister) scheduler.Scheduler {
+	capacity := int32(10) // TODO: config
+
+	stateAccessor := newStateBuilder(logging.FromContext(ctx), lister, capacity)
+	autoscaler := NewAutoscaler(ctx, namespace, name, stateAccessor, 10*time.Second) // TODO: config
+
+	go autoscaler.Start(ctx)
+
+	return NewStatefulSetScheduler(ctx, namespace, name, lister, stateAccessor, autoscaler)
+}
+
+// StatefulSetScheduler is a scheduler placing VPod into statefulset-managed set of pods
 type StatefulSetScheduler struct {
 	logger            *zap.SugaredLogger
 	statefulSetName   string
 	statefulSetClient clientappsv1.StatefulSetInterface
 	vpodLister        scheduler.VPodLister
 	lock              sync.Locker
+	stateAccessor     stateAccessor
+	autoscaler        Autoscaler
 
-	// replicas is the (cached) number of statefulset replicas
+	// replicas is the (cached) number of statefulset replicas.
 	replicas int32
 
-	// capacity is the total number of virtual replicas available per pod.
-	capacity int32
-
 	// pending tracks the number of virtual replicas that haven't been scheduled yet.
 	// Used by the autoscaler
 	pending map[types.NamespacedName]int32
-
-	// enq is used to trigger the reconciliation of the named object
-	enq func(types.NamespacedName)
 }
 
-func NewStatefulSetScheduler(ctx context.Context, namespace, name string, lister scheduler.VPodLister, enq func(types.NamespacedName)) scheduler.Scheduler {
+func NewStatefulSetScheduler(ctx context.Context, namespace, name string, lister scheduler.VPodLister, stateAccessor stateAccessor, autoscaler Autoscaler) scheduler.Scheduler {
 	scheduler := &StatefulSetScheduler{
 		logger:            logging.FromContext(ctx),
 		statefulSetName:   name,
 		statefulSetClient: kubeclient.Get(ctx).AppsV1().StatefulSets(namespace),
 		vpodLister:        lister,
-		capacity:          10, // TODO: from annotation on statefulset, configmap?
-		enq:               enq,
 		pending:           make(map[types.NamespacedName]int32),
 		lock:              new(sync.Mutex),
+		stateAccessor:     stateAccessor,
+		autoscaler:        autoscaler,
 	}
 
 	// Monitor our statefulset
@@ -77,10 +86,6 @@ func NewStatefulSetScheduler(ctx context.Context, namespace, name string, lister
 		Handler:    controller.HandleAll(scheduler.updateStatefulset),
 	})
 
-	// Start the autoscaler. Eventually we may want to rely on HPA + custom metrics
-	// in particular to handle fluctuation
-	go scheduler.autoscale(ctx)
-
 	return scheduler
 }
 
@@ -91,25 +96,18 @@ func (s *StatefulSetScheduler) Schedule(vpod scheduler.VPod) ([]duckv1alpha1.Pla
 	logger := s.logger.With("key", zap.Any("key", vpod.GetKey()))
 	logger.Info("scheduling")
 
-	// Get the current placement state
+	// Get the current placements state
 	// Quite an expensive operation but safe and simple.
-	snapshot, err := s.snapshot(logger)
+	state, err := s.stateAccessor.State()
 	if err != nil {
 		logger.Info("error while refreshing scheduler state (will retry)", zap.Error(err))
 		return nil, err
 	}
 
-	if s.replicas == 0 {
-		logger.Info("scheduling failed (no replicas available)")
-		s.pending[vpod.GetKey()] = vpod.GetVReplicas()
-		return make([]duckv1alpha1.Placement, 0), scheduler.ErrNoReplicas
-	}
-
-	// Remove placements from scaled-down pods
-	placements := s.reschedule(vpod.GetPlacements())
+	placements := vpod.GetPlacements()
 
 	// The scheduler
-	// - allocates as many vreplicas as possible to the same pod(s) (for placed vpod)
+	// - allocates as many vreplicas as possible to the same pod(s) (TODO: HA)
 	// - allocates remaining vreplicas to new pods
 
 	// Exact number of vreplicas => do nothing
@@ -126,17 +124,26 @@ func (s *StatefulSetScheduler) Schedule(vpod scheduler.VPod) ([]duckv1alpha1.Pla
 	if tr > vpod.GetVReplicas() {
 		logger.Infow("scaling down", zap.Int32("vreplicas", tr), zap.Int32("new vreplicas", vpod.GetVReplicas()))
 		placements := s.removeReplicas(tr-vpod.GetVReplicas(), placements)
+
+		// Do not trigger the autoscaler to avoid unnecessary churn
+
 		return placements, nil
 	}
 
 	// Need more => scale up
 	logger.Infow("scaling up", zap.Int32("vreplicas", tr), zap.Int32("new vreplicas", vpod.GetVReplicas()))
-	placements, left := s.addReplicas(snapshot, vpod.GetVReplicas()-tr, placements)
+	placements, left := s.addReplicas(state, vpod.GetVReplicas()-tr, placements)
 	if left > 0 {
 		// Give time for the autoscaler to do its job
 		logger.Info("scheduling failed (not enough pod replicas)", zap.Any("placement", placements))
 
 		s.pending[vpod.GetKey()] = left
+
+		// Trigger the autoscaler
+		if s.autoscaler != nil {
+			s.autoscaler.Autoscale(s.pendingVReplicas())
+		}
+
 		return placements, scheduler.ErrNotEnoughReplicas
 	}
 
@@ -145,28 +152,6 @@ func (s *StatefulSetScheduler) Schedule(vpod scheduler.VPod) ([]duckv1alpha1.Pla
 	return placements, nil
 }
 
-func (s *StatefulSetScheduler) reschedule(placements []duckv1alpha1.Placement) []duckv1alpha1.Placement {
-	if s.needRescheduling(placements) {
-		newPlacements := make([]duckv1alpha1.Placement, 0, len(placements))
-		for i := 0; i < len(placements); i++ {
-			if ordinalFromPodName(placements[i].PodName) >= s.replicas {
-				newPlacements = append(newPlacements, placements[i])
-			}
-		}
-		return newPlacements
-	}
-	return placements
-}
-
-func (s *StatefulSetScheduler) needRescheduling(placements []duckv1alpha1.Placement) bool {
-	for i := 0; i < len(placements); i++ {
-		if ordinalFromPodName(placements[i].PodName) >= s.replicas {
-			return true
-		}
-	}
-	return false
-}
-
 func (s *StatefulSetScheduler) removeReplicas(diff int32, placements []duckv1alpha1.Placement) []duckv1alpha1.Placement {
 	newPlacements := make([]duckv1alpha1.Placement, 0, len(placements))
 	for i := 0; i < len(placements); i++ {
@@ -184,7 +169,7 @@ func (s *StatefulSetScheduler) removeReplicas(diff int32, placements []duckv1alp
 	return newPlacements
 }
 
-func (s *StatefulSetScheduler) addReplicas(snapshot *Snapshot, diff int32, placements []duckv1alpha1.Placement) ([]duckv1alpha1.Placement, int32) {
+func (s *StatefulSetScheduler) addReplicas(state *state, diff int32, placements []duckv1alpha1.Placement) ([]duckv1alpha1.Placement, int32) {
 	// Pod affinity algorithm: prefer adding replicas to existing pods before considering other replicas
 	// In the future, we might want to spread replicas across pods in different regions.
 	newPlacements := make([]duckv1alpha1.Placement, 0, len(placements))
@@ -192,17 +177,19 @@ func (s *StatefulSetScheduler) addReplicas(snapshot *Snapshot, diff int32, place
 	// Add to existing
 	for i := 0; i < len(placements); i++ {
 		podName := placements[i].PodName
+		ordinal := ordinalFromPodName(podName)
 
 		// Is there space in PodName?
-		if diff >= 0 && snapshot.free[podName] > 0 {
-			allocation := integer.Int32Min(snapshot.free[podName], diff)
+		f := state.Free(ordinal)
+		if diff >= 0 && f > 0 {
+			allocation := integer.Int32Min(f, diff)
 			newPlacements = append(newPlacements, duckv1alpha1.Placement{
 				PodName:   podName,
 				VReplicas: placements[i].VReplicas + allocation,
 			})
 
 			diff -= allocation
-			snapshot.free[podName] -= allocation
+			state.SetFree(ordinal, f-allocation)
 		} else {
 			newPlacements = append(newPlacements, placements[i])
 		}
@@ -210,18 +197,17 @@ func (s *StatefulSetScheduler) addReplicas(snapshot *Snapshot, diff int32, place
 
 	if diff > 0 {
 		// Needs to allocate replicas to additional pods
-		for i := int32(0); i < s.replicas; i++ {
-			podName := podNameFromOrdinal(s.statefulSetName, i)
-
-			if snapshot.free[podName] > 0 {
-				allocation := integer.Int32Min(snapshot.free[podName], diff)
+		for ordinal := int32(0); ordinal < s.replicas; ordinal++ {
+			f := state.Free(ordinal)
+			if f > 0 {
+				allocation := integer.Int32Min(f, diff)
 				newPlacements = append(newPlacements, duckv1alpha1.Placement{
-					PodName:   podName,
+					PodName:   podNameFromOrdinal(s.statefulSetName, ordinal),
 					VReplicas: allocation,
 				})
 
 				diff -= allocation
-				snapshot.free[podName] -= allocation
+				state.SetFree(ordinal, f-allocation)
 			}
 
 			if diff == 0 {
@@ -233,7 +219,15 @@ func (s *StatefulSetScheduler) addReplicas(snapshot *Snapshot, diff int32, place
 	return newPlacements, diff
 }
 
-// --- StatefulSet monitoring helpers
+// pendingReplicas returns the total number of vreplicas
+// that haven't been scheduled yet
+func (s *StatefulSetScheduler) pendingVReplicas() int32 {
+	t := int32(0)
+	for _, v := range s.pending {
+		t += v
+	}
+	return t
+}
 
 func (s *StatefulSetScheduler) updateStatefulset(obj interface{}) {
 	statefulset, ok := obj.(*appsv1.StatefulSet)
@@ -244,34 +238,11 @@ func (s *StatefulSetScheduler) updateStatefulset(obj interface{}) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	if s.replicas < statefulset.Status.Replicas {
-		// Less replicas: evict VPods scheduled on those replicas.
-		s.evict()
-
+	if statefulset.Spec.Replicas == nil {
+		s.replicas = 1
+	} else {
+		s.replicas = *statefulset.Spec.Replicas
 	}
-	// else if s.replicas > statefulset.Status.Replicas {
-	//   noop: controller will retry to schedule pending vreplicas
-	// }
-
-	s.replicas = statefulset.Status.Replicas
 
 	s.logger.Infow("updated statefulset replicas", zap.Int32("replicas", s.replicas))
 }
-
-func (s *StatefulSetScheduler) evict() {
-	vpods, err := s.vpodLister()
-	if err != nil {
-		s.logger.Fatalw("failed to list schedulable objects", zap.Error(err))
-		return
-	}
-
-	for _, vpod := range vpods {
-		ps := vpod.GetPlacements()
-		for i := 0; i < len(ps); i++ {
-			if ordinalFromPodName(ps[i].PodName) >= s.replicas {
-				s.enq(vpod.GetKey())
-				break
-			}
-		}
-	}
-}
diff --git a/pkg/common/scheduler/statefulset/snapshot.go b/pkg/common/scheduler/statefulset/snapshot.go
deleted file mode 100644
index 58632ba4c..000000000
--- a/pkg/common/scheduler/statefulset/snapshot.go
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package statefulset
-
-import "go.uber.org/zap"
-
-// Snapshot provides some information about the current scheduling of all vpods
-type Snapshot struct {
-	// free tracks the free capacity of each pod.
-	free map[string]int32
-
-	// lastOrdinal is the ordinal index corresponding to the last statefulset replica
-	// with placed vpods.
-	lastOrdinal int32
-}
-
-func (s *StatefulSetScheduler) snapshot(logger *zap.SugaredLogger) (*Snapshot, error) {
-	vpods, err := s.vpodLister()
-	if err != nil {
-		return nil, err
-	}
-
-	free := make(map[string]int32)
-	for i := int32(0); i < s.replicas; i++ {
-		podName := podNameFromOrdinal(s.statefulSetName, i)
-		free[podName] = s.capacity
-	}
-
-	last := int32(0)
-
-	for _, vpod := range vpods {
-		ps := vpod.GetPlacements()
-
-		for i := 0; i < len(ps); i++ {
-			podName := ps[i].PodName
-			vreplicas := ps[i].VReplicas
-
-			free[podName] -= vreplicas
-
-			if free[podName] < 0 {
-				// Over committed. The autoscaler will fix it.
-				logger.Infow("pod is overcommitted", zap.String("podName", podName))
-			}
-
-			ordinal := ordinalFromPodName(podName)
-			if ordinal > last {
-				last = ordinal
-			}
-		}
-	}
-	return &Snapshot{free: free, lastOrdinal: last}, nil
-}
diff --git a/pkg/common/scheduler/statefulset/state.go b/pkg/common/scheduler/statefulset/state.go
new file mode 100644
index 000000000..055011cc4
--- /dev/null
+++ b/pkg/common/scheduler/statefulset/state.go
@@ -0,0 +1,130 @@
+/*
+Copyright 2020 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package statefulset
+
+import (
+	"go.uber.org/zap"
+	"knative.dev/eventing-kafka/pkg/common/scheduler"
+)
+
+type stateAccessor interface {
+	// State returns the current state about placed vpods get
+	State() (*state, error)
+}
+
+// state provides information about the current scheduling of all vpods
+// It is used by for the scheduler and the autoscaler
+type state struct {
+	// free tracks the free capacity of each pod.
+	free []int32
+
+	// lastOrdinal is the ordinal index corresponding to the last statefulset replica
+	// with placed vpods.
+	lastOrdinal int32
+
+	// Pod capacity. Used
+	capacity int32
+}
+
+// Free safely returns the free capacity at the given ordinal
+func (s *state) Free(ordinal int32) int32 {
+	if int32(len(s.free)) <= ordinal {
+		return s.capacity
+	}
+	return s.free[ordinal]
+}
+
+// SetFree safely sets the free capacity at the given ordinal
+func (s *state) SetFree(ordinal int32, value int32) {
+	s.free = grow(s.free, ordinal, s.capacity)
+	s.free[int(ordinal)] = value
+}
+
+// freeCapacity returns the number of vreplicas that can be used,
+// up to the last ordinal
+func (s *state) freeCapacity() int32 {
+	t := int32(0)
+	for i := int32(0); i <= s.lastOrdinal; i++ {
+		t += s.free[i]
+	}
+	return t
+}
+
+// stateBuilder reconstruct the state from scratch, by listing vpods
+type stateBuilder struct {
+	logger     *zap.SugaredLogger
+	vpodLister scheduler.VPodLister
+	capacity   int32
+}
+
+// newStateBuilder returns a StateAccessor recreating the state from scratch each time it is requested
+func newStateBuilder(logger *zap.SugaredLogger, lister scheduler.VPodLister, podCapacity int32) stateAccessor {
+	return &stateBuilder{
+		logger:     logger,
+		vpodLister: lister,
+		capacity:   podCapacity,
+	}
+}
+
+func (s *stateBuilder) State() (*state, error) {
+	vpods, err := s.vpodLister()
+	if err != nil {
+		return nil, err
+	}
+
+	free := make([]int32, 0, 256)
+	last := int32(-1)
+
+	for _, vpod := range vpods {
+		ps := vpod.GetPlacements()
+
+		for i := 0; i < len(ps); i++ {
+			podName := ps[i].PodName
+			vreplicas := ps[i].VReplicas
+
+			ordinal := ordinalFromPodName(podName)
+
+			free = grow(free, ordinal, s.capacity)
+
+			free[ordinal] -= vreplicas
+
+			if free[ordinal] < 0 {
+				// Over committed. The autoscaler will fix it.
+				s.logger.Infow("pod is overcommitted", zap.String("podName", podName))
+			}
+
+			if ordinal > last && free[ordinal] != s.capacity {
+				last = ordinal
+			}
+		}
+	}
+	return &state{free: free, lastOrdinal: last, capacity: s.capacity}, nil
+}
+
+func grow(slice []int32, ordinal int32, def int32) []int32 {
+	l := int32(len(slice))
+	diff := ordinal - l + 1
+
+	if diff <= 0 {
+		return slice
+	}
+
+	for i := int32(0); i < diff; i++ {
+		slice = append(slice, def)
+	}
+	return slice
+}
diff --git a/pkg/source/reconciler/mtsource/controller.go b/pkg/source/reconciler/mtsource/controller.go
index 7584ccf95..bfabb631a 100644
--- a/pkg/source/reconciler/mtsource/controller.go
+++ b/pkg/source/reconciler/mtsource/controller.go
@@ -55,7 +55,7 @@ func NewController(
 	// Use a different set of conditions
 	sourcesv1beta1.RegisterAlternateKafkaConditionSet(sourcesv1beta1.KafkaMTSourceCondSet)
 
-	c.scheduler = scheduler.NewStatefulSetScheduler(ctx, system.Namespace(), mtadapterName, c.vpodLister, impl.EnqueueKey)
+	c.scheduler = scheduler.NewScheduler(ctx, system.Namespace(), mtadapterName, c.vpodLister)
 
 	logging.FromContext(ctx).Info("Setting up kafka event handlers")
 
diff --git a/vendor/knative.dev/eventing/pkg/adapter/v2/main.go b/vendor/knative.dev/eventing/pkg/adapter/v2/main.go
index 9354d7db0..b4b76988e 100644
--- a/vendor/knative.dev/eventing/pkg/adapter/v2/main.go
+++ b/vendor/knative.dev/eventing/pkg/adapter/v2/main.go
@@ -33,7 +33,6 @@ import (
 	kubeclient "knative.dev/pkg/client/injection/kube/client"
 	"knative.dev/pkg/controller"
 	"knative.dev/pkg/injection"
-	"knative.dev/pkg/injection/sharedmain"
 	"knative.dev/pkg/leaderelection"
 	"knative.dev/pkg/logging"
 	"knative.dev/pkg/metrics"
@@ -203,7 +202,7 @@ func MainWithInformers(ctx context.Context, component string, env EnvConfigAcces
 func ConstructEnvOrDie(ector EnvConfigConstructor) EnvConfigAccessor {
 	env := ector()
 	if err := envconfig.Process("", env); err != nil {
-		log.Fatalf("Error processing env var: %s", err)
+		log.Panicf("Error processing env var: %s", err)
 	}
 	return env
 }
@@ -217,7 +216,7 @@ func SetupInformers(ctx context.Context, logger *zap.SugaredLogger) (context.Con
 		logger.Infof("Registering %d informers", len(injection.Default.GetInformers()))
 		logger.Infof("Registering %d ducks", len(injection.Default.GetDucks()))
 
-		cfg := sharedmain.ParseAndGetConfigOrDie()
+		cfg := injection.ParseAndGetRESTConfigOrDie()
 		return injection.Default.SetupInformers(ctx, cfg)
 	}
 	return ctx, nil
diff --git a/vendor/knative.dev/eventing/pkg/adapter/v2/main_message.go b/vendor/knative.dev/eventing/pkg/adapter/v2/main_message.go
index 854cbb553..b798967b4 100644
--- a/vendor/knative.dev/eventing/pkg/adapter/v2/main_message.go
+++ b/vendor/knative.dev/eventing/pkg/adapter/v2/main_message.go
@@ -42,9 +42,7 @@ type MessageAdapter interface {
 type MessageAdapterConstructor func(ctx context.Context, env EnvConfigAccessor, adapter *kncloudevents.HTTPMessageSender, reporter source.StatsReporter) MessageAdapter
 
 func MainMessageAdapter(component string, ector EnvConfigConstructor, ctor MessageAdapterConstructor) {
-	ctx := signals.NewContext()
-
-	MainMessageAdapterWithContext(ctx, component, ector, ctor)
+	MainMessageAdapterWithContext(signals.NewContext(), component, ector, ctor)
 }
 
 func MainMessageAdapterWithContext(ctx context.Context, component string, ector EnvConfigConstructor, ctor MessageAdapterConstructor) {
diff --git a/vendor/knative.dev/pkg/configmap/informer/informed_watcher.go b/vendor/knative.dev/pkg/configmap/informer/informed_watcher.go
index 2231ee52f..1ae7ef8d4 100644
--- a/vendor/knative.dev/pkg/configmap/informer/informed_watcher.go
+++ b/vendor/knative.dev/pkg/configmap/informer/informed_watcher.go
@@ -107,7 +107,8 @@ var _ configmap.Watcher = (*InformedWatcher)(nil)
 // Asserts that InformedWatcher implements DefaultingWatcher.
 var _ configmap.DefaultingWatcher = (*InformedWatcher)(nil)
 
-// WatchWithDefault implements DefaultingWatcher.
+// WatchWithDefault implements DefaultingWatcher. Adding a default for the configMap being watched means that when
+// Start is called, Start will not wait for the add event from the API server.
 func (i *InformedWatcher) WatchWithDefault(cm corev1.ConfigMap, o ...configmap.Observer) {
 	i.defaults[cm.Name] = &cm
 
@@ -126,31 +127,59 @@ func (i *InformedWatcher) WatchWithDefault(cm corev1.ConfigMap, o ...configmap.O
 	i.Watch(cm.Name, o...)
 }
 
-// Start implements Watcher.
-func (i *InformedWatcher) Start(stopCh <-chan struct{}) error {
-	// Pretend that all the defaulted ConfigMaps were just created. This is done before we start
-	// the informer to ensure that if a defaulted ConfigMap does exist, then the real value is
-	// processed after the default one.
+func (i *InformedWatcher) triggerAddEventForDefaultedConfigMaps(addConfigMapEvent func(obj interface{})) {
 	i.ForEach(func(k string, _ []configmap.Observer) error {
 		if def, ok := i.defaults[k]; ok {
-			i.addConfigMapEvent(def)
+			addConfigMapEvent(def)
 		}
 		return nil
 	})
+}
+
+func (i *InformedWatcher) getConfigMapNames() []string {
+	var configMaps []string
+	i.ForEach(func(k string, _ []configmap.Observer) error {
+		configMaps = append(configMaps, k)
+		return nil
+	})
+	return configMaps
+}
 
-	if err := i.registerCallbackAndStartInformer(stopCh); err != nil {
+// Start implements Watcher. Start will wait for all watched resources to exist and for the add event handler to be
+// invoked at least once for each before continuing or for the stopCh to be signalled, whichever happens first. If
+// the watched resource is defaulted, Start will invoke the add event handler directly and will not wait for a further
+// add event from the API server.
+func (i *InformedWatcher) Start(stopCh <-chan struct{}) error {
+	// using the synced callback wrapper around the add event handler will allow the caller
+	// to wait for the add event to be processed for all configmaps
+	s := newSyncedCallback(i.getConfigMapNames(), i.addConfigMapEvent)
+	addConfigMapEvent := func(obj interface{}) {
+		configMap := obj.(*corev1.ConfigMap)
+		s.Call(obj, configMap.Name)
+	}
+	// Pretend that all the defaulted ConfigMaps were just created. This is done before we start
+	// the informer to ensure that if a defaulted ConfigMap does exist, then the real value is
+	// processed after the default one.
+	i.triggerAddEventForDefaultedConfigMaps(addConfigMapEvent)
+
+	if err := i.registerCallbackAndStartInformer(addConfigMapEvent, stopCh); err != nil {
 		return err
 	}
 
-	// Wait until it has been synced (WITHOUT holing the mutex, so callbacks happen)
+	// Wait until the shared informer has been synced (WITHOUT holing the mutex, so callbacks happen)
 	if ok := cache.WaitForCacheSync(stopCh, i.informer.Informer().HasSynced); !ok {
 		return errors.New("error waiting for ConfigMap informer to sync")
 	}
 
-	return i.checkObservedResourcesExist()
+	if err := i.checkObservedResourcesExist(); err != nil {
+		return err
+	}
+
+	// Wait until all config maps have been at least initially processed
+	return s.WaitForAllKeys(stopCh)
 }
 
-func (i *InformedWatcher) registerCallbackAndStartInformer(stopCh <-chan struct{}) error {
+func (i *InformedWatcher) registerCallbackAndStartInformer(addConfigMapEvent func(obj interface{}), stopCh <-chan struct{}) error {
 	i.Lock()
 	defer i.Unlock()
 	if i.started {
@@ -159,13 +188,14 @@ func (i *InformedWatcher) registerCallbackAndStartInformer(stopCh <-chan struct{
 	i.started = true
 
 	i.informer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
-		AddFunc:    i.addConfigMapEvent,
+		AddFunc:    addConfigMapEvent,
 		UpdateFunc: i.updateConfigMapEvent,
 		DeleteFunc: i.deleteConfigMapEvent,
 	})
 
 	// Start the shared informer factory (non-blocking).
 	i.sif.Start(stopCh)
+
 	return nil
 }
 
diff --git a/vendor/knative.dev/pkg/configmap/informer/synced_callback.go b/vendor/knative.dev/pkg/configmap/informer/synced_callback.go
new file mode 100644
index 000000000..355c512b5
--- /dev/null
+++ b/vendor/knative.dev/pkg/configmap/informer/synced_callback.go
@@ -0,0 +1,112 @@
+/*
+Copyright 2021 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package informer
+
+import (
+	"sync"
+
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/apimachinery/pkg/util/wait"
+)
+
+// namedWaitGroup is used to increment and decrement a WaitGroup by name
+type namedWaitGroup struct {
+	waitGroup sync.WaitGroup
+	keys      sets.String
+	mu        sync.Mutex
+}
+
+// newNamedWaitGroup returns an instantiated namedWaitGroup.
+func newNamedWaitGroup() *namedWaitGroup {
+	return &namedWaitGroup{
+		keys: sets.NewString(),
+	}
+}
+
+// Add will add the key to the list of keys being tracked and increment the wait group.
+// If the key has already been added, the wait group will not be incremented again.
+func (n *namedWaitGroup) Add(key string) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	if !n.keys.Has(key) {
+		n.keys.Insert(key)
+		n.waitGroup.Add(1)
+	}
+}
+
+// Done will decrement the counter if the key is present in the tracked keys. If it is not present
+// it will be ignored.
+func (n *namedWaitGroup) Done(key string) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	if n.keys.Has(key) {
+		n.keys.Delete(key)
+		n.waitGroup.Done()
+	}
+}
+
+// Wait will wait for the underlying waitGroup to complete.
+func (n *namedWaitGroup) Wait() {
+	n.waitGroup.Wait()
+}
+
+// syncedCallback can be used to wait for a callback to be called at least once for a list of keys.
+type syncedCallback struct {
+	// namedWaitGroup will block until the callback has been called for all tracked entities
+	namedWaitGroup *namedWaitGroup
+
+	// callback is the callback that is intended to be called at least once for each key
+	// being tracked via WaitGroup
+	callback func(obj interface{})
+}
+
+// newSyncedCallback will return a syncedCallback that will track the provided keys.
+func newSyncedCallback(keys []string, callback func(obj interface{})) *syncedCallback {
+	s := &syncedCallback{
+		callback:       callback,
+		namedWaitGroup: newNamedWaitGroup(),
+	}
+	for _, key := range keys {
+		s.namedWaitGroup.Add(key)
+	}
+	return s
+}
+
+// Event is intended to be a wrapper for the actual event handler; this wrapper will signal via
+// the wait group that the event handler has been called at least once for the key.
+func (s *syncedCallback) Call(obj interface{}, key string) {
+	s.callback(obj)
+	s.namedWaitGroup.Done(key)
+}
+
+// WaitForAllKeys will block until s.Call has been called for all the keys we are tracking or the stop signal is
+// received.
+func (s *syncedCallback) WaitForAllKeys(stopCh <-chan struct{}) error {
+	c := make(chan struct{})
+	go func() {
+		defer close(c)
+		s.namedWaitGroup.Wait()
+	}()
+	select {
+	case <-c:
+		return nil
+	case <-stopCh:
+		return wait.ErrWaitTimeout
+	}
+}
diff --git a/vendor/modules.txt b/vendor/modules.txt
index 886acaa5d..085e9bded 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -1080,7 +1080,7 @@ k8s.io/utils/buffer
 k8s.io/utils/integer
 k8s.io/utils/pointer
 k8s.io/utils/trace
-# knative.dev/eventing v0.21.1-0.20210302004649-f22066824960
+# knative.dev/eventing v0.21.1-0.20210302231422-0f2612ef8aba
 ## explicit
 knative.dev/eventing/pkg/adapter/v2
 knative.dev/eventing/pkg/adapter/v2/test
@@ -1219,7 +1219,7 @@ knative.dev/eventing/test/test_images/recordevents
 # knative.dev/hack v0.0.0-20210203173706-8368e1f6eacf
 ## explicit
 knative.dev/hack
-# knative.dev/pkg v0.0.0-20210226182947-9039dc189ced
+# knative.dev/pkg v0.0.0-20210302195815-2bc944bb68d0
 ## explicit
 knative.dev/pkg/apis
 knative.dev/pkg/apis/duck


Test patch:
diff --git a/pkg/common/scheduler/statefulset/autoscaler_test.go b/pkg/common/scheduler/statefulset/autoscaler_test.go
new file mode 100644
index 000000000..0f1f3c9f0
--- /dev/null
+++ b/pkg/common/scheduler/statefulset/autoscaler_test.go
@@ -0,0 +1,272 @@
+/*
+Copyright 2020 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package statefulset
+
+import (
+	"testing"
+	"time"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	gtesting "k8s.io/client-go/testing"
+
+	kubeclient "knative.dev/pkg/client/injection/kube/client/fake"
+	_ "knative.dev/pkg/client/injection/kube/informers/apps/v1/statefulset/fake"
+	"knative.dev/pkg/logging"
+
+	duckv1alpha1 "knative.dev/eventing-kafka/pkg/apis/duck/v1alpha1"
+	"knative.dev/eventing-kafka/pkg/common/scheduler"
+	tscheduler "knative.dev/eventing-kafka/pkg/common/scheduler/testing"
+)
+
+const (
+	testNs = "test-ns"
+)
+
+func TestAutoscaler(t *testing.T) {
+	testCases := []struct {
+		name         string
+		replicas     int32
+		vpods        []scheduler.VPod
+		pendings     int32
+		scaleDown    bool
+		wantReplicas int32
+	}{
+		{
+			name:     "no replicas, no placements, no pending",
+			replicas: int32(0),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 0, nil),
+			},
+			pendings:     int32(0),
+			wantReplicas: int32(0),
+		},
+		{
+			name:     "no replicas, no placements, with pending",
+			replicas: int32(0),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 5, nil),
+			},
+			pendings:     int32(5),
+			wantReplicas: int32(1),
+		},
+		{
+			name:     "no replicas, with placements, no pending",
+			replicas: int32(0),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(0),
+			wantReplicas: int32(2),
+		},
+		{
+			name:     "no replicas, with placements, with pending, enough capacity",
+			replicas: int32(0),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(3),
+			wantReplicas: int32(2),
+		},
+		{
+			name:     "no replicas, with placements, with pending, not enough capacity",
+			replicas: int32(0),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(8),
+			wantReplicas: int32(3),
+		},
+		{
+			name:     "with replicas, no placements, no pending, scale down",
+			replicas: int32(3),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 0, nil),
+			},
+			pendings:     int32(0),
+			scaleDown:    true,
+			wantReplicas: int32(0),
+		},
+		{
+			name:     "with replicas, no placements, with pending, scale down",
+			replicas: int32(3),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 5, nil),
+			},
+			pendings:     int32(5),
+			scaleDown:    true,
+			wantReplicas: int32(1),
+		},
+		{
+			name:     "with replicas, no placements, with pending, scale down disabled",
+			replicas: int32(3),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 5, nil),
+			},
+			pendings:     int32(5),
+			scaleDown:    false,
+			wantReplicas: int32(3),
+		},
+		{
+			name:     "with replicas, no placements, with pending, scale up",
+			replicas: int32(3),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 5, nil),
+			},
+			pendings:     int32(40),
+			wantReplicas: int32(4),
+		},
+		{
+			name:     "with replicas, with placements, no pending, no change",
+			replicas: int32(2),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(0),
+			wantReplicas: int32(2),
+		},
+		{
+			name:     "with replicas, with placements, no pending, scale down",
+			replicas: int32(5),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(0),
+			scaleDown:    true,
+			wantReplicas: int32(2),
+		},
+		{
+			name:     "with replicas, with placements, with pending, enough capacity",
+			replicas: int32(2),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(3),
+			wantReplicas: int32(2),
+		},
+		{
+			name:     "with replicas, with placements, with pending, not enough capacity",
+			replicas: int32(2),
+			vpods: []scheduler.VPod{
+				tscheduler.NewVPod(testNs, "vpod-1", 15, []duckv1alpha1.Placement{
+					{PodName: "pod-0", VReplicas: int32(8)},
+					{PodName: "pod-1", VReplicas: int32(7)}}),
+			},
+			pendings:     int32(8),
+			wantReplicas: int32(3),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, _ := setupFakeContext(t)
+
+			vpodClient := tscheduler.NewVPodClient()
+			stateAccessor := newStateBuilder(logging.FromContext(ctx), vpodClient.List, 10)
+
+			sfsClient := kubeclient.Get(ctx).AppsV1().StatefulSets(testNs)
+			_, err := sfsClient.Create(ctx, makeStatefulset(ctx, testNs, sfsName, tc.replicas), metav1.CreateOptions{})
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+
+			autoscaler := NewAutoscaler(ctx, testNs, sfsName, stateAccessor, 10*time.Second).(*autoscaler)
+
+			for _, vpod := range tc.vpods {
+				vpodClient.Append(vpod)
+			}
+
+			err = autoscaler.doautoscale(ctx, tc.scaleDown, tc.pendings)
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+
+			scale, err := sfsClient.GetScale(ctx, sfsName, metav1.GetOptions{})
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+			if scale.Spec.Replicas != tc.wantReplicas {
+				t.Errorf("unexpected number of replicas, got %d, want %d", scale.Spec.Replicas, tc.wantReplicas)
+			}
+
+		})
+	}
+}
+
+func TestAutoscalerScaleDownToZero(t *testing.T) {
+	ctx, cancel := setupFakeContext(t)
+
+	afterUpdate := make(chan bool)
+	kubeclient.Get(ctx).PrependReactor("update", "statefulsets", func(action gtesting.Action) (handled bool, ret runtime.Object, err error) {
+		if action.GetSubresource() == "scale" {
+			afterUpdate <- true
+		}
+		return false, nil, nil
+	})
+
+	vpodClient := tscheduler.NewVPodClient()
+	stateAccessor := newStateBuilder(logging.FromContext(ctx), vpodClient.List, 10)
+
+	sfsClient := kubeclient.Get(ctx).AppsV1().StatefulSets(testNs)
+	_, err := sfsClient.Create(ctx, makeStatefulset(ctx, testNs, sfsName, 10), metav1.CreateOptions{})
+	if err != nil {
+		t.Fatal("unexpected error", err)
+	}
+
+	autoscaler := NewAutoscaler(ctx, testNs, sfsName, stateAccessor, 2*time.Second).(*autoscaler)
+
+	done := make(chan bool)
+	go func() {
+		autoscaler.Start(ctx)
+		done <- true
+	}()
+
+	select {
+	case <-afterUpdate:
+	case <-time.After(4 * time.Second):
+		t.Fatal("timeout waiting for scale subresource to be updated")
+
+	}
+
+	sfs, err := sfsClient.Get(ctx, sfsName, metav1.GetOptions{})
+	if err != nil {
+		t.Fatal("unexpected error", err)
+	}
+	if *sfs.Spec.Replicas != 0 {
+		t.Errorf("unexpected number of replicas, got %d, want 0", *sfs.Spec.Replicas)
+	}
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(1 * time.Second):
+		t.Fatal("timeout waiting for autoscaler to stop")
+	}
+}
diff --git a/pkg/common/scheduler/statefulset/scheduler_test.go b/pkg/common/scheduler/statefulset/scheduler_test.go
index d1c062930..55cae4584 100644
--- a/pkg/common/scheduler/statefulset/scheduler_test.go
+++ b/pkg/common/scheduler/statefulset/scheduler_test.go
@@ -18,14 +18,13 @@ package statefulset
 
 import (
 	"context"
-	"errors"
-	"fmt"
 	"reflect"
 	"testing"
 	"time"
 
 	appsv1 "k8s.io/api/apps/v1"
 	autoscalingv1 "k8s.io/api/autoscaling/v1"
+	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	gtesting "k8s.io/client-go/testing"
@@ -33,6 +32,7 @@ import (
 	kubeclient "knative.dev/pkg/client/injection/kube/client/fake"
 	_ "knative.dev/pkg/client/injection/kube/informers/apps/v1/statefulset/fake"
 	"knative.dev/pkg/controller"
+	"knative.dev/pkg/logging"
 	rectesting "knative.dev/pkg/reconciler/testing"
 
 	duckv1alpha1 "knative.dev/eventing-kafka/pkg/apis/duck/v1alpha1"
@@ -41,70 +41,104 @@ import (
 )
 
 const (
-	sfsName = "statefulset-name"
+	sfsName       = "statefulset-name"
+	vpodName      = "source-name"
+	vpodNamespace = "source-namespace"
 )
 
 func TestStatefulsetScheduler(t *testing.T) {
-	ctx, cancel := setupFakeContext(t)
-
 	testCases := []struct {
 		name       string
-		vpods      []int32
+		vreplicas  int32
 		replicas   int32
 		placements []duckv1alpha1.Placement
+		expected   []duckv1alpha1.Placement
 		err        error
 	}{
 		{
-			name:       "no replicas",
-			vpods:      []int32{1},
-			replicas:   int32(0),
-			err:        scheduler.ErrNoReplicas,
-			placements: []duckv1alpha1.Placement{},
+			name:      "no replicas, no vreplicas",
+			vreplicas: 0,
+			replicas:  int32(0),
+			expected:  nil,
+		},
+		{
+			name:      "no replicas, 1 vreplicas, fail.",
+			vreplicas: 1,
+			replicas:  int32(0),
+			err:       scheduler.ErrNotEnoughReplicas,
+			expected:  []duckv1alpha1.Placement{},
+		},
+		{
+			name:      "one replica, one vreplicas",
+			vreplicas: 1,
+			replicas:  int32(1),
+			expected:  []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 1}},
 		},
 		{
-			name:       "one vpod, one vreplicas",
-			vpods:      []int32{1},
-			replicas:   int32(1),
-			placements: []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 1}},
+			name:      "one replica, 3 vreplicas",
+			vreplicas: 3,
+			replicas:  int32(1),
+			expected:  []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 3}},
 		},
 		{
-			name:       "one vpod, 3 vreplicas",
-			vpods:      []int32{3},
-			replicas:   int32(1),
-			placements: []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 3}},
+			name:      "one replica, 15 vreplicas, unschedulable",
+			vreplicas: 15,
+			replicas:  int32(1),
+			err:       scheduler.ErrNotEnoughReplicas,
+			expected:  []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 10}},
 		},
 		{
-			name:       "one vpod, 15 vreplicas, partial scheduling",
-			vpods:      []int32{15},
-			replicas:   int32(1),
-			err:        scheduler.ErrNotEnoughReplicas,
-			placements: []duckv1alpha1.Placement{{PodName: "statefulset-name-0", VReplicas: 10}},
+			name:      "two replicas, 15 vreplicas, scheduled",
+			vreplicas: 15,
+			replicas:  int32(2),
+			expected: []duckv1alpha1.Placement{
+				{PodName: "statefulset-name-0", VReplicas: 10},
+				{PodName: "statefulset-name-1", VReplicas: 5},
+			},
 		},
 		{
-			name:     "one vpod, 15 vreplicas, full scheduling",
-			vpods:    []int32{15},
-			replicas: int32(2),
+			name:      "two replicas, 15 vreplicas, already scheduled",
+			vreplicas: 15,
+			replicas:  int32(2),
 			placements: []duckv1alpha1.Placement{
 				{PodName: "statefulset-name-0", VReplicas: 10},
 				{PodName: "statefulset-name-1", VReplicas: 5},
 			},
+			expected: []duckv1alpha1.Placement{
+				{PodName: "statefulset-name-0", VReplicas: 10},
+				{PodName: "statefulset-name-1", VReplicas: 5},
+			},
+		},
+		{
+			name:      "two replicas, 15 vreplicas, too much scheduled (scale down)",
+			vreplicas: 15,
+			replicas:  int32(2),
+			placements: []duckv1alpha1.Placement{
+				{PodName: "statefulset-name-0", VReplicas: 10},
+				{PodName: "statefulset-name-1", VReplicas: 10},
+			},
+			expected: []duckv1alpha1.Placement{
+				{PodName: "statefulset-name-0", VReplicas: 5},
+				{PodName: "statefulset-name-1", VReplicas: 10},
+			},
 		},
 	}
 
-	for i, tc := range testCases {
+	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
+			ctx, _ := setupFakeContext(t)
+
 			vpodClient := tscheduler.NewVPodClient()
-			sfsNs := fmt.Sprint("ns-", i)
 
-			_, err := kubeclient.Get(ctx).AppsV1().StatefulSets(sfsNs).Create(ctx, makeStatefulset(ctx, sfsNs, sfsName, tc.replicas), metav1.CreateOptions{})
+			_, err := kubeclient.Get(ctx).AppsV1().StatefulSets(testNs).Create(ctx, makeStatefulset(ctx, testNs, sfsName, tc.replicas), metav1.CreateOptions{})
 			if err != nil {
 				t.Fatal("unexpected error", err)
 			}
+			sa := newStateBuilder(logging.FromContext(ctx), vpodClient.List, 10)
+			s := NewStatefulSetScheduler(ctx, testNs, sfsName, vpodClient.List, sa, nil).(*StatefulSetScheduler)
 
-			s := NewStatefulSetScheduler(ctx, sfsNs, sfsName, vpodClient.List, nil).(*StatefulSetScheduler)
-
-			// Give some time for the informer to notify the scheduler
-			time.Sleep(500 * time.Millisecond)
+			// Give some time for the informer to notify the scheduler and set the number of replicas
+			time.Sleep(200 * time.Millisecond)
 
 			func() {
 				s.lock.Lock()
@@ -114,29 +148,23 @@ func TestStatefulsetScheduler(t *testing.T) {
 				}
 			}()
 
-			for i, vreplicas := range tc.vpods {
-				vpodName := fmt.Sprint("vpod-name-", i)
-				vpodNamespace := fmt.Sprint("vpod-ns-", i)
-
-				vpod := vpodClient.Create(vpodNamespace, vpodName, vreplicas, nil)
-				placements, err := s.Schedule(vpod)
+			vpod := vpodClient.Create(vpodNamespace, vpodName, tc.vreplicas, tc.placements)
+			placements, err := s.Schedule(vpod)
 
-				if tc.err == nil && err != nil {
-					t.Fatal("unexpected error", err)
-				}
+			if tc.err == nil && err != nil {
+				t.Fatal("unexpected error", err)
+			}
 
-				if tc.err != nil && err == nil {
-					t.Fatal("expected error, got none")
-				}
+			if tc.err != nil && err == nil {
+				t.Fatal("expected error, got none")
+			}
 
-				if !reflect.DeepEqual(placements, tc.placements) {
-					t.Errorf("got %v, want %v", placements, tc.placements)
-				}
+			if !reflect.DeepEqual(placements, tc.expected) {
+				t.Errorf("got %v, want %v", placements, tc.expected)
 			}
+
 		})
 	}
-	cancel()
-	<-ctx.Done()
 }
 
 func makeStatefulset(ctx context.Context, ns, name string, replicas int32) *appsv1.StatefulSet {
@@ -153,18 +181,6 @@ func makeStatefulset(ctx context.Context, ns, name string, replicas int32) *apps
 		},
 	}
 
-	kc := kubeclient.Get(ctx)
-	kc.PrependReactor("get", "statefulsets", func(action gtesting.Action) (handled bool, ret runtime.Object, err error) {
-		if action.GetSubresource() == "scale" {
-			return true, &autoscalingv1.Scale{
-				Spec: autoscalingv1.ScaleSpec{
-					Replicas: obj.Status.Replicas,
-				},
-			}, nil
-		}
-		return false, nil, nil
-	})
-
 	return obj
 }
 
@@ -176,9 +192,67 @@ func setupFakeContext(t *testing.T) (context.Context, context.CancelFunc) {
 	}
 
 	kc := kubeclient.Get(ctx)
+	kc.PrependReactor("create", "statefulsets", func(action gtesting.Action) (handled bool, ret runtime.Object, err error) {
+		createAction := action.(gtesting.CreateActionImpl)
+		sfs := createAction.GetObject().(*appsv1.StatefulSet)
+		scale := &autoscalingv1.Scale{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      sfs.Name,
+				Namespace: sfs.Namespace,
+			},
+			Spec: autoscalingv1.ScaleSpec{
+				Replicas: func() int32 {
+					if sfs.Spec.Replicas == nil {
+						return 1
+					}
+					return *sfs.Spec.Replicas
+				}(),
+			},
+		}
+		kc.Tracker().Add(scale)
+		return false, nil, nil
+	})
+
 	kc.PrependReactor("get", "statefulsets", func(action gtesting.Action) (handled bool, ret runtime.Object, err error) {
+		getAction := action.(gtesting.GetAction)
+		if action.GetSubresource() == "scale" {
+			scale, err := kc.Tracker().Get(autoscalingv1.SchemeGroupVersion.WithResource("scales"), getAction.GetNamespace(), getAction.GetName())
+			return true, scale, err
+
+		}
+		return false, nil, nil
+	})
+
+	kc.PrependReactor("update", "statefulsets", func(action gtesting.Action) (handled bool, ret runtime.Object, err error) {
+		updateAction := action.(gtesting.UpdateActionImpl)
 		if action.GetSubresource() == "scale" {
-			return true, nil, errors.New("object not found")
+			scale := updateAction.GetObject().(*autoscalingv1.Scale)
+
+			err := kc.Tracker().Update(autoscalingv1.SchemeGroupVersion.WithResource("scales"), scale, scale.GetNamespace())
+			if err != nil {
+				return true, nil, err
+			}
+
+			meta, err := meta.Accessor(updateAction.GetObject())
+			if err != nil {
+				return true, nil, err
+			}
+
+			obj, err := kc.Tracker().Get(appsv1.SchemeGroupVersion.WithResource("statefulsets"), meta.GetNamespace(), meta.GetName())
+			if err != nil {
+				return true, nil, err
+			}
+
+			sfs := obj.(*appsv1.StatefulSet)
+			sfs.Spec.Replicas = &scale.Spec.Replicas
+
+			err = kc.Tracker().Update(appsv1.SchemeGroupVersion.WithResource("statefulsets"), sfs, sfs.GetNamespace())
+			if err != nil {
+				return true, nil, err
+			}
+
+			return true, scale, nil
+
 		}
 		return false, nil, nil
 	})
diff --git a/pkg/common/scheduler/statefulset/snapshot_test.go b/pkg/common/scheduler/statefulset/snapshot_test.go
deleted file mode 100644
index 73170a534..000000000
--- a/pkg/common/scheduler/statefulset/snapshot_test.go
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package statefulset
-
-import (
-	"fmt"
-	"reflect"
-	"testing"
-	"time"
-
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-
-	"knative.dev/pkg/logging"
-
-	duckv1alpha1 "knative.dev/eventing-kafka/pkg/apis/duck/v1alpha1"
-	tscheduler "knative.dev/eventing-kafka/pkg/common/scheduler/testing"
-	kubeclient "knative.dev/pkg/client/injection/kube/client/fake"
-)
-
-func TestSnapshot(t *testing.T) {
-	ctx, cancel := setupFakeContext(t)
-
-	testCases := []struct {
-		name     string
-		vpods    [][]duckv1alpha1.Placement
-		replicas int32
-		expected Snapshot
-		err      error
-	}{
-		{
-			name:     "no vpods, no replicas",
-			vpods:    [][]duckv1alpha1.Placement{},
-			replicas: int32(0),
-			expected: Snapshot{free: map[string]int32{}, lastOrdinal: 0},
-		},
-		{
-			name:     "no vpods, one replica",
-			vpods:    [][]duckv1alpha1.Placement{},
-			replicas: int32(1),
-			expected: Snapshot{free: map[string]int32{"statefulset-name-0": int32(10)}, lastOrdinal: 0},
-		},
-		{
-			name:     "one vpods, one replica",
-			vpods:    [][]duckv1alpha1.Placement{{{PodName: "statefulset-name-0", VReplicas: 1}}},
-			replicas: int32(1),
-			expected: Snapshot{free: map[string]int32{"statefulset-name-0": int32(9)}, lastOrdinal: 0},
-		},
-		{
-			name: "many vpods, many replicas",
-			vpods: [][]duckv1alpha1.Placement{
-				{{PodName: "statefulset-name-0", VReplicas: 1}, {PodName: "statefulset-name-2", VReplicas: 5}},
-				{{PodName: "statefulset-name-1", VReplicas: 2}},
-				{{PodName: "statefulset-name-1", VReplicas: 3}, {PodName: "statefulset-name-0", VReplicas: 1}},
-			},
-			replicas: int32(4),
-			expected: Snapshot{free: map[string]int32{
-				"statefulset-name-0": int32(8),
-				"statefulset-name-1": int32(5),
-				"statefulset-name-2": int32(5),
-				"statefulset-name-3": int32(10)}, lastOrdinal: 2},
-		},
-	}
-
-	for i, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			vpodClient := tscheduler.NewVPodClient()
-			sfsNs := fmt.Sprint("ns-", i)
-
-			_, err := kubeclient.Get(ctx).AppsV1().StatefulSets(sfsNs).Create(ctx, makeStatefulset(ctx, sfsNs, sfsName, tc.replicas), metav1.CreateOptions{})
-			if err != nil {
-				t.Fatal("unexpected error", err)
-			}
-
-			s := NewStatefulSetScheduler(ctx, sfsNs, sfsName, vpodClient.List, nil).(*StatefulSetScheduler)
-
-			// Give some time for the informer to notify the scheduler
-			time.Sleep(500 * time.Millisecond)
-
-			func() {
-				s.lock.Lock()
-				defer s.lock.Unlock()
-				if s.replicas != tc.replicas {
-					t.Fatalf("expected number of statefulset replica to be %d (got %d)", tc.replicas, s.replicas)
-				}
-			}()
-
-			for i, placements := range tc.vpods {
-				vpodName := fmt.Sprint("vpod-name-", i)
-				vpodNamespace := fmt.Sprint("vpod-ns-", i)
-
-				vpodClient.Create(vpodNamespace, vpodName, 1, placements)
-
-			}
-
-			shot, err := s.snapshot(logging.FromContext(ctx))
-			if err != nil {
-				t.Fatal("unexpected error", err)
-			}
-
-			if !reflect.DeepEqual(*shot, tc.expected) {
-				t.Errorf("got %v, want %v", shot, tc.expected)
-			}
-		})
-	}
-
-	cancel()
-	<-ctx.Done()
-}
diff --git a/pkg/common/scheduler/statefulset/state_test.go b/pkg/common/scheduler/statefulset/state_test.go
new file mode 100644
index 000000000..1da07fdf5
--- /dev/null
+++ b/pkg/common/scheduler/statefulset/state_test.go
@@ -0,0 +1,98 @@
+/*
+Copyright 2020 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package statefulset
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	logtesting "knative.dev/pkg/logging/testing"
+
+	duckv1alpha1 "knative.dev/eventing-kafka/pkg/apis/duck/v1alpha1"
+	tscheduler "knative.dev/eventing-kafka/pkg/common/scheduler/testing"
+)
+
+func TestStateBuilder(t *testing.T) {
+	testCases := []struct {
+		name     string
+		vpods    [][]duckv1alpha1.Placement
+		expected state
+		freec    int32
+		err      error
+	}{
+		{
+			name:     "no vpods",
+			vpods:    [][]duckv1alpha1.Placement{},
+			expected: state{capacity: 10, free: []int32{}, lastOrdinal: -1},
+			freec:    int32(0),
+		},
+		{
+			name:     "one vpods",
+			vpods:    [][]duckv1alpha1.Placement{{{PodName: "statefulset-name-0", VReplicas: 1}}},
+			expected: state{capacity: 10, free: []int32{int32(9)}, lastOrdinal: 0},
+			freec:    int32(9),
+		},
+		{
+			name: "many vpods, no gaps",
+			vpods: [][]duckv1alpha1.Placement{
+				{{PodName: "statefulset-name-0", VReplicas: 1}, {PodName: "statefulset-name-2", VReplicas: 5}},
+				{{PodName: "statefulset-name-1", VReplicas: 2}},
+				{{PodName: "statefulset-name-1", VReplicas: 3}, {PodName: "statefulset-name-0", VReplicas: 1}},
+			},
+			expected: state{capacity: 10, free: []int32{int32(8), int32(5), int32(5)}, lastOrdinal: 2},
+			freec:    int32(18),
+		},
+		{
+			name: "many vpods, with gaps",
+			vpods: [][]duckv1alpha1.Placement{
+				{{PodName: "statefulset-name-0", VReplicas: 1}, {PodName: "statefulset-name-2", VReplicas: 5}},
+				{{PodName: "statefulset-name-1", VReplicas: 0}},
+				{{PodName: "statefulset-name-1", VReplicas: 0}, {PodName: "statefulset-name-3", VReplicas: 0}},
+			},
+			expected: state{capacity: 10, free: []int32{int32(9), int32(10), int32(5), int32(10)}, lastOrdinal: 2},
+			freec:    int32(24),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			vpodClient := tscheduler.NewVPodClient()
+
+			for i, placements := range tc.vpods {
+				vpodName := fmt.Sprint("vpod-name-", i)
+				vpodNamespace := fmt.Sprint("vpod-ns-", i)
+
+				vpodClient.Create(vpodNamespace, vpodName, 1, placements)
+			}
+
+			stateBuilder := newStateBuilder(logtesting.TestLogger(t), vpodClient.List, int32(10))
+			state, err := stateBuilder.State()
+			if err != nil {
+				t.Fatal("unexpected error", err)
+			}
+
+			if !reflect.DeepEqual(*state, tc.expected) {
+				t.Errorf("unexpected state, got %v, want %v", state, tc.expected)
+			}
+
+			if state.freeCapacity() != tc.freec {
+				t.Errorf("unexpected free capacity, got %d, want %d", state.freeCapacity(), tc.freec)
+			}
+		})
+	}
+}
diff --git a/pkg/common/scheduler/testing/client.go b/pkg/common/scheduler/testing/client.go
index 2d51251c8..d72c54bf3 100644
--- a/pkg/common/scheduler/testing/client.go
+++ b/pkg/common/scheduler/testing/client.go
@@ -35,13 +35,19 @@ func NewVPodClient() *VPodClient {
 }
 
 func (s *VPodClient) Create(ns, name string, vreplicas int32, placements []duckv1alpha1.Placement) scheduler.VPod {
-	vpod := newVPod(ns, name, vreplicas, placements)
+	vpod := NewVPod(ns, name, vreplicas, placements)
 	s.store.lock.Lock()
 	defer s.store.lock.Unlock()
 	s.store.vpods = append(s.store.vpods, vpod)
 	return vpod
 }
 
+func (s *VPodClient) Append(vpod scheduler.VPod) {
+	s.store.lock.Lock()
+	defer s.store.lock.Unlock()
+	s.store.vpods = append(s.store.vpods, vpod)
+}
+
 func (s *VPodClient) List() ([]scheduler.VPod, error) {
 	return s.lister()
 }
diff --git a/pkg/common/scheduler/testing/vpod.go b/pkg/common/scheduler/testing/vpod.go
index 4751fa46a..c8a33fa8c 100644
--- a/pkg/common/scheduler/testing/vpod.go
+++ b/pkg/common/scheduler/testing/vpod.go
@@ -27,7 +27,7 @@ type sampleVPod struct {
 	placements []duckv1alpha1.Placement
 }
 
-func newVPod(ns, name string, vreplicas int32, placements []duckv1alpha1.Placement) *sampleVPod {
+func NewVPod(ns, name string, vreplicas int32, placements []duckv1alpha1.Placement) *sampleVPod {
 	return &sampleVPod{
 		key: types.NamespacedName{
 			Namespace: ns,
diff --git a/vendor/knative.dev/eventing/pkg/adapter/v2/test/test_client.go b/vendor/knative.dev/eventing/pkg/adapter/v2/test/test_client.go
index 395c66c33..0f27d0533 100644
--- a/vendor/knative.dev/eventing/pkg/adapter/v2/test/test_client.go
+++ b/vendor/knative.dev/eventing/pkg/adapter/v2/test/test_client.go
@@ -28,13 +28,33 @@ import (
 )
 
 type TestCloudEventsClient struct {
-	lock  sync.Mutex
-	sent  []cloudevents.Event
-	delay time.Duration
+	lock          sync.Mutex
+	sent          []cloudevents.Event
+	delay         time.Duration
+	resultSend    []protocol.Result
+	resultRequest []struct {
+		event  *event.Event
+		result protocol.Result
+	}
 }
 
 var _ cloudevents.Client = (*TestCloudEventsClient)(nil)
 
+// Send_AppendResult will enqueue a response for the following Send call.
+// For testing.
+func (c *TestCloudEventsClient) Send_AppendResult(r protocol.Result) {
+	c.resultSend = append(c.resultSend, r)
+}
+
+// Request_AppendResult will enqueue a response for the following Request call.
+// For testing.
+func (c *TestCloudEventsClient) Request_AppendResult(e *event.Event, r protocol.Result) {
+	c.resultRequest = append(c.resultRequest, struct {
+		event  *event.Event
+		result protocol.Result
+	}{event: e, result: r})
+}
+
 func (c *TestCloudEventsClient) Send(ctx context.Context, out event.Event) protocol.Result {
 	if c.delay > 0 {
 		time.Sleep(c.delay)
@@ -43,6 +63,11 @@ func (c *TestCloudEventsClient) Send(ctx context.Context, out event.Event) proto
 	defer c.lock.Unlock()
 	// TODO: improve later.
 	c.sent = append(c.sent, out)
+	if len(c.resultSend) != 0 {
+		resp := c.resultSend[0]
+		c.resultSend = c.resultSend[1:]
+		return resp
+	}
 	return http.NewResult(200, "%w", protocol.ResultACK)
 }