openbilibili/vendor/google.golang.org/grpc/stress/metrics_client/main.go

/*
 *
 * Copyright 2016 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package main

import (
	"context"
	"flag"
	"fmt"
	"io"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	metricspb "google.golang.org/grpc/stress/grpc_testing"
)

var (
	metricsServerAddress = flag.String("metrics_server_address", "", "The metrics server addresses in the fomrat <hostname>:<port>")
	totalOnly            = flag.Bool("total_only", false, "If true, this prints only the total value of all gauges")
)

func printMetrics(client metricspb.MetricsServiceClient, totalOnly bool) {
	stream, err := client.GetAllGauges(context.Background(), &metricspb.EmptyMessage{})
	if err != nil {
		grpclog.Fatalf("failed to call GetAllGuages: %v", err)
	}

	var (
		overallQPS int64
		rpcStatus  error
	)
	for {
		gaugeResponse, err := stream.Recv()
		if err != nil {
			rpcStatus = err
			break
		}
		if _, ok := gaugeResponse.GetValue().(*metricspb.GaugeResponse_LongValue); !ok {
			panic(fmt.Sprintf("gauge %s is not a long value", gaugeResponse.Name))
		}
		v := gaugeResponse.GetLongValue()
		if !totalOnly {
			grpclog.Infof("%s: %d", gaugeResponse.Name, v)
		}
		overallQPS += v
	}
	if rpcStatus != io.EOF {
		grpclog.Fatalf("failed to finish server streaming: %v", rpcStatus)
	}
	grpclog.Infof("overall qps: %d", overallQPS)
}

func main() {
	flag.Parse()
	if *metricsServerAddress == "" {
		grpclog.Fatalf("Metrics server address is empty.")
	}

	conn, err := grpc.Dial(*metricsServerAddress, grpc.WithInsecure())
	if err != nil {
		grpclog.Fatalf("cannot connect to metrics server: %v", err)
	}
	defer conn.Close()

	c := metricspb.NewMetricsServiceClient(conn)
	printMetrics(c, *totalOnly)
}