Repository: knative/pkg
Pull request #1521 (2020-07-19 22:06:08)

Harden and reaarange the test for drain handler
/assign mattmoor

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

This pull-request has been approved by: *<a href="https://github.com/knative/pkg/pull/1521#" title="Author self-approved">vagababov</a>*

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

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:

- ~~[network/OWNERS](https://github.com/knative/pkg/blob/master/network/OWNERS)~~ [vagababov]

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":[]} -->
vagababov: ```
>rt -count=50
PASS
ok      knative.dev/pkg/network/handlers        11.396s

```
`rt == go test -race`

Patch:
diff --git a/network/handlers/drain.go b/network/handlers/drain.go
index 7bdeffb9d..b85516088 100644
--- a/network/handlers/drain.go
+++ b/network/handlers/drain.go
@@ -105,6 +105,5 @@ func (d *Drainer) reset() {
 func (d *Drainer) draining() bool {
 	d.RLock()
 	defer d.RUnlock()
-
 	return d.timer != nil
 }


Test patch:
diff --git a/network/handlers/drain_test.go b/network/handlers/drain_test.go
index 769ed0902..e73331b34 100644
--- a/network/handlers/drain_test.go
+++ b/network/handlers/drain_test.go
@@ -17,7 +17,6 @@ limitations under the License.
 package handlers
 
 import (
-	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -49,109 +48,96 @@ func TestDrainMechanics(t *testing.T) {
 	drainer.ServeHTTP(w, req)
 	drainer.ServeHTTP(w, req)
 
-	// Check for 200 OK
+	// Check for 200 OK.
 	resp := httptest.NewRecorder()
 	drainer.ServeHTTP(resp, probe)
 	if got, want := resp.Code, http.StatusOK; got != want {
-		t.Errorf("probe status = %d, wanted %d", got, want)
+		t.Errorf("Probe status = %d, wanted %d", got, want)
 	}
 
 	// Start to drain, and cancel the context when it returns.
-	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
 	go func() {
+		defer close(done)
 		drainer.Drain()
-		cancel()
 	}()
 
 	select {
 	case <-time.After(40 * time.Millisecond):
 		// Drain is blocking.
-	case <-ctx.Done():
+	case <-done:
 		t.Error("Drain terminated prematurely.")
 	}
 	// Now send a request to reset things.
 	drainer.ServeHTTP(w, req)
 
-	// Check for 400 shutting down
+	// Check for 503 as a probe response when shutting down.
 	resp = httptest.NewRecorder()
 	drainer.ServeHTTP(resp, probe)
 	if got, want := resp.Code, http.StatusServiceUnavailable; got != want {
-		t.Errorf("probe status = %d, wanted %d", got, want)
+		t.Errorf("Probe status = %d, wanted %d", got, want)
 	}
 
-	select {
-	case <-time.After(40 * time.Millisecond):
-		// Drain is blocking.
-	case <-ctx.Done():
-		t.Error("Drain terminated prematurely.")
-	}
-	// Now send a request to reset things.
-	drainer.ServeHTTP(w, req)
-
-	select {
-	case <-time.After(40 * time.Millisecond):
-		// Drain is blocking.
-	case <-ctx.Done():
-		t.Error("Drain terminated prematurely.")
-	}
-	// Now send a request to reset things.
-	drainer.ServeHTTP(w, req)
-
-	select {
-	case <-time.After(40 * time.Millisecond):
-		// Drain is blocking.
-	case <-ctx.Done():
-		t.Error("Drain terminated prematurely.")
+	for i := 0; i < 3; i++ {
+		select {
+		case <-time.After(40 * time.Millisecond):
+			// Drain is blocking.
+		case <-done:
+			t.Error("Drain terminated prematurely.")
+		}
+		// For the last one we don't want to reset the drain timer.
+		if i < 2 {
+			drainer.ServeHTTP(w, req)
+		}
 	}
 	// Probing does not reset the clock.
-	// Check for 500 shutting down
+	// Check for 503 on a probe when shutting down.
 	resp = httptest.NewRecorder()
 	drainer.ServeHTTP(resp, probe)
 	if got, want := resp.Code, http.StatusServiceUnavailable; got != want {
-		t.Errorf("probe status = %d, wanted %d", got, want)
+		t.Errorf("Probe status = %d, wanted %d", got, want)
 	}
 
 	// Big finish, test that multiple invocations of Drain all block.
-	ctx1, cancel1 := context.WithCancel(context.Background())
+	done1 := make(chan struct{})
 	go func() {
+		defer close(done1)
 		drainer.Drain()
-		cancel1()
 	}()
-	ctx2, cancel2 := context.WithCancel(context.Background())
+	done2 := make(chan struct{})
 	go func() {
+		defer close(done2)
 		drainer.Drain()
-		cancel2()
 	}()
-	ctx3, cancel3 := context.WithCancel(context.Background())
+	done3 := make(chan struct{})
 	go func() {
+		defer close(done3)
 		drainer.Drain()
-		cancel3()
 	}()
 
 	select {
-	case <-time.After(70 * time.Millisecond):
+	case <-time.After(80 * time.Millisecond):
 		t.Error("Timed out waiting for Drain to return.")
-
-	case <-ctx.Done():
-	case <-ctx1.Done():
-	case <-ctx2.Done():
-	case <-ctx3.Done():
+	case <-done:
+	case <-done1:
+	case <-done2:
+	case <-done3:
 		// Once the first context is cancelled, check that all of them are cancelled.
 	}
 
 	// Check that a 4th and final one after things complete finishes instantly.
-	ctx4, cancel4 := context.WithCancel(context.Background())
+	done4 := make(chan struct{})
 	go func() {
+		defer close(done4)
 		drainer.Drain()
-		cancel4()
 	}()
 
-	// Give the rest a short window to complete.
-	time.Sleep(time.Millisecond)
+	// Give the test a short window to launch and execute the go routine.
+	time.Sleep(5 * time.Millisecond)
 
-	for idx, ictx := range []context.Context{ctx, ctx1, ctx2, ctx3, ctx4} {
+	for idx, dch := range []chan struct{}{done, done1, done2, done3, done4} {
 		select {
-		case <-ictx.Done():
+		case <-dch:
 			// Should be done.
 		default:
 			t.Errorf("Drain[%d] did not complete.", idx)