benchmark/server/main.go

/*
 *
 * Copyright 2017 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 provides a server used for benchmarking.  It launches a server
which is listening on port 50051.  An example to start the server can be found
at:

	go run benchmark/server/main.go -test_name=grpc_test

After starting the server, the client can be run separately and used to test
qps and latency.
*/
package main

import (
	"flag"
	"fmt"
	"net"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/benchmark"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/internal/syscall"
)

var (
	port     = flag.String("port", "50051", "Localhost port to listen on.")
	testName = flag.String("test_name", "", "Name of the test used for creating profiles.")

	logger = grpclog.Component("benchmark")
)

func main() {
	flag.Parse()
	if *testName == "" {
		logger.Fatal("-test_name not set")
	}
	lis, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		logger.Fatalf("Failed to listen: %v", err)
	}
	defer lis.Close()

	cf, err := os.Create("/tmp/" + *testName + ".cpu")
	if err != nil {
		logger.Fatalf("Failed to create file: %v", err)
	}
	defer cf.Close()
	pprof.StartCPUProfile(cf)
	cpuBeg := syscall.GetCPUTime()
	// Launch server in a separate goroutine.
	stop := benchmark.StartServer(benchmark.ServerInfo{Type: "protobuf", Listener: lis},
		grpc.WriteBufferSize(128*1024),
		grpc.ReadBufferSize(128*1024),
	)
	// Wait on OS terminate signal.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	<-ch
	cpu := time.Duration(syscall.GetCPUTime() - cpuBeg)
	stop()
	pprof.StopCPUProfile()
	mf, err := os.Create("/tmp/" + *testName + ".mem")
	if err != nil {
		logger.Fatalf("Failed to create file: %v", err)
	}
	defer mf.Close()
	runtime.GC() // materialize all statistics
	if err := pprof.WriteHeapProfile(mf); err != nil {
		logger.Fatalf("Failed to write memory profile: %v", err)
	}
	fmt.Println("Server CPU utilization:", cpu)
	fmt.Println("Server CPU profile:", cf.Name())
	fmt.Println("Server Mem Profile:", mf.Name())
}