Method TestE2ECallMetricsStreaming

orca/call_metrics_test.go:151–287

TestE2ECallMetricsStreaming tests the injection of custom backend metrics from the server application for a streaming RPC, and verifies that expected load reports are received at the client.

(t *testing.T)

// from the server application for a streaming RPC, and verifies that expected
// load reports are received at the client.
func (s) TestE2ECallMetricsStreaming(t *testing.T) {
	tests := []struct {
		desc          string
		injectMetrics bool
		wantProto     *v3orcapb.OrcaLoadReport
	}{
		{
			desc:          "with custom backend metrics",
			injectMetrics: true,
			wantProto: &v3orcapb.OrcaLoadReport{
				CpuUtilization: 1.0,
				MemUtilization: 0.5,
				RequestCost:    map[string]float64{"queryCost": 0.25},
				Utilization:    map[string]float64{"queueSize": 0.75},
			},
		},
		{
			desc:          "with no custom backend metrics",
			injectMetrics: false,
		},
	}

	for _, test := range tests {
		t.Run(test.desc, func(t *testing.T) {
			// A server option to enable reporting of per-call backend metrics.
			smr := orca.NewServerMetricsRecorder()
			callMetricsServerOption := orca.CallMetricsServerOption(smr)
			smr.SetCPUUtilization(1.0)

			// An interceptor which injects custom backend metrics, added only
			// when the injectMetrics field in the test is set.
			injectingInterceptor := func(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
				recorder := orca.CallMetricsRecorderFromContext(ss.Context())
				if recorder == nil {
					err := errors.New("Failed to retrieve per-RPC custom metrics recorder from the RPC context")
					t.Error(err)
					return err
				}
				recorder.SetMemoryUtilization(0.5)
				// This value will be overwritten by a write to the same metric
				// from the server handler.
				recorder.SetNamedUtilization("queueSize", 1.0)
				return handler(srv, ss)
			}

			// A stub server whose streaming handler injects custom metrics, if
			// the injectMetrics field in the test is set. It overwrites one of
			// the values injected above, by the interceptor.
			srv := stubserver.StubServer{
				FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
					if test.injectMetrics {
						recorder := orca.CallMetricsRecorderFromContext(stream.Context())
						if recorder == nil {
							err := errors.New("Failed to retrieve per-RPC custom metrics recorder from the RPC context")
							t.Error(err)
							return err
						}
						recorder.SetRequestCost("queryCost", 0.25)