github.com/martinohmann/rfoutlet@v1.2.1-0.20220707195255-8a66aa411105/internal/websocket/hub.go (about)

     1  // Package websocket provides the http handler, client and hub for managing the
     2  // websocket connections of clients using rfoutlet. This is both used to react
     3  // on commands of a single client as well as to broadcast updates to all
     4  // clients so that state changes are immediately visible to everybody.
     5  package websocket
     6  
     7  // Hub acts as a central registry for connected websocket clients and can be
     8  // used to broadcast messages to everyone.
     9  type Hub struct {
    10  	clients    map[*client]struct{}
    11  	register   chan *client
    12  	unregister chan *client
    13  	broadcast  chan []byte
    14  	send       chan clientMsg
    15  }
    16  
    17  // NewHub creates a new hub for handling communicating between connected
    18  // websocket clients.
    19  func NewHub() *Hub {
    20  	return &Hub{
    21  		clients:    make(map[*client]struct{}),
    22  		register:   make(chan *client),
    23  		unregister: make(chan *client),
    24  		broadcast:  make(chan []byte),
    25  		send:       make(chan clientMsg),
    26  	}
    27  }
    28  
    29  // Run runs the control loop. If stopCh is closed, the hub will disconnect all
    30  // clients and stop the control loop.
    31  func (h *Hub) Run(stopCh <-chan struct{}) {
    32  	for {
    33  		select {
    34  		case client := <-h.register:
    35  			h.clients[client] = struct{}{}
    36  			log.WithField("uuid", client.uuid).Info("new client registered")
    37  		case client := <-h.unregister:
    38  			if _, ok := h.clients[client]; ok {
    39  				h.unregisterClient(client)
    40  			}
    41  		case cm := <-h.send:
    42  			if _, ok := h.clients[cm.client]; ok {
    43  				h.sendClientMessage(cm.client, cm.msg)
    44  			}
    45  		case msg := <-h.broadcast:
    46  			log.WithField("length", len(msg)).Debug("broadcasting message")
    47  			for client := range h.clients {
    48  				h.sendClientMessage(client, msg)
    49  			}
    50  		case <-stopCh:
    51  			log.Infof("shutting down hub")
    52  			for client := range h.clients {
    53  				h.unregisterClient(client)
    54  			}
    55  			return
    56  		}
    57  	}
    58  }
    59  
    60  // Broadcast broadcasts msg to all connected clients.
    61  func (h *Hub) Broadcast(msg []byte) {
    62  	h.broadcast <- msg
    63  }
    64  
    65  // Send sends a message to a specific client.
    66  func (h *Hub) Send(client *client, msg []byte) {
    67  	h.send <- clientMsg{client, msg}
    68  }
    69  
    70  func (h *Hub) unregisterClient(client *client) {
    71  	close(client.send)
    72  	delete(h.clients, client)
    73  	log.WithField("uuid", client.uuid).Info("client unregistered")
    74  }
    75  
    76  func (h *Hub) sendClientMessage(client *client, msg []byte) {
    77  	select {
    78  	case client.send <- msg:
    79  	default:
    80  		h.unregister <- client
    81  	}
    82  }
    83  
    84  // clientMsg is a wrapper type for a message destined for a specific client.
    85  type clientMsg struct {
    86  	client *client
    87  	msg    []byte
    88  }