Golang Conn.SetPingHandler方法代码示例

func connectToAgent(timeout int, onReady func(*websocket.Conn)) (*websocket.Conn, error) {
	var ws *websocket.Conn
	var err error

	t := 0
	for {
		if t > timeout {
			return nil, errors.New("Connection to Agent : timeout reached")
		}

		ws, err = newClient()
		if err == nil {
			break
		}
		time.Sleep(1 * time.Second)
		t++
	}

	ready := false
	h := func(message string) error {
		err := ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Second))
		if err != nil {
			return err
		}
		if !ready {
			ready = true
			onReady(ws)
		}
		return nil
	}
	ws.SetPingHandler(h)

	return ws, nil
}

func AddWSSink(receivedChan chan []byte, port string, path string) (*websocket.Conn, <-chan struct{}) {
	connectionDroppedChannel := make(chan struct{}, 1)

	var ws *websocket.Conn

	ip, _ := localip.LocalIP()
	fullURL := "ws://" + ip + ":" + port + path

	Eventually(func() error {
		var err error
		ws, _, err = websocket.DefaultDialer.Dial(fullURL, http.Header{})
		return err
	}, 5, 1).ShouldNot(HaveOccurred(), fmt.Sprintf("Unable to connect to server at %s.", fullURL))

	ws.SetPingHandler(func(message string) error {
		ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
		return nil
	})

	go func() {
		for {
			_, data, err := ws.ReadMessage()
			if err != nil {
				close(connectionDroppedChannel)
				close(receivedChan)
				return
			}

			receivedChan <- data
		}

	}()

	return ws, connectionDroppedChannel
}

/* reader reads input from the web socket connection and sends it to the master.
 * reader runs in its own goroutine.
 * If the connection closes, the reader goroutine terminates.
 */
func reader(conn *websocket.Conn, ch chan readerResult) {
	conn.SetPongHandler(func(appData string) error {
		ch <- readerResult{conn, 10, nil, nil}
		return nil
	})

	conn.SetPingHandler(func(appData string) error {
		ch <- readerResult{conn, 9, nil, nil}
		return nil
	})

	for {
		messageType, data, err := conn.ReadMessage()
		readerResult := readerResult{conn, messageType, data, err}
		ch <- readerResult
		if err != nil {
			return
		}
	}
}

func wsSetPingHandler(t *WSTunnelServer, ws *websocket.Conn, rs *remoteServer) {
	// timeout handler sends a close message, waits a few seconds, then kills the socket
	timeout := func() {
		ws.WriteControl(websocket.CloseMessage, nil, time.Now().Add(1*time.Second))
		time.Sleep(5 * time.Second)
		rs.log.Info("WS closing due to ping timeout", "ws", wsp(ws))
		ws.Close()
	}
	// timeout timer
	timer := time.AfterFunc(t.WSTimeout, timeout)
	// ping handler resets last ping time
	ph := func(message string) error {
		timer.Reset(t.WSTimeout)
		ws.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(t.WSTimeout/3))
		// update lastActivity
		rs.lastActivity = time.Now()
		return nil
	}
	ws.SetPingHandler(ph)
}

// Reads incoming data from the websocket and forwards it to stdout.
func readWS(ws *websocket.Conn, query *regexp.Regexp, done chan struct{}) {
	defer func() {
		done <- struct{}{}
	}()
	ws.SetPingHandler(func(string) error {
		ws.SetWriteDeadline(time.Now().Add(writeTimeout))
		return ws.WriteMessage(websocket.PongMessage, []byte{})
	})
	for {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			return
		}
		var log LogMessage
		err = json.Unmarshal(msg, &log)
		if err == nil {
			if query == nil || query.MatchString(log.Message) {
				logrus.Printf("%s - %s", log.Timestamp, log.Message)
			}
		} else {
			logrus.StandardLogger().Out.Write(msg)
		}
	}
}

func addWSSink(receivedChan chan []byte, url string) (chan struct{}, <-chan struct{}, func()) {
	stopKeepAlive := make(chan struct{})
	connectionDropped := make(chan struct{})

	var ws *websocket.Conn
	tryToConnect := func() error {
		var err error
		ws, _, err = websocket.DefaultDialer.Dial(url, http.Header{})
		return err
	}
	Eventually(tryToConnect, 5).ShouldNot(HaveOccurred())

	ws.SetPingHandler(func(message string) error {
		select {
		case <-stopKeepAlive:
			return nil
		default:
			return ws.WriteControl(websocket.PongMessage, []byte(message), time.Time{})
		}
	})

	go func() {
		defer close(connectionDropped)
		for {
			_, data, err := ws.ReadMessage()
			if err != nil {
				return
			}
			receivedChan <- data
		}
	}()

	return stopKeepAlive, connectionDropped, func() {
		ws.Close()
	}
}

		keepAliveCompleted = make(chan struct{})
		testServer = httptest.NewServer(makeTestHandler(keepAliveCompleted))
		var err error
		wsClient, _, err = websocket.DefaultDialer.Dial(httpToWs(testServer.URL), nil)
		Expect(err).NotTo(HaveOccurred())
	})

	AfterEach(func() {
		wsClient.Close()
		testServer.Close()
	})

	It("sends pings to the client", func() {
		var pingCount int32
		wsClient.SetPingHandler(func(string) error {
			atomic.AddInt32(&pingCount, 1)
			return nil
		})

		go wsClient.ReadMessage()
		Eventually(func() int32 { return atomic.LoadInt32(&pingCount) }).ShouldNot(BeZero())
	})

	It("doesn't close the client when it responds with pong frames", func() {
		go wsClient.ReadMessage()
		// default ping handler responds with pong frames
		Consistently(keepAliveCompleted).ShouldNot(BeClosed())
	})

	It("closes the client when it doesn't respond with pong frames", func() {
		wsClient.SetPingHandler(func(string) error { return nil })
		go wsClient.ReadMessage()

func CmdLoop() {
	const (
		// Time allowed to write a message to the peer.
		writeWait = 10 * time.Second

		// Time allowed to resond to a ping event from peer
		pingWait = 55 * time.Second
	)

	var (
		// Connection object
		ws *websocket.Conn
	)

	bigobjectURL := url.URL{
		Scheme: "ws",
		Host:   net.JoinHostPort(host, port),
		Path:   "exec",
	}

	var err error
	dialer := websocket.DefaultDialer
	if ws, _, err = dialer.Dial(bigobjectURL.String(), nil); err != nil {
		if !quiet {
			log.Fatalln(err)
		} else {
			os.Exit(1)
		}
	}
	ws.SetPingHandler(nil)

	input, next := MakeReadline()
	ticker := time.Tick(pingWait)

	defer func() {
		if err := recover(); err != nil {
			ws.Close()
			if !quiet {
				log.Fatalln(err)
			} else {
				os.Exit(1)
			}
		}
	}()

	for {
		select {
		case stmt, ok := <-input:
			if !ok {
				return
			}
			now := time.Now()
			switch parseAction(stmt) {
			default:
				req := &api.RPCRequest{Stmt: stmt}
				if err := ws.WriteJSON(req); err != nil {
					if !quiet {
						fmt.Println(err)
					}
				} else {
					var resp interface{}
					if err := ws.ReadJSON(&resp); err != nil {
						panic(err)
					}
					text, _ := json.MarshalIndent(resp, "", "    ")
					fmt.Println(string(text))
				}
				break
			case GET_RESULT_NOW:
				req := &api.RPCRequest{Stmt: stmt, Opts: &api.RPCOpts{Handle: handle}}
				if err := ws.WriteJSON(req); err != nil {
					if !quiet {
						fmt.Println(err)
					}
				} else if !handle {
					var idx int64 = 1
					for idx > 0 {
						var resp map[string]interface{}
						if err := ws.ReadJSON(&resp); err != nil {
							panic(err)
						}
						if thing, ok := resp["Content"]; ok && thing != nil {
							payload := thing.(map[string]interface{})
							data := payload["content"].([]interface{})
							idx = int64(payload["index"].(float64))
							for _, thing := range data {
								text, _ := json.MarshalIndent(thing, "", "    ")
								fmt.Println(string(text))
							}
						} else {
							text, _ := json.MarshalIndent(resp, "", "    ")
							fmt.Println(string(text))
							break
						}
					}
				} else {
					var resp interface{}
					if err := ws.ReadJSON(&resp); err != nil {
						panic(err)
					}
//.........这里部分代码省略.........