Repository: knative-sandbox/eventing-kafka
Pull request #496 (2021-04-01 18:33:03)

[v0.20] [Consolidated Channel] Backport of #487 
<!-- Please include the 'why' behind your changes if no issue exists -->

## Proposed Changes

- 🐛 Fix bug of probing the dispatcher for subscription readiness after dispatcher pods change.
- Backport of #487 

<!--
If this change has user-visible impact, follow the instructions below.
Examples include:

- 🎁 Add new feature
- 🐛 Fix bug
- 🧽 Update or clean up current behavior
- 🗑️ Remove feature or internal logic

Otherwise delete the rest of this template.
-->

**Release Note**

<!--
🗒️ If this change has user-visible impact, write a release note in the block
below. Include the string "action required" if additional action is required of
users switching to the new release, for example in case of a breaking change.

Write as if you are speaking to users, not other Knative contributors. If this
change has no user-visible impact, no release-note is needed.
-->

```release-note
🐛 A bug fix was back ported to the consolidated channel so that subscriptions becomes eventually ready even when dispatchers scale up or down. 

```

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 **phamilton** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @phamilton` 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/release-0.20/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":["phamilton"]} -->
knative-prow-robot: @devguyio: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-knative-sandbox-eventing-kafka-integration-test-mt-source | 3ac4aff8fdca63f8e8059d5f693c311a3c293a77 | [link](https://prow.knative.dev/view/gs/knative-prow/pr-logs/pull/knative-sandbox_eventing-kafka/496/pull-knative-sandbox-eventing-kafka-integration-test-mt-source/1377690498147291136) | `/test pull-knative-sandbox-eventing-kafka-integration-test-mt-source`

[Full PR test history](https://gubernator.knative.dev/pr/knative-sandbox_eventing-kafka/496). [Your PR dashboard](https://gubernator.knative.dev/pr/devguyio).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
devguyio: please ignore failures in the `pull-knative-sandbox-eventing-kafka-integration-test-mt-source` since it was meant from 0.21.0 onward.  
knative-metrics-robot: The following is the coverage report on the affected files.
Say `/test pull-knative-sandbox-eventing-kafka-go-coverage` to re-run this coverage report

File | Old Coverage | New Coverage | Delta
---- |:------------:|:------------:|:-----:
[pkg/channel/consolidated/reconciler/controller/kafkachannel.go](https://storage.cloud.google.com/knative-prow/pr-logs/pull/knative-sandbox_eventing-kafka/496/pull-knative-sandbox-eventing-kafka-go-coverage/1377690498235371520/artifacts/line-cov.html#file1) | 55.3% | 55.9% | 0.6

codecov[bot]: # [Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/496?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`release-0.20@c14fb2b`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `n/a`.

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

```diff
@@               Coverage Diff               @@
##             release-0.20     #496   +/-   ##
===============================================
  Coverage                ?   75.67%           
===============================================
  Files                   ?      116           
  Lines                   ?     4642           
  Branches                ?        0           
===============================================
  Hits                    ?     3513           
  Misses                  ?      929           
  Partials                ?      200           
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/496?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/496?src=pr&el=footer). Last update [c14fb2b...3ac4aff](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/496?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

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

This pull-request has been approved by: *<a href="https://github.com/knative-sandbox/eventing-kafka/pull/496#pullrequestreview-626588863" title="LGTM">lberk</a>*

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

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-sandbox/eventing-kafka/blob/release-0.20/OWNERS)~~ [lberk]

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":[]} -->

Patch:
diff --git a/pkg/channel/consolidated/reconciler/controller/controller.go b/pkg/channel/consolidated/reconciler/controller/controller.go
index b1b9842ae..a06d01c9b 100644
--- a/pkg/channel/consolidated/reconciler/controller/controller.go
+++ b/pkg/channel/consolidated/reconciler/controller/controller.go
@@ -22,7 +22,6 @@ import (
 	"github.com/kelseyhightower/envconfig"
 	"go.uber.org/zap"
 	corev1 "k8s.io/api/core/v1"
-	v1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
@@ -30,6 +29,7 @@ import (
 
 	"knative.dev/eventing-kafka/pkg/apis/messaging/v1beta1"
 	"knative.dev/eventing-kafka/pkg/channel/consolidated/status"
+	kafkamessagingv1beta1 "knative.dev/eventing-kafka/pkg/client/informers/externalversions/messaging/v1beta1"
 	kafkaChannelClient "knative.dev/eventing-kafka/pkg/client/injection/client"
 	"knative.dev/eventing-kafka/pkg/client/injection/informers/messaging/v1beta1/kafkachannel"
 	kafkaChannelReconciler "knative.dev/eventing-kafka/pkg/client/injection/reconciler/messaging/v1beta1/kafkachannel"
@@ -110,7 +110,7 @@ func NewController(
 	statusProber.Start(ctx.Done())
 	// Get and Watch the Kakfa config map and dynamically update Kafka configuration.
 	if _, err := kubeclient.Get(ctx).CoreV1().ConfigMaps(system.Namespace()).Get(ctx, "config-kafka", metav1.GetOptions{}); err == nil {
-		cmw.Watch("config-kafka", func(configMap *v1.ConfigMap) {
+		cmw.Watch("config-kafka", func(configMap *corev1.ConfigMap) {
 			r.updateKafkaConfig(ctx, configMap)
 		})
 	} else if !apierrors.IsNotFound(err) {
@@ -158,17 +158,22 @@ func NewController(
 		),
 		Handler: cache.ResourceEventHandlerFuncs{
 			// Cancel probing when a Pod is deleted
-			DeleteFunc: func(obj interface{}) {
-				pod, ok := obj.(*corev1.Pod)
-				if ok && pod != nil {
-					logger.Debugw("Dispatcher pod deleted. Canceling pod probing.",
-						zap.String("pod", pod.GetName()))
-					statusProber.CancelPodProbing(*pod)
-					impl.GlobalResync(kafkaChannelInformer.Informer())
-				}
-			},
+			DeleteFunc: getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Delete"),
+			AddFunc:    getPodInformerEventHandler(ctx, logger, statusProber, impl, kafkaChannelInformer, "Add"),
 		},
 	})
 
 	return impl
 }
+
+func getPodInformerEventHandler(ctx context.Context, logger *zap.SugaredLogger, statusProber *status.Prober, impl *controller.Impl, kafkaChannelInformer kafkamessagingv1beta1.KafkaChannelInformer, handlerType string) func(obj interface{}) {
+	return func(obj interface{}) {
+		pod, ok := obj.(*corev1.Pod)
+		if ok && pod != nil {
+			logger.Debugw("%s pods. Refreshing pod probing.", handlerType,
+				zap.String("pod", pod.GetName()))
+			statusProber.RefreshPodProbing(ctx)
+			impl.GlobalResync(kafkaChannelInformer.Informer())
+		}
+	}
+}
diff --git a/pkg/channel/consolidated/reconciler/controller/kafkachannel.go b/pkg/channel/consolidated/reconciler/controller/kafkachannel.go
index 808512fd7..a43a6d8a1 100644
--- a/pkg/channel/consolidated/reconciler/controller/kafkachannel.go
+++ b/pkg/channel/consolidated/reconciler/controller/kafkachannel.go
@@ -163,6 +163,7 @@ func (r *Reconciler) ReconcileKind(ctx context.Context, kc *v1beta1.KafkaChannel
 		kc.Status.MarkConfigFailed("InvalidConfiguration", "Unable to build Kafka admin client for channel %s: %v", kc.Name, err)
 		return err
 	}
+	defer kafkaClusterAdmin.Close()
 
 	kc.Status.MarkConfigTrue()
 
@@ -226,7 +227,6 @@ func (r *Reconciler) ReconcileKind(ctx context.Context, kc *v1beta1.KafkaChannel
 	// Reconcile the k8s service representing the actual Channel. It points to the Dispatcher service via ExternalName
 	svc, err := r.reconcileChannelService(ctx, dispatcherNamespace, kc)
 	if err != nil {
-
 		return err
 	}
 	kc.Status.MarkChannelServiceTrue()
@@ -234,15 +234,9 @@ func (r *Reconciler) ReconcileKind(ctx context.Context, kc *v1beta1.KafkaChannel
 		Scheme: "http",
 		Host:   network.GetServiceHostname(svc.Name, svc.Namespace),
 	})
-	err = r.setupSubscriptionStatusWatcher(ctx, kc)
+	err = r.reconcileSubscribers(ctx, kc)
 	if err != nil {
-		logger.Errorw("error setting up some subscription status watchers", zap.Error(err))
-	}
-	// close the connection
-	err = kafkaClusterAdmin.Close()
-	if err != nil {
-		logger.Errorw("Error closing the connection", zap.Error(err))
-		return err
+		return fmt.Errorf("error reconciling subscribers %v", err)
 	}
 
 	// Ok, so now the Dispatcher Deployment & Service have been created, we're golden since the
@@ -250,10 +244,9 @@ func (r *Reconciler) ReconcileKind(ctx context.Context, kc *v1beta1.KafkaChannel
 	return newReconciledNormal(kc.Namespace, kc.Name)
 }
 
-func (r *Reconciler) setupSubscriptionStatusWatcher(ctx context.Context, ch *v1beta1.KafkaChannel) error {
+func (r *Reconciler) reconcileSubscribers(ctx context.Context, ch *v1beta1.KafkaChannel) error {
 	after := ch.DeepCopy()
 	after.Status.Subscribers = make([]v1.SubscriberStatus, 0)
-
 	for _, s := range ch.Spec.Subscribers {
 		if r, _ := r.statusManager.IsReady(ctx, *ch, s); r {
 			logging.FromContext(ctx).Debugw("marking subscription", zap.Any("subscription", s))
diff --git a/pkg/channel/consolidated/reconciler/controller/lister.go b/pkg/channel/consolidated/reconciler/controller/lister.go
index db39a2b25..07fbc6816 100644
--- a/pkg/channel/consolidated/reconciler/controller/lister.go
+++ b/pkg/channel/consolidated/reconciler/controller/lister.go
@@ -70,7 +70,6 @@ func (t *DispatcherPodsLister) ListProbeTargets(ctx context.Context, kc v1beta1.
 	return &status.ProbeTarget{
 		PodIPs:  sets.NewString(readyIPs...),
 		PodPort: "8081",
-		Port:    "8081",
 		URL:     u,
 	}, nil
 }
diff --git a/pkg/channel/consolidated/status/status.go b/pkg/channel/consolidated/status/status.go
index b42f7bf0f..4c21e3253 100644
--- a/pkg/channel/consolidated/status/status.go
+++ b/pkg/channel/consolidated/status/status.go
@@ -63,10 +63,10 @@ type targetState struct {
 	pendingCount atomic.Int32
 	// readyCount is the number of pods that have the subscription ready
 	readyPartitions sets.Int
-	initialCount    int
+	probedPods      sets.String
 	lastAccessed    time.Time
-
-	cancel func()
+	ready           bool
+	cancel          func()
 }
 
 // podState represents the probing state of a Pod (for a specific subscription)
@@ -97,7 +97,6 @@ type workItem struct {
 type ProbeTarget struct {
 	PodIPs  sets.String
 	PodPort string
-	Port    string
 	URL     *url.URL
 }
 
@@ -107,7 +106,7 @@ type ProbeTargetLister interface {
 	ListProbeTargets(ctx context.Context, ch messagingv1beta1.KafkaChannel) (*ProbeTarget, error)
 }
 
-// Manager provides a way to check if an Ingress is ready
+// Manager provides a way to check if a Subscription is ready
 type Manager interface {
 	IsReady(ctx context.Context, ch messagingv1beta1.KafkaChannel, sub eventingduckv1.SubscriberSpec) (bool, error)
 	CancelProbing(sub eventingduckv1.SubscriberSpec)
@@ -160,40 +159,52 @@ func NewProber(
 }
 
 func (m *Prober) checkReadiness(state *targetState) bool {
-	consumers := int32(state.initialCount)
+	state.readyLock.Lock()
+	defer state.readyLock.Unlock()
 	partitions := state.ch.Spec.NumPartitions
 	m.logger.Debugw("Checking subscription readiness",
-		zap.Any("initial probed consumers", consumers),
+		zap.Any("subscription", state.sub.UID),
+		zap.Any("channel", state.ch.Name),
+		zap.Any("pod ips", state.probedPods),
 		zap.Any("channel partitions", partitions),
 		zap.Any("ready partitions", state.readyPartitions.List()),
 	)
-	return state.readyPartitions.Len() == int(partitions)
+	if !state.ready {
+		state.ready = state.readyPartitions.Len() == int(partitions)
+	}
+	return state.ready
 }
 
 func (m *Prober) IsReady(ctx context.Context, ch messagingv1beta1.KafkaChannel, sub eventingduckv1.SubscriberSpec) (bool, error) {
-	subscriptionKey := sub.UID
 	logger := logging.FromContext(ctx)
 
-	if ready, ok := func() (bool, bool) {
+	// Get the probe targets
+	target, err := m.targetLister.ListProbeTargets(ctx, ch)
+	if err != nil {
+		logger.Errorw("Error listing probe targets", zap.Error(err),
+			zap.Any("subscription", sub.UID))
+		return false, err
+	}
+	// get the state while locking for very short scope
+	state, ok := func() (*targetState, bool) {
 		m.mu.Lock()
 		defer m.mu.Unlock()
-		if state, ok := m.targetStates[subscriptionKey]; ok {
-			if state.sub.Generation == sub.Generation {
-				state.lastAccessed = time.Now()
-				logger.Debugw("Subscription is cached. Checking readiness",
-					zap.Any("subscription", sub.UID))
-				return m.checkReadiness(state), true
-			}
-
-			// Cancel the polling for the outdated version
-			state.cancel()
-			delete(m.targetStates, subscriptionKey)
+		s, o := m.targetStates[sub.UID]
+		return s, o
+	}()
+	if ok {
+		if !isOutdatedTargetState(state, sub, target.PodIPs) {
+			return m.checkReadiness(state), nil
 		}
-		return false, false
-	}(); ok {
-		return ready, nil
+		m.ejectStateUnsafe(sub)
 	}
+	m.probeTarget(ctx, ch, sub, target)
+	return false, nil
+}
 
+func (m *Prober) probeTarget(ctx context.Context, ch messagingv1beta1.KafkaChannel, sub eventingduckv1.SubscriberSpec, target *ProbeTarget) {
+	subscriptionKey := sub.UID
+	logger := logging.FromContext(ctx)
 	subCtx, cancel := context.WithCancel(context.Background())
 	subscriptionState := &targetState{
 		sub:          sub,
@@ -202,14 +213,7 @@ func (m *Prober) IsReady(ctx context.Context, ch messagingv1beta1.KafkaChannel,
 		cancel:       cancel,
 	}
 
-	// Get the probe targets and group them by IP
-	target, err := m.targetLister.ListProbeTargets(ctx, ch)
-	if err != nil {
-		logger.Errorw("Error listing probe targets", zap.Error(err),
-			zap.Any("subscription", sub.UID))
-		return false, err
-	}
-
+	// Group the probe targets by IP
 	workItems := make(map[string][]*workItem)
 	for ip := range target.PodIPs {
 		workItems[ip] = append(workItems[ip], &workItem{
@@ -221,7 +225,7 @@ func (m *Prober) IsReady(ctx context.Context, ch messagingv1beta1.KafkaChannel,
 		})
 	}
 
-	subscriptionState.initialCount = target.PodIPs.Len()
+	subscriptionState.probedPods = target.PodIPs
 	subscriptionState.pendingCount.Store(int32(len(workItems)))
 	subscriptionState.readyPartitions = sets.Int{}
 
@@ -281,7 +285,6 @@ func (m *Prober) IsReady(ctx context.Context, ch messagingv1beta1.KafkaChannel,
 		defer m.mu.Unlock()
 		m.targetStates[subscriptionKey] = subscriptionState
 	}()
-	return false, nil
 }
 
 // Start starts the Manager background operations
@@ -317,6 +320,11 @@ func (m *Prober) Start(done <-chan struct{}) chan struct{} {
 func (m *Prober) CancelProbing(sub eventingduckv1.SubscriberSpec) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+	m.ejectStateUnsafe(sub)
+}
+
+// ejectStateUnsafe ejects a state from Cache, it's not safe for concurrent access and is meant for internal use only under proper locking.
+func (m *Prober) ejectStateUnsafe(sub eventingduckv1.SubscriberSpec) {
 	if state, ok := m.targetStates[sub.UID]; ok {
 		m.logger.Debugw("Canceling state", zap.Any("subscription", sub))
 		state.cancel()
@@ -335,6 +343,35 @@ func (m *Prober) CancelPodProbing(pod corev1.Pod) {
 	}
 }
 
+// RefreshPodProbing lists probe targets and invalidates any in-flight (non-ready) states whose initial probed targets changed from the
+// current ones.
+func (m *Prober) RefreshPodProbing(ctx context.Context) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	logger := logging.FromContext(ctx)
+	for _, state := range m.targetStates {
+		if !m.checkReadiness(state) {
+			// This is an in-flight state
+			sub := state.sub
+			ch := state.ch
+			// Get the probe targets
+			target, err := m.targetLister.ListProbeTargets(ctx, ch)
+			if err != nil {
+				logger.Errorw("Error listing probe targets", zap.Error(err),
+					zap.Any("subscription", sub.UID))
+				return
+			}
+			m.ejectStateUnsafe(sub)
+			func() {
+				// probeTarget requires an unlocked mutex.
+				m.mu.Unlock()
+				defer m.mu.Lock()
+				m.probeTarget(ctx, ch, sub, target)
+			}()
+		}
+	}
+}
+
 // processWorkItem processes a single work item from workQueue.
 // It returns false when there is no more items to process, true otherwise.
 func (m *Prober) processWorkItem() bool {
@@ -386,8 +423,10 @@ func (m *Prober) processWorkItem() bool {
 	if err != nil || !ok {
 		// In case of error, enqueue for retry
 		m.workQueue.AddRateLimited(obj)
-		item.logger.Debugw("Probing of %s failed, IP: %s:%s, ready: %t, error: %v (depth: %d)",
-			item.url, item.podIP, item.podPort, ok, err, m.workQueue.Len())
+		item.logger.Debugw("Probing failed",
+			zap.Any("url", item.url), zap.Any("IP", item.podIP),
+			zap.Any("port", item.podPort), zap.Bool("ready", ok), zap.Error(err),
+			zap.Int("depth", m.workQueue.Len()))
 	} else {
 		m.onProbingSuccess(item.targetStates, item.podState)
 	}
@@ -399,7 +438,7 @@ func (m *Prober) onProbingSuccess(subscriptionState *targetState, podState *podS
 	if podState.pendingCount.Dec() == 0 {
 		// Unlock the goroutine blocked on <-podCtx.Done()
 		podState.cancel()
-		// This is the last pod being successfully probed, the subscription is ready
+		// This is the last pod being successfully probed, the subscription might ready
 		if m.checkReadiness(subscriptionState) {
 			subscriptionState.cancel()
 			m.readyCallback(subscriptionState.ch, subscriptionState.sub)
@@ -414,13 +453,13 @@ func (m *Prober) onProbingCancellation(subscriptionState *targetState, podState
 			// Probing succeeded, nothing to do
 			return
 		}
-
 		// Attempt to set pendingCount to 0.
 		if podState.pendingCount.CAS(pendingCount, 0) {
-			// This is the last pod being successfully probed, the subscription is ready
+			// This is the last pod being successfully probed, the subscription might be ready
 			if subscriptionState.pendingCount.Dec() == 0 {
-				subscriptionState.cancel()
-				m.readyCallback(subscriptionState.ch, subscriptionState.sub)
+				if m.checkReadiness(subscriptionState) {
+					m.readyCallback(subscriptionState.ch, subscriptionState.sub)
+				}
 			}
 			return
 		}
@@ -430,7 +469,8 @@ func (m *Prober) onProbingCancellation(subscriptionState *targetState, podState
 func (m *Prober) probeVerifier(item *workItem) prober.Verifier {
 	return func(r *http.Response, b []byte) (bool, error) {
 		m.logger.Debugw("Verifying response", zap.Int("status code", r.StatusCode),
-			zap.ByteString("body", b))
+			zap.ByteString("body", b), zap.Any("subscription", item.targetStates.sub.UID),
+			zap.Any("channel", item.targetStates.ch))
 		switch r.StatusCode {
 		case http.StatusOK:
 			var subscriptions = make(map[string][]int)
@@ -467,6 +507,14 @@ func (m *Prober) probeVerifier(item *workItem) prober.Verifier {
 	}
 }
 
+// A target state is outdated if the generation is different or if the target IPs change before it becomes
+// ready.
+func isOutdatedTargetState(state *targetState, sub eventingduckv1.SubscriberSpec, podIPs sets.String) bool {
+	state.readyLock.RLock()
+	defer state.readyLock.RUnlock()
+	return state.sub.Generation != sub.Generation || (!state.ready && !state.probedPods.Equal(podIPs))
+}
+
 // deepCopy copies a URL into a new one
 func deepCopy(in *url.URL) *url.URL {
 	// Safe to ignore the error since this is a deep copy


Test patch:
diff --git a/pkg/channel/consolidated/status/status_test.go b/pkg/channel/consolidated/status/status_test.go
index c9d0efe60..4d99781f3 100644
--- a/pkg/channel/consolidated/status/status_test.go
+++ b/pkg/channel/consolidated/status/status_test.go
@@ -39,49 +39,12 @@ import (
 	"knative.dev/pkg/apis"
 )
 
-var (
-	channelTemplate = &v1beta1.KafkaChannel{
-		ObjectMeta: metav1.ObjectMeta{
-			Namespace: "default",
-			Name:      "chan4prober",
-		},
-		Spec: v1beta1.KafkaChannelSpec{
-			NumPartitions:     1,
-			ReplicationFactor: 1,
-		},
-	}
-	subscriptionTemplate = eventingduckv1.SubscriberSpec{
-		UID:           types.UID("90713ffd-f527-42bf-b158-57630b68ebe2"),
-		Generation:    1,
-		SubscriberURI: getURL("http://subscr.ns.local"),
-	}
-)
-
-const dispatcherReadySubHeader = "K-Subscriber-Status"
-
-func getURL(s string) *apis.URL {
-	u, _ := apis.ParseURL(s)
-	return u
+var channelObjectMeta = metav1.ObjectMeta{
+	Namespace: "default",
+	Name:      "chan4prober",
 }
 
-func handleProbe(t *testing.T) func(http.ResponseWriter, *http.Request) {
-	return func(w http.ResponseWriter, r *http.Request) {
-		channelRefName := channelTemplate.ObjectMeta.Name
-		channelRefNamespace := channelTemplate.ObjectMeta.Namespace
-		var subscriptions = map[string][]int{
-			string(subscriptionTemplate.UID): {0},
-		}
-		w.Header().Set(dispatcherReadySubHeader, channelRefName)
-		jsonResult, err := json.Marshal(subscriptions)
-		if err != nil {
-			t.Fatalf("Error marshalling json for sub-status channelref: %s/%s, %v", channelRefNamespace, channelRefName, err)
-		}
-		_, err = w.Write(jsonResult)
-		if err != nil {
-			t.Fatalf("Error writing jsonResult to serveHTTP writer: %v", err)
-		}
-	}
-}
+const dispatcherReadySubHeader = "K-Subscriber-Status"
 
 type ReadyPair struct {
 	c v1beta1.KafkaChannel
@@ -89,46 +52,58 @@ type ReadyPair struct {
 }
 
 func TestProbeSinglePod(t *testing.T) {
-	var succeed atomic.Bool
-
-	ch := channelTemplate.DeepCopy()
-	sub := subscriptionTemplate.DeepCopy()
+	ch := getChannel(1)
+	sub := getSubscription()
+	var subscriptions = map[string][]int{
+		string(sub.UID): {0},
+	}
 
-	probeHandler := http.HandlerFunc(handleProbe(t))
+	// This should be called when we want the dispatcher to return a successful result
+	successHandler := http.HandlerFunc(readyJSONHandler(t, subscriptions))
 
 	// Probes only succeed if succeed is true
+	var succeed atomic.Bool
+
+	// This is a latch channel that will lock the handler goroutine until we drain it
 	probeRequests := make(chan *http.Request)
-	finalHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler := func(w http.ResponseWriter, r *http.Request) {
 		probeRequests <- r
 		if !succeed.Load() {
 			w.WriteHeader(http.StatusNotFound)
 			return
 		}
 
-		probeHandler.ServeHTTP(w, r)
-	})
+		successHandler.ServeHTTP(w, r)
+	}
 
-	ts := httptest.NewServer(finalHandler)
+	ts := getDispatcherServer(handler)
 	defer ts.Close()
-	tsURL, err := url.Parse(ts.URL)
-	if err != nil {
-		t.Fatalf("Failed to parse URL %q: %v", ts.URL, err)
-	}
-	port, err := strconv.Atoi(tsURL.Port())
-	if err != nil {
-		t.Fatalf("Failed to parse port %q: %v", tsURL.Port(), err)
+
+	lister := fakeProbeTargetLister{
+		target: getTargetLister(t, ts.URL),
 	}
-	hostname := tsURL.Hostname()
 
-	ready := make(chan *ReadyPair)
+	prober, ready := getProber(t, &lister)
+
+	done := make(chan struct{})
+	cancelled := prober.Start(done)
+	defer func() {
+		close(done)
+		<-cancelled
+	}()
+
+	assertEventuallyReady(t, prober, ch, sub, ready, &succeed, &probeRequests)
+}
+
+func TestProbeListerFail(t *testing.T) {
+	ch := getChannel(1)
+	sub := getSubscription()
 
+	ready := make(chan *ReadyPair)
+	defer close(ready)
 	prober := NewProber(
 		zaptest.NewLogger(t).Sugar(),
-		fakeProbeTargetLister{
-			PodIPs:  sets.NewString(hostname),
-			PodPort: strconv.Itoa(port),
-			URL:     tsURL,
-		},
+		notFoundLister{},
 		func(c v1beta1.KafkaChannel, s eventingduckv1.SubscriberSpec) {
 			ready <- &ReadyPair{
 				c,
@@ -136,6 +111,73 @@ func TestProbeSinglePod(t *testing.T) {
 			}
 		})
 
+	// If we can't list, this  must fail and return false
+	ok, err := prober.IsReady(context.Background(), *ch, *sub)
+	if err == nil {
+		t.Fatal("IsReady returned unexpected success")
+	}
+	if ok {
+		t.Fatal("IsReady() returned true")
+	}
+}
+
+func TestSucceedAfterRefreshPodProbing(t *testing.T) {
+	// We have a channel with three partitions
+	ch := getChannel(3)
+	sub := getSubscription()
+
+	// Dispatcher D1 will return only one ready partition
+	var subsD1 = map[string][]int{
+		string(sub.UID): {0},
+	}
+
+	// Dispatcher D2 will return the three ready partitions
+	var subsD2 = map[string][]int{
+		string(sub.UID): {0, 1, 2},
+	}
+
+	// The success handler for dispatcher D1, will return one partition only
+	successHandlerD1 := http.HandlerFunc(readyJSONHandler(t, subsD1))
+
+	// This is a latch channel that will lock the handler goroutine until we drain it
+	probeRequestsD1 := make(chan *http.Request)
+
+	handlerD1 := func(w http.ResponseWriter, r *http.Request) {
+		probeRequestsD1 <- r
+		successHandlerD1.ServeHTTP(w, r)
+	}
+
+	serverD1 := getDispatcherServer(handlerD1)
+	defer serverD1.Close()
+
+	// Probes only succeed if succeed is true
+	var succeed atomic.Bool
+
+	// The success handler for dispatcher D2, will return all three needed partitions
+	probeHandlerD2 := http.HandlerFunc(readyJSONHandler(t, subsD2))
+
+	// This is a latch channel that will lock the handler goroutine until we drain it
+	probeRequestsD2 := make(chan *http.Request)
+	handlerD2 := func(w http.ResponseWriter, r *http.Request) {
+		probeRequestsD2 <- r
+		if !succeed.Load() {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+
+		probeHandlerD2.ServeHTTP(w, r)
+	}
+
+	serverD2 := getDispatcherServer(handlerD2)
+	defer serverD2.Close()
+
+	// Initially, lister points to d1
+	lister := fakeProbeTargetLister{
+		target: getTargetLister(t, serverD1.URL),
+	}
+
+	prober, ready := getProber(t, &lister)
+
 	done := make(chan struct{})
 	cancelled := prober.Start(done)
 	defer func() {
@@ -143,6 +185,23 @@ func TestProbeSinglePod(t *testing.T) {
 		<-cancelled
 	}()
 
+	// Assert we're not ready.
+	assertNeverReady(t, prober, ch, sub, ready, &probeRequestsD1)
+
+	// Switch to new dispatcher
+	lister.target = getTargetLister(t, serverD2.URL)
+
+	// Assert we're still not ready
+	assertNeverReady(t, prober, ch, sub, ready, &probeRequestsD1)
+
+	// Refresh the pod probing, now the prober should probe the new disptacher
+	prober.RefreshPodProbing(context.Background())
+
+	// Assert that probing will be successful eventually
+	assertEventuallyReady(t, prober, ch, sub, ready, &succeed, &probeRequestsD2)
+}
+
+func assertNeverReady(t *testing.T, prober *Prober, ch *messagingv1beta1.KafkaChannel, sub *eventingduckv1.SubscriberSpec, ready chan *ReadyPair, probeRequests *chan *http.Request) {
 	// The first call to IsReady must succeed and return false
 	ok, err := prober.IsReady(context.Background(), *ch, *sub)
 	if err != nil {
@@ -152,21 +211,33 @@ func TestProbeSinglePod(t *testing.T) {
 		t.Fatal("IsReady() returned true")
 	}
 
+	// Just drain the requests in the channel to not block the handler
+	go func() {
+		for range *probeRequests {
+		}
+	}()
+
 	select {
 	case <-ready:
-		// Since succeed is still false and we don't return 200, the prober shouldn't be ready
+		// Prober shouldn't be ready
 		t.Fatal("Prober shouldn't be ready")
 	case <-time.After(1 * time.Second):
 		// Not ideal but it gives time to the prober to write to ready
 		break
 	}
+}
+
+func assertEventuallyReady(t *testing.T, prober *Prober, ch *messagingv1beta1.KafkaChannel, sub *eventingduckv1.SubscriberSpec, ready chan *ReadyPair, succeed *atomic.Bool, probeRequests *chan *http.Request) {
+
+	// Since succeed is still false the prober shouldn't be ready
+	assertNeverReady(t, prober, ch, sub, ready, probeRequests)
 
-	// Make probes to hostB succeed
+	// Make probes succeed
 	succeed.Store(true)
 
 	// Just drain the requests in the channel to not block the handler
 	go func() {
-		for range probeRequests {
+		for range *probeRequests {
 		}
 	}()
 
@@ -178,37 +249,88 @@ func TestProbeSinglePod(t *testing.T) {
 	}
 }
 
-func TestProbeListerFail(t *testing.T) {
-	ch := channelTemplate.DeepCopy()
-	sub := subscriptionTemplate.DeepCopy()
+func getDispatcherServer(handler func(w http.ResponseWriter, r *http.Request)) *httptest.Server {
+	ts := httptest.NewServer(http.HandlerFunc(handler))
+	return ts
+}
+
+func getChannel(partitionsNum int32) *v1beta1.KafkaChannel {
+	return (&v1beta1.KafkaChannel{
+		ObjectMeta: channelObjectMeta,
+		Spec: v1beta1.KafkaChannelSpec{
+			NumPartitions:     partitionsNum,
+			ReplicationFactor: 1,
+		},
+	}).DeepCopy()
+}
+
+func getSubscription() *eventingduckv1.SubscriberSpec {
+	return (&eventingduckv1.SubscriberSpec{
+		UID:           types.UID("90713ffd-f527-42bf-b158-57630b68ebe2"),
+		Generation:    1,
+		SubscriberURI: getURL("http://subscr.ns.local"),
+	}).DeepCopy()
+}
 
+func getURL(s string) *apis.URL {
+	u, _ := apis.ParseURL(s)
+	return u
+}
+
+// readyJSONHandler is a factory for a handler which responds with a JSON of the ready subscriptions
+func readyJSONHandler(t *testing.T, subscriptions map[string][]int) func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		channelRefName := channelObjectMeta.Name
+		channelRefNamespace := channelObjectMeta.Namespace
+		w.Header().Set(dispatcherReadySubHeader, channelRefName)
+		jsonResult, err := json.Marshal(subscriptions)
+		if err != nil {
+			t.Fatalf("Error marshalling json for sub-status channelref: %s/%s, %v", channelRefNamespace, channelRefName, err)
+		}
+		_, err = w.Write(jsonResult)
+		if err != nil {
+			t.Fatalf("Error writing jsonResult to serveHTTP writer: %v", err)
+		}
+	}
+}
+
+func getTargetLister(t *testing.T, dURL string) *ProbeTarget {
+	tsURL, err := url.Parse(dURL)
+	if err != nil {
+		t.Fatalf("Failed to parse URL %q: %v", dURL, err)
+	}
+	port, err := strconv.Atoi(tsURL.Port())
+	if err != nil {
+		t.Fatalf("Failed to parse port %q: %v", tsURL.Port(), err)
+	}
+	hostname := tsURL.Hostname()
+	return &ProbeTarget{
+		PodIPs:  sets.NewString(hostname),
+		PodPort: strconv.Itoa(port),
+		URL:     tsURL,
+	}
+}
+
+func getProber(t *testing.T, lister ProbeTargetLister) (*Prober, chan *ReadyPair) {
 	ready := make(chan *ReadyPair)
-	defer close(ready)
 	prober := NewProber(
 		zaptest.NewLogger(t).Sugar(),
-		notFoundLister{},
+		lister,
 		func(c v1beta1.KafkaChannel, s eventingduckv1.SubscriberSpec) {
 			ready <- &ReadyPair{
 				c,
 				s,
 			}
 		})
-
-	// If we can't list, this  must fail and return false
-	ok, err := prober.IsReady(context.Background(), *ch, *sub)
-	if err == nil {
-		t.Fatal("IsReady returned unexpected success")
-	}
-	if ok {
-		t.Fatal("IsReady() returned true")
-	}
+	return prober, ready
 }
 
-type fakeProbeTargetLister ProbeTarget
+type fakeProbeTargetLister struct {
+	target *ProbeTarget
+}
 
 func (l fakeProbeTargetLister) ListProbeTargets(ctx context.Context, kc messagingv1beta1.KafkaChannel) (*ProbeTarget, error) {
-	t := ProbeTarget(l)
-	return &t, nil
+	return l.target, nil
 }
 
 type notFoundLister struct{}