ceph-csi/vendor/google.golang.org/grpc/test/gracefulstop_test.go

216 lines
5.6 KiB
Go
Raw Normal View History

2018-01-09 18:57:14 +00:00
/*
*
* Copyright 2017 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package test
import (
"fmt"
"net"
"sync"
"testing"
"time"
"golang.org/x/net/context"
"google.golang.org/grpc"
"google.golang.org/grpc/test/leakcheck"
testpb "google.golang.org/grpc/test/grpc_testing"
)
type delayListener struct {
net.Listener
closeCalled chan struct{}
acceptCalled chan struct{}
allowCloseCh chan struct{}
cc *delayConn
dialed bool
}
func (d *delayListener) Accept() (net.Conn, error) {
select {
case <-d.acceptCalled:
// On the second call, block until closed, then return an error.
<-d.closeCalled
<-d.allowCloseCh
return nil, fmt.Errorf("listener is closed")
default:
close(d.acceptCalled)
return d.Listener.Accept()
}
}
func (d *delayListener) allowClose() {
close(d.allowCloseCh)
}
func (d *delayListener) Close() error {
close(d.closeCalled)
go func() {
<-d.allowCloseCh
d.Listener.Close()
}()
return nil
}
func (d *delayListener) allowClientRead() {
d.cc.allowRead()
}
func (d *delayListener) Dial(to time.Duration) (net.Conn, error) {
if d.dialed {
// Only hand out one connection (net.Dial can return more even after the
// listener is closed). This is not thread-safe, but Dial should never be
// called concurrently in this environment.
return nil, fmt.Errorf("no more conns")
}
d.dialed = true
c, err := net.DialTimeout("tcp", d.Listener.Addr().String(), to)
if err != nil {
return nil, err
}
d.cc = &delayConn{Conn: c, blockRead: make(chan struct{})}
return d.cc, nil
}
func (d *delayListener) clientWriteCalledChan() <-chan struct{} {
return d.cc.writeCalledChan()
}
type delayConn struct {
net.Conn
blockRead chan struct{}
mu sync.Mutex
writeCalled chan struct{}
}
func (d *delayConn) writeCalledChan() <-chan struct{} {
d.mu.Lock()
defer d.mu.Unlock()
d.writeCalled = make(chan struct{})
return d.writeCalled
}
func (d *delayConn) allowRead() {
close(d.blockRead)
}
func (d *delayConn) Read(b []byte) (n int, err error) {
<-d.blockRead
return d.Conn.Read(b)
}
func (d *delayConn) Write(b []byte) (n int, err error) {
d.mu.Lock()
if d.writeCalled != nil {
close(d.writeCalled)
d.writeCalled = nil
}
d.mu.Unlock()
return d.Conn.Write(b)
}
func TestGracefulStop(t *testing.T) {
defer leakcheck.Check(t)
// This test ensures GracefulStop cannot race and break RPCs on new
// connections created after GracefulStop was called but before
// listener.Accept() returns a "closing" error.
//
// Steps of this test:
// 1. Start Server
// 2. GracefulStop() Server after listener's Accept is called, but don't
// allow Accept() to exit when Close() is called on it.
// 3. Create a new connection to the server after listener.Close() is called.
// Server will want to send a GoAway on the new conn, but we delay client
// reads until 5.
// 4. Send an RPC on the new connection.
// 5. Allow the client to read the GoAway. The RPC should complete
// successfully.
lis, err := net.Listen("tcp", "localhost:0")
if err != nil {
t.Fatalf("Error listenening: %v", err)
}
dlis := &delayListener{
Listener: lis,
acceptCalled: make(chan struct{}),
closeCalled: make(chan struct{}),
allowCloseCh: make(chan struct{}),
}
d := func(_ string, to time.Duration) (net.Conn, error) { return dlis.Dial(to) }
ss := &stubServer{
emptyCall: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
return &testpb.Empty{}, nil
},
}
s := grpc.NewServer()
testpb.RegisterTestServiceServer(s, ss)
// 1. Start Server
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
s.Serve(dlis)
wg.Done()
}()
// 2. GracefulStop() Server after listener's Accept is called, but don't
// allow Accept() to exit when Close() is called on it.
<-dlis.acceptCalled
wg.Add(1)
go func() {
s.GracefulStop()
wg.Done()
}()
// 3. Create a new connection to the server after listener.Close() is called.
// Server will want to send a GoAway on the new conn, but we delay it
// until 5.
<-dlis.closeCalled // Block until GracefulStop calls dlis.Close()
// Now dial. The listener's Accept method will return a valid connection,
// even though GracefulStop has closed the listener.
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
cc, err := grpc.DialContext(ctx, "", grpc.WithInsecure(), grpc.WithBlock(), grpc.WithDialer(d))
if err != nil {
t.Fatalf("grpc.Dial(%q) = %v", lis.Addr().String(), err)
}
cancel()
client := testpb.NewTestServiceClient(cc)
defer cc.Close()
dlis.allowClose()
wcch := dlis.clientWriteCalledChan()
go func() {
// 5. Allow the client to read the GoAway. The RPC should complete
// successfully.
<-wcch
dlis.allowClientRead()
}()
// 4. Send an RPC on the new connection.
// The server would send a GOAWAY first, but we are delaying the server's
// writes for now until the client writes more than the preface.
ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
t.Fatalf("EmptyCall() = %v; want <nil>", err)
}
// 5. happens above, then we finish the call.
cancel()
wg.Wait()
}