Repository: grpc/grpc-go
File: interop/test_utils.go (Go)

Code before:

// UseSharedChannel reuses the provided currentChannel.
func UseSharedChannel(currentChannel *grpc.ClientConn) (*grpc.ClientConn, testgrpc.TestServiceClient) {
	client := testgrpc.NewTestServiceClient(currentChannel)
	return currentChannel, client
}

func doOneSoakIteration(
	ctx context.Context,
	client testgrpc.TestServiceClient,
	soakRequestSize int,
	soakResponseSize int,
	copts []grpc.CallOption) (SoakIterationResult, error) {
	start := time.Now()
	var err error
	// Do a large-unary RPC.
	// Create the request payload.
	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, soakRequestSize)
	req := &testpb.SimpleRequest{
		ResponseType: testpb.PayloadType_COMPRESSABLE,
		ResponseSize: int32(soakResponseSize),
		Payload:      pl,
	}
	// Perform the GRPC call.
	var reply *testpb.SimpleResponse
	reply, err = client.UnaryCall(ctx, req, copts...)
	if err != nil {
		err = fmt.Errorf("/TestService/UnaryCall RPC failed: %s", err)
		return SoakIterationResult{}, err
	}
	// Validate response.
	t := reply.GetPayload().GetType()
	s := len(reply.GetPayload().GetBody())
	if t != testpb.PayloadType_COMPRESSABLE || s != soakResponseSize {
		err = fmt.Errorf("got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, soakResponseSize)
		return SoakIterationResult{}, err
	}
	// Calculate latency and return result.
	latency := time.Since(start).Milliseconds()
	return SoakIterationResult{
		LatencyMs: latency,
		Status:    "OK",
	}, nil
}

func executeSoakTestInThread(
	ctx context.Context,
	soakIterationsPerThread int,
	startNs int64,
	minTimeBetweenRPCs time.Duration,
	soakRequestSize int,

Review comment (suggestion):
Nit: use type inference instead:

```suggestion
	reply, err := client.UnaryCall(ctx, req, copts...)
```

Code after:
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		if initialMetadata, ok := md[initialMetadataKey]; ok {
			header := metadata.Pairs(initialMetadataKey, initialMetadata[0])
			grpc.SendHeader(ctx, header)
		}
		if trailingMetadata, ok := md[trailingMetadataKey]; ok {
			trailer := metadata.Pairs(trailingMetadataKey, trailingMetadata[0])
			grpc.SetTrailer(ctx, trailer)
		}
	}
	if st != nil && st.Code != 0 {
		return nil, status.Error(codes.Code(st.Code), st.Message)
	}
	pl, err := serverNewPayload(in.GetResponseType(), in.GetResponseSize())
	if err != nil {
		return nil, err
	}
	if r, orcaData := orca.CallMetricsRecorderFromContext(ctx), in.GetOrcaPerQueryReport(); r != nil && orcaData != nil {
		// Transfer the request's per-Call ORCA data to the call metrics
		// recorder in the context, if present.
		setORCAMetrics(r, orcaData)
	}
	return &testpb.SimpleResponse{
		Payload: pl,
	}, nil
}

func setORCAMetrics(r orca.ServerMetricsRecorder, orcaData *testpb.TestOrcaReport) {
	r.SetCPUUtilization(orcaData.CpuUtilization)
	r.SetMemoryUtilization(orcaData.MemoryUtilization)
	if rq, ok := r.(orca.CallMetricsRecorder); ok {
		for k, v := range orcaData.RequestCost {
			rq.SetRequestCost(k, v)
		}
	}
	for k, v := range orcaData.Utilization {
		r.SetNamedUtilization(k, v)
	}
}

func (s *testServer) StreamingOutputCall(args *testpb.StreamingOutputCallRequest, stream testgrpc.TestService_StreamingOutputCallServer) error {
	cs := args.GetResponseParameters()
	for _, c := range cs {
		if us := c.GetIntervalUs(); us > 0 {
			time.Sleep(time.Duration(us) * time.Microsecond)
		}
		pl, err := serverNewPayload(args.GetResponseType(), c.GetSize())
		if err != nil {
			return err
		}
		if err := stream.Send(&testpb.StreamingOutputCallResponse{