Live: RunStream improvements – backoff, better layout, tests (#33029)

2021-04-17 01:17:08 +03:00
parent 1ed724ea6e
commit 7d5a46ffda
7 changed files with 434 additions and 298 deletions
@@ -0,0 +1,289 @@
+package runstream
+
+import (
+	"context"
+	"errors"
+	"math"
+	"sync"
+	"time"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+var (
+	logger = log.New("live.runstream")
+)
+
+//go:generate mockgen -destination=mock.go -package=runstream github.com/grafana/grafana/pkg/services/live/runstream StreamPacketSender,PresenceGetter,StreamRunner
+
+type StreamPacketSender interface {
+	Send(channel string, packet *backend.StreamPacket) error
+}
+
+type PresenceGetter interface {
+	GetNumSubscribers(channel string) (int, error)
+}
+
+type StreamRunner interface {
+	RunStream(ctx context.Context, request *backend.RunStreamRequest, sender backend.StreamPacketSender) error
+}
+
+type streamSender struct {
+	channel      string
+	packetSender StreamPacketSender
+}
+
+func newStreamSender(channel string, packetSender StreamPacketSender) *streamSender {
+	return &streamSender{
+		channel:      channel,
+		packetSender: packetSender,
+	}
+}
+
+func (p *streamSender) Send(packet *backend.StreamPacket) error {
+	return p.packetSender.Send(p.channel, packet)
+}
+
+// Manager manages streams from Grafana to plugins (i.e. RunStream method).
+type Manager struct {
+	mu             sync.RWMutex
+	streams        map[string]chan struct{}
+	presenceGetter PresenceGetter
+	packetSender   StreamPacketSender
+	registerCh     chan submitRequest
+	closedCh       chan struct{}
+	checkInterval  time.Duration
+	maxChecks      int
+}
+
+// ManagerOption modifies Manager behavior (used for tests for example).
+type ManagerOption func(*Manager)
+
+// WithCheckConfig allows setting custom check rules.
+func WithCheckConfig(interval time.Duration, maxChecks int) ManagerOption {
+	return func(sm *Manager) {
+		sm.checkInterval = interval
+		sm.maxChecks = maxChecks
+	}
+}
+
+const (
+	defaultCheckInterval = 5 * time.Second
+	defaultMaxChecks     = 3
+)
+
+// NewManager creates new Manager.
+func NewManager(packetSender StreamPacketSender, presenceGetter PresenceGetter, opts ...ManagerOption) *Manager {
+	sm := &Manager{
+		streams:        make(map[string]chan struct{}),
+		packetSender:   packetSender,
+		presenceGetter: presenceGetter,
+		registerCh:     make(chan submitRequest),
+		closedCh:       make(chan struct{}),
+		checkInterval:  defaultCheckInterval,
+		maxChecks:      defaultMaxChecks,
+	}
+	for _, opt := range opts {
+		opt(sm)
+	}
+	return sm
+}
+
+func (s *Manager) stopStream(sr streamRequest, cancelFn func()) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	closeCh, ok := s.streams[sr.Channel]
+	if !ok {
+		return
+	}
+	delete(s.streams, sr.Channel)
+	cancelFn()
+	close(closeCh)
+}
+
+func (s *Manager) watchStream(ctx context.Context, cancelFn func(), sr streamRequest) {
+	numNoSubscribersChecks := 0
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(s.checkInterval):
+			numSubscribers, err := s.presenceGetter.GetNumSubscribers(sr.Channel)
+			if err != nil {
+				logger.Error("Error checking num subscribers", "channel", sr.Channel, "path", sr.Path)
+				continue
+			}
+			if numSubscribers > 0 {
+				// reset counter since channel has active subscribers.
+				numNoSubscribersChecks = 0
+				continue
+			}
+			numNoSubscribersChecks++
+			if numNoSubscribersChecks >= s.maxChecks {
+				logger.Debug("Stop stream since no active subscribers", "channel", sr.Channel, "path", sr.Path)
+				s.stopStream(sr, cancelFn)
+				return
+			}
+		}
+	}
+}
+
+const streamDurationThreshold = 100 * time.Millisecond
+const coolDownDelay = 100 * time.Millisecond
+const maxDelay = 5 * time.Second
+
+func getDelay(numErrors int) time.Duration {
+	if numErrors == 0 {
+		return 0
+	}
+	delay := coolDownDelay * time.Duration(math.Pow(2, float64(numErrors)))
+	if delay > maxDelay {
+		return maxDelay
+	}
+	return delay
+}
+
+// run stream until context canceled or stream finished without an error.
+func (s *Manager) runStream(ctx context.Context, cancelFn func(), sr streamRequest) {
+	var numFastErrors int
+	var delay time.Duration
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+		startTime := time.Now()
+		err := sr.StreamRunner.RunStream(
+			ctx,
+			&backend.RunStreamRequest{
+				PluginContext: sr.PluginContext,
+				Path:          sr.Path,
+			},
+			newStreamSender(sr.Channel, s.packetSender),
+		)
+		if err != nil {
+			if errors.Is(ctx.Err(), context.Canceled) {
+				logger.Debug("Stream cleanly finished", "path", sr.Path)
+				return
+			}
+			// Best effort to cool down re-establishment process. We don't have a
+			// nice way to understand whether we really need to wait here - so relying
+			// on RunStream duration time.
+			if time.Since(startTime) < streamDurationThreshold {
+				if delay < maxDelay {
+					// Due to not calling getDelay after we have delay larger than maxDelay
+					// we avoid possible float overflow errors while calculating delay duration
+					// based on numFastErrors.
+					delay = getDelay(numFastErrors)
+				}
+				numFastErrors++
+			} else {
+				// Assuming that stream successfully started.
+				delay = 0
+				numFastErrors = 0
+			}
+			logger.Error("Error running stream, re-establishing", "path", sr.Path, "error", err, "wait", delay)
+			time.Sleep(delay)
+			continue
+		}
+		logger.Debug("Stream finished without error, stopping it", "path", sr.Path)
+		s.stopStream(sr, cancelFn)
+		return
+	}
+}
+
+var errClosed = errors.New("stream manager closed")
+
+func (s *Manager) registerStream(ctx context.Context, sr submitRequest) {
+	s.mu.Lock()
+	if closeCh, ok := s.streams[sr.streamRequest.Channel]; ok {
+		s.mu.Unlock()
+		sr.responseCh <- submitResponse{Result: submitResult{StreamExists: true, CloseNotify: closeCh}}
+		return
+	}
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	closeCh := make(chan struct{})
+	s.streams[sr.streamRequest.Channel] = closeCh
+	s.mu.Unlock()
+	sr.responseCh <- submitResponse{Result: submitResult{StreamExists: false, CloseNotify: closeCh}}
+	go s.watchStream(ctx, cancel, sr.streamRequest)
+	s.runStream(ctx, cancel, sr.streamRequest)
+}
+
+// Run Manager till context canceled.
+func (s *Manager) Run(ctx context.Context) error {
+	for {
+		select {
+		case sr := <-s.registerCh:
+			go s.registerStream(ctx, sr)
+		case <-ctx.Done():
+			close(s.closedCh)
+			return ctx.Err()
+		}
+	}
+}
+
+type streamRequest struct {
+	Channel       string
+	Path          string
+	PluginContext backend.PluginContext
+	StreamRunner  StreamRunner
+}
+
+type submitRequest struct {
+	responseCh    chan submitResponse
+	streamRequest streamRequest
+}
+
+type submitResult struct {
+	// StreamExists tells whether stream have been already opened.
+	StreamExists bool
+	// CloseNotify will be closed as soon as stream cleanly exited.
+	CloseNotify chan struct{}
+}
+
+type submitResponse struct {
+	Error  error
+	Result submitResult
+}
+
+// SubmitStream submits stream handler in Manager to manage.
+// The stream will be opened and kept till channel has active subscribers.
+func (s *Manager) SubmitStream(ctx context.Context, channel string, path string, pCtx backend.PluginContext, streamRunner StreamRunner) (*submitResult, error) {
+	req := submitRequest{
+		responseCh: make(chan submitResponse, 1),
+		streamRequest: streamRequest{
+			Channel:       channel,
+			Path:          path,
+			PluginContext: pCtx,
+			StreamRunner:  streamRunner,
+		},
+	}
+
+	// Send submit request.
+	select {
+	case s.registerCh <- req:
+	case <-s.closedCh:
+		close(s.registerCh)
+		return nil, errClosed
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+
+	// Wait for submit response.
+	select {
+	case resp := <-req.responseCh:
+		if resp.Error != nil {
+			return nil, resp.Error
+		}
+		return &resp.Result, nil
+	case <-s.closedCh:
+		return nil, errClosed
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
@@ -0,0 +1,197 @@
+package runstream
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+)
+
+// wait until channel closed with timeout.
+func waitWithTimeout(tb testing.TB, ch chan struct{}, timeout time.Duration) {
+	select {
+	case <-ch:
+	case <-time.After(timeout):
+		tb.Fatal("timeout")
+	}
+}
+
+func TestStreamManager_Run(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mockPacketSender := NewMockStreamPacketSender(mockCtrl)
+	mockPresenceGetter := NewMockPresenceGetter(mockCtrl)
+
+	manager := NewManager(mockPacketSender, mockPresenceGetter)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		cancel()
+	}()
+
+	err := manager.Run(ctx)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestStreamManager_SubmitStream_Send(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mockPacketSender := NewMockStreamPacketSender(mockCtrl)
+	mockPresenceGetter := NewMockPresenceGetter(mockCtrl)
+
+	manager := NewManager(mockPacketSender, mockPresenceGetter)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		_ = manager.Run(ctx)
+	}()
+
+	startedCh := make(chan struct{})
+	doneCh := make(chan struct{})
+
+	mockPacketSender.EXPECT().Send("test", gomock.Any()).Times(1)
+
+	mockStreamRunner := NewMockStreamRunner(mockCtrl)
+	mockStreamRunner.EXPECT().RunStream(
+		gomock.Any(), gomock.Any(), gomock.Any(),
+	).DoAndReturn(func(ctx context.Context, req *backend.RunStreamRequest, sender backend.StreamPacketSender) error {
+		require.Equal(t, "test", req.Path)
+		close(startedCh)
+		err := sender.Send(&backend.StreamPacket{
+			Data: []byte("test"),
+		})
+		require.NoError(t, err)
+		<-ctx.Done()
+		close(doneCh)
+		return ctx.Err()
+	}).Times(1)
+
+	result, err := manager.SubmitStream(context.Background(), "test", "test", backend.PluginContext{}, mockStreamRunner)
+	require.NoError(t, err)
+	require.False(t, result.StreamExists)
+
+	// try submit the same.
+	result, err = manager.SubmitStream(context.Background(), "test", "test", backend.PluginContext{}, mockStreamRunner)
+	require.NoError(t, err)
+	require.True(t, result.StreamExists)
+
+	waitWithTimeout(t, startedCh, time.Second)
+	require.Len(t, manager.streams, 1)
+	cancel()
+	waitWithTimeout(t, doneCh, time.Second)
+}
+
+func TestStreamManager_SubmitStream_CloseNoSubscribers(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mockPacketSender := NewMockStreamPacketSender(mockCtrl)
+	mockPresenceGetter := NewMockPresenceGetter(mockCtrl)
+
+	// Create manager with very fast num subscribers checks.
+	manager := NewManager(
+		mockPacketSender,
+		mockPresenceGetter,
+		WithCheckConfig(10*time.Millisecond, 3),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		_ = manager.Run(ctx)
+	}()
+
+	startedCh := make(chan struct{})
+	doneCh := make(chan struct{})
+
+	mockPresenceGetter.EXPECT().GetNumSubscribers("test").Return(0, nil).Times(3)
+
+	mockStreamRunner := NewMockStreamRunner(mockCtrl)
+	mockStreamRunner.EXPECT().RunStream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, req *backend.RunStreamRequest, sender backend.StreamPacketSender) error {
+		close(startedCh)
+		<-ctx.Done()
+		close(doneCh)
+		return ctx.Err()
+	}).Times(1)
+
+	_, err := manager.SubmitStream(context.Background(), "test", "test", backend.PluginContext{}, mockStreamRunner)
+	require.NoError(t, err)
+
+	waitWithTimeout(t, startedCh, time.Second)
+	waitWithTimeout(t, doneCh, time.Second)
+	require.Len(t, manager.streams, 0)
+}
+
+func TestStreamManager_SubmitStream_ErrorRestartsRunStream(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mockPacketSender := NewMockStreamPacketSender(mockCtrl)
+	mockPresenceGetter := NewMockPresenceGetter(mockCtrl)
+
+	manager := NewManager(mockPacketSender, mockPresenceGetter)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		_ = manager.Run(ctx)
+	}()
+
+	numErrors := 3
+	currentErrors := 0
+
+	mockStreamRunner := NewMockStreamRunner(mockCtrl)
+	mockStreamRunner.EXPECT().RunStream(
+		gomock.Any(), gomock.Any(), gomock.Any(),
+	).DoAndReturn(func(ctx context.Context, req *backend.RunStreamRequest, sender backend.StreamPacketSender) error {
+		if currentErrors >= numErrors {
+			return nil
+		}
+		currentErrors++
+		return errors.New("boom")
+	}).Times(numErrors + 1)
+
+	result, err := manager.SubmitStream(context.Background(), "test", "test", backend.PluginContext{}, mockStreamRunner)
+	require.NoError(t, err)
+	require.False(t, result.StreamExists)
+
+	waitWithTimeout(t, result.CloseNotify, time.Second)
+}
+
+func TestStreamManager_SubmitStream_NilErrorStopsRunStream(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mockPacketSender := NewMockStreamPacketSender(mockCtrl)
+	mockPresenceGetter := NewMockPresenceGetter(mockCtrl)
+
+	manager := NewManager(mockPacketSender, mockPresenceGetter)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		_ = manager.Run(ctx)
+	}()
+
+	mockStreamRunner := NewMockStreamRunner(mockCtrl)
+	mockStreamRunner.EXPECT().RunStream(
+		gomock.Any(), gomock.Any(), gomock.Any(),
+	).DoAndReturn(func(ctx context.Context, req *backend.RunStreamRequest, sender backend.StreamPacketSender) error {
+		return nil
+	}).Times(1)
+
+	result, err := manager.SubmitStream(context.Background(), "test", "test", backend.PluginContext{}, mockStreamRunner)
+	require.NoError(t, err)
+	require.False(t, result.StreamExists)
+	waitWithTimeout(t, result.CloseNotify, time.Second)
+}
@@ -0,0 +1,125 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/grafana/grafana/pkg/services/live/runstream (interfaces: StreamPacketSender,PresenceGetter,StreamRunner)
+
+// Package runstream is a generated GoMock package.
+package runstream
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+	backend "github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+// MockStreamPacketSender is a mock of StreamPacketSender interface.
+type MockStreamPacketSender struct {
+	ctrl     *gomock.Controller
+	recorder *MockStreamPacketSenderMockRecorder
+}
+
+// MockStreamPacketSenderMockRecorder is the mock recorder for MockStreamPacketSender.
+type MockStreamPacketSenderMockRecorder struct {
+	mock *MockStreamPacketSender
+}
+
+// NewMockStreamPacketSender creates a new mock instance.
+func NewMockStreamPacketSender(ctrl *gomock.Controller) *MockStreamPacketSender {
+	mock := &MockStreamPacketSender{ctrl: ctrl}
+	mock.recorder = &MockStreamPacketSenderMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockStreamPacketSender) EXPECT() *MockStreamPacketSenderMockRecorder {
+	return m.recorder
+}
+
+// Send mocks base method.
+func (m *MockStreamPacketSender) Send(arg0 string, arg1 *backend.StreamPacket) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Send", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Send indicates an expected call of Send.
+func (mr *MockStreamPacketSenderMockRecorder) Send(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MockStreamPacketSender)(nil).Send), arg0, arg1)
+}
+
+// MockPresenceGetter is a mock of PresenceGetter interface.
+type MockPresenceGetter struct {
+	ctrl     *gomock.Controller
+	recorder *MockPresenceGetterMockRecorder
+}
+
+// MockPresenceGetterMockRecorder is the mock recorder for MockPresenceGetter.
+type MockPresenceGetterMockRecorder struct {
+	mock *MockPresenceGetter
+}
+
+// NewMockPresenceGetter creates a new mock instance.
+func NewMockPresenceGetter(ctrl *gomock.Controller) *MockPresenceGetter {
+	mock := &MockPresenceGetter{ctrl: ctrl}
+	mock.recorder = &MockPresenceGetterMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockPresenceGetter) EXPECT() *MockPresenceGetterMockRecorder {
+	return m.recorder
+}
+
+// GetNumSubscribers mocks base method.
+func (m *MockPresenceGetter) GetNumSubscribers(arg0 string) (int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetNumSubscribers", arg0)
+	ret0, _ := ret[0].(int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetNumSubscribers indicates an expected call of GetNumSubscribers.
+func (mr *MockPresenceGetterMockRecorder) GetNumSubscribers(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNumSubscribers", reflect.TypeOf((*MockPresenceGetter)(nil).GetNumSubscribers), arg0)
+}
+
+// MockStreamRunner is a mock of StreamRunner interface.
+type MockStreamRunner struct {
+	ctrl     *gomock.Controller
+	recorder *MockStreamRunnerMockRecorder
+}
+
+// MockStreamRunnerMockRecorder is the mock recorder for MockStreamRunner.
+type MockStreamRunnerMockRecorder struct {
+	mock *MockStreamRunner
+}
+
+// NewMockStreamRunner creates a new mock instance.
+func NewMockStreamRunner(ctrl *gomock.Controller) *MockStreamRunner {
+	mock := &MockStreamRunner{ctrl: ctrl}
+	mock.recorder = &MockStreamRunnerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockStreamRunner) EXPECT() *MockStreamRunnerMockRecorder {
+	return m.recorder
+}
+
+// RunStream mocks base method.
+func (m *MockStreamRunner) RunStream(arg0 context.Context, arg1 *backend.RunStreamRequest, arg2 backend.StreamPacketSender) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RunStream", arg0, arg1, arg2)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RunStream indicates an expected call of RunStream.
+func (mr *MockStreamRunnerMockRecorder) RunStream(arg0, arg1, arg2 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RunStream", reflect.TypeOf((*MockStreamRunner)(nil).RunStream), arg0, arg1, arg2)
+}