Golang websocket.Conn類代碼示例

// new websocket connection wrapper
func NewWSConnection(wsConn *websocket.Conn) *WSConnection {
	return &WSConnection{
		id:        wsConn.RemoteAddr().String(),
		wsConn:    wsConn,
		writeChan: make(chan WSResponse, WriteChanBufferSize), // buffered. we keep track when its full
		quitChan:  make(chan struct{}),
	}
}

// new websocket connection wrapper
func NewWSConnection(baseConn *websocket.Conn, funcMap map[string]*RPCFunc, evsw *events.EventSwitch) *WSConnection {
	wsc := &WSConnection{
		id:        baseConn.RemoteAddr().String(),
		baseConn:  baseConn,
		writeChan: make(chan RPCResponse, writeChanCapacity), // error when full.
		funcMap:   funcMap,
		evsw:      evsw,
	}
	wsc.QuitService = *NewQuitService(log, "WSConnection", wsc)
	return wsc
}

// unsubscribe from an event
func unsubscribe(t *testing.T, con *websocket.Conn, eventid string) {
	err := con.WriteJSON(rpctypes.RPCRequest{
		JSONRPC: "2.0",
		ID:      "",
		Method:  "unsubscribe",
		Params:  []interface{}{eventid},
	})
	if err != nil {
		t.Fatal(err)
	}
}

func writer(ws *websocket.Conn, lastMod time.Time) {
	lastError := ""
	pingTicker := time.NewTicker(pingPeriod)
	fileTicker := time.NewTicker(filePeriod)
	defer func() {
		pingTicker.Stop()
		fileTicker.Stop()
		ws.Close()
	}()
	for {
		select {
		case <-fileTicker.C:
			var p []byte
			var err error

			p, lastMod, err = readFileIfModified(lastMod)

			if err != nil {
				if s := err.Error(); s != lastError {
					lastError = s
					p = []byte(lastError)
				}
			} else {
				lastError = ""
			}

			if p != nil {
				ws.SetWriteDeadline(time.Now().Add(writeWait))
				if err := ws.WriteMessage(websocket.TextMessage, p); err != nil {
					return
				}
			}
		case <-pingTicker.C:
			ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func reader(ws *websocket.Conn) {
	defer ws.Close()
	ws.SetReadLimit(512)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			break
		}
	}
}

// wait for an event; do things that might trigger events, and check them when they are received
// the check function takes an event id and the byte slice read off the ws
func waitForEvent(t *testing.T, con *websocket.Conn, eventid string, dieOnTimeout bool, f func(), check func(string, []byte) error) {
	// go routine to wait for webscoket msg
	goodCh := make(chan []byte)
	errCh := make(chan error)
	quitCh := make(chan struct{})
	defer close(quitCh)

	// Read message
	go func() {
		for {
			_, p, err := con.ReadMessage()
			if err != nil {
				errCh <- err
				break
			} else {
				// if the event id isnt what we're waiting on
				// ignore it
				var response ctypes.Response
				var err error
				wire.ReadJSON(&response, p, &err)
				if err != nil {
					errCh <- err
					break
				}
				event, ok := response.Result.(*ctypes.ResultEvent)
				if ok && event.Event == eventid {
					goodCh <- p
					break
				}
			}
		}
	}()

	// do stuff (transactions)
	f()

	// wait for an event or timeout
	timeout := time.NewTimer(10 * time.Second)
	select {
	case <-timeout.C:
		if dieOnTimeout {
			con.Close()
			t.Fatalf("%s event was not received in time", eventid)
		}
		// else that's great, we didn't hear the event
		// and we shouldn't have
	case p := <-goodCh:
		if dieOnTimeout {
			// message was received and expected
			// run the check
			err := check(eventid, p)
			if err != nil {
				t.Fatal(err)
				panic(err) // Show the stack trace.
			}
		} else {
			con.Close()
			t.Fatalf("%s event was not expected", eventid)
		}
	case err := <-errCh:
		t.Fatal(err)
		panic(err) // Show the stack trace.
	}
}