go websocketテストケース

前回websocketサービスを書き終わった後、最大どれだけの接続に耐えられるかをテストしたいと思って、試してみたらバグが出てきました

close(ch)
chdata

package main

import (
	"golang.org/x/net/websocket"
	"log"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var (
	origin = "https://baidu.com"
	url    = "wss://*"
	start  = make(chan struct{})
	//  
	done int32
	//   = msgNums + 1
	msgNums = 50
	wg      sync.WaitGroup
)

func Worker(id int) {
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		log.Println("worker ", id, " dial fail:", err)
		return
	}

	//  
	start

	send := 0

	defer func() {
		ws.Close()
		log.Printf("worker %3d done, send:%3d 
", id, send)
		wg.Done()
	}()

	for {
		msg := []byte("{\"msg\":\"worker " + strconv.Itoa(id) + "\"}")
		_, err := ws.Write(msg)
		if err != nil {
			return
		}
		send++
		//  
		if send > msgNums {
			//  
			atomic.AddInt32(&done, 1)
			return
		}

		time.Sleep(time.Second)
	}
}

func main() {
	//  
	for i := range [70][0]int{} {
		go Worker(i)
		wg.Add(1)
	}
	//  
	close(start)
	//  
	wg.Wait()
	//  
	log.Println("done:", done)

	//  ， 
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		panic(err)
	}
	for range [5][0]int{} {
		msg := []byte("end")
		_, err := ws.Write(msg)
		if err != nil {
			panic(err)
		}
		time.Sleep(time.Millisecond * 500)
	}
}

package main

import (
	"golang.org/x/net/websocket"
	"io/ioutil"
	"log"
	"sync"
)

//  ，  err， 
func Read(ws *websocket.Conn) (data []byte, ok bool) {
	r, err := ws.NewFrameReader()
	if err != nil {
		return
	}
	fr, err := ws.HandleFrame(r)
	if err != nil {
		return
	}
	if fr == nil {
		return
	}
	data, err = ioutil.ReadAll(fr)
	if err != nil {
		return
	}
	return data, true
}

var (
	origin   = "https://baidu.com"
	url      = "wss://*"
	m        sync.Map
	msgNums  = 51
	readChan = make(chan string, 2000)
)

func Worker() {
	ws, err := websocket.Dial(url, "", origin)
	if err != nil {
		log.Println("worker dial fail:", err)
		return
	}

	exit := make(chan struct{})

	go func() {
		for {
			select {
			case s := readChan:
				value, ok := m.Load(s)
				if ok {
					m.Store(s, value.(int)+1)
				} else {
					m.Store(s, 1)
				}
			case exit:
				return
			}
		}
	}()

	for {
		data, ok := Read(ws)
		if ok {
			s := string(data)
			if s == "end" {
				close(exit)
				return
			}
			readChan  s
		}
	}
}

func main() {
	defer func() {
		num := 0
		all := 0
		m.Range(func(key, value interface{}) bool {
			log.Printf("%20s : %3d", key, value)
			num++
			if value == msgNums {
				all++
			}
			return true
		})
		//  worker 
		log.Println("total worker: ", num)
		//  
		log.Println("all recv num: ", all)
	}()
	Worker()
}