Re-subscribe all triggers when a capability LOOPP restarts

Author: bolekk

pkg/capabilities/registry/base.go

@@ -152,6 +152,9 @@ func (r *baseRegistry) List(_ context.Context) ([]capabilities.BaseCapability, e
 }
 
 func (r *baseRegistry) Add(ctx context.Context, c capabilities.BaseCapability) error {
+	if c == nil {
+		return errors.New("cannot add a nil capability to the registry")
+	}
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	info, err := c.Info(ctx)
@@ -175,11 +178,11 @@ func (r *baseRegistry) Add(ctx context.Context, c capabilities.BaseCapability) e
 		var ac atomicBaseCapability
 		switch info.CapabilityType {
 		case capabilities.CapabilityTypeTrigger:
-			ac = &atomicTriggerCapability{}
+			ac = newAtomicTriggerCapability()
 		case capabilities.CapabilityTypeAction, capabilities.CapabilityTypeConsensus, capabilities.CapabilityTypeTarget:
 			ac = &atomicExecuteCapability{}
 		case capabilities.CapabilityTypeCombined:
-			ac = &atomicExecuteAndTriggerCapability{}
+			ac = newAtomicExecuteAndTriggerCapability()
 		default:
 			return fmt.Errorf("unknown capability type: %s", info.CapabilityType)
 		}
@@ -206,23 +209,161 @@ func (r *baseRegistry) Remove(_ context.Context, id string) error {
 	return nil
 }
 
+type triggerRegistration struct {
+	request capabilities.TriggerRegistrationRequest
+	outCh   chan capabilities.TriggerResponse
+	cancel  context.CancelFunc // used to shut down the forwarding goroutine when the trigger is unregistered
+}
+
+// Caches all trigger registrations and replays them when the underlying capability is updated.
+// Owns channels passed to the higher layer (Engine or Don2Don) and goroutines forwarding events
+// from the underlying capability.
+type triggerRegistrationManager struct {
+	mu     sync.RWMutex
+	regs   map[string]*triggerRegistration
+	bindMu sync.Mutex
+}
+
+func newTriggerRegistrationManager() *triggerRegistrationManager {
+	return &triggerRegistrationManager{
+		regs: make(map[string]*triggerRegistration),
+	}
+}
+
+func (m *triggerRegistrationManager) register(ctx context.Context, loadExec func() (capabilities.TriggerExecutable, error), req capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
+	m.bindMu.Lock()
+	defer m.bindMu.Unlock()
+	exec, err := loadExec()
+	if err != nil {
+		return nil, err
+	}
+	in, err := exec.RegisterTrigger(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	return m.upsertRegistration(req, nil, in), nil
+}
+
+func (m *triggerRegistrationManager) unregister(ctx context.Context, exec capabilities.TriggerExecutable, req capabilities.TriggerRegistrationRequest) error {
+	m.bindMu.Lock()
+	defer m.bindMu.Unlock()
+	var out chan capabilities.TriggerResponse
+	m.mu.Lock()
+	if reg, ok := m.regs[req.TriggerID]; ok {
+		if reg.cancel != nil {
+			reg.cancel()
+		}
+		out = reg.outCh
+		delete(m.regs, req.TriggerID)
+	}
+	m.mu.Unlock()
+
+	if out != nil {
+		close(out)
+	}
+	return exec.UnregisterTrigger(ctx, req)
+}
+
+func (m *triggerRegistrationManager) upsertRegistration(req capabilities.TriggerRegistrationRequest, outCh chan capabilities.TriggerResponse, in <-chan capabilities.TriggerResponse) chan capabilities.TriggerResponse {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	regInMap, ok := m.regs[req.TriggerID]
+	if !ok {
+		if outCh == nil {
+			outCh = make(chan capabilities.TriggerResponse)
+		}
+		regInMap = &triggerRegistration{
+			request: req,
+			outCh:   outCh,
+		}
+		m.regs[req.TriggerID] = regInMap
+	} else {
+		regInMap.request = req
+		if outCh != nil {
+			regInMap.outCh = outCh
+		}
+		if regInMap.cancel != nil {
+			regInMap.cancel() // shuts down the previous forwarding goroutine
+		}
+	}
+	ctxForward, cancel := context.WithCancel(context.Background())
+	regInMap.cancel = cancel
+	go forwardTriggerResponses(ctxForward, in, regInMap.outCh)
+	return regInMap.outCh
+}
+
+func (m *triggerRegistrationManager) rebind(newExec capabilities.TriggerExecutable) error {
+	m.bindMu.Lock()
+	defer m.bindMu.Unlock()
+	m.mu.RLock()
+	regs := make([]*triggerRegistration, 0, len(m.regs))
+	for _, reg := range m.regs {
+		regs = append(regs, &triggerRegistration{
+			request: reg.request,
+			outCh:   reg.outCh,
+			cancel:  reg.cancel,
+		})
+	}
+	m.mu.RUnlock()
+
+	for _, reg := range regs {
+		if reg.cancel != nil {
+			reg.cancel()
+		}
+		if newExec != nil {
+			in, err := newExec.RegisterTrigger(context.Background(), reg.request)
+			if err != nil {
+				return fmt.Errorf("failed to re-register trigger %s: %w", reg.request.TriggerID, err)
+			}
+			_ = m.upsertRegistration(reg.request, reg.outCh, in)
+		}
+	}
+	return nil
+}
+
+func forwardTriggerResponses(ctx context.Context, in <-chan capabilities.TriggerResponse, out chan<- capabilities.TriggerResponse) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case resp, ok := <-in:
+			if !ok {
+				return
+			}
+			select {
+			case <-ctx.Done():
+				return
+			case out <- resp:
+			}
+		}
+	}
+}
+
 var _ capabilities.TriggerCapability = &atomicTriggerCapability{}
 
 type atomicTriggerCapability struct {
 	atomic.Pointer[capabilities.TriggerCapability]
+	registrations *triggerRegistrationManager
+}
+
+func newAtomicTriggerCapability() *atomicTriggerCapability {
+	return &atomicTriggerCapability{
+		registrations: newTriggerRegistrationManager(),
+	}
 }
 
 func (a *atomicTriggerCapability) Update(c capabilities.BaseCapability) error {
 	if c == nil {
 		a.Store(nil)
-		return nil
+		return a.registrations.rebind(nil)
 	}
 	tc, ok := c.(capabilities.TriggerCapability)
 	if !ok {
 		return errors.New("trigger capability does not satisfy TriggerCapability interface")
 	}
 	a.Store(&tc)
-	return nil
+	return a.registrations.rebind(tc)
 }
 
 func (a *atomicTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
@@ -245,19 +386,22 @@ func (a *atomicTriggerCapability) GetState() connectivity.State {
 }
 
 func (a *atomicTriggerCapability) RegisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
-	c := a.Load()
-	if c == nil {
-		return nil, errors.New("capability unavailable")
+	loadExec := func() (capabilities.TriggerExecutable, error) {
+		c := a.Load()
+		if c == nil {
+			return nil, errors.New("capability unavailable")
+		}
+		return *c, nil
 	}
-	return (*c).RegisterTrigger(ctx, request)
+	return a.registrations.register(ctx, loadExec, request)
 }
 
 func (a *atomicTriggerCapability) UnregisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) error {
 	c := a.Load()
 	if c == nil {
 		return errors.New("capability unavailable")
 	}
-	return (*c).UnregisterTrigger(ctx, request)
+	return a.registrations.unregister(ctx, *c, request)
 }
 
 var _ capabilities.ExecutableCapability = &atomicExecuteCapability{}
@@ -326,19 +470,26 @@ var _ capabilities.ExecutableAndTriggerCapability = &atomicExecuteAndTriggerCapa
 
 type atomicExecuteAndTriggerCapability struct {
 	atomic.Pointer[capabilities.ExecutableAndTriggerCapability]
+	registrations *triggerRegistrationManager
+}
+
+func newAtomicExecuteAndTriggerCapability() *atomicExecuteAndTriggerCapability {
+	return &atomicExecuteAndTriggerCapability{
+		registrations: newTriggerRegistrationManager(),
+	}
 }
 
 func (a *atomicExecuteAndTriggerCapability) Update(c capabilities.BaseCapability) error {
 	if c == nil {
 		a.Store(nil)
-		return nil
+		return a.registrations.rebind(nil)
 	}
 	tc, ok := c.(capabilities.ExecutableAndTriggerCapability)
 	if !ok {
 		return errors.New("target capability does not satisfy ExecutableAndTriggerCapability interface")
 	}
 	a.Store(&tc)
-	return nil
+	return a.registrations.rebind(tc)
 }
 
 func (a *atomicExecuteAndTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
@@ -361,19 +512,22 @@ func (a *atomicExecuteAndTriggerCapability) GetState() connectivity.State {
 }
 
 func (a *atomicExecuteAndTriggerCapability) RegisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
-	c := a.Load()
-	if c == nil {
-		return nil, errors.New("capability unavailable")
+	loadExec := func() (capabilities.TriggerExecutable, error) {
+		c := a.Load()
+		if c == nil {
+			return nil, errors.New("capability unavailable")
+		}
+		return *c, nil
 	}
-	return (*c).RegisterTrigger(ctx, request)
+	return a.registrations.register(ctx, loadExec, request)
 }
 
 func (a *atomicExecuteAndTriggerCapability) UnregisterTrigger(ctx context.Context, request capabilities.TriggerRegistrationRequest) error {
 	c := a.Load()
 	if c == nil {
 		return errors.New("capability unavailable")
 	}
-	return (*c).UnregisterTrigger(ctx, request)
+	return a.registrations.unregister(ctx, *c, request)
 }
 
 func (a *atomicExecuteAndTriggerCapability) RegisterToWorkflow(ctx context.Context, request capabilities.RegisterToWorkflowRequest) error {