github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/websocket-wrapper/ws_server.go (about)

     1  package ww
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sync"
     7  	"sync/atomic"
     8  
     9  	"github.com/gorilla/websocket"
    10  )
    11  
    12  // WSHandle means a websocket connection
    13  type WSHandle uint64
    14  
    15  // wsServer save websocket server connection info
    16  type wsServer struct {
    17  	// the pair of map save relate info,but the key of connMap is the value of connHandleMap
    18  	connMap       map[*websocket.Conn]WSHandle
    19  	connHandleMap map[WSHandle]*websocket.Conn
    20  	connChanMap   map[WSHandle]chan *WSMessage
    21  	// the rwlock for concurrent safety
    22  	rwLock sync.RWMutex
    23  }
    24  
    25  var (
    26  	// connHandle keep the conn has different handle value
    27  	connHandle uint64
    28  	// wss keep server's info
    29  	wss wsServer
    30  )
    31  
    32  func init() {
    33  	// init the wss fields
    34  	wss = wsServer{}
    35  	wss.rwLock.Lock()
    36  	defer wss.rwLock.Unlock()
    37  	if wss.connMap == nil {
    38  		wss.connMap = make(map[*websocket.Conn]WSHandle)
    39  	}
    40  	if wss.connHandleMap == nil {
    41  		wss.connHandleMap = make(map[WSHandle]*websocket.Conn)
    42  	}
    43  	if wss.connChanMap == nil {
    44  		wss.connChanMap = make(map[WSHandle]chan *WSMessage)
    45  	}
    46  	// add the other value for initialize
    47  }
    48  
    49  // new a WSHandle as the index of conn
    50  func (wss *wsServer) newWSHandle() WSHandle {
    51  	u := atomic.AddUint64(&connHandle, 1)
    52  	return WSHandle(u)
    53  }
    54  
    55  // save a conn
    56  func (wss *wsServer) saveConn(c *websocket.Conn) WSHandle {
    57  	hd := wss.newWSHandle()
    58  	wss.rwLock.Lock()
    59  	defer wss.rwLock.Unlock()
    60  	wss.connMap[c] = hd
    61  	wss.connHandleMap[hd] = c
    62  	return hd
    63  }
    64  
    65  // remove a conn
    66  func (wss *wsServer) removeConnByIndex(hd WSHandle) {
    67  	wss.rwLock.Lock()
    68  	defer wss.rwLock.Unlock()
    69  	if k, ok := wss.connHandleMap[hd]; ok {
    70  		delete(wss.connMap, k)
    71  		delete(wss.connHandleMap, hd)
    72  	}
    73  }
    74  
    75  // remove a conn by conn
    76  func (wss *wsServer) removeConn(conn *websocket.Conn) {
    77  	wss.rwLock.Lock()
    78  	defer wss.rwLock.Unlock()
    79  	if k, ok := wss.connMap[conn]; ok {
    80  		delete(wss.connHandleMap, k)
    81  		delete(wss.connMap, conn)
    82  	}
    83  }
    84  
    85  // save chan
    86  func (wss *wsServer) saveChan(hd WSHandle, c chan *WSMessage) {
    87  	wss.rwLock.Lock()
    88  	defer wss.rwLock.Unlock()
    89  	wss.connChanMap[hd] = c
    90  }
    91  
    92  // remove chan
    93  func (wss *wsServer) removeChan(hd WSHandle) {
    94  	wss.rwLock.Lock()
    95  	defer wss.rwLock.Unlock()
    96  	delete(wss.connChanMap, hd)
    97  }
    98  
    99  // get chan
   100  func (wss *wsServer) getChan(hd WSHandle) (chan *WSMessage, error) {
   101  	wss.rwLock.RLock()
   102  	defer wss.rwLock.RUnlock()
   103  	c, ok := wss.connChanMap[hd]
   104  	if !ok {
   105  		errorMsg := fmt.Sprintf("handle:%d not found chan", hd)
   106  		return nil, errors.New(errorMsg)
   107  	}
   108  	return c, nil
   109  }
   110  
   111  // SendMessge ,send a message to chan
   112  func (h WSHandle) SendMessage(msg *WSMessage) (err error) {
   113  	if c, err := wss.getChan(h); err != nil {
   114  		fmt.Println("SendMessge error:", err.Error())
   115  		return err
   116  	} else {
   117  		defer func() {
   118  			if e := recover(); e != nil {
   119  				Error("SendMessge recover error:%+v", e)
   120  				err = errors.New("SendMessge erros,write to chan fail")
   121  			}
   122  		}()
   123  		c <- msg
   124  	}
   125  	return nil
   126  }