fix(service): keep subscriptions made before Start, drop publishes before it

pkg/ysk/ysk_test.go

@@ -25,10 +25,7 @@ func setup(t *testing.T) (*service.EventServiceWS, *service.YSKService, func())
 
 	ctx := context.Background()
 	go s.Start(&ctx)
-	// EventServiceWS exposes no readiness signal and Publish dereferences the
-	// context it installs in Start, so the boot wait cannot be made conditional.
-	// Subscribing before Start would also be wiped by its subscriberChannels reset.
-	time.Sleep(1 * time.Second)
+	<-wsService.Ready()
 
 	return wsService, yskService, func() {
 		repository.Close()

service/action_service_websocket.go

@@ -14,17 +14,44 @@ import (
 type ActionServiceWS struct {
 	typeService *ActionTypeService
 
-	ctx   *context.Context
+	// mutex guards every field below: Start runs in a goroutine of its own,
+	// Trigger and Subscribe come from the routes
 	mutex sync.Mutex
+	ctx   *context.Context
 	stop  chan struct{}
 
 	inboundChannel     chan model.Action
 	subscriberChannels map[string]map[string][]chan model.Action
+
+	// ready: closed once Start has installed its state; Trigger delivers from then on
+	ready     chan struct{}
+	readyOnce sync.Once
+}
+
+// Ready is closed once Start runs: actions triggered from then on are delivered.
+func (s *ActionServiceWS) Ready() <-chan struct{} {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	return s.readyLocked()
+}
+
+func (s *ActionServiceWS) readyLocked() chan struct{} {
+	if s.ready == nil {
+		s.ready = make(chan struct{})
+	}
+	return s.ready
 }
 
 func (s *ActionServiceWS) Trigger(action model.Action) {
-	if s.inboundChannel == nil {
-		logger.Error("error when triggering action via websocket", zap.Error(ErrInboundChannelNotFound))
+	// Start runs in a goroutine, sets both under the lock; before that: drop (no nil-ctx panic)
+	inboundChannel, ctx := func() (chan model.Action, *context.Context) {
+		s.mutex.Lock()
+		defer s.mutex.Unlock()
+		return s.inboundChannel, s.ctx
+	}()
+	if inboundChannel == nil || ctx == nil {
+		logger.Warn("action triggered before the bus started: dropped", zap.String("sourceID", action.SourceID), zap.String("name", action.Name))
+		return
 	}
 
 	if action.Timestamp == 0 {
@@ -33,16 +60,17 @@ func (s *ActionServiceWS) Trigger(action model.Action) {
 
 	// TODO - ensure properties are valid for action type
 
+	// Waits for the dispatch to take it, never long (it waits for no
+	// subscriber), as EventServiceWS.Publish does: the route triggers from a
+	// goroutine of its own.
 	select {
-	case s.inboundChannel <- action:
+	case inboundChannel <- action:
 
-	case <-(*s.ctx).Done():
-		if err := (*s.ctx).Err(); err != nil {
+	case <-(*ctx).Done():
+		if err := (*ctx).Err(); err != nil {
 			logger.Info(err.Error())
 		}
 		return
-
-	default: // drop action if no one is listening
 	}
 }
 
@@ -69,6 +97,9 @@ func (s *ActionServiceWS) Subscribe(sourceID string, names []string) (chan model
 		}
 	}
 
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
 	if s.subscriberChannels == nil {
 		s.subscriberChannels = make(map[string]map[string][]chan model.Action)
 	}
@@ -118,14 +149,27 @@ func (s *ActionServiceWS) Unsubscribe(sourceID string, name string, c chan model
 }
 
 func (s *ActionServiceWS) Start(ctx *context.Context) {
-	s.ctx = ctx
-	s.mutex = sync.Mutex{}
+	inboundChannel := func() chan model.Action {
+		s.mutex.Lock()
+		defer s.mutex.Unlock()
+
+		s.ctx = ctx
 
-	s.inboundChannel = make(chan model.Action)
-	s.subscriberChannels = make(map[string]map[string][]chan model.Action)
-	s.stop = make(chan struct{})
+		s.inboundChannel = make(chan model.Action)
+		// keep subscriptions made before Start (see EventServiceWS.Start)
+		if s.subscriberChannels == nil {
+			s.subscriberChannels = make(map[string]map[string][]chan model.Action)
+		}
+		s.stop = make(chan struct{})
+		ready := s.readyLocked()
+		s.readyOnce.Do(func() { close(ready) })
+		return s.inboundChannel
+	}()
 
 	defer func() {
+		s.mutex.Lock()
+		defer s.mutex.Unlock()
+
 		if s.subscriberChannels != nil {
 			for sourceID, source := range s.subscriberChannels {
 				for actionName, subscribers := range source {
@@ -159,60 +203,54 @@ func (s *ActionServiceWS) Start(ctx *context.Context) {
 	for {
 		select {
 
-		case <-(*s.ctx).Done():
+		case <-(*ctx).Done():
 			return
 
-		case action, ok := <-s.inboundChannel:
+		case action, ok := <-inboundChannel:
 			if !ok {
 				return
 			}
 
-			if s.subscriberChannels == nil {
-				continue
-			}
+			// under the lock Subscribe and Unsubscribe change the lists with (see
+			// EventServiceWS.Start)
+			func() {
+				s.mutex.Lock()
+				defer s.mutex.Unlock()
 
-			if s.subscriberChannels[action.SourceID] == nil {
-				continue
-			}
-
-			if s.subscriberChannels[action.SourceID][action.Name] == nil {
-				continue
-			}
-
-			for _, c := range s.subscriberChannels[action.SourceID][action.Name] {
-				select {
-				case c <- action:
-				case <-(*s.ctx).Done():
+				if s.subscriberChannels == nil || s.subscriberChannels[action.SourceID] == nil {
 					return
-				default: // drop event if no one is listening
-					continue
 				}
-			}
 
-		case <-ticker.C:
-			if s.subscriberChannels == nil {
-				continue
-			}
+				for _, c := range s.subscriberChannels[action.SourceID][action.Name] {
+					select {
+					case c <- action:
+					default: // drop action if no one is listening
+					}
+				}
+			}()
 
+		case <-ticker.C:
 			heartbeat := model.Action{
 				SourceID:  common.MessageBusSourceID,
 				Name:      common.MessageBusHeartbeatName,
 				Timestamp: time.Now().Unix(),
 			}
 
-			for _, source := range s.subscriberChannels {
-				for _, subscribers := range source {
-					for _, subscriber := range subscribers {
-						select {
-						case subscriber <- heartbeat:
-						case <-(*s.ctx).Done():
-							return
-						default: // drop event if no one is listening
-							continue
+			func() {
+				s.mutex.Lock()
+				defer s.mutex.Unlock()
+
+				for _, source := range s.subscriberChannels {
+					for _, subscribers := range source {
+						for _, subscriber := range subscribers {
+							select {
+							case subscriber <- heartbeat:
+							default: // drop event if no one is listening
+							}
 						}
 					}
 				}
-			}
+			}()
 		}
 	}
 }

service/action_service_websocket_test.go

@@ -0,0 +1,109 @@
+package service
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/ReCasaOS/CasaOS-Common/utils/logger"
+	"github.com/ReCasaOS/CasaOS-MessageBus/model"
+	"github.com/ReCasaOS/CasaOS-MessageBus/repository"
+	"go.uber.org/goleak"
+	"gotest.tools/assert"
+)
+
+func newActionService(t *testing.T) *ActionServiceWS {
+	t.Helper()
+	logger.LogInitConsoleOnly()
+	repository, err := repository.NewDatabaseRepositoryInMemory()
+	assert.NilError(t, err)
+	t.Cleanup(func() { repository.Close() })
+	typeService := NewActionTypeService(&repository)
+	_, err = typeService.RegisterActionType(model.ActionType{SourceID: "app-management", Name: "app:install"})
+	assert.NilError(t, err)
+	return NewActionServiceWS(typeService)
+}
+
+func receiveAction(t *testing.T, channel chan model.Action) model.Action {
+	t.Helper()
+	select {
+	case action := <-channel:
+		return action
+	case <-time.After(5 * time.Second):
+		t.Fatal("no action arrived")
+		return model.Action{}
+	}
+}
+
+// A subscription made before Start must survive Start (see the event service).
+func TestActionSubscriptionBeforeStartSurvivesStart(t *testing.T) {
+	// a cleanup: runs after the repository's (registered later, run first)
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+	ws := newActionService(t)
+
+	channel, err := ws.Subscribe("app-management", []string{"app:install"})
+	assert.NilError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go ws.Start(&ctx)
+	<-ws.Ready()
+
+	within(t, "Trigger after Start", func() {
+		ws.Trigger(model.Action{SourceID: "app-management", Name: "app:install"})
+	})
+	assert.Equal(t, receiveAction(t, channel).Name, "app:install")
+}
+
+// A Trigger before Start is dropped; it used to dereference the nil context.
+func TestTriggerBeforeStartIsDropped(t *testing.T) {
+	ws := newActionService(t)
+	within(t, "Trigger before Start", func() {
+		ws.Trigger(model.Action{SourceID: "app-management", Name: "app:install"})
+	})
+}
+
+// Subscribe and Trigger from the routes while Start installs its state, and
+// while it dispatches: -race finds an unguarded field.
+func TestActionSubscribeAndTriggerRaceWithStart(t *testing.T) {
+	// a cleanup: runs after the repository's (registered later, run first)
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+	ws := newActionService(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	stopped := make(chan struct{})
+
+	var wg sync.WaitGroup
+	channels := make(chan chan model.Action, 8)
+	for range 8 {
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			channel, err := ws.Subscribe("app-management", []string{"app:install"})
+			assert.Check(t, err)
+			channels <- channel
+		}()
+		go func() {
+			defer wg.Done()
+			ws.Trigger(model.Action{SourceID: "app-management", Name: "app:install"})
+		}()
+	}
+	go func() {
+		ws.Start(&ctx)
+		close(stopped)
+	}()
+	wg.Wait()
+	close(channels)
+
+	// every subscription made concurrently with Start is kept
+	<-ws.Ready()
+	ws.Trigger(model.Action{SourceID: "app-management", Name: "app:install"})
+	for channel := range channels {
+		// buffered (1): the last action, or an earlier one if it was still waiting
+		assert.Equal(t, receiveAction(t, channel).Name, "app:install")
+	}
+
+	cancel()
+	<-stopped
+}