wsbench.go

package main

import (
	"crypto/tls"
	"flag"
	"log"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"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")

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{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}

	waiter := &sync.WaitGroup{}
	waiter.Add(1)

	for i := 0; i < *conns; i++ {
		time.Sleep(10 * time.Millisecond)
		go func() {
			c, _, err := dialer.Dial(u.String(), nil)
			if err != nil {
				log.Fatal("dial:", err)
			}
			defer c.Close()

			waiter.Wait()
		}()
	}

	<-interrupt
	waiter.Done()
}