Live: RunStream improvements – backoff, better layout, tests (#33029)
This commit is contained in:
Alexander Emelin
@@ -1,165 +0,0 @@
|
||||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: github.com/grafana/grafana/pkg/services/live/features (interfaces: StreamPacketSender,PresenceGetter,PluginContextGetter,StreamRunner)
|
||||
|
||||
// Package features is a generated GoMock package.
|
||||
package features
|
||||
|
||||
import (
|
||||
context "context"
|
||||
reflect "reflect"
|
||||
|
||||
gomock "github.com/golang/mock/gomock"
|
||||
backend "github.com/grafana/grafana-plugin-sdk-go/backend"
|
||||
models "github.com/grafana/grafana/pkg/models"
|
||||
)
|
||||
|
||||
// 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)
|
||||
}
|
||||
|
||||
// MockPluginContextGetter is a mock of PluginContextGetter interface.
|
||||
type MockPluginContextGetter struct {
|
||||
ctrl *gomock.Controller
|
||||
recorder *MockPluginContextGetterMockRecorder
|
||||
}
|
||||
|
||||
// MockPluginContextGetterMockRecorder is the mock recorder for MockPluginContextGetter.
|
||||
type MockPluginContextGetterMockRecorder struct {
|
||||
mock *MockPluginContextGetter
|
||||
}
|
||||
|
||||
// NewMockPluginContextGetter creates a new mock instance.
|
||||
func NewMockPluginContextGetter(ctrl *gomock.Controller) *MockPluginContextGetter {
|
||||
mock := &MockPluginContextGetter{ctrl: ctrl}
|
||||
mock.recorder = &MockPluginContextGetterMockRecorder{mock}
|
||||
return mock
|
||||
}
|
||||
|
||||
// EXPECT returns an object that allows the caller to indicate expected use.
|
||||
func (m *MockPluginContextGetter) EXPECT() *MockPluginContextGetterMockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// GetPluginContext mocks base method.
|
||||
func (m *MockPluginContextGetter) GetPluginContext(arg0 *models.SignedInUser, arg1, arg2 string) (backend.PluginContext, bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetPluginContext", arg0, arg1, arg2)
|
||||
ret0, _ := ret[0].(backend.PluginContext)
|
||||
ret1, _ := ret[1].(bool)
|
||||
ret2, _ := ret[2].(error)
|
||||
return ret0, ret1, ret2
|
||||
}
|
||||
|
||||
// GetPluginContext indicates an expected call of GetPluginContext.
|
||||
func (mr *MockPluginContextGetterMockRecorder) GetPluginContext(arg0, arg1, arg2 interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPluginContext", reflect.TypeOf((*MockPluginContextGetter)(nil).GetPluginContext), arg0, arg1, arg2)
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
@@ -3,62 +3,36 @@ package features
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/grafana/grafana/pkg/services/live/runstream"
|
||||
|
||||
"github.com/centrifugal/centrifuge"
|
||||
"github.com/grafana/grafana-plugin-sdk-go/backend"
|
||||
"github.com/grafana/grafana/pkg/models"
|
||||
)
|
||||
|
||||
//go:generate mockgen -destination=mock.go -package=features github.com/grafana/grafana/pkg/services/live/features StreamPacketSender,PresenceGetter,PluginContextGetter,StreamRunner
|
||||
|
||||
type StreamPacketSender interface {
|
||||
Send(channel string, packet *backend.StreamPacket) error
|
||||
}
|
||||
|
||||
type PresenceGetter interface {
|
||||
GetNumSubscribers(channel string) (int, error)
|
||||
}
|
||||
//go:generate mockgen -destination=plugin_mock.go -package=features github.com/grafana/grafana/pkg/services/live/features PluginContextGetter
|
||||
|
||||
type PluginContextGetter interface {
|
||||
GetPluginContext(user *models.SignedInUser, pluginID string, datasourceUID string) (backend.PluginContext, bool, 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)
|
||||
}
|
||||
|
||||
// PluginRunner can handle streaming operations for channels belonging to plugins.
|
||||
type PluginRunner struct {
|
||||
pluginID string
|
||||
datasourceUID string
|
||||
pluginContextGetter PluginContextGetter
|
||||
handler backend.StreamHandler
|
||||
streamManager *StreamManager
|
||||
runStreamManager *runstream.Manager
|
||||
}
|
||||
|
||||
// NewPluginRunner creates new PluginRunner.
|
||||
func NewPluginRunner(pluginID string, datasourceUID string, streamManager *StreamManager, pluginContextGetter PluginContextGetter, handler backend.StreamHandler) *PluginRunner {
|
||||
func NewPluginRunner(pluginID string, datasourceUID string, runStreamManager *runstream.Manager, pluginContextGetter PluginContextGetter, handler backend.StreamHandler) *PluginRunner {
|
||||
return &PluginRunner{
|
||||
pluginID: pluginID,
|
||||
datasourceUID: datasourceUID,
|
||||
pluginContextGetter: pluginContextGetter,
|
||||
handler: handler,
|
||||
streamManager: streamManager,
|
||||
runStreamManager: runStreamManager,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,7 +42,7 @@ func (m *PluginRunner) GetHandlerForPath(path string) (models.ChannelHandler, er
|
||||
path: path,
|
||||
pluginID: m.pluginID,
|
||||
datasourceUID: m.datasourceUID,
|
||||
streamManager: m.streamManager,
|
||||
runStreamManager: m.runStreamManager,
|
||||
handler: m.handler,
|
||||
pluginContextGetter: m.pluginContextGetter,
|
||||
}, nil
|
||||
@@ -79,7 +53,7 @@ type PluginPathRunner struct {
|
||||
path string
|
||||
pluginID string
|
||||
datasourceUID string
|
||||
streamManager *StreamManager
|
||||
runStreamManager *runstream.Manager
|
||||
handler backend.StreamHandler
|
||||
pluginContextGetter PluginContextGetter
|
||||
}
|
||||
@@ -108,7 +82,7 @@ func (r *PluginPathRunner) OnSubscribe(ctx context.Context, user *models.SignedI
|
||||
}
|
||||
|
||||
if resp.UseRunStream {
|
||||
submitResult, err := r.streamManager.SubmitStream(ctx, e.Channel, r.path, pCtx, r.handler)
|
||||
submitResult, err := r.runStreamManager.SubmitStream(ctx, e.Channel, r.path, pCtx, r.handler)
|
||||
if err != nil {
|
||||
logger.Error("Error submitting stream to manager", "error", err, "path", r.path)
|
||||
return models.SubscribeReply{}, 0, centrifuge.ErrorInternal
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: github.com/grafana/grafana/pkg/services/live/features (interfaces: PluginContextGetter)
|
||||
|
||||
// Package features is a generated GoMock package.
|
||||
package features
|
||||
|
||||
import (
|
||||
reflect "reflect"
|
||||
|
||||
gomock "github.com/golang/mock/gomock"
|
||||
backend "github.com/grafana/grafana-plugin-sdk-go/backend"
|
||||
models "github.com/grafana/grafana/pkg/models"
|
||||
)
|
||||
|
||||
// MockPluginContextGetter is a mock of PluginContextGetter interface.
|
||||
type MockPluginContextGetter struct {
|
||||
ctrl *gomock.Controller
|
||||
recorder *MockPluginContextGetterMockRecorder
|
||||
}
|
||||
|
||||
// MockPluginContextGetterMockRecorder is the mock recorder for MockPluginContextGetter.
|
||||
type MockPluginContextGetterMockRecorder struct {
|
||||
mock *MockPluginContextGetter
|
||||
}
|
||||
|
||||
// NewMockPluginContextGetter creates a new mock instance.
|
||||
func NewMockPluginContextGetter(ctrl *gomock.Controller) *MockPluginContextGetter {
|
||||
mock := &MockPluginContextGetter{ctrl: ctrl}
|
||||
mock.recorder = &MockPluginContextGetterMockRecorder{mock}
|
||||
return mock
|
||||
}
|
||||
|
||||
// EXPECT returns an object that allows the caller to indicate expected use.
|
||||
func (m *MockPluginContextGetter) EXPECT() *MockPluginContextGetterMockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// GetPluginContext mocks base method.
|
||||
func (m *MockPluginContextGetter) GetPluginContext(arg0 *models.SignedInUser, arg1, arg2 string) (backend.PluginContext, bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetPluginContext", arg0, arg1, arg2)
|
||||
ret0, _ := ret[0].(backend.PluginContext)
|
||||
ret1, _ := ret[1].(bool)
|
||||
ret2, _ := ret[2].(error)
|
||||
return ret0, ret1, ret2
|
||||
}
|
||||
|
||||
// GetPluginContext indicates an expected call of GetPluginContext.
|
||||
func (mr *MockPluginContextGetterMockRecorder) GetPluginContext(arg0, arg1, arg2 interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPluginContext", reflect.TypeOf((*MockPluginContextGetter)(nil).GetPluginContext), arg0, arg1, arg2)
|
||||
}
|
||||
@@ -1,207 +0,0 @@
|
||||
package features
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/grafana/grafana-plugin-sdk-go/backend"
|
||||
)
|
||||
|
||||
// StreamManager manages streams from Grafana to plugins.
|
||||
type StreamManager struct {
|
||||
mu sync.RWMutex
|
||||
streams map[string]struct{}
|
||||
presenceGetter PresenceGetter
|
||||
packetSender StreamPacketSender
|
||||
registerCh chan submitRequest
|
||||
closedCh chan struct{}
|
||||
checkInterval time.Duration
|
||||
maxChecks int
|
||||
}
|
||||
|
||||
// StreamManagerOption modifies StreamManager behavior (used for tests for example).
|
||||
type StreamManagerOption func(*StreamManager)
|
||||
|
||||
// WithCheckConfig allows setting custom check rules.
|
||||
func WithCheckConfig(interval time.Duration, maxChecks int) StreamManagerOption {
|
||||
return func(sm *StreamManager) {
|
||||
sm.checkInterval = interval
|
||||
sm.maxChecks = maxChecks
|
||||
}
|
||||
}
|
||||
|
||||
const (
|
||||
defaultCheckInterval = 5 * time.Second
|
||||
defaultMaxChecks = 3
|
||||
)
|
||||
|
||||
// NewStreamManager creates new StreamManager.
|
||||
func NewStreamManager(packetSender StreamPacketSender, presenceGetter PresenceGetter, opts ...StreamManagerOption) *StreamManager {
|
||||
sm := &StreamManager{
|
||||
streams: make(map[string]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 *StreamManager) stopStream(sr streamRequest, cancelFn func()) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.streams, sr.Channel)
|
||||
cancelFn()
|
||||
}
|
||||
|
||||
func (s *StreamManager) 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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// run stream until context canceled.
|
||||
func (s *StreamManager) runStream(ctx context.Context, sr streamRequest) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
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
|
||||
}
|
||||
logger.Error("Error running stream, retrying", "path", sr.Path, "error", err)
|
||||
continue
|
||||
}
|
||||
logger.Warn("Stream finished without error?", "path", sr.Path)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
var errClosed = errors.New("stream manager closed")
|
||||
|
||||
func (s *StreamManager) registerStream(ctx context.Context, sr submitRequest) {
|
||||
s.mu.Lock()
|
||||
if _, ok := s.streams[sr.streamRequest.Channel]; ok {
|
||||
s.mu.Unlock()
|
||||
sr.responseCh <- submitResponse{Result: submitResult{StreamExists: true}}
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
s.streams[sr.streamRequest.Channel] = struct{}{}
|
||||
s.mu.Unlock()
|
||||
sr.responseCh <- submitResponse{Result: submitResult{StreamExists: false}}
|
||||
go s.watchStream(ctx, cancel, sr.streamRequest)
|
||||
s.runStream(ctx, sr.streamRequest)
|
||||
}
|
||||
|
||||
// Run StreamManager till context canceled.
|
||||
func (s *StreamManager) 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 bool
|
||||
}
|
||||
|
||||
type submitResponse struct {
|
||||
Error error
|
||||
Result submitResult
|
||||
}
|
||||
|
||||
// SubmitStream submits stream handler in StreamManager to manage.
|
||||
// The stream will be opened and kept till channel has active subscribers.
|
||||
func (s *StreamManager) 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()
|
||||
}
|
||||
}
|
||||
@@ -1,131 +0,0 @@
|
||||
package features
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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 := NewStreamManager(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 := NewStreamManager(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(),
|
||||
).Do(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)
|
||||
|
||||
manager := NewStreamManager(
|
||||
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()).Do(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)
|
||||
}
|
||||
Reference in New Issue
Block a user