origin/network/websocketclient.go

package network

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"runtime/debug"

	"github.com/duanhf2012/origin/service"
	"github.com/duanhf2012/origin/sysmodule"

	"github.com/gorilla/websocket"

	"time"
)

//IWebsocketClient ...
type IWebsocketClient interface {
	Init(slf IWebsocketClient, strurl, strProxyPath string, timeoutsec time.Duration) error
	Start() error
	WriteMessage(msg []byte) error
	OnDisconnect() error
	OnConnected() error
	OnReadMessage(msg []byte) error
	ReConnect()
}

//WebsocketClient ...
type WebsocketClient struct {
	WsDailer   *websocket.Dialer
	conn       *websocket.Conn
	url        string
	state      int //0未连接状态   1正在重连   2连接状态
	bwritemsg  chan []byte
	closer     chan bool
	slf        IWebsocketClient
	timeoutsec time.Duration

	bRun bool
	ping []byte
}

const (
	MAX_WRITE_MSG = 10240
)

//Init ...
func (ws *WebsocketClient) Init(slf IWebsocketClient, strurl, strProxyPath string, timeoutsec time.Duration) error {

	ws.timeoutsec = timeoutsec
	ws.slf = slf
	if strProxyPath != "" {
		proxy := func(_ *http.Request) (*url.URL, error) {
			return url.Parse(strProxyPath)
		}

		if timeoutsec > 0 {
			tosec := timeoutsec * time.Second
			ws.WsDailer = &websocket.Dialer{Proxy: proxy, HandshakeTimeout: tosec}
		} else {
			ws.WsDailer = &websocket.Dialer{Proxy: proxy}
		}
	} else {
		if timeoutsec > 0 {
			tosec := timeoutsec * time.Second
			ws.WsDailer = &websocket.Dialer{HandshakeTimeout: tosec}
		} else {
			ws.WsDailer = &websocket.Dialer{}
		}
	}

	ws.url = strurl
	ws.ping = []byte(`ping`)
	return nil
}

func (ws *WebsocketClient) SetPing(ping string) {
	ws.ping = []byte(ping)
}

//OnRun ...
func (ws *WebsocketClient) OnRun() error {
	defer func() {
		if r := recover(); r != nil {
			coreInfo := string(debug.Stack())
			coreInfo += "\n" + fmt.Sprintf("Core WebsocketClient url is %s. Core information is %v\n", ws.url, r)
			service.GetLogger().Printf(service.LEVER_FATAL, coreInfo)
			go ws.OnRun()
		}
	}()

	for {
		if ws.bRun == false {
			break
		}

		if ws.state == 0 {
			time.Sleep(2 * time.Second)
			ws.StartConnect()
		} else if ws.state == 1 {
			ws.state = 0
			close(ws.closer)
			ws.conn.Close()
			ws.slf.OnDisconnect()
		} else if ws.state == 2 {
			ws.conn.SetReadDeadline(time.Now().Add(ws.timeoutsec * time.Second))
			_, message, err := ws.conn.ReadMessage()

			if err != nil {
				service.GetLogger().Printf(service.LEVER_WARN, "websocket client is disconnect [%s],information is %v", ws.url, err)
				ws.conn.Close()
				ws.state = 0
				close(ws.closer)
				ws.slf.OnDisconnect()
				continue
			}

			ws.slf.OnReadMessage(message)
		}
	}

	return nil
}

//StartConnect ...
func (ws *WebsocketClient) StartConnect() error {

	var err error
	ws.conn, _, err = ws.WsDailer.Dial(ws.url, nil)
	service.GetLogger().Printf(sysmodule.LEVER_INFO, "connecting %s, %+v\n", ws.url, err)
	if err != nil {
		return err
	}
	ws.closer = make(chan bool)
	ws.bwritemsg = make(chan []byte, MAX_WRITE_MSG)
	ws.state = 2
	ws.slf.OnConnected()

	return nil
}

//Start ...
func (ws *WebsocketClient) Start() error {

	if ws.bRun == false {
		ws.bRun = true
		ws.state = 0
		go ws.OnRun()
		go ws.writeMsg()
	}
	return nil
}

//触发
func (ws *WebsocketClient) writeMsg() error {
	//dump处理
	defer func() {
		if r := recover(); r != nil {
			coreInfo := string(debug.Stack())
			coreInfo += "\n" + fmt.Sprintf("Core WebsocketClient url is %s. Core information is %v\n", ws.url, r)
			service.GetLogger().Printf(service.LEVER_FATAL, coreInfo)
			go ws.writeMsg()
		}
	}()

	timerC := time.NewTicker(time.Second * 5).C
	for {
		if ws.bRun == false {
			break
		}

		if ws.state == 0 {
			time.Sleep(1 * time.Second)
			continue
		}
		select {
		case _, ok := <-ws.closer:
			if ok == false {
				break
			}

		case <-timerC:
			if ws.state == 2 {
				err := ws.WriteMessage(ws.ping)
				if err != nil {
					service.GetLogger().Printf(service.LEVER_WARN, "websocket client is disconnect [%s],information is %v", ws.url, err)
					ws.state = 0
					ws.conn.Close()
					ws.slf.OnDisconnect()
				}
			}
		case msg, ok := <-ws.bwritemsg:
			if ok == true && ws.state == 2 {
				ws.conn.SetWriteDeadline(time.Now().Add(ws.timeoutsec * time.Second))
				err := ws.conn.WriteMessage(websocket.TextMessage, msg)
				if err != nil {
					service.GetLogger().Printf(service.LEVER_WARN, "websocket client is disconnect [%s],information is %v", ws.url, err)
					ws.state = 0
					ws.conn.Close()
					ws.slf.OnDisconnect()
				}
			}
		}
	}

	return nil
}

//ReConnect ...
func (ws *WebsocketClient) ReConnect() {
	ws.state = 1
}

//WriteMessage ...
func (ws *WebsocketClient) WriteMessage(msg []byte) error {
	if ws.closer == nil || ws.bwritemsg == nil {
		service.GetLogger().Printf(service.LEVER_WARN, "WriteMessage data fail,websocket client is disconnect.")
		return errors.New("riteMessage data fail,websocket client is disconnect.")
	}
	select {
	case <-ws.closer:
		service.GetLogger().Printf(service.LEVER_WARN, "WriteMessage data fail,websocket client is disconnect.")
		return errors.New("riteMessage data fail,websocket client is disconnect.")
	default:
		if len(ws.bwritemsg) < MAX_WRITE_MSG {
			ws.bwritemsg <- msg
		} else {
			service.GetLogger().Printf(service.LEVER_ERROR, "WriteMessage data fail,bwriteMsg is overload.")
			return errors.New("WriteMessage data fail,bwriteMsg is overload.")
		}

	}

	return nil
}

//OnDisconnect ...
func (ws *WebsocketClient) OnDisconnect() error {

	return nil
}

//OnConnected ...
func (ws *WebsocketClient) OnConnected() error {

	return nil
}

//OnReadMessage 触发
func (ws *WebsocketClient) OnReadMessage(msg []byte) error {

	return nil
}

//Stop ...
func (ws *WebsocketClient) Stop() {
	ws.bRun = false
}