Live: persisting last message in cache for broadcast scope (#32938)

pkg/models/live.go

@@ -79,3 +79,22 @@ type DashboardActivityChannel interface {
 	// someone is subscribed to the `grafana/dashboards/gitops` channel
 	HasGitOpsObserver() bool
 }
+
+type LiveMessage struct {
+	Id        int64
+	OrgId     int64
+	Channel   string
+	Data      json.RawMessage
+	Published time.Time
+}
+
+type SaveLiveMessageQuery struct {
+	OrgId   int64
+	Channel string
+	Data    json.RawMessage
+}
+
+type GetLiveMessageQuery struct {
+	OrgId   int64
+	Channel string
+}

pkg/services/live/database/migrations.go

@@ -0,0 +1,24 @@
+package database
+
+import "github.com/grafana/grafana/pkg/services/sqlstore/migrator"
+
+// For now disable migration. For now we are using local cache as storage to evaluate ideas.
+// This will be turned on soon though.
+func AddLiveChannelMigrations(mg *migrator.Migrator) {
+	//liveMessage := migrator.Table{
+	//	Name: "live_message",
+	//	Columns: []*migrator.Column{
+	//		{Name: "id", Type: migrator.DB_BigInt, Nullable: false, IsPrimaryKey: true, IsAutoIncrement: true},
+	//		{Name: "org_id", Type: migrator.DB_BigInt, Nullable: false},
+	//		{Name: "channel", Type: migrator.DB_NVarchar, Length: 189, Nullable: false},
+	//		{Name: "data", Type: migrator.DB_Text, Nullable: false},
+	//		{Name: "published", Type: migrator.DB_DateTime, Nullable: false},
+	//	},
+	//	Indices: []*migrator.Index{
+	//		{Cols: []string{"org_id", "channel"}, Type: migrator.UniqueIndex},
+	//	},
+	//}
+	//
+	//mg.AddMigration("create live message table", migrator.NewAddTableMigration(liveMessage))
+	//mg.AddMigration("add index live_message.org_id_channel_unique", migrator.NewAddIndexMigration(liveMessage, liveMessage.Indices[0]))
+}

pkg/services/live/database/storage.go

@@ -0,0 +1,66 @@
+package database
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/localcache"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+)
+
+type Storage struct {
+	store *sqlstore.SQLStore
+	cache *localcache.CacheService
+}
+
+func NewStorage(store *sqlstore.SQLStore, cache *localcache.CacheService) *Storage {
+	return &Storage{store: store, cache: cache}
+}
+
+func getLiveMessageCacheKey(orgID int64, channel string) string {
+	return fmt.Sprintf("live_message_%d_%s", orgID, channel)
+}
+
+func (s *Storage) SaveLiveMessage(query *models.SaveLiveMessageQuery) error {
+	// Come back to saving into database after evaluating database structure.
+	//err := s.store.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+	//	params := []interface{}{query.OrgId, query.Channel, query.Data, time.Now()}
+	//	upsertSQL := s.store.Dialect.UpsertSQL(
+	//		"live_message",
+	//		[]string{"org_id", "channel"},
+	//		[]string{"org_id", "channel", "data", "published"})
+	//	_, err := sess.SQL(upsertSQL, params...).Query()
+	//	return err
+	//})
+	// return err
+	s.cache.Set(getLiveMessageCacheKey(query.OrgId, query.Channel), models.LiveMessage{
+		Id:        0, // Not used actually.
+		OrgId:     query.OrgId,
+		Channel:   query.Channel,
+		Data:      query.Data,
+		Published: time.Now(),
+	}, 0)
+	return nil
+}
+
+func (s *Storage) GetLiveMessage(query *models.GetLiveMessageQuery) (models.LiveMessage, bool, error) {
+	// Come back to saving into database after evaluating database structure.
+	//var msg models.LiveMessage
+	//var exists bool
+	//err := s.store.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
+	//	var err error
+	//	exists, err = sess.Where("org_id=? AND channel=?", query.OrgId, query.Channel).Get(&msg)
+	//	return err
+	//})
+	//return msg, exists, err
+	m, ok := s.cache.Get(getLiveMessageCacheKey(query.OrgId, query.Channel))
+	if !ok {
+		return models.LiveMessage{}, false, nil
+	}
+	msg, ok := m.(models.LiveMessage)
+	if !ok {
+		return models.LiveMessage{}, false, fmt.Errorf("unexpected live message type in cache: %T", m)
+	}
+	return msg, true, nil
+}

pkg/services/live/database/tests/setup.go

@@ -0,0 +1,44 @@
+package tests
+
+import (
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/localcache"
+	"github.com/grafana/grafana/pkg/registry"
+	"github.com/grafana/grafana/pkg/services/live"
+	"github.com/grafana/grafana/pkg/services/live/database"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+// SetupTestStorage initializes a storage to used by the integration tests.
+// This is required to properly register and execute migrations.
+func SetupTestStorage(t *testing.T) *database.Storage {
+	cfg := setting.NewCfg()
+	// Live is disabled by default and only if it's enabled its database migrations run
+	// and the related database tables are created.
+	cfg.FeatureToggles = map[string]bool{"live": true}
+
+	gLive := live.NewGrafanaLive()
+	gLive.Cfg = cfg
+
+	// Hook for initialising the service after the Cfg is populated
+	// so that database migrations will run.
+	overrideServiceFunc := func(descriptor registry.Descriptor) (*registry.Descriptor, bool) {
+		if _, ok := descriptor.Instance.(*live.GrafanaLive); ok {
+			return &registry.Descriptor{
+				Name:         descriptor.Name,
+				Instance:     gLive,
+				InitPriority: descriptor.InitPriority,
+			}, true
+		}
+		return nil, false
+	}
+	registry.RegisterOverride(overrideServiceFunc)
+
+	// Now we can use sql.Store.
+	sqlStore := sqlstore.InitTestDB(t)
+	localCache := localcache.New(time.Hour, time.Hour)
+	return database.NewStorage(sqlStore, localCache)
+}

pkg/services/live/database/tests/storage_test.go

@@ -0,0 +1,61 @@
+// +build integration
+
+package tests
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/models"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLiveMessage(t *testing.T) {
+	storage := SetupTestStorage(t)
+
+	getQuery := &models.GetLiveMessageQuery{
+		OrgId:   1,
+		Channel: "test_channel",
+	}
+	_, ok, err := storage.GetLiveMessage(getQuery)
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	saveQuery := &models.SaveLiveMessageQuery{
+		OrgId:   1,
+		Channel: "test_channel",
+		Data:    []byte(`{}`),
+	}
+	err = storage.SaveLiveMessage(saveQuery)
+	require.NoError(t, err)
+
+	msg, ok, err := storage.GetLiveMessage(getQuery)
+	require.NoError(t, err)
+	require.True(t, ok)
+	require.Equal(t, int64(1), msg.OrgId)
+	require.Equal(t, "test_channel", msg.Channel)
+	require.Equal(t, json.RawMessage(`{}`), msg.Data)
+	require.NotZero(t, msg.Published)
+
+	// try saving again, should be replaced.
+	saveQuery2 := &models.SaveLiveMessageQuery{
+		OrgId:   1,
+		Channel: "test_channel",
+		Data:    []byte(`{"input": "hello"}`),
+	}
+	err = storage.SaveLiveMessage(saveQuery2)
+	require.NoError(t, err)
+
+	getQuery2 := &models.GetLiveMessageQuery{
+		OrgId:   1,
+		Channel: "test_channel",
+	}
+	msg2, ok, err := storage.GetLiveMessage(getQuery2)
+	require.NoError(t, err)
+	require.True(t, ok)
+	require.Equal(t, int64(1), msg2.OrgId)
+	require.Equal(t, "test_channel", msg2.Channel)
+	require.Equal(t, json.RawMessage(`{"input": "hello"}`), msg2.Data)
+	require.NotZero(t, msg2.Published)
+}

pkg/services/live/features/broadcast.go

@@ -2,41 +2,68 @@ package features
 
 import (
 	"context"
-	"time"
 
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/models"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-
-	"github.com/grafana/grafana/pkg/models"
 )
 
 var (
 	logger = log.New("live.features") // scoped to all features?
 )
 
+//go:generate mockgen -destination=broadcast_mock.go -package=features github.com/grafana/grafana/pkg/services/live/features LiveMessageStore
+
+type LiveMessageStore interface {
+	SaveLiveMessage(query *models.SaveLiveMessageQuery) error
+	GetLiveMessage(query *models.GetLiveMessageQuery) (models.LiveMessage, bool, error)
+}
+
 // BroadcastRunner will simply broadcast all events to `grafana/broadcast/*` channels
 // This assumes that data is a JSON object
-type BroadcastRunner struct{}
+type BroadcastRunner struct {
+	liveMessageStore LiveMessageStore
+}
+
+func NewBroadcastRunner(liveMessageStore LiveMessageStore) *BroadcastRunner {
+	return &BroadcastRunner{liveMessageStore: liveMessageStore}
+}
 
 // GetHandlerForPath called on init
-func (b *BroadcastRunner) GetHandlerForPath(path string) (models.ChannelHandler, error) {
+func (b *BroadcastRunner) GetHandlerForPath(_ string) (models.ChannelHandler, error) {
 	return b, nil // all dashboards share the same handler
 }
 
 // OnSubscribe will let anyone connect to the path
-func (b *BroadcastRunner) OnSubscribe(ctx context.Context, _ *models.SignedInUser, e models.SubscribeEvent) (models.SubscribeReply, backend.SubscribeStreamStatus, error) {
-	return models.SubscribeReply{
+func (b *BroadcastRunner) OnSubscribe(_ context.Context, u *models.SignedInUser, e models.SubscribeEvent) (models.SubscribeReply, backend.SubscribeStreamStatus, error) {
+	reply := models.SubscribeReply{
 		Presence:  true,
 		JoinLeave: true,
-		Recover:   true, // loads the saved value from history
-	}, backend.SubscribeStreamStatusOK, nil
+	}
+	query := &models.GetLiveMessageQuery{
+		OrgId:   u.OrgId,
+		Channel: e.Channel,
+	}
+	msg, ok, err := b.liveMessageStore.GetLiveMessage(query)
+	if err != nil {
+		return models.SubscribeReply{}, 0, err
+	}
+	if ok {
+		reply.Data = msg.Data
+	}
+	return reply, backend.SubscribeStreamStatusOK, nil
 }
 
 // OnPublish is called when a client wants to broadcast on the websocket
-func (b *BroadcastRunner) OnPublish(ctx context.Context, _ *models.SignedInUser, e models.PublishEvent) (models.PublishReply, backend.PublishStreamStatus, error) {
-	return models.PublishReply{
-		HistorySize: 1, // The last message is saved for 10 min.
-		HistoryTTL:  10 * time.Minute,
-	}, backend.PublishStreamStatusOK, nil
+func (b *BroadcastRunner) OnPublish(_ context.Context, u *models.SignedInUser, e models.PublishEvent) (models.PublishReply, backend.PublishStreamStatus, error) {
+	query := &models.SaveLiveMessageQuery{
+		OrgId:   u.OrgId,
+		Channel: e.Channel,
+		Data:    e.Data,
+	}
+	if err := b.liveMessageStore.SaveLiveMessage(query); err != nil {
+		return models.PublishReply{}, 0, err
+	}
+	return models.PublishReply{}, backend.PublishStreamStatusOK, nil
 }

pkg/services/live/features/broadcast_mock.go

@@ -0,0 +1,65 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/grafana/grafana/pkg/services/live/features (interfaces: LiveMessageStore)
+
+// Package features is a generated GoMock package.
+package features
+
+import (
+	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+	models "github.com/grafana/grafana/pkg/models"
+)
+
+// MockLiveMessageStore is a mock of LiveMessageStore interface.
+type MockLiveMessageStore struct {
+	ctrl     *gomock.Controller
+	recorder *MockLiveMessageStoreMockRecorder
+}
+
+// MockLiveMessageStoreMockRecorder is the mock recorder for MockLiveMessageStore.
+type MockLiveMessageStoreMockRecorder struct {
+	mock *MockLiveMessageStore
+}
+
+// NewMockLiveMessageStore creates a new mock instance.
+func NewMockLiveMessageStore(ctrl *gomock.Controller) *MockLiveMessageStore {
+	mock := &MockLiveMessageStore{ctrl: ctrl}
+	mock.recorder = &MockLiveMessageStoreMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockLiveMessageStore) EXPECT() *MockLiveMessageStoreMockRecorder {
+	return m.recorder
+}
+
+// GetLiveMessage mocks base method.
+func (m *MockLiveMessageStore) GetLiveMessage(arg0 *models.GetLiveMessageQuery) (models.LiveMessage, bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetLiveMessage", arg0)
+	ret0, _ := ret[0].(models.LiveMessage)
+	ret1, _ := ret[1].(bool)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
+}
+
+// GetLiveMessage indicates an expected call of GetLiveMessage.
+func (mr *MockLiveMessageStoreMockRecorder) GetLiveMessage(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLiveMessage", reflect.TypeOf((*MockLiveMessageStore)(nil).GetLiveMessage), arg0)
+}
+
+// SaveLiveMessage mocks base method.
+func (m *MockLiveMessageStore) SaveLiveMessage(arg0 *models.SaveLiveMessageQuery) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "SaveLiveMessage", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SaveLiveMessage indicates an expected call of SaveLiveMessage.
+func (mr *MockLiveMessageStoreMockRecorder) SaveLiveMessage(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveLiveMessage", reflect.TypeOf((*MockLiveMessageStore)(nil).SaveLiveMessage), arg0)
+}

pkg/services/live/features/broadcast_test.go

@@ -0,0 +1,85 @@
+package features
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewBroadcastRunner(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	d := NewMockLiveMessageStore(mockCtrl)
+	br := NewBroadcastRunner(d)
+	require.NotNil(t, br)
+}
+
+func TestBroadcastRunner_OnSubscribe(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mockDispatcher := NewMockLiveMessageStore(mockCtrl)
+
+	channel := "stream/channel/test"
+	data := json.RawMessage(`{}`)
+
+	mockDispatcher.EXPECT().GetLiveMessage(&models.GetLiveMessageQuery{
+		OrgId:   1,
+		Channel: channel,
+	}).DoAndReturn(func(query *models.GetLiveMessageQuery) (models.LiveMessage, bool, error) {
+		return models.LiveMessage{
+			Data: data,
+		}, true, nil
+	}).Times(1)
+
+	br := NewBroadcastRunner(mockDispatcher)
+	require.NotNil(t, br)
+	handler, err := br.GetHandlerForPath("test")
+	require.NoError(t, err)
+	reply, status, err := handler.OnSubscribe(
+		context.Background(),
+		&models.SignedInUser{OrgId: 1, UserId: 2},
+		models.SubscribeEvent{Channel: channel, Path: "test"},
+	)
+	require.NoError(t, err)
+	require.Equal(t, backend.SubscribeStreamStatusOK, status)
+	require.Equal(t, data, reply.Data)
+	require.True(t, reply.Presence)
+	require.True(t, reply.JoinLeave)
+	require.False(t, reply.Recover)
+}
+
+func TestBroadcastRunner_OnPublish(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mockDispatcher := NewMockLiveMessageStore(mockCtrl)
+
+	channel := "stream/channel/test"
+	data := json.RawMessage(`{}`)
+	var orgID int64 = 1
+
+	mockDispatcher.EXPECT().SaveLiveMessage(&models.SaveLiveMessageQuery{
+		OrgId:   orgID,
+		Channel: channel,
+		Data:    data,
+	}).DoAndReturn(func(query *models.SaveLiveMessageQuery) error {
+		return nil
+	}).Times(1)
+
+	br := NewBroadcastRunner(mockDispatcher)
+	require.NotNil(t, br)
+	handler, err := br.GetHandlerForPath("test")
+	require.NoError(t, err)
+	reply, status, err := handler.OnPublish(
+		context.Background(),
+		&models.SignedInUser{OrgId: 1, UserId: 2},
+		models.PublishEvent{Channel: channel, Path: "test", Data: data},
+	)
+	require.NoError(t, err)
+	require.Equal(t, backend.PublishStreamStatusOK, status)
+	require.Nil(t, reply.Data)
+}

pkg/services/live/live.go

@@ -8,6 +8,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
+
 	"github.com/grafana/grafana/pkg/api/dtos"
 	"github.com/grafana/grafana/pkg/api/response"
 	"github.com/grafana/grafana/pkg/api/routing"
@@ -18,12 +20,15 @@ import (
 	"github.com/grafana/grafana/pkg/plugins/plugincontext"
 	"github.com/grafana/grafana/pkg/registry"
 	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/live/database"
 	"github.com/grafana/grafana/pkg/services/live/demultiplexer"
 	"github.com/grafana/grafana/pkg/services/live/features"
 	"github.com/grafana/grafana/pkg/services/live/livecontext"
 	"github.com/grafana/grafana/pkg/services/live/managedstream"
 	"github.com/grafana/grafana/pkg/services/live/pushws"
 	"github.com/grafana/grafana/pkg/services/live/runstream"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/cloudwatch"
 	"github.com/grafana/grafana/pkg/util"
@@ -40,13 +45,17 @@ var (
 )
 
 func init() {
-	registry.RegisterServiceWithPriority(&GrafanaLive{
+	registry.RegisterServiceWithPriority(NewGrafanaLive(), registry.Low)
+}
+
+func NewGrafanaLive() *GrafanaLive {
+	return &GrafanaLive{
 		channels:   make(map[string]models.ChannelHandler),
 		channelsMu: sync.RWMutex{},
 		GrafanaScope: CoreGrafanaScope{
 			Features: make(map[string]models.ChannelHandlerFactory),
 		},
-	}, registry.Low)
+	}
 }
 
 // CoreGrafanaScope list of core features
@@ -68,7 +77,9 @@ type GrafanaLive struct {
 	RouteRegister         routing.RouteRegister    `inject:""`
 	LogsService           *cloudwatch.LogsService  `inject:""`
 	PluginManager         *manager.PluginManager   `inject:""`
+	CacheService          *localcache.CacheService `inject:""`
 	DatasourceCache       datasources.CacheService `inject:""`
+	SQLStore              *sqlstore.SQLStore       `inject:""`
 
 	node *centrifuge.Node
 
@@ -87,6 +98,7 @@ type GrafanaLive struct {
 
 	contextGetter    *pluginContextGetter
 	runStreamManager *runstream.Manager
+	storage          *database.Storage
 }
 
 func (g *GrafanaLive) getStreamPlugin(pluginID string) (backend.StreamHandler, error) {
@@ -101,6 +113,15 @@ func (g *GrafanaLive) getStreamPlugin(pluginID string) (backend.StreamHandler, e
 	return streamHandler, nil
 }
 
+// AddMigration defines database migrations.
+// This is an implementation of registry.DatabaseMigrator.
+func (g *GrafanaLive) AddMigration(mg *migrator.Migrator) {
+	if !g.IsEnabled() {
+		return
+	}
+	database.AddLiveChannelMigrations(mg)
+}
+
 func (g *GrafanaLive) Run(ctx context.Context) error {
 	if g.runStreamManager != nil {
 		// Only run stream manager if GrafanaLive properly initialized.
@@ -146,9 +167,10 @@ func (g *GrafanaLive) Init() error {
 		Publisher:   g.Publish,
 		ClientCount: g.ClientCount,
 	}
+	g.storage = database.NewStorage(g.SQLStore, g.CacheService)
 	g.GrafanaScope.Dashboards = dash
 	g.GrafanaScope.Features["dashboard"] = dash
-	g.GrafanaScope.Features["broadcast"] = &features.BroadcastRunner{}
+	g.GrafanaScope.Features["broadcast"] = features.NewBroadcastRunner(g.storage)
 
 	g.ManagedStreamRunner = managedstream.NewRunner(g.Publish)
 
@@ -477,7 +499,10 @@ func (g *GrafanaLive) ClientCount(channel string) (int, error) {
 
 // IsEnabled returns true if the Grafana Live feature is enabled.
 func (g *GrafanaLive) IsEnabled() bool {
-	return g != nil && g.Cfg.IsLiveEnabled()
+	if g == nil || g.Cfg == nil {
+		return false
+	}
+	return g.Cfg.IsLiveEnabled()
 }
 
 func (g *GrafanaLive) HandleHTTPPublish(ctx *models.ReqContext, cmd dtos.LivePublishCmd) response.Response {

pkg/services/sqlstore/migrator/migrator.go

@@ -1,6 +1,7 @@
 package migrator
 
 import (
+	"fmt"
 	"time"
 
 	_ "github.com/go-sql-driver/mysql"
@@ -117,7 +118,7 @@ func (mg *Migrator) Start() error {
 			return err
 		})
 		if err != nil {
-			return errutil.Wrap("migration failed", err)
+			return errutil.Wrap(fmt.Sprintf("migration failed (id = %s)", m.Id()), err)
 		}
 	}