From bb6e22c4e01618c4e0c999ff20a034e70c6e26ad Mon Sep 17 00:00:00 2001
From: prometheus <prometheustar@outlook.com>
Date: Thu, 15 Feb 2024 23:45:11 +0800
Subject: [PATCH] benchmark

---
 benchmark/main.go | 16 ++++++++++++----
 generate.go       |  5 ++++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/benchmark/main.go b/benchmark/main.go
index eef6bb9..585bc3c 100644
--- a/benchmark/main.go
+++ b/benchmark/main.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"flag"
 	"fmt"
 	"github.com/bytedance/sonic"
 	"github.com/gorilla/websocket"
@@ -33,7 +34,7 @@ import (
 )
 
 var (
-	benchmarkMode   = "deliver" // deliver / groupDeliver
+	benchmarkMode   = "deliver" // deliver/group
 	mockUsers       = 2000
 	eachUserSend    = 100
 	mockNetUsers    []*NetUser
@@ -68,8 +69,15 @@ func init() {
 }
 
 func main() {
-	// benchmark()
-	benchmarkGroup()
+	flag.StringVar(&benchmarkMode, "mode", "deliver", "benchmark mode: deliver/group")
+	flag.IntVar(&mockUsers, "users", 10, "mock users")
+	flag.IntVar(&eachUserSend, "send", 10, "each user send msg count")
+
+	if benchmarkMode == "group" {
+		benchmarkGroup()
+	} else {
+		benchmark()
+	}
 }
 
 func records(ctx context.Context) {
@@ -335,7 +343,7 @@ func listen(ctx context.Context, conn *websocket.Conn) {
 		}
 		header := payload.Header
 
-		if benchmarkMode == "groupDeliver" {
+		if benchmarkMode == "group" {
 			atomic.AddInt64(&receiverCounter, 1)
 		}
 
diff --git a/generate.go b/generate.go
index b6e80e3..0d9cbe8 100644
--- a/generate.go
+++ b/generate.go
@@ -3,7 +3,10 @@ package main
 import "os"
 
 //go:generate go run generate.go
-//go:generate protoc --go_out=./api/gen --go-grpc_out=./api/gen ./api/*.proto
+//go:generate protoc --go_out=./api/gen --go-grpc_out=./api/gen ./api/postal.proto
+//go:generate protoc --go_out=./api/gen --go-grpc_out=./api/gen ./api/auth.proto
+//go:generate protoc --go_out=./api/gen --go-grpc_out=./api/gen ./api/chat.proto
+//go:generate protoc --go_out=./api/gen --go-grpc_out=./api/gen ./api/mahjong.proto
 //go:generate pbjs -t static-module -w es6 -o api/genjs/postal.js api/postal.proto --no-service
 //go:generate pbjs -t static-module -w es6 -o api/genjs/auth.js api/auth.proto --no-service
 //go:generate pbjs -t static-module -w es6 -o api/genjs/chat.js api/chat.proto --no-service