Method TestGracefulClose

internal/transport/transport_test.go:1007–1075

TestGracefulClose ensures that GracefulClose allows in-flight streams to proceed until they complete naturally, while not allowing creation of new streams during this window.

(t *testing.T)

// proceed until they complete naturally, while not allowing creation of new
// streams during this window.
func (s) TestGracefulClose(t *testing.T) {
	leakcheck.SetTrackingBufferPool(t)
	server, ct, cancel := setUp(t, 0, pingpong)
	defer cancel()
	defer func() {
		// Stop the server's listener to make the server's goroutines terminate
		// (after the last active stream is done).
		server.lis.Close()
		// Check for goroutine leaks (i.e. GracefulClose with an active stream
		// doesn't eventually close the connection when that stream completes).
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		leakcheck.CheckGoroutines(ctx, t)
		leakcheck.CheckTrackingBufferPool()
		// Correctly clean up the server
		server.stop()
	}()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
	defer cancel()

	// Create a stream that will exist for this whole test and confirm basic
	// functionality.
	s, err := ct.NewStream(ctx, &CallHdr{}, nil)
	if err != nil {
		t.Fatalf("NewStream(_, _) = _, %v, want _, <nil>", err)
	}
	msg := make([]byte, 1024)
	outgoingHeader := make([]byte, 5)
	outgoingHeader[0] = byte(0)
	binary.BigEndian.PutUint32(outgoingHeader[1:], uint32(len(msg)))
	incomingHeader := make([]byte, 5)
	if err := s.Write(outgoingHeader, newBufferSlice(msg), &WriteOptions{}); err != nil {
		t.Fatalf("Error while writing: %v", err)
	}
	if _, err := s.readTo(incomingHeader); err != nil {
		t.Fatalf("Error while reading: %v", err)
	}
	sz := binary.BigEndian.Uint32(incomingHeader[1:])
	recvMsg := make([]byte, int(sz))
	if _, err := s.readTo(recvMsg); err != nil {
		t.Fatalf("Error while reading: %v", err)
	}

	// Gracefully close the transport, which should not affect the existing
	// stream.
	ct.GracefulClose()

	var wg sync.WaitGroup
	// Expect errors creating new streams because the client transport has been
	// gracefully closed.
	for i := 0; i < 200; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := ct.NewStream(ctx, &CallHdr{}, nil)
			if err != nil && err.(*NewStreamError).Err == ErrConnClosing && err.(*NewStreamError).AllowTransparentRetry {
				return
			}