gocommon/wshandler/room.go

package wshandler

import (
	"context"

	"github.com/gorilla/websocket"
	"repositories.action2quare.com/ayo/gocommon/logger"
)

type room struct {
	inChan      chan *wsconn
	outChan     chan *wsconn
	messageChan chan *UpstreamMessage
	name        string
}

func makeRoom(name string) *room {
	return &room{
		inChan:      make(chan *wsconn, 10),
		outChan:     make(chan *wsconn, 10),
		messageChan: make(chan *UpstreamMessage, 100),
		name:        name,
	}
}

func (r *room) broadcast(msg *UpstreamMessage) {
	r.messageChan <- msg
}

func (r *room) in(conn *wsconn) {
	r.inChan <- conn
}

func (r *room) out(conn *wsconn) {
	r.outChan <- conn
}

func (r *room) start(ctx context.Context) {
	go func(ctx context.Context) {
		conns := make(map[string]*wsconn)
		normal := false
		for !normal {
			normal = r.loop(ctx, &conns)
		}
	}(ctx)
}

func (r *room) loop(ctx context.Context, conns *map[string]*wsconn) (normalEnd bool) {
	defer func() {
		s := recover()
		if s != nil {
			logger.Error(s)
			normalEnd = false
		}
	}()

	a, b, c := []byte(`{"alias":"`), []byte(`","body":"`), []byte(`"}`)

	for {
		select {
		case <-ctx.Done():
			return true

		case conn := <-r.inChan:
			(*conns)[conn.alias] = conn

		case conn := <-r.outChan:
			delete((*conns), conn.alias)

		case msg := <-r.messageChan:
			for _, conn := range *conns {
				writer, _ := conn.NextWriter(websocket.TextMessage)
				writer.Write(a)
				writer.Write([]byte(msg.Alias))
				writer.Write(b)
				writer.Write(msg.Body)
				writer.Write(c)
				writer.Close()
			}
		}
	}
}