snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

benchmark/worker/benchmark_client.go

@@ -0,0 +1,428 @@
+/*
+ *
+ * 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"
+	"math"
+	rand "math/rand/v2"
+	"runtime"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/benchmark"
+	"google.golang.org/grpc/benchmark/stats"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/syscall"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/testdata"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+
+	_ "google.golang.org/grpc/xds" // To install the xds resolvers and balancers.
+)
+
+var caFile = flag.String("ca_file", "", "The file containing the CA root cert file")
+
+type lockingHistogram struct {
+	mu        sync.Mutex
+	histogram *stats.Histogram
+}
+
+func (h *lockingHistogram) add(value int64) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.histogram.Add(value)
+}
+
+// swap sets h.histogram to o and returns its old value.
+func (h *lockingHistogram) swap(o *stats.Histogram) *stats.Histogram {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	old := h.histogram
+	h.histogram = o
+	return old
+}
+
+func (h *lockingHistogram) mergeInto(merged *stats.Histogram) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	merged.Merge(h.histogram)
+}
+
+type benchmarkClient struct {
+	closeConns        func()
+	lastResetTime     time.Time
+	histogramOptions  stats.HistogramOptions
+	lockingHistograms []lockingHistogram
+	rusageLastReset   *syscall.Rusage
+}
+
+func printClientConfig(config *testpb.ClientConfig) {
+	// Some config options are ignored:
+	// - client type:
+	//     will always create sync client
+	// - async client threads.
+	// - core list
+	logger.Infof(" * client type: %v (ignored, always creates sync client)", config.ClientType)
+	logger.Infof(" * async client threads: %v (ignored)", config.AsyncClientThreads)
+	// TODO: use cores specified by CoreList when setting list of cores is supported in go.
+	logger.Infof(" * core list: %v (ignored)", config.CoreList)
+
+	logger.Infof(" - security params: %v", config.SecurityParams)
+	logger.Infof(" - core limit: %v", config.CoreLimit)
+	logger.Infof(" - payload config: %v", config.PayloadConfig)
+	logger.Infof(" - rpcs per chann: %v", config.OutstandingRpcsPerChannel)
+	logger.Infof(" - channel number: %v", config.ClientChannels)
+	logger.Infof(" - load params: %v", config.LoadParams)
+	logger.Infof(" - rpc type: %v", config.RpcType)
+	logger.Infof(" - histogram params: %v", config.HistogramParams)
+	logger.Infof(" - server targets: %v", config.ServerTargets)
+}
+
+func setupClientEnv(config *testpb.ClientConfig) {
+	// Use all cpu cores available on machine by default.
+	// TODO: Revisit this for the optimal default setup.
+	if config.CoreLimit > 0 {
+		runtime.GOMAXPROCS(int(config.CoreLimit))
+	} else {
+		runtime.GOMAXPROCS(runtime.NumCPU())
+	}
+}
+
+// createConns creates connections according to given config.
+// It returns the connections and corresponding function to close them.
+// It returns non-nil error if there is anything wrong.
+func createConns(config *testpb.ClientConfig) ([]*grpc.ClientConn, func(), error) {
+	opts := []grpc.DialOption{
+		grpc.WithWriteBufferSize(128 * 1024),
+		grpc.WithReadBufferSize(128 * 1024),
+	}
+
+	// Sanity check for client type.
+	switch config.ClientType {
+	case testpb.ClientType_SYNC_CLIENT:
+	case testpb.ClientType_ASYNC_CLIENT:
+	default:
+		return nil, nil, status.Errorf(codes.InvalidArgument, "unknown client type: %v", config.ClientType)
+	}
+
+	// Check and set security options.
+	if config.SecurityParams != nil {
+		if *caFile == "" {
+			*caFile = testdata.Path("ca.pem")
+		}
+		creds, err := credentials.NewClientTLSFromFile(*caFile, config.SecurityParams.ServerHostOverride)
+		if err != nil {
+			return nil, nil, status.Errorf(codes.InvalidArgument, "failed to create TLS credentials: %v", err)
+		}
+		opts = append(opts, grpc.WithTransportCredentials(creds))
+	} else {
+		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+
+	// Use byteBufCodec if it is required.
+	if config.PayloadConfig != nil {
+		switch config.PayloadConfig.Payload.(type) {
+		case *testpb.PayloadConfig_BytebufParams:
+			opts = append(opts, grpc.WithDefaultCallOptions(grpc.CallCustomCodec(byteBufCodec{})))
+		case *testpb.PayloadConfig_SimpleParams:
+		default:
+			return nil, nil, status.Errorf(codes.InvalidArgument, "unknown payload config: %v", config.PayloadConfig)
+		}
+	}
+
+	// Create connections.
+	connCount := int(config.ClientChannels)
+	conns := make([]*grpc.ClientConn, connCount)
+	for connIndex := 0; connIndex < connCount; connIndex++ {
+		conns[connIndex] = benchmark.NewClientConn(config.ServerTargets[connIndex%len(config.ServerTargets)], opts...)
+	}
+
+	return conns, func() {
+		for _, conn := range conns {
+			conn.Close()
+		}
+	}, nil
+}
+
+func performRPCs(ctx context.Context, config *testpb.ClientConfig, conns []*grpc.ClientConn, bc *benchmarkClient) error {
+	// Read payload size and type from config.
+	var (
+		payloadReqSize, payloadRespSize int
+		payloadType                     string
+	)
+	if config.PayloadConfig != nil {
+		switch c := config.PayloadConfig.Payload.(type) {
+		case *testpb.PayloadConfig_BytebufParams:
+			payloadReqSize = int(c.BytebufParams.ReqSize)
+			payloadRespSize = int(c.BytebufParams.RespSize)
+			payloadType = "bytebuf"
+		case *testpb.PayloadConfig_SimpleParams:
+			payloadReqSize = int(c.SimpleParams.ReqSize)
+			payloadRespSize = int(c.SimpleParams.RespSize)
+			payloadType = "protobuf"
+		default:
+			return status.Errorf(codes.InvalidArgument, "unknown payload config: %v", config.PayloadConfig)
+		}
+	}
+
+	// If set, perform an open loop, if not perform a closed loop. An open loop
+	// asynchronously starts RPCs based on random start times derived from a
+	// Poisson distribution. A closed loop performs RPCs in a blocking manner,
+	// and runs the next RPC after the previous RPC completes and returns.
+	var poissonLambda *float64
+	switch t := config.LoadParams.Load.(type) {
+	case *testpb.LoadParams_ClosedLoop:
+	case *testpb.LoadParams_Poisson:
+		if t.Poisson == nil {
+			return status.Errorf(codes.InvalidArgument, "poisson is nil, needs to be set")
+		}
+		if t.Poisson.OfferedLoad <= 0 {
+			return status.Errorf(codes.InvalidArgument, "poisson.offered is <= 0: %v, needs to be >0", t.Poisson.OfferedLoad)
+		}
+		poissonLambda = &t.Poisson.OfferedLoad
+	default:
+		return status.Errorf(codes.InvalidArgument, "unknown load params: %v", config.LoadParams)
+	}
+
+	rpcCountPerConn := int(config.OutstandingRpcsPerChannel)
+
+	switch config.RpcType {
+	case testpb.RpcType_UNARY:
+		bc.unaryLoop(ctx, conns, rpcCountPerConn, payloadReqSize, payloadRespSize, poissonLambda)
+	case testpb.RpcType_STREAMING:
+		bc.streamingLoop(ctx, conns, rpcCountPerConn, payloadReqSize, payloadRespSize, payloadType, poissonLambda)
+	default:
+		return status.Errorf(codes.InvalidArgument, "unknown rpc type: %v", config.RpcType)
+	}
+
+	return nil
+}
+
+func startBenchmarkClient(ctx context.Context, config *testpb.ClientConfig) (*benchmarkClient, error) {
+	printClientConfig(config)
+
+	// Set running environment like how many cores to use.
+	setupClientEnv(config)
+
+	conns, closeConns, err := createConns(config)
+	if err != nil {
+		return nil, err
+	}
+
+	rpcCountPerConn := int(config.OutstandingRpcsPerChannel)
+	bc := &benchmarkClient{
+		histogramOptions: stats.HistogramOptions{
+			NumBuckets:     int(math.Log(config.HistogramParams.MaxPossible)/math.Log(1+config.HistogramParams.Resolution)) + 1,
+			GrowthFactor:   config.HistogramParams.Resolution,
+			BaseBucketSize: (1 + config.HistogramParams.Resolution),
+			MinValue:       0,
+		},
+		lockingHistograms: make([]lockingHistogram, rpcCountPerConn*len(conns)),
+
+		lastResetTime:   time.Now(),
+		closeConns:      closeConns,
+		rusageLastReset: syscall.GetRusage(),
+	}
+
+	if err = performRPCs(ctx, config, conns, bc); err != nil {
+		// Close all connections if performRPCs failed.
+		closeConns()
+		return nil, err
+	}
+
+	return bc, nil
+}
+
+func (bc *benchmarkClient) unaryLoop(ctx context.Context, conns []*grpc.ClientConn, rpcCountPerConn int, reqSize int, respSize int, poissonLambda *float64) {
+	for ic, conn := range conns {
+		client := testgrpc.NewBenchmarkServiceClient(conn)
+		// For each connection, create rpcCountPerConn goroutines to do rpc.
+		for j := range rpcCountPerConn {
+			// Create histogram for each goroutine.
+			idx := ic*rpcCountPerConn + j
+			bc.lockingHistograms[idx].histogram = stats.NewHistogram(bc.histogramOptions)
+			// Start goroutine on the created mutex and histogram.
+			go func(idx int) {
+				// TODO: do warm up if necessary.
+				// Now relying on worker client to reserve time to do warm up.
+				// The worker client needs to wait for some time after client is created,
+				// before starting benchmark.
+				if poissonLambda == nil { // Closed loop.
+					for {
+						if ctx.Err() != nil {
+							break
+						}
+						start := time.Now()
+						if err := benchmark.DoUnaryCall(client, reqSize, respSize); err != nil {
+							continue
+						}
+						elapse := time.Since(start)
+						bc.lockingHistograms[idx].add(int64(elapse))
+					}
+				} else { // Open loop.
+					timeBetweenRPCs := time.Duration((rand.ExpFloat64() / *poissonLambda) * float64(time.Second))
+					time.AfterFunc(timeBetweenRPCs, func() {
+						bc.poissonUnary(client, idx, reqSize, respSize, *poissonLambda)
+					})
+				}
+			}(idx)
+		}
+	}
+}
+
+func (bc *benchmarkClient) streamingLoop(ctx context.Context, conns []*grpc.ClientConn, rpcCountPerConn int, reqSize int, respSize int, payloadType string, poissonLambda *float64) {
+	var doRPC func(testgrpc.BenchmarkService_StreamingCallClient, int, int) error
+	if payloadType == "bytebuf" {
+		doRPC = benchmark.DoByteBufStreamingRoundTrip
+	} else {
+		doRPC = benchmark.DoStreamingRoundTrip
+	}
+	for ic, conn := range conns {
+		// For each connection, create rpcCountPerConn goroutines to do rpc.
+		for j := 0; j < rpcCountPerConn; j++ {
+			c := testgrpc.NewBenchmarkServiceClient(conn)
+			stream, err := c.StreamingCall(context.Background())
+			if err != nil {
+				logger.Fatalf("%v.StreamingCall(_) = _, %v", c, err)
+			}
+			idx := ic*rpcCountPerConn + j
+			bc.lockingHistograms[idx].histogram = stats.NewHistogram(bc.histogramOptions)
+			if poissonLambda == nil { // Closed loop.
+				// Start goroutine on the created mutex and histogram.
+				go func(idx int) {
+					// TODO: do warm up if necessary.
+					// Now relying on worker client to reserve time to do warm up.
+					// The worker client needs to wait for some time after client is created,
+					// before starting benchmark.
+					for {
+						start := time.Now()
+						if err := doRPC(stream, reqSize, respSize); err != nil {
+							return
+						}
+						elapse := time.Since(start)
+						bc.lockingHistograms[idx].add(int64(elapse))
+						if ctx.Err() != nil {
+							return
+						}
+					}
+				}(idx)
+			} else { // Open loop.
+				timeBetweenRPCs := time.Duration((rand.ExpFloat64() / *poissonLambda) * float64(time.Second))
+				time.AfterFunc(timeBetweenRPCs, func() {
+					bc.poissonStreaming(stream, idx, reqSize, respSize, *poissonLambda, doRPC)
+				})
+			}
+		}
+	}
+}
+
+func (bc *benchmarkClient) poissonUnary(client testgrpc.BenchmarkServiceClient, idx int, reqSize int, respSize int, lambda float64) {
+	go func() {
+		start := time.Now()
+		if err := benchmark.DoUnaryCall(client, reqSize, respSize); err != nil {
+			return
+		}
+		elapse := time.Since(start)
+		bc.lockingHistograms[idx].add(int64(elapse))
+	}()
+	timeBetweenRPCs := time.Duration((rand.ExpFloat64() / lambda) * float64(time.Second))
+	time.AfterFunc(timeBetweenRPCs, func() {
+		bc.poissonUnary(client, idx, reqSize, respSize, lambda)
+	})
+}
+
+func (bc *benchmarkClient) poissonStreaming(stream testgrpc.BenchmarkService_StreamingCallClient, idx int, reqSize int, respSize int, lambda float64, doRPC func(testgrpc.BenchmarkService_StreamingCallClient, int, int) error) {
+	go func() {
+		start := time.Now()
+
+		if err := doRPC(stream, reqSize, respSize); err != nil {
+			return
+		}
+		elapse := time.Since(start)
+		bc.lockingHistograms[idx].add(int64(elapse))
+	}()
+	timeBetweenRPCs := time.Duration((rand.ExpFloat64() / lambda) * float64(time.Second))
+	time.AfterFunc(timeBetweenRPCs, func() {
+		bc.poissonStreaming(stream, idx, reqSize, respSize, lambda, doRPC)
+	})
+}
+
+// getStats returns the stats for benchmark client.
+// It resets lastResetTime and all histograms if argument reset is true.
+func (bc *benchmarkClient) getStats(reset bool) *testpb.ClientStats {
+	var wallTimeElapsed, uTimeElapsed, sTimeElapsed float64
+	mergedHistogram := stats.NewHistogram(bc.histogramOptions)
+
+	if reset {
+		// Merging histogram may take some time.
+		// Put all histograms aside and merge later.
+		toMerge := make([]*stats.Histogram, len(bc.lockingHistograms))
+		for i := range bc.lockingHistograms {
+			toMerge[i] = bc.lockingHistograms[i].swap(stats.NewHistogram(bc.histogramOptions))
+		}
+
+		for i := 0; i < len(toMerge); i++ {
+			mergedHistogram.Merge(toMerge[i])
+		}
+
+		wallTimeElapsed = time.Since(bc.lastResetTime).Seconds()
+		latestRusage := syscall.GetRusage()
+		uTimeElapsed, sTimeElapsed = syscall.CPUTimeDiff(bc.rusageLastReset, latestRusage)
+
+		bc.rusageLastReset = latestRusage
+		bc.lastResetTime = time.Now()
+	} else {
+		// Merge only, not reset.
+		for i := range bc.lockingHistograms {
+			bc.lockingHistograms[i].mergeInto(mergedHistogram)
+		}
+
+		wallTimeElapsed = time.Since(bc.lastResetTime).Seconds()
+		uTimeElapsed, sTimeElapsed = syscall.CPUTimeDiff(bc.rusageLastReset, syscall.GetRusage())
+	}
+
+	b := make([]uint32, len(mergedHistogram.Buckets))
+	for i, v := range mergedHistogram.Buckets {
+		b[i] = uint32(v.Count)
+	}
+	return &testpb.ClientStats{
+		Latencies: &testpb.HistogramData{
+			Bucket:       b,
+			MinSeen:      float64(mergedHistogram.Min),
+			MaxSeen:      float64(mergedHistogram.Max),
+			Sum:          float64(mergedHistogram.Sum),
+			SumOfSquares: float64(mergedHistogram.SumOfSquares),
+			Count:        float64(mergedHistogram.Count),
+		},
+		TimeElapsed: wallTimeElapsed,
+		TimeUser:    uTimeElapsed,
+		TimeSystem:  sTimeElapsed,
+	}
+}
+
+func (bc *benchmarkClient) shutdown() {
+	bc.closeConns()
+}

benchmark/worker/benchmark_server.go

@@ -0,0 +1,187 @@
+/*
+ *
+ * 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 (
+	"flag"
+	"fmt"
+	"net"
+	"runtime"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/benchmark"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/internal/syscall"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/testdata"
+)
+
+var (
+	certFile = flag.String("tls_cert_file", "", "The TLS cert file")
+	keyFile  = flag.String("tls_key_file", "", "The TLS key file")
+)
+
+type benchmarkServer struct {
+	port      int
+	cores     int
+	closeFunc func()
+
+	mu              sync.Mutex
+	lastResetTime   time.Time
+	rusageLastReset *syscall.Rusage
+}
+
+func printServerConfig(config *testpb.ServerConfig) {
+	// Some config options are ignored:
+	// - server type:
+	//     will always start sync server
+	// - async server threads
+	// - core list
+	logger.Infof(" * server type: %v (ignored, always starts sync server)", config.ServerType)
+	logger.Infof(" * async server threads: %v (ignored)", config.AsyncServerThreads)
+	// TODO: use cores specified by CoreList when setting list of cores is supported in go.
+	logger.Infof(" * core list: %v (ignored)", config.CoreList)
+
+	logger.Infof(" - security params: %v", config.SecurityParams)
+	logger.Infof(" - core limit: %v", config.CoreLimit)
+	logger.Infof(" - port: %v", config.Port)
+	logger.Infof(" - payload config: %v", config.PayloadConfig)
+}
+
+func startBenchmarkServer(config *testpb.ServerConfig, serverPort int) (*benchmarkServer, error) {
+	printServerConfig(config)
+
+	// Use all cpu cores available on machine by default.
+	// TODO: Revisit this for the optimal default setup.
+	numOfCores := runtime.NumCPU()
+	if config.CoreLimit > 0 {
+		numOfCores = int(config.CoreLimit)
+	}
+	runtime.GOMAXPROCS(numOfCores)
+
+	opts := []grpc.ServerOption{
+		grpc.WriteBufferSize(128 * 1024),
+		grpc.ReadBufferSize(128 * 1024),
+	}
+
+	// Sanity check for server type.
+	switch config.ServerType {
+	case testpb.ServerType_SYNC_SERVER:
+	case testpb.ServerType_ASYNC_SERVER:
+	case testpb.ServerType_ASYNC_GENERIC_SERVER:
+	default:
+		return nil, status.Errorf(codes.InvalidArgument, "unknown server type: %v", config.ServerType)
+	}
+
+	// Set security options.
+	if config.SecurityParams != nil {
+		if *certFile == "" {
+			*certFile = testdata.Path("server1.pem")
+		}
+		if *keyFile == "" {
+			*keyFile = testdata.Path("server1.key")
+		}
+		creds, err := credentials.NewServerTLSFromFile(*certFile, *keyFile)
+		if err != nil {
+			logger.Fatalf("failed to generate credentials: %v", err)
+		}
+		opts = append(opts, grpc.Creds(creds))
+	}
+
+	// Priority: config.Port > serverPort > default (0).
+	port := int(config.Port)
+	if port == 0 {
+		port = serverPort
+	}
+	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+	if err != nil {
+		logger.Fatalf("Failed to listen: %v", err)
+	}
+	addr := lis.Addr().String()
+
+	// Create different benchmark server according to config.
+	var closeFunc func()
+	if config.PayloadConfig != nil {
+		switch payload := config.PayloadConfig.Payload.(type) {
+		case *testpb.PayloadConfig_BytebufParams:
+			opts = append(opts, grpc.CustomCodec(byteBufCodec{}))
+			closeFunc = benchmark.StartServer(benchmark.ServerInfo{
+				Type:     "bytebuf",
+				Metadata: payload.BytebufParams.RespSize,
+				Listener: lis,
+			}, opts...)
+		case *testpb.PayloadConfig_SimpleParams:
+			closeFunc = benchmark.StartServer(benchmark.ServerInfo{
+				Type:     "protobuf",
+				Listener: lis,
+			}, opts...)
+		case *testpb.PayloadConfig_ComplexParams:
+			return nil, status.Errorf(codes.Unimplemented, "unsupported payload config: %v", config.PayloadConfig)
+		default:
+			return nil, status.Errorf(codes.InvalidArgument, "unknown payload config: %v", config.PayloadConfig)
+		}
+	} else {
+		// Start protobuf server if payload config is nil.
+		closeFunc = benchmark.StartServer(benchmark.ServerInfo{
+			Type:     "protobuf",
+			Listener: lis,
+		}, opts...)
+	}
+
+	logger.Infof("benchmark server listening at %v", addr)
+	addrSplitted := strings.Split(addr, ":")
+	p, err := strconv.Atoi(addrSplitted[len(addrSplitted)-1])
+	if err != nil {
+		logger.Fatalf("failed to get port number from server address: %v", err)
+	}
+
+	return &benchmarkServer{
+		port:            p,
+		cores:           numOfCores,
+		closeFunc:       closeFunc,
+		lastResetTime:   time.Now(),
+		rusageLastReset: syscall.GetRusage(),
+	}, nil
+}
+
+// getStats returns the stats for benchmark server.
+// It resets lastResetTime if argument reset is true.
+func (bs *benchmarkServer) getStats(reset bool) *testpb.ServerStats {
+	bs.mu.Lock()
+	defer bs.mu.Unlock()
+	wallTimeElapsed := time.Since(bs.lastResetTime).Seconds()
+	rusageLatest := syscall.GetRusage()
+	uTimeElapsed, sTimeElapsed := syscall.CPUTimeDiff(bs.rusageLastReset, rusageLatest)
+
+	if reset {
+		bs.lastResetTime = time.Now()
+		bs.rusageLastReset = rusageLatest
+	}
+	return &testpb.ServerStats{
+		TimeElapsed: wallTimeElapsed,
+		TimeUser:    uTimeElapsed,
+		TimeSystem:  sTimeElapsed,
+	}
+}

benchmark/worker/main.go

@@ -0,0 +1,239 @@
+/*
+ *
+ * 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.
+ *
+ */
+
+// Binary worker implements the benchmark worker that can turn into a benchmark
+// client or server.
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	_ "net/http/pprof"
+	"runtime"
+	"strconv"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/status"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+var (
+	driverPort    = flag.Int("driver_port", 10000, "port for communication with driver")
+	serverPort    = flag.Int("server_port", 0, "port for benchmark server if not specified by server config message")
+	pprofPort     = flag.Int("pprof_port", -1, "Port for pprof debug server to listen on. Pprof server doesn't start if unset")
+	blockProfRate = flag.Int("block_prof_rate", 0, "fraction of goroutine blocking events to report in blocking profile")
+
+	logger = grpclog.Component("benchmark")
+)
+
+type byteBufCodec struct {
+}
+
+func (byteBufCodec) Marshal(v any) ([]byte, error) {
+	b, ok := v.(*[]byte)
+	if !ok {
+		return nil, fmt.Errorf("failed to marshal: %v is not type of *[]byte", v)
+	}
+	return *b, nil
+}
+
+func (byteBufCodec) Unmarshal(data []byte, v any) error {
+	b, ok := v.(*[]byte)
+	if !ok {
+		return fmt.Errorf("failed to marshal: %v is not type of *[]byte", v)
+	}
+	*b = data
+	return nil
+}
+
+func (byteBufCodec) String() string {
+	return "bytebuffer"
+}
+
+// workerServer implements WorkerService rpc handlers.
+// It can create benchmarkServer or benchmarkClient on demand.
+type workerServer struct {
+	testgrpc.UnimplementedWorkerServiceServer
+	stop       chan<- bool
+	serverPort int
+}
+
+func (s *workerServer) RunServer(stream testgrpc.WorkerService_RunServerServer) error {
+	var bs *benchmarkServer
+	defer func() {
+		// Close benchmark server when stream ends.
+		logger.Infof("closing benchmark server")
+		if bs != nil {
+			bs.closeFunc()
+		}
+	}()
+	for {
+		in, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+
+		var out *testpb.ServerStatus
+		switch argtype := in.Argtype.(type) {
+		case *testpb.ServerArgs_Setup:
+			logger.Infof("server setup received:")
+			if bs != nil {
+				logger.Infof("server setup received when server already exists, closing the existing server")
+				bs.closeFunc()
+			}
+			bs, err = startBenchmarkServer(argtype.Setup, s.serverPort)
+			if err != nil {
+				return err
+			}
+			out = &testpb.ServerStatus{
+				Stats: bs.getStats(false),
+				Port:  int32(bs.port),
+				Cores: int32(bs.cores),
+			}
+
+		case *testpb.ServerArgs_Mark:
+			logger.Infof("server mark received:")
+			logger.Infof(" - %v", argtype)
+			if bs == nil {
+				return status.Error(codes.InvalidArgument, "server does not exist when mark received")
+			}
+			out = &testpb.ServerStatus{
+				Stats: bs.getStats(argtype.Mark.Reset_),
+				Port:  int32(bs.port),
+				Cores: int32(bs.cores),
+			}
+		}
+
+		if err := stream.Send(out); err != nil {
+			return err
+		}
+	}
+}
+
+func (s *workerServer) RunClient(stream testgrpc.WorkerService_RunClientServer) error {
+	var bc *benchmarkClient
+	ctx, cancel := context.WithCancel(stream.Context())
+	defer func() {
+		cancel()
+		// Shut down benchmark client when stream ends.
+		logger.Infof("shutting down benchmark client")
+		if bc != nil {
+			bc.shutdown()
+		}
+	}()
+	for {
+		in, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+
+		var out *testpb.ClientStatus
+		switch t := in.Argtype.(type) {
+		case *testpb.ClientArgs_Setup:
+			logger.Infof("client setup received:")
+			if bc != nil {
+				logger.Infof("client setup received when client already exists, shutting down the existing client")
+				bc.shutdown()
+			}
+			bc, err = startBenchmarkClient(ctx, t.Setup)
+			if err != nil {
+				return err
+			}
+			out = &testpb.ClientStatus{
+				Stats: bc.getStats(false),
+			}
+
+		case *testpb.ClientArgs_Mark:
+			logger.Infof("client mark received:")
+			logger.Infof(" - %v", t)
+			if bc == nil {
+				return status.Error(codes.InvalidArgument, "client does not exist when mark received")
+			}
+			out = &testpb.ClientStatus{
+				Stats: bc.getStats(t.Mark.Reset_),
+			}
+		}
+
+		if err := stream.Send(out); err != nil {
+			return err
+		}
+	}
+}
+
+func (s *workerServer) CoreCount(context.Context, *testpb.CoreRequest) (*testpb.CoreResponse, error) {
+	logger.Infof("core count: %v", runtime.NumCPU())
+	return &testpb.CoreResponse{Cores: int32(runtime.NumCPU())}, nil
+}
+
+func (s *workerServer) QuitWorker(context.Context, *testpb.Void) (*testpb.Void, error) {
+	logger.Infof("quitting worker")
+	s.stop <- true
+	return &testpb.Void{}, nil
+}
+
+func main() {
+	grpc.EnableTracing = false
+
+	flag.Parse()
+	lis, err := net.Listen("tcp", ":"+strconv.Itoa(*driverPort))
+	if err != nil {
+		logger.Fatalf("failed to listen: %v", err)
+	}
+	logger.Infof("worker listening at port %v", *driverPort)
+
+	s := grpc.NewServer()
+	stop := make(chan bool)
+	testgrpc.RegisterWorkerServiceServer(s, &workerServer{
+		stop:       stop,
+		serverPort: *serverPort,
+	})
+
+	go func() {
+		<-stop
+		// Wait for 1 second before stopping the server to make sure the return value of QuitWorker is sent to client.
+		// TODO revise this once server graceful stop is supported in gRPC.
+		time.Sleep(time.Second)
+		s.Stop()
+	}()
+
+	runtime.SetBlockProfileRate(*blockProfRate)
+
+	if *pprofPort >= 0 {
+		go func() {
+			logger.Infoln("Starting pprof server on port " + strconv.Itoa(*pprofPort))
+			logger.Infoln(http.ListenAndServe("localhost:"+strconv.Itoa(*pprofPort), nil))
+		}()
+	}
+
+	s.Serve(lis)
+}