fix TestServerCredsDispatch and stats test race (#1554)
diff --git a/stats/stats_test.go b/stats/stats_test.go
index d66485f..358909c 100644
--- a/stats/stats_test.go
+++ b/stats/stats_test.go
@@ -813,7 +813,9 @@
err: err,
}
+ h.mu.Lock()
checkConnStats(t, h.gotConn)
+ h.mu.Unlock()
checkServerStats(t, h.gotRPC, expect, checkFuncs)
}
diff --git a/test/end2end_test.go b/test/end2end_test.go
index dd0c343..0add563 100644
--- a/test/end2end_test.go
+++ b/test/end2end_test.go
@@ -4145,21 +4145,22 @@
}
type serverDispatchCred struct {
- ready chan struct{}
- rawConn net.Conn
+ rawConnCh chan net.Conn
}
func newServerDispatchCred() *serverDispatchCred {
return &serverDispatchCred{
- ready: make(chan struct{}),
+ rawConnCh: make(chan net.Conn, 1),
}
}
func (c *serverDispatchCred) ClientHandshake(ctx context.Context, addr string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
return rawConn, nil, nil
}
func (c *serverDispatchCred) ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
- c.rawConn = rawConn
- close(c.ready)
+ select {
+ case c.rawConnCh <- rawConn:
+ default:
+ }
return nil, nil, credentials.ErrConnDispatched
}
func (c *serverDispatchCred) Info() credentials.ProtocolInfo {
@@ -4172,8 +4173,7 @@
return nil
}
func (c *serverDispatchCred) getRawConn() net.Conn {
- <-c.ready
- return c.rawConn
+ return <-c.rawConnCh
}
func TestServerCredsDispatch(t *testing.T) {
@@ -4192,8 +4192,12 @@
}
defer cc.Close()
+ rawConn := cred.getRawConn()
+ // Give grpc a chance to see the error and potentially close the connection.
+ // And check that connection is not closed after that.
+ time.Sleep(100 * time.Millisecond)
// Check rawConn is not closed.
- if n, err := cred.getRawConn().Write([]byte{0}); n <= 0 || err != nil {
+ if n, err := rawConn.Write([]byte{0}); n <= 0 || err != nil {
t.Errorf("Read() = %v, %v; want n>0, <nil>", n, err)
}
}
