From 1af2206bb4233558032e41b2361ac9de4ecc05b5 Mon Sep 17 00:00:00 2001
From: Arjan Bal <arjansbal@google.com>
Date: Tue, 17 Dec 2024 15:09:36 +0530
Subject: [PATCH] Move ejection tracking to sc wrapper

---
 .../balancer/outlierdetection/balancer.go     | 20 +------------------
 .../outlierdetection/subconn_wrapper.go       |  8 ++++++++
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/xds/internal/balancer/outlierdetection/balancer.go b/xds/internal/balancer/outlierdetection/balancer.go
index d2e08db637c9..e607f0785a6d 100644
--- a/xds/internal/balancer/outlierdetection/balancer.go
+++ b/xds/internal/balancer/outlierdetection/balancer.go
@@ -593,29 +593,11 @@ func (b *outlierDetectionBalancer) Target() string {
 func (b *outlierDetectionBalancer) handleSubConnUpdate(u *scUpdate) {
 	scw := u.scw
 	scw.clearHealthListener()
-
-	if scw.listener == nil {
-		return
-	}
-
-	// If the health listener is being used for ejection, forward the
-	// connectivity updates unconditionally.
-	if scw.healthListenerEnabled {
-		b.child.updateSubConnState(scw, u.state)
-		return
-	}
-
-	// Raw connectivity listener is being used for ejection.
-	if scw.ejected {
-		return
-	}
 	b.child.updateSubConnState(scw, u.state)
 }
 
 func (b *outlierDetectionBalancer) handleSubConnHealthUpdate(u *scHealthUpdate) {
-	if !u.scw.ejected {
-		b.child.updateSubConnHealthState(u.scw, u.state)
-	}
+	b.child.updateSubConnHealthState(u.scw, u.state)
 }
 
 // handleEjectedUpdate handles any SubConns that get ejected/unejected, and
diff --git a/xds/internal/balancer/outlierdetection/subconn_wrapper.go b/xds/internal/balancer/outlierdetection/subconn_wrapper.go
index 9f1ec959a80b..7c0a2c117bdf 100644
--- a/xds/internal/balancer/outlierdetection/subconn_wrapper.go
+++ b/xds/internal/balancer/outlierdetection/subconn_wrapper.go
@@ -140,6 +140,9 @@ func (scw *subConnWrapper) RegisterHealthListener(listener func(balancer.SubConn
 // sends updates the health listener.
 func (scw *subConnWrapper) updateSubConnHealthState(scs balancer.SubConnState) {
 	scw.latestHealthState = scs
+	if scw.ejected {
+		return
+	}
 	scw.mu.Lock()
 	defer scw.mu.Unlock()
 	if scw.healthListener != nil {
@@ -151,6 +154,11 @@ func (scw *subConnWrapper) updateSubConnHealthState(scs balancer.SubConnState) {
 // unejection and updates the raw connectivity listener.
 func (scw *subConnWrapper) updateSubConnConnectivityState(scs balancer.SubConnState) {
 	scw.latestRawConnectivityState = scs
+	// If the raw connectivity listener is used for ejection, and the SubConn is
+	// ejected, don't send the update.
+	if scw.ejected && !scw.healthListenerEnabled {
+		return
+	}
 	if scw.listener != nil {
 		scw.listener(scs)
 	}