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{
		HandshakeTimeout: 1 * time.Second,
		TLSClientConfig:  &tls.Config{InsecureSkipVerify: true},
	}

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

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

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

	<-interrupt
	waiter.Done()
}