github.com/oarkflow/sio@v0.0.6/hub.go (about)

     1  package sio
     2  
     3  import (
     4  	"slices"
     5  	"sync"
     6  )
     7  
     8  type hub struct {
     9  	sockets          map[string]*Socket
    10  	rooms            map[string]*room
    11  	shutdownCh       chan bool
    12  	socketList       chan []*Socket
    13  	addCh            chan *Socket
    14  	delCh            chan *Socket
    15  	joinRoomCh       chan *joinRequest
    16  	leaveRoomCh      chan *leaveRequest
    17  	roomMsgCh        chan *RoomMsg
    18  	broomcastCh      chan *RoomMsg // for passing data from the backend
    19  	broadcastCh      chan *BroadcastMsg
    20  	bbroadcastCh     chan *BroadcastMsg
    21  	multihomeEnabled bool
    22  	multihomeBackend Adapter
    23  	l                sync.RWMutex
    24  }
    25  
    26  type room struct {
    27  	name    string
    28  	sockets map[string]*Socket
    29  	l       sync.RWMutex
    30  }
    31  
    32  type joinRequest struct {
    33  	roomName string
    34  	socket   *Socket
    35  }
    36  
    37  type leaveRequest struct {
    38  	roomName string
    39  	socket   *Socket
    40  }
    41  
    42  // RoomMsg represents an event to be dispatched to a room of sockets
    43  type RoomMsg struct {
    44  	RoomName  string
    45  	Except    []string
    46  	EventName string
    47  	Data      any
    48  }
    49  
    50  // BroadcastMsg represents an event to be dispatched to all Sockets on the Server
    51  type BroadcastMsg struct {
    52  	EventName string
    53  	Except    []string
    54  	Data      any
    55  }
    56  
    57  func (h *hub) addSocket(s *Socket) {
    58  	h.addCh <- s
    59  }
    60  
    61  func (h *hub) removeSocket(s *Socket) {
    62  	h.delCh <- s
    63  }
    64  
    65  func (h *hub) joinRoom(j *joinRequest) {
    66  	h.joinRoomCh <- j
    67  }
    68  
    69  func (h *hub) leaveRoom(l *leaveRequest) {
    70  	h.leaveRoomCh <- l
    71  }
    72  
    73  func (h *hub) toRoom(msg *RoomMsg) {
    74  	h.roomMsgCh <- msg
    75  }
    76  
    77  func (h *hub) broadcast(b *BroadcastMsg) {
    78  	h.broadcastCh <- b
    79  }
    80  
    81  func (h *hub) setMultihomeBackend(b Adapter) {
    82  	if h.multihomeEnabled {
    83  		return // can't have two backends... yet
    84  	}
    85  
    86  	h.multihomeBackend = b
    87  	h.multihomeEnabled = true
    88  
    89  	h.multihomeBackend.Init()
    90  
    91  	go h.multihomeBackend.BroadcastFromBackend(h.bbroadcastCh)
    92  	go h.multihomeBackend.RoomcastFromBackend(h.broomcastCh)
    93  }
    94  
    95  func (h *hub) listen() {
    96  	for {
    97  		select {
    98  		case c := <-h.addCh:
    99  			h.l.Lock()
   100  			h.sockets[c.ID()] = c
   101  			h.l.Unlock()
   102  		case c := <-h.delCh:
   103  			delete(h.sockets, c.ID())
   104  		case c := <-h.joinRoomCh:
   105  			if _, exists := h.rooms[c.roomName]; !exists { // make the room if it doesn't exist
   106  				h.rooms[c.roomName] = &room{name: c.roomName, sockets: make(map[string]*Socket)}
   107  			}
   108  			h.rooms[c.roomName].l.Lock()
   109  			h.rooms[c.roomName].sockets[c.socket.ID()] = c.socket
   110  			h.rooms[c.roomName].l.Unlock()
   111  		case c := <-h.leaveRoomCh:
   112  			if room, exists := h.rooms[c.roomName]; exists {
   113  				room.l.Lock()
   114  				delete(room.sockets, c.socket.ID())
   115  				room.l.Unlock()
   116  				if len(room.sockets) == 0 { // room is empty, delete it
   117  					delete(h.rooms, c.roomName)
   118  				}
   119  			}
   120  		case c := <-h.roomMsgCh:
   121  			if room, exists := h.rooms[c.RoomName]; exists {
   122  				room.l.Lock()
   123  				for _, s := range room.sockets {
   124  					if len(c.Except) > 0 {
   125  						if !slices.Contains(c.Except, s.ID()) {
   126  							s.Emit(c.EventName, c.Data)
   127  						}
   128  					} else {
   129  						s.Emit(c.EventName, c.Data)
   130  					}
   131  				}
   132  				room.l.Unlock()
   133  			}
   134  			if h.multihomeEnabled { // the room may exist on the other end
   135  				go h.multihomeBackend.RoomcastToBackend(c)
   136  			}
   137  		case c := <-h.broomcastCh:
   138  			if room, exists := h.rooms[c.RoomName]; exists {
   139  				room.l.Lock()
   140  				for _, s := range room.sockets {
   141  					if len(c.Except) > 0 {
   142  						if !slices.Contains(c.Except, s.ID()) {
   143  							s.Emit(c.EventName, c.Data)
   144  						}
   145  					} else {
   146  						s.Emit(c.EventName, c.Data)
   147  					}
   148  				}
   149  				room.l.Unlock()
   150  			}
   151  		case c := <-h.broadcastCh:
   152  			h.l.Lock()
   153  			for _, s := range h.sockets {
   154  				if len(c.Except) > 0 {
   155  					if !slices.Contains(c.Except, s.ID()) {
   156  						s.Emit(c.EventName, c.Data)
   157  					}
   158  				} else {
   159  					s.Emit(c.EventName, c.Data)
   160  				}
   161  			}
   162  			h.l.Unlock()
   163  			if h.multihomeEnabled {
   164  				go h.multihomeBackend.BroadcastToBackend(c)
   165  			}
   166  		case c := <-h.bbroadcastCh:
   167  			h.l.Lock()
   168  			for _, s := range h.sockets {
   169  				if len(c.Except) > 0 {
   170  					if !slices.Contains(c.Except, s.ID()) {
   171  						s.Emit(c.EventName, c.Data)
   172  					}
   173  				} else {
   174  					s.Emit(c.EventName, c.Data)
   175  				}
   176  			}
   177  			h.l.Unlock()
   178  		case _ = <-h.shutdownCh:
   179  			var socketList []*Socket
   180  			h.l.Lock()
   181  			for _, s := range h.sockets {
   182  				socketList = append(socketList, s)
   183  			}
   184  			h.l.Unlock()
   185  			h.socketList <- socketList
   186  		}
   187  	}
   188  }
   189  
   190  func newHub() *hub {
   191  	h := &hub{
   192  		shutdownCh:       make(chan bool),
   193  		socketList:       make(chan []*Socket),
   194  		sockets:          make(map[string]*Socket),
   195  		rooms:            make(map[string]*room),
   196  		addCh:            make(chan *Socket),
   197  		delCh:            make(chan *Socket),
   198  		joinRoomCh:       make(chan *joinRequest),
   199  		leaveRoomCh:      make(chan *leaveRequest),
   200  		roomMsgCh:        make(chan *RoomMsg),
   201  		broomcastCh:      make(chan *RoomMsg),
   202  		broadcastCh:      make(chan *BroadcastMsg),
   203  		bbroadcastCh:     make(chan *BroadcastMsg),
   204  		multihomeEnabled: false,
   205  	}
   206  
   207  	go h.listen()
   208  
   209  	return h
   210  }