From 766f70d9cb8f977e9fc607080b46c8f85ba124af Mon Sep 17 00:00:00 2001
From: Dave Protasowski <dprotaso@gmail.com>
Date: Wed, 17 Aug 2022 19:38:48 -0400
Subject: [PATCH] fix double close of channel (#2574)

---
 network/handlers/drain.go      | 47 ++++++++++++++++------------------
 network/handlers/drain_test.go |  3 +++
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/network/handlers/drain.go b/network/handlers/drain.go
index 5252016f0..5eee2ea42 100644
--- a/network/handlers/drain.go
+++ b/network/handlers/drain.go
@@ -73,8 +73,11 @@ type Drainer struct {
 	// timer is used to orchestrate the drain.
 	timer timer
 
-	// used to synchronize callers of Drain and Reset
-	ch chan struct{}
+	// used to synchronize callers of Drain
+	drainCh chan struct{}
+
+	// used to synchronize Drain and Reset
+	resetCh chan struct{}
 
 	// HealthCheckUAPrefixes are the additional user agent prefixes that trigger the
 	// drainer's health check
@@ -118,8 +121,8 @@ func (d *Drainer) Drain() {
 	ch := func() chan struct{} {
 		d.Lock()
 		defer d.Unlock()
-		if d.ch != nil {
-			return d.ch
+		if d.drainCh != nil {
+			return d.drainCh
 		}
 
 		if d.QuietPeriod <= 0 {
@@ -127,32 +130,26 @@ func (d *Drainer) Drain() {
 		}
 
 		timer := newTimer(d.QuietPeriod)
-		ch := make(chan struct{})
+		drainCh := make(chan struct{})
+		resetCh := make(chan struct{})
 
 		go func() {
 			select {
-			case <-ch:
-				// closed by reset
+			case <-resetCh:
 			case <-timer.tickChan():
-				close(ch)
 			}
+			close(drainCh)
 		}()
 
-		d.ch = ch
+		d.drainCh = drainCh
+		d.resetCh = resetCh
 		d.timer = timer
-		return ch
+		return drainCh
 	}()
 
 	<-ch
 }
 
-func drainTimer(tc <-chan time.Time) {
-	select {
-	case <-tc:
-	default:
-	}
-}
-
 // isHealthcheckRequest validates if the request has a user agent that is for healthcheck
 func (d *Drainer) isHealthCheckRequest(r *http.Request) bool {
 	if network.IsKubeletProbe(r) {
@@ -175,16 +172,16 @@ func (d *Drainer) Reset() {
 	defer d.Unlock()
 
 	if d.timer != nil {
-		if d.timer.Stop() {
-			d.timer = nil
-		} else {
-			drainTimer(d.timer.tickChan())
-		}
+		d.timer.Stop()
+		d.timer = nil
 	}
 
-	if d.ch != nil {
-		close(d.ch)
-		d.ch = nil
+	if d.resetCh != nil {
+		close(d.resetCh)
+		d.resetCh = nil
+	}
+	if d.drainCh != nil {
+		d.drainCh = nil
 	}
 }
 
diff --git a/network/handlers/drain_test.go b/network/handlers/drain_test.go
index 22c55706b..082b8100d 100644
--- a/network/handlers/drain_test.go
+++ b/network/handlers/drain_test.go
@@ -542,4 +542,7 @@ func TestReset(t *testing.T) {
 	if diff > 50*time.Millisecond {
 		t.Error("expected to drain to wait QuietPeriod time after reset")
 	}
+
+	// Calling reset after a drain should succeed
+	d.Reset()
 }