refactor(http): refactoring http server

pkg/api/live/conn.go

@@ -32,12 +32,14 @@ var upgrader = websocket.Upgrader{
 }
 
 type connection struct {
+	hub  *hub
 	ws   *websocket.Conn
 	send chan []byte
 }
 
-func newConnection(ws *websocket.Conn) *connection {
+func newConnection(ws *websocket.Conn, hub *hub) *connection {
 	return &connection{
+		hub:  hub,
 		send: make(chan []byte, 256),
 		ws:   ws,
 	}
@@ -45,7 +47,7 @@ func newConnection(ws *websocket.Conn) *connection {
 
 func (c *connection) readPump() {
 	defer func() {
-		h.unregister <- c
+		c.hub.unregister <- c
 		c.ws.Close()
 	}()
 
@@ -81,9 +83,9 @@ func (c *connection) handleMessage(message []byte) {
 
 	switch msgType {
 	case "subscribe":
-		h.subChannel <- &streamSubscription{name: streamName, conn: c}
+		c.hub.subChannel <- &streamSubscription{name: streamName, conn: c}
 	case "unsubscribe":
-		h.subChannel <- &streamSubscription{name: streamName, conn: c, remove: true}
+		c.hub.subChannel <- &streamSubscription{name: streamName, conn: c, remove: true}
 	}
 
 }

pkg/api/live/hub.go

@@ -1,6 +1,8 @@
 package live
 
 import (
+	"context"
+
 	"github.com/grafana/grafana/pkg/api/dtos"
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/log"
@@ -23,22 +25,26 @@ type streamSubscription struct {
 	remove bool
 }
 
-var h = hub{
-	connections:   make(map[*connection]bool),
-	streams:       make(map[string]map[*connection]bool),
-	register:      make(chan *connection),
-	unregister:    make(chan *connection),
-	streamChannel: make(chan *dtos.StreamMessage),
-	subChannel:    make(chan *streamSubscription),
-	log:           log.New("live.hub"),
+func newHub() *hub {
+	return &hub{
+		connections:   make(map[*connection]bool),
+		streams:       make(map[string]map[*connection]bool),
+		register:      make(chan *connection),
+		unregister:    make(chan *connection),
+		streamChannel: make(chan *dtos.StreamMessage),
+		subChannel:    make(chan *streamSubscription),
+		log:           log.New("stream.hub"),
+	}
 }
 
 func (h *hub) removeConnection() {
 }
 
-func (h *hub) run() {
+func (h *hub) run(ctx context.Context) {
 	for {
 		select {
+		case <-ctx.Done():
+			return
 		case c := <-h.register:
 			h.connections[c] = true
 			h.log.Info("New connection", "total", len(h.connections))
@@ -49,7 +55,7 @@ func (h *hub) run() {
 				delete(h.connections, c)
 				close(c.send)
 			}
-		// hand stream subscriptions
+			// hand stream subscriptions
 		case sub := <-h.subChannel:
 			h.log.Info("Subscribing", "channel", sub.name, "remove", sub.remove)
 			subscribers, exists := h.streams[sub.name]

pkg/api/live/live.go

@@ -1,40 +0,0 @@
-package live
-
-import (
-	"net/http"
-
-	"github.com/grafana/grafana/pkg/api/dtos"
-	"github.com/grafana/grafana/pkg/log"
-	"github.com/grafana/grafana/pkg/middleware"
-)
-
-type LiveConn struct {
-	log log.Logger
-}
-
-func New() *LiveConn {
-	go h.run()
-
-	return &LiveConn{log: log.New("live.server")}
-}
-
-func (lc *LiveConn) Serve(w http.ResponseWriter, r *http.Request) {
-	lc.log.Info("Upgrading to WebSocket")
-
-	ws, err := upgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Error(3, "Live: Failed to upgrade connection to WebSocket", err)
-		return
-	}
-
-	c := newConnection(ws)
-	h.register <- c
-
-	go c.writePump()
-	c.readPump()
-}
-
-func (lc *LiveConn) PushToStream(c *middleware.Context, message dtos.StreamMessage) {
-	h.streamChannel <- &message
-	c.JsonOK("Message recevived")
-}

pkg/api/live/stream_manager.go

@@ -1,6 +1,8 @@
 package live
 
 import (
+	"context"
+	"net/http"
 	"sync"
 
 	"github.com/grafana/grafana/pkg/components/simplejson"
@@ -8,21 +10,48 @@ import (
 	m "github.com/grafana/grafana/pkg/models"
 )
 
-type StreamManagerImpl struct {
+type StreamManager struct {
 	log           log.Logger
 	streams       map[string]*Stream
 	streamRWMutex *sync.RWMutex
+	hub           *hub
 }
 
-func NewStreamManager() m.StreamManager {
-	return &StreamManagerImpl{
-		log:           log.New("live.stream.manager"),
+func NewStreamManager() *StreamManager {
+	return &StreamManager{
+		hub:           newHub(),
+		log:           log.New("stream.manager"),
 		streams:       make(map[string]*Stream),
 		streamRWMutex: &sync.RWMutex{},
 	}
 }
 
-func (s *StreamManagerImpl) GetStreamList() m.StreamList {
+func (sm *StreamManager) Run(context context.Context) {
+	log.Info("Initializing Stream Manager")
+
+	go func() {
+		sm.hub.run(context)
+		log.Info("Stopped Stream Manager")
+	}()
+}
+
+func (sm *StreamManager) Serve(w http.ResponseWriter, r *http.Request) {
+	sm.log.Info("Upgrading to WebSocket")
+
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		sm.log.Error("Failed to upgrade connection to WebSocket", "error", err)
+		return
+	}
+
+	c := newConnection(ws, sm.hub)
+	sm.hub.register <- c
+
+	go c.writePump()
+	c.readPump()
+}
+
+func (s *StreamManager) GetStreamList() m.StreamList {
 	list := make(m.StreamList, 0)
 
 	for _, stream := range s.streams {
@@ -34,7 +63,7 @@ func (s *StreamManagerImpl) GetStreamList() m.StreamList {
 	return list
 }
 
-func (s *StreamManagerImpl) Push(packet *m.StreamPacket) {
+func (s *StreamManager) Push(packet *m.StreamPacket) {
 	stream, exist := s.streams[packet.Stream]
 
 	if !exist {