docs: improve of websocket
continuous-integration/drone the build was successful Details

This commit is contained in:
Geno 2021-06-01 18:08:29 +02:00
parent 1edf36b7aa
commit de11bae835
2 changed files with 142 additions and 110 deletions

View File

@ -1,7 +1,10 @@
package ws package ws
const ( const (
// BodyError Message Body map typ for errors
BodyError = "error" BodyError = "error"
BodySet = "set" // BodySet Message Body map typ for set values
BodyGet = "get" BodySet = "set"
// BodyGet Message Body map typ for get values
BodyGet = "get"
) )

View File

@ -8,12 +8,14 @@ import (
"github.com/bdlm/log" "github.com/bdlm/log"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/google/uuid"
"golang.org/x/time/rate" "golang.org/x/time/rate"
"nhooyr.io/websocket" "nhooyr.io/websocket"
"nhooyr.io/websocket/wsjson" "nhooyr.io/websocket/wsjson"
) )
// WebsocketEndpoint to handle Request
type WebsocketEndpoint struct { type WebsocketEndpoint struct {
// publishLimiter controls the rate limit applied to the publish endpoint. // publishLimiter controls the rate limit applied to the publish endpoint.
// //
@ -24,25 +26,52 @@ type WebsocketEndpoint struct {
Subscribers map[*Subscriber]struct{} Subscribers map[*Subscriber]struct{}
// Message Handler // Message Handler
handlers map[string]MessageHandleFunc handlers map[string]MessageHandleFunc
// DefaultMessageHandler if no other handler for MessageType found
DefaultMessageHandler MessageHandleFunc DefaultMessageHandler MessageHandleFunc
OnOpen SubscriberEventFunc // Run Function on open connection by subscriper
OnClose SubscriberEventFunc OnOpen SubscriberEventFunc
// Run Function on close connection to subscriper
OnClose SubscriberEventFunc
} }
// MessageHandleFunc for handling messages // Subscriber of websocket endpoint
type MessageHandleFunc func(ctx context.Context, msg *Message) type Subscriber struct {
out chan *Message
closeSlow func()
}
// SubscriberEventFunc for handling connection state of Subsriber
type SubscriberEventFunc func(s *Subscriber, msg chan<- *Message) type SubscriberEventFunc func(s *Subscriber, msg chan<- *Message)
// Message on websocket // Message on websocket
type Message struct { type Message struct {
Type string `json:"type"` Type string `json:"type"`
ID *uuid.UUID `json:"id,omitempty"`
ReplyID *uuid.UUID `json:"reply_id,omitempty"`
Body map[string]interface{} `json:"body"` Body map[string]interface{} `json:"body"`
Reply chan<- *Message `json:"-"`
Subscriber *Subscriber `json:"-"` Subscriber *Subscriber `json:"-"`
} }
// Reply to Message
func (m *Message) Reply(msg *Message) {
if m == nil || m.Subscriber == nil {
return
}
if m.ID != nil {
msg.ReplyID = m.ID
if msg.ID == nil {
id := uuid.New()
msg.ID = &id
}
}
m.Subscriber.out <- msg
}
// MessageHandleFunc for handling messages
type MessageHandleFunc func(ctx context.Context, msg *Message)
// NewEndpoint - create an empty websocket
func NewEndpoint() *WebsocketEndpoint { func NewEndpoint() *WebsocketEndpoint {
return &WebsocketEndpoint{ return &WebsocketEndpoint{
publishLimiter: rate.NewLimiter(rate.Every(time.Millisecond*100), 8), publishLimiter: rate.NewLimiter(rate.Every(time.Millisecond*100), 8),
@ -50,109 +79,8 @@ func NewEndpoint() *WebsocketEndpoint {
handlers: make(map[string]MessageHandleFunc), handlers: make(map[string]MessageHandleFunc),
} }
} }
func (this *WebsocketEndpoint) Handler(ctx *gin.Context) {
c, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
InsecureSkipVerify: true,
})
if err != nil {
ctx.JSON(http.StatusBadRequest, false)
return
}
defer c.Close(websocket.StatusInternalError, "")
err = this.addSubscriber(ctx, c)
if websocket.CloseStatus(err) == websocket.StatusNormalClosure ||
websocket.CloseStatus(err) == websocket.StatusGoingAway {
return
}
log.Errorf("subscriber stopped: %s", err)
}
func (this *WebsocketEndpoint) AddMessageHandler(typ string, f MessageHandleFunc) {
this.handlers[typ] = f
}
type Subscriber struct {
out chan *Message
closeSlow func()
}
func (this *WebsocketEndpoint) readWorker(ctx context.Context, c *websocket.Conn, s *Subscriber) error {
for {
var msg Message
err := wsjson.Read(ctx, c, &msg)
if err != nil {
return err
}
log.WithField("type", msg.Type).Debug("receive")
msg.Subscriber = s
msg.Reply = s.out
if handler, ok := this.handlers[msg.Type]; ok {
handler(ctx, &msg)
} else if this.DefaultMessageHandler != nil {
this.DefaultMessageHandler(ctx, &msg)
}
}
}
func (this *WebsocketEndpoint) addSubscriber(ctxGin *gin.Context, c *websocket.Conn) error {
s := &Subscriber{
out: make(chan *Message, 10),
closeSlow: func() {
c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages")
},
}
this.subscribersMu.Lock()
this.Subscribers[s] = struct{}{}
this.subscribersMu.Unlock()
defer func() {
this.subscribersMu.Lock()
delete(this.Subscribers, s)
this.subscribersMu.Unlock()
if this.OnClose != nil {
this.OnClose(s, s.out)
}
log.Debug("websocket closed")
}()
if this.OnOpen != nil {
this.OnOpen(s, s.out)
}
ctx := ctxGin.Request.Context()
go func() {
err := this.readWorker(ctx, c, s)
if websocket.CloseStatus(err) == websocket.StatusNormalClosure ||
websocket.CloseStatus(err) == websocket.StatusGoingAway {
return
}
log.Errorf("websocket reading error: %s", err)
}()
log.Debug("websocket started")
for {
select {
case msg := <-s.out:
err := writeTimeout(ctx, time.Second*5, c, msg)
if err != nil {
return err
}
case <-ctx.Done():
return ctx.Err()
}
}
}
func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg *Message) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
return wsjson.Write(ctx, c, msg)
}
// Broadcast Message to all subscriber (exclude sender of Message)
func (this *WebsocketEndpoint) Broadcast(msg *Message) { func (this *WebsocketEndpoint) Broadcast(msg *Message) {
this.subscribersMu.Lock() this.subscribersMu.Lock()
defer this.subscribersMu.Unlock() defer this.subscribersMu.Unlock()
@ -170,3 +98,104 @@ func (this *WebsocketEndpoint) Broadcast(msg *Message) {
} }
} }
} }
// AddMessageHandler - add websocket message handler
func (we *WebsocketEndpoint) AddMessageHandler(typ string, f MessageHandleFunc) {
we.handlers[typ] = f
}
// Handler - to register in gin webservice
func (we *WebsocketEndpoint) Handler(ctx *gin.Context) {
c, err := websocket.Accept(ctx.Writer, ctx.Request, nil)
if err != nil {
ctx.JSON(http.StatusBadRequest, false)
return
}
defer c.Close(websocket.StatusInternalError, "")
err = we.addSubscriber(ctx, c)
if websocket.CloseStatus(err) == websocket.StatusNormalClosure ||
websocket.CloseStatus(err) == websocket.StatusGoingAway {
return
}
log.Errorf("subscriber stopped: %s", err)
}
// addSubscriber and startup of websocket endpoint
func (we *WebsocketEndpoint) addSubscriber(ctxGin *gin.Context, c *websocket.Conn) error {
s := &Subscriber{
out: make(chan *Message, 10),
closeSlow: func() {
c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages")
},
}
we.subscribersMu.Lock()
we.Subscribers[s] = struct{}{}
we.subscribersMu.Unlock()
defer func() {
we.subscribersMu.Lock()
delete(we.Subscribers, s)
we.subscribersMu.Unlock()
if we.OnClose != nil {
we.OnClose(s, s.out)
}
log.Debug("websocket closed")
}()
if we.OnOpen != nil {
we.OnOpen(s, s.out)
}
ctx := ctxGin.Request.Context()
go func() {
err := we.readWorker(ctx, c, s)
if websocket.CloseStatus(err) == websocket.StatusNormalClosure ||
websocket.CloseStatus(err) == websocket.StatusGoingAway {
return
}
log.Errorf("websocket reading error: %s", err)
}()
log.Debug("websocket started")
for {
select {
case msg := <-s.out:
err := writeTimeout(ctx, time.Second*5, c, msg)
if err != nil {
return err
}
case <-ctx.Done():
return ctx.Err()
}
}
}
// readWorker of subscriber
func (we *WebsocketEndpoint) readWorker(ctx context.Context, c *websocket.Conn, s *Subscriber) error {
for {
var msg Message
err := wsjson.Read(ctx, c, &msg)
if err != nil {
return err
}
log.WithField("type", msg.Type).Debug("receive")
msg.Subscriber = s
if handler, ok := we.handlers[msg.Type]; ok {
handler(ctx, &msg)
} else if we.DefaultMessageHandler != nil {
we.DefaultMessageHandler(ctx, &msg)
}
}
}
// writeTimeout send message to subscriber with timeout
func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg *Message) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
return wsjson.Write(ctx, c, msg)
}