main.go

package main

import (
	"flag"
	"net"
	"net/url"
	"os/exec"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"go.coder.com/crand"
	"go.uber.org/zap"
)

var dialer = net.Dialer{
	Timeout:   time.Second * 10,
	KeepAlive: time.Second * 10,
}

func main() {
	log, _ := zap.NewDevelopment()
	var (
		wsURL   = flag.String("ws-url", "", "")
		clients = flag.Int("clients", 1, "number of dials")
	)
	flag.Parse()

	u, err := url.Parse(*wsURL)
	if err != nil {
		log.Error("failed to parse ws-url",
			zap.Error(err),
		)
		return
	}
	log.Info("starting all connections",
		zap.Stringer("ws-url", u),
	)

	query := u.Query()

	var ws sync.WaitGroup
	for i := 0; i < *clients; i++ {
		id := crand.MustString(10)
		log := log.With(
			zap.Int("i", i),
			zap.String("id", id),
		)
		query.Set("id", id)
		u.RawQuery = query.Encode()
		websocket.DefaultDialer.NetDial = func(network, addr string) (net.Conn, error) {
			return dialer.Dial(network, addr)
		}
		c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
		if err != nil {
			log.Error("failed to dial",
				zap.Error(err),
			)
			continue
		}
		log.Info("dial succeeded",
			zap.String("url", u.String()),
		)

		pih := c.PingHandler()
		c.SetPingHandler(func(d string) error {
			log.Info("got ping")
			return pih(d)
		})
		poh := c.PongHandler()
		c.SetPongHandler(func(d string) error {
			log.Info("got pong")
			return poh(d)
		})

		go func() {
			defer c.Close()

			t := time.NewTicker(time.Second * 10)
			defer t.Stop()
			for range t.C {
				deadline := time.Now().Add(time.Second * 20)
				err := c.WriteControl(websocket.PingMessage, nil, deadline)
				if err != nil {
					log.Error("failed to write ping",
						zap.Error(err),
					)
					return
				}
				log.Info("pinged")
			}
		}()

		ws.Add(1)
		go func() {
			defer c.Close()
			defer exec.Command("noti").Run()
			defer ws.Done()

			_, _, err := c.NextReader()
			log.Error("failed to get next reader",
				zap.Error(err),
			)
		}()
		time.Sleep(time.Millisecond * 50)
	}
	log.Info("initialized all connections")
	ws.Wait()
}