gocommon/wshandler/room.go

package wshandler

import (
	"bytes"
	"context"
	"encoding/json"

	"github.com/gorilla/websocket"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"repositories.action2quare.com/ayo/gocommon/logger"
)

type room struct {
	inChan      chan *wsconn
	outChan     chan primitive.ObjectID
	messageChan chan *UpstreamMessage
	name        string
	destroyChan chan<- string
	sendMsgChan chan<- send_msg_queue_elem
}

// 만약 destroyChan가 nil이면 room이 비어도 파괴되지 않는다. 영구 유지되는 room
func makeRoom(name string, destroyChan chan<- string, sendMsgChan chan<- send_msg_queue_elem) *room {
	return &room{
		inChan:      make(chan *wsconn, 10),
		outChan:     make(chan primitive.ObjectID, 10),
		messageChan: make(chan *UpstreamMessage, 100),
		name:        name,
		destroyChan: destroyChan,
		sendMsgChan: sendMsgChan,
	}
}

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

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

func (r *room) out(accid primitive.ObjectID) *room {
	r.outChan <- accid
	return r
}

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
		}
	}()

	tag := "#" + r.name
	for {
		select {
		case <-ctx.Done():
			return true

		case conn := <-r.inChan:
			(*conns)[conn.sender.Accid.Hex()] = conn

		case accid := <-r.outChan:
			delete((*conns), accid.Hex())
			if len(*conns) == 0 && r.destroyChan != nil {
				r.destroyChan <- r.name
				return true
			}

		case msg := <-r.messageChan:
			ds := DownstreamMessage{
				Alias: msg.Alias,
				Body:  msg.Body,
				Tag:   append(msg.Tag, tag),
			}

			buff := new(bytes.Buffer)
			enc := json.NewEncoder(buff)
			enc.SetEscapeHTML(false)
			enc.Encode(ds)

			for _, conn := range *conns {
				pmsg, err := websocket.NewPreparedMessage(websocket.TextMessage, buff.Bytes())
				if err != nil {
					logger.Println("websocket.NewPreparedMessage failed :", err)
				} else {
					r.sendMsgChan <- send_msg_queue_elem{
						to:   conn.Conn,
						pmsg: pmsg,
					}
				}
			}
		}
	}
}