Repository: grpc/grpc-go
File: test/gracefulstop_test.go (Go)

Code before:
	d.dialed = true
	return (&net.Dialer{}).DialContext(ctx, "tcp", d.Listener.Addr().String())
}

// TestGracefulStop ensures GracefulStop causes new connections to fail.
//
// 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 should close this connection immediately, before handshaking.
//  4. Send an RPC on the new connection.  Should see Unavailable error
//     because the ClientConn is in transient failure.
func (s) TestGracefulStop(t *testing.T) {
	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(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) }

	ss := &stubserver.StubServer{
		Listener: dlis,
		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
			_, err := stream.Recv()
			if err != nil {
				return err
			}
			return stream.Send(&testpb.StreamingOutputCallResponse{})
		},
		S: grpc.NewServer(),
	}
	stubserver.StartTestService(t, ss)

	// 1. Start Server
	wg := sync.WaitGroup{}

	// 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() {
		ss.S.GracefulStop()
		wg.Done()
	}()

Review comment (suggestion):
While you are here, could you please rename this `d` as `dialer` and move it a line right before it is used. Currently it is only used on line 153, when it is passed to `grpc.WithContextDialer` in `grpc.DialContext`.

Code after:
	return (&net.Dialer{}).DialContext(ctx, "tcp", d.Listener.Addr().String())
}

// TestGracefulStop ensures GracefulStop causes new connections to fail.
//
// 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 should close this connection immediately, before handshaking.
//  4. Send an RPC on the new connection.  Should see Unavailable error
//     because the ClientConn is in transient failure.
func (s) TestGracefulStop(t *testing.T) {
	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{}),
	}

	ss := &stubserver.StubServer{
		Listener: dlis,
		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
			if _, err := stream.Recv(); err != nil {
				return err
			}
			return stream.Send(&testpb.StreamingOutputCallResponse{})
		},
		S: grpc.NewServer(),
	}
	// 1. Start Server and start serving by calling Serve().
	stubserver.StartTestService(t, ss)

	// 2. Call GracefulStop from a goroutine. It will trigger Close on the
	// listener, but the listener will not actually close until a connection
	// is accepted.
	gracefulStopDone := make(chan struct{})
	<-dlis.acceptCalled
	go func() {
		ss.S.GracefulStop()
		close(gracefulStopDone)
	}()

	// 3. Create a new connection to the server after listener.Close() is called.
	// Server should close this connection immediately, before handshaking.