Fix flaky TestCloseSemantics, use separate context per call

Currently, we share the same context for the whole test, but we should
instead be creating a separate context per call, which will help avoid
timeouts in the test.

close_test.go

@@ -275,8 +275,6 @@ func TestCloseStress(t *testing.T) {
 
 type closeSemanticsTest struct {
 	*testing.T
-
-	ctx      context.Context
 	isolated bool
 }
 
@@ -301,15 +299,16 @@ func (t *closeSemanticsTest) withNewClient(f func(ch *Channel)) {
 }
 
 func (t *closeSemanticsTest) startCall(from *Channel, to *Channel, method string) (*OutboundCall, error) {
+	ctx, _ := NewContext(time.Second)
 	var call *OutboundCall
 	var err error
 	toPeer := to.PeerInfo()
 	if t.isolated {
 		sc := from.GetSubChannel(toPeer.ServiceName, Isolated)
 		sc.Peers().Add(toPeer.HostPort)
-		call, err = sc.BeginCall(t.ctx, method, nil)
+		call, err = sc.BeginCall(ctx, method, nil)
 	} else {
-		call, err = from.BeginCall(t.ctx, toPeer.HostPort, toPeer.ServiceName, method, nil)
+		call, err = from.BeginCall(ctx, toPeer.HostPort, toPeer.ServiceName, method, nil)
 	}
 	return call, err
 }
@@ -340,7 +339,7 @@ func (t *closeSemanticsTest) callStream(from *Channel, to *Channel) <-chan struc
 	return c
 }
 
-func (t *closeSemanticsTest) runTest(ctx context.Context) {
+func (t *closeSemanticsTest) runTest() {
 	s1, s1C := t.makeServer("s1")
 	s2, s2C := t.makeServer("s2")
 
@@ -359,6 +358,7 @@ func (t *closeSemanticsTest) runTest(ctx context.Context) {
 	// Close s1, should no longer be able to call it.
 	s1.Close()
 	assert.Equal(t, ChannelStartClose, s1.State())
+
 	t.withNewClient(func(ch *Channel) {
 		assert.Error(t, t.call(ch, s1), "closed channel should not accept incoming calls")
 		require.NoError(t, t.call(ch, s2),
@@ -396,28 +396,20 @@ func TestCloseSemantics(t *testing.T) {
 	defer goroutines.VerifyNoLeaks(t, nil)
 	defer testutils.SetTimeout(t, 2*time.Second)()
 
-	ctx, cancel := NewContext(time.Second)
-	defer cancel()
-
-	ct := &closeSemanticsTest{t, ctx, false /* isolated */}
-	ct.runTest(ctx)
+	ct := &closeSemanticsTest{t, false /* isolated */}
+	ct.runTest()
 }
 
 func TestCloseSemanticsIsolated(t *testing.T) {
 	// We defer the check as we want it to run after the SetTimeout clears the timeout.
 	defer goroutines.VerifyNoLeaks(t, nil)
 	defer testutils.SetTimeout(t, 2*time.Second)()
 
-	ctx, cancel := NewContext(time.Second)
-	defer cancel()
-
-	ct := &closeSemanticsTest{t, ctx, true /* isolated */}
-	ct.runTest(ctx)
+	ct := &closeSemanticsTest{t, true /* isolated */}
+	ct.runTest()
 }
 
 func TestCloseSingleChannel(t *testing.T) {
-	ctx, cancel := NewContext(time.Second)
-	defer cancel()
 	ch := testutils.NewServer(t, nil)
 
 	var connected sync.WaitGroup
@@ -437,6 +429,9 @@ func TestCloseSingleChannel(t *testing.T) {
 		connected.Add(1)
 		completed.Add(1)
 		go func() {
+			ctx, cancel := NewContext(time.Second)
+			defer cancel()
+
 			peerInfo := ch.PeerInfo()
 			_, _, _, err := raw.Call(ctx, ch, peerInfo.HostPort, peerInfo.ServiceName, "echo", nil, nil)
 			assert.NoError(t, err, "Call failed")
@@ -458,9 +453,6 @@ func TestCloseSingleChannel(t *testing.T) {
 }
 
 func TestCloseOneSide(t *testing.T) {
-	ctx, cancel := NewContext(time.Second)
-	defer cancel()
-
 	ch1 := testutils.NewServer(t, &testutils.ChannelOpts{ServiceName: "client"})
 	ch2 := testutils.NewServer(t, &testutils.ChannelOpts{ServiceName: "server"})
 
@@ -477,6 +469,8 @@ func TestCloseOneSide(t *testing.T) {
 	})
 
 	go func() {
+		ctx, cancel := NewContext(time.Second)
+		defer cancel()
 		ch2Peer := ch2.PeerInfo()
 		_, _, _, err := raw.Call(ctx, ch1, ch2Peer.HostPort, ch2Peer.ServiceName, "echo", nil, nil)
 		assert.NoError(t, err, "Call failed")