Live: using organisation IDs throughout stack (#33714)
This commit is contained in:
@@ -21,7 +21,7 @@ var (
|
||||
// Runner keeps ManagedStream per streamID.
|
||||
type Runner struct {
|
||||
mu sync.RWMutex
|
||||
streams map[string]*ManagedStream
|
||||
streams map[int64]map[string]*ManagedStream
|
||||
publisher models.ChannelPublisher
|
||||
}
|
||||
|
||||
@@ -29,16 +29,19 @@ type Runner struct {
|
||||
func NewRunner(publisher models.ChannelPublisher) *Runner {
|
||||
return &Runner{
|
||||
publisher: publisher,
|
||||
streams: map[string]*ManagedStream{},
|
||||
streams: map[int64]map[string]*ManagedStream{},
|
||||
}
|
||||
}
|
||||
|
||||
// Streams returns a map of active managed streams (per streamID).
|
||||
func (r *Runner) Streams() map[string]*ManagedStream {
|
||||
func (r *Runner) Streams(orgID int64) map[string]*ManagedStream {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
streams := make(map[string]*ManagedStream, len(r.streams))
|
||||
for k, v := range r.streams {
|
||||
if _, ok := r.streams[orgID]; !ok {
|
||||
return map[string]*ManagedStream{}
|
||||
}
|
||||
streams := make(map[string]*ManagedStream, len(r.streams[orgID]))
|
||||
for k, v := range r.streams[orgID] {
|
||||
streams[k] = v
|
||||
}
|
||||
return streams
|
||||
@@ -46,13 +49,17 @@ func (r *Runner) Streams() map[string]*ManagedStream {
|
||||
|
||||
// GetOrCreateStream -- for now this will create new manager for each key.
|
||||
// Eventually, the stream behavior will need to be configured explicitly
|
||||
func (r *Runner) GetOrCreateStream(streamID string) (*ManagedStream, error) {
|
||||
func (r *Runner) GetOrCreateStream(orgID int64, streamID string) (*ManagedStream, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
s, ok := r.streams[streamID]
|
||||
_, ok := r.streams[orgID]
|
||||
if !ok {
|
||||
r.streams[orgID] = map[string]*ManagedStream{}
|
||||
}
|
||||
s, ok := r.streams[orgID][streamID]
|
||||
if !ok {
|
||||
s = NewManagedStream(streamID, r.publisher)
|
||||
r.streams[streamID] = s
|
||||
r.streams[orgID][streamID] = s
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
@@ -62,7 +69,7 @@ type ManagedStream struct {
|
||||
mu sync.RWMutex
|
||||
id string
|
||||
start time.Time
|
||||
last map[string]json.RawMessage
|
||||
last map[int64]map[string]json.RawMessage
|
||||
publisher models.ChannelPublisher
|
||||
}
|
||||
|
||||
@@ -71,18 +78,22 @@ func NewManagedStream(id string, publisher models.ChannelPublisher) *ManagedStre
|
||||
return &ManagedStream{
|
||||
id: id,
|
||||
start: time.Now(),
|
||||
last: map[string]json.RawMessage{},
|
||||
last: map[int64]map[string]json.RawMessage{},
|
||||
publisher: publisher,
|
||||
}
|
||||
}
|
||||
|
||||
// ListChannels returns info for the UI about this stream.
|
||||
func (s *ManagedStream) ListChannels(prefix string) []util.DynMap {
|
||||
func (s *ManagedStream) ListChannels(orgID int64, prefix string) []util.DynMap {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
info := make([]util.DynMap, 0, len(s.last))
|
||||
for k, v := range s.last {
|
||||
if _, ok := s.last[orgID]; !ok {
|
||||
return []util.DynMap{}
|
||||
}
|
||||
|
||||
info := make([]util.DynMap, 0, len(s.last[orgID]))
|
||||
for k, v := range s.last[orgID] {
|
||||
ch := util.DynMap{}
|
||||
ch["channel"] = prefix + k
|
||||
ch["data"] = v
|
||||
@@ -93,7 +104,7 @@ func (s *ManagedStream) ListChannels(prefix string) []util.DynMap {
|
||||
|
||||
// Push sends frame to the stream and saves it for later retrieval by subscribers.
|
||||
// unstableSchema flag can be set to disable schema caching for a path.
|
||||
func (s *ManagedStream) Push(path string, frame *data.Frame, unstableSchema bool) error {
|
||||
func (s *ManagedStream) Push(orgID int64, path string, frame *data.Frame, unstableSchema bool) error {
|
||||
// Keep schema + data for last packet.
|
||||
frameJSON, err := data.FrameToJSON(frame, true, true)
|
||||
if err != nil {
|
||||
@@ -105,8 +116,11 @@ func (s *ManagedStream) Push(path string, frame *data.Frame, unstableSchema bool
|
||||
// If schema is stable we can safely cache it, and only send values if
|
||||
// stream already has schema cached.
|
||||
s.mu.Lock()
|
||||
_, exists := s.last[path]
|
||||
s.last[path] = frameJSON
|
||||
if _, ok := s.last[orgID]; !ok {
|
||||
s.last[orgID] = map[string]json.RawMessage{}
|
||||
}
|
||||
_, exists := s.last[orgID][path]
|
||||
s.last[orgID][path] = frameJSON
|
||||
s.mu.Unlock()
|
||||
|
||||
// When the packet already exits, only send the data.
|
||||
@@ -125,20 +139,26 @@ func (s *ManagedStream) Push(path string, frame *data.Frame, unstableSchema bool
|
||||
// And we don't want to cache schema for unstable case. But we still need to
|
||||
// set path to a map to make stream visible in UI stream select widget.
|
||||
s.mu.Lock()
|
||||
s.last[path] = nil
|
||||
if _, ok := s.last[orgID]; ok {
|
||||
s.last[orgID][path] = nil
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
// The channel this will be posted into.
|
||||
channel := live.Channel{Scope: live.ScopeStream, Namespace: s.id, Path: path}.String()
|
||||
logger.Debug("Publish data to channel", "channel", channel, "dataLength", len(frameJSON))
|
||||
return s.publisher(channel, frameJSON)
|
||||
return s.publisher(orgID, channel, frameJSON)
|
||||
}
|
||||
|
||||
// getLastPacket retrieves schema for a channel.
|
||||
func (s *ManagedStream) getLastPacket(path string) (json.RawMessage, bool) {
|
||||
func (s *ManagedStream) getLastPacket(orgId int64, path string) (json.RawMessage, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
schema, ok := s.last[path]
|
||||
_, ok := s.last[orgId]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
schema, ok := s.last[orgId][path]
|
||||
return schema, ok && schema != nil
|
||||
}
|
||||
|
||||
@@ -146,23 +166,23 @@ func (s *ManagedStream) GetHandlerForPath(_ string) (models.ChannelHandler, erro
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *ManagedStream) OnSubscribe(_ context.Context, _ *models.SignedInUser, e models.SubscribeEvent) (models.SubscribeReply, backend.SubscribeStreamStatus, error) {
|
||||
func (s *ManagedStream) OnSubscribe(_ context.Context, u *models.SignedInUser, e models.SubscribeEvent) (models.SubscribeReply, backend.SubscribeStreamStatus, error) {
|
||||
reply := models.SubscribeReply{}
|
||||
packet, ok := s.getLastPacket(e.Path)
|
||||
packet, ok := s.getLastPacket(u.OrgId, e.Path)
|
||||
if ok {
|
||||
reply.Data = packet
|
||||
}
|
||||
return reply, backend.SubscribeStreamStatusOK, nil
|
||||
}
|
||||
|
||||
func (s *ManagedStream) OnPublish(_ context.Context, _ *models.SignedInUser, evt models.PublishEvent) (models.PublishReply, backend.PublishStreamStatus, error) {
|
||||
func (s *ManagedStream) OnPublish(_ context.Context, u *models.SignedInUser, evt models.PublishEvent) (models.PublishReply, backend.PublishStreamStatus, error) {
|
||||
var frame data.Frame
|
||||
err := json.Unmarshal(evt.Data, &frame)
|
||||
if err != nil {
|
||||
// Stream scope only deals with data frames.
|
||||
return models.PublishReply{}, 0, err
|
||||
}
|
||||
err = s.Push(evt.Path, &frame, true)
|
||||
err = s.Push(u.OrgId, evt.Path, &frame, true)
|
||||
if err != nil {
|
||||
// Stream scope only deals with data frames.
|
||||
return models.PublishReply{}, 0, err
|
||||
|
||||
@@ -7,35 +7,46 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var noopPublisher = func(p string, b []byte) error {
|
||||
type testPublisher struct {
|
||||
orgID int64
|
||||
t *testing.T
|
||||
}
|
||||
|
||||
func (p *testPublisher) publish(orgID int64, _ string, _ []byte) error {
|
||||
require.Equal(p.t, p.orgID, orgID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestNewManagedStream(t *testing.T) {
|
||||
c := NewManagedStream("a", noopPublisher)
|
||||
publisher := &testPublisher{orgID: 1, t: t}
|
||||
c := NewManagedStream("a", publisher.publish)
|
||||
require.NotNil(t, c)
|
||||
}
|
||||
|
||||
func TestManagedStream_GetLastPacket_UnstableSchema(t *testing.T) {
|
||||
c := NewManagedStream("a", noopPublisher)
|
||||
_, ok := c.getLastPacket("test")
|
||||
var orgID int64 = 1
|
||||
publisher := &testPublisher{orgID: orgID, t: t}
|
||||
c := NewManagedStream("a", publisher.publish)
|
||||
_, ok := c.getLastPacket(orgID, "test")
|
||||
require.False(t, ok)
|
||||
err := c.Push("test", data.NewFrame("hello"), true)
|
||||
err := c.Push(orgID, "test", data.NewFrame("hello"), true)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok = c.getLastPacket("test")
|
||||
_, ok = c.getLastPacket(orgID, "test")
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestManagedStream_GetLastPacket(t *testing.T) {
|
||||
c := NewManagedStream("a", noopPublisher)
|
||||
_, ok := c.getLastPacket("test")
|
||||
var orgID int64 = 1
|
||||
publisher := &testPublisher{orgID: orgID, t: t}
|
||||
c := NewManagedStream("a", publisher.publish)
|
||||
_, ok := c.getLastPacket(orgID, "test")
|
||||
require.False(t, ok)
|
||||
err := c.Push("test", data.NewFrame("hello"), false)
|
||||
err := c.Push(orgID, "test", data.NewFrame("hello"), false)
|
||||
require.NoError(t, err)
|
||||
|
||||
s, ok := c.getLastPacket("test")
|
||||
s, ok := c.getLastPacket(orgID, "test")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, `{"schema":{"name":"hello","fields":[]},"data":{"values":[]}}`, string(s))
|
||||
|
||||
Reference in New Issue
Block a user