websocket/examples: enhance the server and client

In the server, change the listen port to 9101 to prevent conflict
with autobahn test suite.

In the client, add command "chat" and "chatbot".
The "chat" command run client chart as before.
The "chatbot" command create client for all known accounts and send
N messages continuously.

2 files changed, 104 insertions, 14 deletions

diff --git a/lib/websocket/examples/cmd/client/main.go b/lib/websocket/examples/cmd/client/main.go
index 53d43416..013b08b1 100644
--- a/lib/websocket/examples/cmd/client/main.go
+++ b/lib/websocket/examples/cmd/client/main.go
@@ -7,12 +7,12 @@
 //
 // To run the client as user ID 1 (Groot),
 //
-//	$ go run . 1
+//	$ go run . chat 1
 //
 // You can open other terminal and run another clients,
 //
-//	$ go run . 2 # or
-//	$ go run . 3
+//	$ go run . chat 2 # or
+//	$ go run . chat 3
 //
 // and start chatting with each others.
 package main
@@ -20,18 +20,25 @@ package main
 import (
 	"bufio"
 	"encoding/json"
+	"flag"
 	"fmt"
 	"log"
 	"net/http"
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/shuLhan/share/lib/websocket"
 	"github.com/shuLhan/share/lib/websocket/examples"
 )
 
+const (
+	cmdChat    = `chat`
+	cmdChatbot = `chatbot`
+)
+
 type ChatClient struct {
 	user *examples.Account
 	conn *websocket.Client
@@ -42,7 +49,7 @@ func NewChatClient(user *examples.Account) (cc *ChatClient) {
 	cc = &ChatClient{
 		user: user,
 		conn: &websocket.Client{
-			Endpoint: "ws://127.0.0.1:9001",
+			Endpoint: `ws://127.0.0.1:9101`,
 			Headers: http.Header{
 				"Key": []string{user.Key},
 			},
@@ -126,14 +133,41 @@ func (cc *ChatClient) handleText(cl *websocket.Client, frame *websocket.Frame) (
 	return nil
 }
 
+func usage() {
+	fmt.Println(`= WebSocket client example
+
+	client <chat | chatbot> <args...>
+
+== USAGE
+
+client ` + cmdChat + ` <id>
+	Connect to chat with others using specific ID: 1, 2, or 3.
+
+client ` + cmdChatbot + ` <N>
+	Connect to the server and sent N messages for each user
+	simultaneously.`)
+}
+
 func main() {
-	log.SetFlags(0)
+	flag.Parse()
 
-	if len(os.Args) <= 1 {
-		log.Printf("client <id>")
+	if len(os.Args) <= 2 {
+		usage()
 		return
 	}
 
+	var cmd = strings.ToLower(flag.Arg(0))
+	switch cmd {
+	case cmdChat:
+		doChat(flag.Arg(1))
+	case cmdChatbot:
+		doChatbot(flag.Arg(1))
+	default:
+		log.Fatalf(`unknown command: %s`, cmd)
+	}
+}
+
+func doChat(userIDStr string) {
 	var (
 		user *examples.Account
 		cc   *ChatClient
@@ -142,7 +176,7 @@ func main() {
 		ok   bool
 	)
 
-	uid, err = strconv.Atoi(os.Args[1])
+	uid, err = strconv.Atoi(userIDStr)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -156,3 +190,60 @@ func main() {
 
 	cc.Start()
 }
+
+func doChatbot(nStr string) {
+	var (
+		wg   sync.WaitGroup
+		user *examples.Account
+		err  error
+		n    int64
+	)
+
+	n, err = strconv.ParseInt(nStr, 10, 64)
+	if err != nil {
+		log.Fatalf(`invalid N: %s`, err)
+	}
+
+	for _, user = range examples.Users {
+		wg.Add(1)
+		go runChatbot(&wg, user, n)
+	}
+	wg.Wait()
+}
+
+func runChatbot(wg *sync.WaitGroup, user *examples.Account, n int64) {
+	var (
+		req = &websocket.Request{
+			Method: http.MethodPost,
+			Target: `/message`,
+		}
+
+		err    error
+		packet []byte
+		x      int64
+	)
+
+	var cc = NewChatClient(user)
+
+	for ; x < n; x++ {
+		req.ID = uint64(time.Now().UnixNano())
+		req.Body = fmt.Sprintf(`#%d Hello from %s at %d`, x, user.Name, req.ID)
+
+		packet, err = json.Marshal(req)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		err = cc.conn.SendText(packet)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}
+	err = cc.conn.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+	wg.Done()
+}
diff --git a/lib/websocket/examples/cmd/server/main.go b/lib/websocket/examples/cmd/server/main.go
index 49733fb4..18a0777c 100644
--- a/lib/websocket/examples/cmd/server/main.go
+++ b/lib/websocket/examples/cmd/server/main.go
@@ -20,11 +20,9 @@ import (
 var server *websocket.Server
 
 func main() {
-	log.SetFlags(0)
-
 	var (
 		opts = &websocket.ServerOptions{
-			Address: ":9001",
+			Address: `:9101`,
 			// Register the authentication handler.
 			HandleAuth:         handleAuth,
 			HandleClientAdd:    handleClientAdd,
@@ -44,7 +42,10 @@ func main() {
 
 	log.Println("server: starting ...")
 
-	server.Start()
+	err = server.Start()
+	if err != nil {
+		log.Fatal(err)
+	}
 }
 
 // handleAuth authenticated the new connection by checking the Header "Key"
@@ -141,8 +142,6 @@ func handlePostMessage(ctx context.Context, req *websocket.Request) (res websock
 		conn   int
 	)
 
-	log.Printf("server: message from %s: %q\n", user.Name, req.Body)
-
 	packet, err = websocket.NewBroadcast(examples.BroadcastMessage, body)
 	if err != nil {
 		res.Code = http.StatusInternalServerError