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

2021-04-30 11:06:33 -07:00
parent c25eab4eda
commit 0c2bcbf2bc
10 changed files with 436 additions and 19 deletions
@@ -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
 }
@@ -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)
+}
@@ -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)
+}