feat(grafana_live): more work pushing data to websocket, #4355

pkg/api/dtos/stream.go

@@ -3,7 +3,11 @@ package dtos
 import "encoding/json"
 
 type StreamMessage struct {
-	Stream     string          `json:"stream"`
-	Metric     string          `json:"metric"`
-	Datapoints [][]json.Number `json:"Datapoints"`
+	Stream string                `json:"stream"`
+	Series []StreamMessageSeries `json:"series"`
+}
+
+type StreamMessageSeries struct {
+	Name       string          `json:"name"`
+	Datapoints [][]json.Number `json:"datapoints"`
 }

pkg/api/live/conn.go

@@ -81,7 +81,10 @@ func (c *connection) handleMessage(message []byte) {
 	switch msgType {
 	case "subscribe":
 		h.subChannel <- &streamSubscription{name: streamName, conn: c}
+	case "unsubscribe":
+		h.subChannel <- &streamSubscription{name: streamName, conn: c, remove: true}
 	}
+
 }
 
 func (c *connection) write(mt int, payload []byte) error {

pkg/api/live/hub.go

@@ -17,8 +17,9 @@ type hub struct {
 }
 
 type streamSubscription struct {
-	conn *connection
-	name string
+	conn   *connection
+	name   string
+	remove bool
 }
 
 var h = hub{
@@ -40,15 +41,24 @@ func (h *hub) run() {
 		case c := <-h.register:
 			h.connections[c] = true
 			log.Info("Live: New connection (Total count: %v)", len(h.connections))
+
 		case c := <-h.unregister:
 			if _, ok := h.connections[c]; ok {
+				log.Info("Live: Closing Connection (Total count: %v)", len(h.connections))
 				delete(h.connections, c)
 				close(c.send)
 			}
 		// hand stream subscriptions
 		case sub := <-h.subChannel:
-			log.Info("Live: Connection subscribing to: %v", sub.name)
+			log.Info("Live: Subscribing to: %v, remove: %v", sub.name, sub.remove)
 			subscribers, exists := h.streams[sub.name]
+
+			// handle unsubscribe
+			if exists && sub.remove {
+				delete(subscribers, sub.conn)
+				continue
+			}
+
 			if !exists {
 				subscribers = make(map[*connection]bool)
 				h.streams[sub.name] = subscribers
@@ -58,13 +68,19 @@ func (h *hub) run() {
 			// handle stream messages
 		case message := <-h.streamChannel:
 			subscribers, exists := h.streams[message.Stream]
-			if !exists {
+			if !exists || len(subscribers) == 0 {
 				log.Info("Live: Message to stream without subscribers: %v", message.Stream)
 				continue
 			}
 
 			messageBytes, _ := simplejson.NewFromAny(message).Encode()
 			for sub := range subscribers {
+				// check if channel is open
+				if _, ok := h.connections[sub]; !ok {
+					delete(subscribers, sub)
+					continue
+				}
+
 				select {
 				case sub.send <- messageBytes:
 				default: