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,180 @@ 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
+}
+
+func newTriggerRegistrationManager() *triggerRegistrationManager {
+	return &triggerRegistrationManager{
+		regs: make(map[string]*triggerRegistration),
+	}
+}
+
+func (m *triggerRegistrationManager) register(ctx context.Context, exec capabilities.TriggerExecutable, req capabilities.TriggerRegistrationRequest) (<-chan capabilities.TriggerResponse, error) {
+	if exec == nil {
+		return nil, errors.New("capability unavailable")
+	}
+	in, err := exec.RegisterTrigger(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	m.mu.Lock()
+	reg, ok := m.regs[req.TriggerID]
+	if !ok {
+		reg = &triggerRegistration{
+			request: req,
+			outCh:   make(chan capabilities.TriggerResponse),
+		}
+		m.regs[req.TriggerID] = reg
+	} else {
+		reg.request = req
+		if reg.cancel != nil {
+			reg.cancel()
+		}
+	}
+	ctxForward, cancel := context.WithCancel(context.Background())
+	reg.cancel = cancel
+	out := reg.outCh
+	m.mu.Unlock()
+
+	go forwardTriggerResponses(ctxForward, in, out)
+
+	return out, nil
+}
+
+func (m *triggerRegistrationManager) unregister(ctx context.Context, exec capabilities.TriggerExecutable, req capabilities.TriggerRegistrationRequest) error {
+	if exec == nil {
+		return errors.New("capability unavailable")
+	}
+
+	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) rebind(oldExec, newExec capabilities.TriggerExecutable) error {
+	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 oldExec != nil {
+			_ = oldExec.UnregisterTrigger(context.Background(), reg.request)
+		}
+		if newExec == nil {
+			continue
+		}
+
+		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.mu.Lock()
+		regInMap, ok := m.regs[reg.request.TriggerID]
+		if !ok {
+			regInMap = &triggerRegistration{
+				request: reg.request,
+				outCh:   reg.outCh,
+			}
+			m.regs[reg.request.TriggerID] = regInMap
+		} else {
+			regInMap.request = reg.request
+			regInMap.outCh = reg.outCh
+			if regInMap.cancel != nil {
+				regInMap.cancel()
+			}
+		}
+		ctxForward, cancel := context.WithCancel(context.Background())
+		regInMap.cancel = cancel
+		out := regInMap.outCh
+		m.mu.Unlock()
+
+		go forwardTriggerResponses(ctxForward, in, out)
+	}
+	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 {
+	var prev capabilities.TriggerExecutable
+	if existing := a.Load(); existing != nil {
+		prev = *existing
+	}
 	if c == nil {
 		a.Store(nil)
-		return nil
+		return a.registrations.rebind(prev, 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(prev, tc)
 }
 
 func (a *atomicTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
@@ -249,15 +409,15 @@ func (a *atomicTriggerCapability) RegisterTrigger(ctx context.Context, request c
 	if c == nil {
 		return nil, errors.New("capability unavailable")
 	}
-	return (*c).RegisterTrigger(ctx, request)
+	return a.registrations.register(ctx, *c, 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 +486,30 @@ 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 {
+	var prev capabilities.TriggerExecutable
+	if existing := a.Load(); existing != nil {
+		prev = *existing
+	}
 	if c == nil {
 		a.Store(nil)
-		return nil
+		return a.registrations.rebind(prev, 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(prev, tc)
 }
 
 func (a *atomicExecuteAndTriggerCapability) Info(ctx context.Context) (capabilities.CapabilityInfo, error) {
@@ -365,15 +536,15 @@ func (a *atomicExecuteAndTriggerCapability) RegisterTrigger(ctx context.Context,
 	if c == nil {
 		return nil, errors.New("capability unavailable")
 	}
-	return (*c).RegisterTrigger(ctx, request)
+	return a.registrations.register(ctx, *c, 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 {