wsbench.go

package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"time"

	"github.com/gorilla/websocket"
)

var dest = flag.String("dest", "wss://localhost:5001/ws", "websocket endpoint")
var conns = flag.Int("conns", 10, "number of connections")
var token = flag.String("token", "", "auth token")

func main() {
	flag.Parse()

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	u, err := url.Parse(*dest)
	if err != nil {
		log.Fatal("url parse:", err)
	}
	log.Printf("connecting to %s", u.String())

	dialer := &websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
	}

	closer := make(chan struct{})

	for i := 0; i < *conns; i++ {
		time.Sleep(10 * time.Millisecond)
		go func(count int) {
			headers := http.Header{}
			if *token != "" {
				headers.Add("Authorization", fmt.Sprintf("Bearer %s", token))
			}
			c, _, err := dialer.Dial(u.String(), headers)
			if err != nil {
				log.Fatalf("failure dialing connection %d: %s", count, err)
			}
			defer c.Close()

			msg := fmt.Sprintf(`{"messageType": "client_id", "value": "wsbench-%d"}`, count)
			err = c.WriteMessage(websocket.TextMessage, []byte(msg))
			if err != nil {
				log.Fatalf("failed to register client %d", count)
			}

			msg = `{"filter": {"match_policy": "include_all", "clauses": [], "actions": {"create": true, "update": true, "delete": true}}}`
			err = c.WriteMessage(websocket.TextMessage, []byte(msg))
			if err != nil {
				log.Fatalf("failed to setup filter for client %d", count)
			}

			<-closer
		}(i + 1)
		if (i+1)%100 == 0 {
			log.Printf("opened %d connections", i+1)
		}
	}

	<-interrupt
	log.Printf("closing connections")
	close(closer)
}