google.golang.org/grpc@v1.72.2/test/stats_test.go (about)

     1  /*
     2   *
     3   * Copyright 2024 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  	"fmt"
    24  	"net"
    25  	"sync"
    26  	"testing"
    27  
    28  	"google.golang.org/grpc"
    29  	"google.golang.org/grpc/credentials/insecure"
    30  	"google.golang.org/grpc/internal/stubserver"
    31  	"google.golang.org/grpc/interop"
    32  	testgrpc "google.golang.org/grpc/interop/grpc_testing"
    33  	testpb "google.golang.org/grpc/interop/grpc_testing"
    34  	"google.golang.org/grpc/peer"
    35  	"google.golang.org/grpc/stats"
    36  )
    37  
    38  // TestPeerForClientStatsHandler configures a stats handler that
    39  // verifies that peer is sent all stats handler callouts instead
    40  // of Begin and PickerUpdated.
    41  func (s) TestPeerForClientStatsHandler(t *testing.T) {
    42  	psh := &peerStatsHandler{}
    43  
    44  	// Stats callouts & peer object population.
    45  	// Note:
    46  	// * Begin stats lack peer info (RPC starts pre-resolution).
    47  	// * PickerUpdated: no peer info (picker lacks transport details).
    48  	expectedCallouts := map[stats.RPCStats]bool{
    49  		&stats.OutPayload{}:    true,
    50  		&stats.InHeader{}:      true,
    51  		&stats.OutHeader{}:     true,
    52  		&stats.InTrailer{}:     true,
    53  		&stats.OutTrailer{}:    true,
    54  		&stats.End{}:           true,
    55  		&stats.Begin{}:         false,
    56  		&stats.PickerUpdated{}: false,
    57  	}
    58  
    59  	// Start server.
    60  	l, err := net.Listen("tcp", "localhost:0")
    61  	if err != nil {
    62  		t.Fatal(err)
    63  	}
    64  	ss := &stubserver.StubServer{
    65  		Listener: l,
    66  		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
    67  			return &testpb.Empty{}, nil
    68  		},
    69  		S: grpc.NewServer(),
    70  	}
    71  	stubserver.StartTestService(t, ss)
    72  	defer ss.S.Stop()
    73  
    74  	// Create client with stats handler and do some calls.
    75  	cc, err := grpc.NewClient(
    76  		l.Addr().String(),
    77  		grpc.WithTransportCredentials(insecure.NewCredentials()),
    78  		grpc.WithStatsHandler(psh))
    79  	if err != nil {
    80  		t.Fatal(err)
    81  	}
    82  	defer cc.Close()
    83  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
    84  	defer cancel()
    85  	client := testgrpc.NewTestServiceClient(cc)
    86  	interop.DoEmptyUnaryCall(ctx, client)
    87  
    88  	psh.mu.Lock()
    89  	pshArgs := psh.args
    90  	psh.mu.Unlock()
    91  
    92  	// Fetch the total unique stats handlers with peer != nil
    93  	uniqueStatsTypes := make(map[string]struct{})
    94  	for _, callbackArgs := range pshArgs {
    95  		key := fmt.Sprintf("%T", callbackArgs.rpcStats)
    96  		if _, exists := uniqueStatsTypes[key]; exists {
    97  			continue
    98  		}
    99  		uniqueStatsTypes[fmt.Sprintf("%T", callbackArgs.rpcStats)] = struct{}{}
   100  	}
   101  	if len(uniqueStatsTypes) != len(expectedCallouts) {
   102  		t.Errorf("Unexpected number of stats handler callouts. Got %v, want %v", len(uniqueStatsTypes), len(expectedCallouts))
   103  	}
   104  
   105  	for _, callbackArgs := range pshArgs {
   106  		expectedPeer, found := expectedCallouts[callbackArgs.rpcStats]
   107  		// In case expectation is set to false and still we got the peer,
   108  		// then it's good to have it. So no need to assert those conditions.
   109  		if found && expectedPeer && callbackArgs.peer != nil {
   110  			continue
   111  		} else if expectedPeer && callbackArgs.peer == nil {
   112  			t.Errorf("peer not populated for: %T", callbackArgs.rpcStats)
   113  		}
   114  	}
   115  }
   116  
   117  type peerStats struct {
   118  	rpcStats stats.RPCStats
   119  	peer     *peer.Peer
   120  }
   121  
   122  type peerStatsHandler struct {
   123  	args []peerStats
   124  	mu   sync.Mutex
   125  }
   126  
   127  func (h *peerStatsHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
   128  	return ctx
   129  }
   130  
   131  func (h *peerStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
   132  	p, _ := peer.FromContext(ctx)
   133  	h.mu.Lock()
   134  	defer h.mu.Unlock()
   135  	h.args = append(h.args, peerStats{rs, p})
   136  }
   137  
   138  func (h *peerStatsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
   139  	return ctx
   140  }
   141  
   142  func (h *peerStatsHandler) HandleConn(context.Context, stats.ConnStats) {}