Live: using organisation IDs throughout stack (#33714)

This commit is contained in:
Alexander Emelin
2021-05-11 22:03:04 +03:00
committed by GitHub
parent 17305ccda3
commit 95a356a840
16 changed files with 288 additions and 88 deletions
+44 -24
View File
@@ -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
+21 -10
View File
@@ -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))