nnet/cpn/client_ws.go

package cpn

import (
	"git.wanbits.io/joe/nnet"
	"github.com/gorilla/websocket"
	"net"
	"time"
)

const (
	WS  = 1
	TCP = 2
)

type ClientOption func(*Client)
type connector func(string, uint64) error

func WithReconn(interval time.Duration) ClientOption {
	return func(c *Client) {
		c.reconnInterval = interval
	}
}

type Client struct {
	*Hub
	addr           string
	id             uint64
	ctr            connector
	reconnInterval time.Duration
}

func NewClient(protocol int, cf *nnet.HubConfig, cb nnet.ISessionCallback, p nnet.IProtocol, opts ...ClientOption) nnet.IHub {
	c := &Client{
		Hub:            newHub(cf, cb, p),
		reconnInterval: 0,
	}

	for _, opt := range opts {
		opt(c)
	}

	c.ctr = c.connectTcp
	if protocol == WS {
		c.ctr = c.connectWs
	}

	return c
}

func NewWsClient(cf *nnet.HubConfig, cb nnet.ISessionCallback, p nnet.IProtocol, opts ...ClientOption) nnet.IHub {
	return NewClient(WS, cf, cb, p, opts...)
}

func (self *Client) NewConnection(addr string, id uint64) error {
	self.addr, self.id = addr, id

	if self.reconnInterval == 0 {
		return self.ctr(addr, id)
	}

	self.StartReconn()
	return nil
}

func (self *Client) StartReconn() {
	go func() {
		var err error
		for {
			err = self.ctr(self.addr, self.id)
			if err == nil {
				return
			}
			time.Sleep(self.reconnInterval)
		}
	}()
}

func (self *Client) connectWs(addr string, id uint64) error {
	conn, _, err := websocket.DefaultDialer.Dial(addr, nil)
	if err != nil {
		return err
	}
	self.wg.Add(1)
	go func() {
		ses := newSession(NewWsConn(conn), self)
		ses.UpdateId(id)
		ses.Do()
		self.wg.Done()
	}()

	return nil
}

func (self *Client) connectTcp(addr string, id uint64) error {
	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	conn, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return err
	}

	self.wg.Add(1)
	go func() {
		ses := newSession(TcpConn{conn}, self)
		ses.UpdateId(id)
		ses.Do()
		self.wg.Done()
	}()
	return nil
}