gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/test/stream_cleanup_test.go (about)

     1  /*
     2   *
     3   * Copyright 2019 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  
    19  package test
    20  
    21  import (
    22  	"context"
    23  	"io"
    24  	"testing"
    25  	"time"
    26  
    27  	grpc "gitee.com/ks-custle/core-gm/grpc"
    28  	"gitee.com/ks-custle/core-gm/grpc/codes"
    29  	"gitee.com/ks-custle/core-gm/grpc/internal/stubserver"
    30  	"gitee.com/ks-custle/core-gm/grpc/status"
    31  	testpb "gitee.com/ks-custle/core-gm/grpc/test/grpc_testing"
    32  )
    33  
    34  func (s) TestStreamCleanup(t *testing.T) {
    35  	const initialWindowSize uint = 70 * 1024 // Must be higher than default 64K, ignored otherwise
    36  	const bodySize = 2 * initialWindowSize   // Something that is not going to fit in a single window
    37  	const callRecvMsgSize uint = 1           // The maximum message size the client can receive
    38  
    39  	ss := &stubserver.StubServer{
    40  		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
    41  			return &testpb.SimpleResponse{Payload: &testpb.Payload{
    42  				Body: make([]byte, bodySize),
    43  			}}, nil
    44  		},
    45  		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
    46  			return &testpb.Empty{}, nil
    47  		},
    48  	}
    49  	if err := ss.Start([]grpc.ServerOption{grpc.MaxConcurrentStreams(1)}, grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(int(callRecvMsgSize))), grpc.WithInitialWindowSize(int32(initialWindowSize))); err != nil {
    50  		t.Fatalf("Error starting endpoint server: %v", err)
    51  	}
    52  	defer ss.Stop()
    53  
    54  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
    55  	defer cancel()
    56  	if _, err := ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{}); status.Code(err) != codes.ResourceExhausted {
    57  		t.Fatalf("should fail with ResourceExhausted, message's body size: %v, maximum message size the client can receive: %v", bodySize, callRecvMsgSize)
    58  	}
    59  	if _, err := ss.Client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
    60  		t.Fatalf("should succeed, err: %v", err)
    61  	}
    62  }
    63  
    64  func (s) TestStreamCleanupAfterSendStatus(t *testing.T) {
    65  	const initialWindowSize uint = 70 * 1024 // Must be higher than default 64K, ignored otherwise
    66  	const bodySize = 2 * initialWindowSize   // Something that is not going to fit in a single window
    67  
    68  	serverReturnedStatus := make(chan struct{})
    69  
    70  	ss := &stubserver.StubServer{
    71  		FullDuplexCallF: func(stream testpb.TestService_FullDuplexCallServer) error {
    72  			defer func() {
    73  				close(serverReturnedStatus)
    74  			}()
    75  			return stream.Send(&testpb.StreamingOutputCallResponse{
    76  				Payload: &testpb.Payload{
    77  					Body: make([]byte, bodySize),
    78  				},
    79  			})
    80  		},
    81  	}
    82  	if err := ss.Start([]grpc.ServerOption{grpc.MaxConcurrentStreams(1)}, grpc.WithInitialWindowSize(int32(initialWindowSize))); err != nil {
    83  		t.Fatalf("Error starting endpoint server: %v", err)
    84  	}
    85  	defer ss.Stop()
    86  
    87  	// This test makes sure we don't delete stream from server transport's
    88  	// activeStreams list too aggressively.
    89  
    90  	// 1. Make a long living stream RPC. So server's activeStream list is not
    91  	// empty.
    92  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    93  	defer cancel()
    94  	stream, err := ss.Client.FullDuplexCall(ctx)
    95  	if err != nil {
    96  		t.Fatalf("FullDuplexCall= _, %v; want _, <nil>", err)
    97  	}
    98  
    99  	// 2. Wait for service handler to return status.
   100  	//
   101  	// This will trigger a stream cleanup code, which will eventually remove
   102  	// this stream from activeStream.
   103  	//
   104  	// But the stream removal won't happen because it's supposed to be done
   105  	// after the status is sent by loopyWriter, and the status send is blocked
   106  	// by flow control.
   107  	<-serverReturnedStatus
   108  
   109  	// 3. GracefulStop (besides sending goaway) checks the number of
   110  	// activeStreams.
   111  	//
   112  	// It will close the connection if there's no active streams. This won't
   113  	// happen because of the pending stream. But if there's a bug in stream
   114  	// cleanup that causes stream to be removed too aggressively, the connection
   115  	// will be closd and the stream will be broken.
   116  	gracefulStopDone := make(chan struct{})
   117  	go func() {
   118  		defer close(gracefulStopDone)
   119  		ss.S.GracefulStop()
   120  	}()
   121  
   122  	// 4. Make sure the stream is not broken.
   123  	if _, err := stream.Recv(); err != nil {
   124  		t.Fatalf("stream.Recv() = _, %v, want _, <nil>", err)
   125  	}
   126  	if _, err := stream.Recv(); err != io.EOF {
   127  		t.Fatalf("stream.Recv() = _, %v, want _, io.EOF", err)
   128  	}
   129  
   130  	timer := time.NewTimer(time.Second)
   131  	select {
   132  	case <-gracefulStopDone:
   133  		timer.Stop()
   134  	case <-timer.C:
   135  		t.Fatalf("s.GracefulStop() didn't finish without 1 second after the last RPC")
   136  	}
   137  }