92 lines
1.5 KiB
Go
92 lines
1.5 KiB
Go
|
package server
|
||
|
|
|
||
|
|
import (
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"github.com/gorilla/websocket"
|
||
|
|
v1 "pkg.bing89.com/travebing/api/v1"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Client struct {
|
||
|
|
plan string
|
||
|
|
id string
|
||
|
|
conn *websocket.Conn
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewClient(plan string, conn *websocket.Conn) *Client {
|
||
|
|
return &Client{
|
||
|
|
id: uuid.NewString(),
|
||
|
|
plan: plan,
|
||
|
|
conn: conn,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) Send(way v1.TravelWay) error {
|
||
|
|
data, err := way.Bytes()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
err = c.conn.WriteMessage(websocket.BinaryMessage, data)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) SendMessage(data []byte) error {
|
||
|
|
err := c.conn.WriteMessage(websocket.BinaryMessage, data)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
type Clients struct {
|
||
|
|
m map[string]*Client
|
||
|
|
lock sync.RWMutex
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewClients() *Clients {
|
||
|
|
return &Clients{
|
||
|
|
m: make(map[string]*Client),
|
||
|
|
lock: sync.RWMutex{},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Clients) Get(plan string) (*Client, bool) {
|
||
|
|
c.lock.Lock()
|
||
|
|
defer c.lock.Unlock()
|
||
|
|
cs, ok := c.m[plan]
|
||
|
|
return cs, ok
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Clients) Notify(way v1.TravelWay) error {
|
||
|
|
data, err := way.Bytes()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
c.lock.Lock()
|
||
|
|
defer c.lock.Unlock()
|
||
|
|
failed := []string{}
|
||
|
|
for id, clt := range c.m {
|
||
|
|
err := clt.SendMessage(data)
|
||
|
|
if err != nil {
|
||
|
|
failed = append(failed, id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
//清理报错的
|
||
|
|
for _, f := range failed {
|
||
|
|
delete(c.m, f)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Clients) Set(clt *Client) {
|
||
|
|
if clt.id == "" {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
c.lock.Lock()
|
||
|
|
defer c.lock.Unlock()
|
||
|
|
c.m[clt.id] = clt
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Clients) Delete(id string) {
|
||
|
|
c.lock.Lock()
|
||
|
|
defer c.lock.Unlock()
|
||
|
|
delete(c.m, id)
|
||
|
|
}
|