1/* 2 * 3 * Copyright 2017 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19package test 20 21import ( 22 "context" 23 "fmt" 24 "net" 25 "sync" 26 "testing" 27 "time" 28 29 "google.golang.org/grpc" 30 "google.golang.org/grpc/codes" 31 "google.golang.org/grpc/status" 32 testpb "google.golang.org/grpc/test/grpc_testing" 33) 34 35type delayListener struct { 36 net.Listener 37 closeCalled chan struct{} 38 acceptCalled chan struct{} 39 allowCloseCh chan struct{} 40 dialed bool 41} 42 43func (d *delayListener) Accept() (net.Conn, error) { 44 select { 45 case <-d.acceptCalled: 46 // On the second call, block until closed, then return an error. 47 <-d.closeCalled 48 <-d.allowCloseCh 49 return nil, fmt.Errorf("listener is closed") 50 default: 51 close(d.acceptCalled) 52 conn, err := d.Listener.Accept() 53 if err != nil { 54 return nil, err 55 } 56 // Allow closing of listener only after accept. 57 // Note: Dial can return successfully, yet Accept 58 // might now have finished. 59 d.allowClose() 60 return conn, nil 61 } 62} 63 64func (d *delayListener) allowClose() { 65 close(d.allowCloseCh) 66} 67func (d *delayListener) Close() error { 68 close(d.closeCalled) 69 go func() { 70 <-d.allowCloseCh 71 d.Listener.Close() 72 }() 73 return nil 74} 75 76func (d *delayListener) Dial(ctx context.Context) (net.Conn, error) { 77 if d.dialed { 78 // Only hand out one connection (net.Dial can return more even after the 79 // listener is closed). This is not thread-safe, but Dial should never be 80 // called concurrently in this environment. 81 return nil, fmt.Errorf("no more conns") 82 } 83 d.dialed = true 84 return (&net.Dialer{}).DialContext(ctx, "tcp", d.Listener.Addr().String()) 85} 86 87func (s) TestGracefulStop(t *testing.T) { 88 // This test ensures GracefulStop causes new connections to fail. 89 // 90 // Steps of this test: 91 // 1. Start Server 92 // 2. GracefulStop() Server after listener's Accept is called, but don't 93 // allow Accept() to exit when Close() is called on it. 94 // 3. Create a new connection to the server after listener.Close() is called. 95 // Server should close this connection immediately, before handshaking. 96 // 4. Send an RPC on the new connection. Should see Unavailable error 97 // because the ClientConn is in transient failure. 98 lis, err := net.Listen("tcp", "localhost:0") 99 if err != nil { 100 t.Fatalf("Error listenening: %v", err) 101 } 102 dlis := &delayListener{ 103 Listener: lis, 104 acceptCalled: make(chan struct{}), 105 closeCalled: make(chan struct{}), 106 allowCloseCh: make(chan struct{}), 107 } 108 d := func(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) } 109 110 ss := &stubServer{ 111 fullDuplexCall: func(stream testpb.TestService_FullDuplexCallServer) error { 112 _, err := stream.Recv() 113 if err != nil { 114 return err 115 } 116 return stream.Send(&testpb.StreamingOutputCallResponse{}) 117 }, 118 } 119 s := grpc.NewServer() 120 testpb.RegisterTestServiceServer(s, ss) 121 122 // 1. Start Server 123 wg := sync.WaitGroup{} 124 wg.Add(1) 125 go func() { 126 s.Serve(dlis) 127 wg.Done() 128 }() 129 130 // 2. GracefulStop() Server after listener's Accept is called, but don't 131 // allow Accept() to exit when Close() is called on it. 132 <-dlis.acceptCalled 133 wg.Add(1) 134 go func() { 135 s.GracefulStop() 136 wg.Done() 137 }() 138 139 // 3. Create a new connection to the server after listener.Close() is called. 140 // Server should close this connection immediately, before handshaking. 141 142 <-dlis.closeCalled // Block until GracefulStop calls dlis.Close() 143 144 // Now dial. The listener's Accept method will return a valid connection, 145 // even though GracefulStop has closed the listener. 146 ctx, dialCancel := context.WithTimeout(context.Background(), 5*time.Second) 147 defer dialCancel() 148 cc, err := grpc.DialContext(ctx, "", grpc.WithInsecure(), grpc.WithContextDialer(d)) 149 if err != nil { 150 t.Fatalf("grpc.DialContext(_, %q, _) = %v", lis.Addr().String(), err) 151 } 152 client := testpb.NewTestServiceClient(cc) 153 defer cc.Close() 154 155 // 4. Send an RPC on the new connection. 156 // The server would send a GOAWAY first, but we are delaying the server's 157 // writes for now until the client writes more than the preface. 158 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) 159 if _, err = client.FullDuplexCall(ctx); err == nil || status.Code(err) != codes.Unavailable { 160 t.Fatalf("FullDuplexCall= _, %v; want _, <status code Unavailable>", err) 161 } 162 cancel() 163 wg.Wait() 164} 165