feat(targets): enhance subscription management with improved stopping and waiting mechanisms

pkg/collector/managers/targets/targets_manager.go

@@ -30,6 +30,10 @@ import (
 	"google.golang.org/grpc"
 )
 
+const (
+	subscriptionReaderStopTimeout = 5 * time.Second
+)
+
 type ManagedTarget struct {
 	sync.RWMutex
 	Name string
@@ -43,6 +47,7 @@ type ManagedTarget struct {
 
 	mu                   *sync.Mutex
 	readersCfn           map[string]context.CancelFunc
+	readersDone          map[string]chan struct{}
 	readerWG             sync.WaitGroup
 	lastError            string // last error message, protected by mu
 	outputs              map[string]struct{}
@@ -77,6 +82,7 @@ func newManagedTarget(name string, cfg *types.TargetConfig, tunServer *tunnel.Se
 		outputs:              make(map[string]struct{}, len(cfg.Outputs)),
 		mu:                   new(sync.Mutex),
 		readersCfn:           make(map[string]context.CancelFunc),
+		readersDone:          make(map[string]chan struct{}),
 		appliedSubscriptions: make([]string, 0, len(cfg.Subscriptions)),
 	}
 	for _, output := range cfg.Outputs {
@@ -375,15 +381,8 @@ func (tm *TargetsManager) apply(name string, cfg *types.TargetConfig) {
 					}
 				}
 				for _, sub := range removed {
-					mt.mu.Lock()
-					cfn, exists := mt.readersCfn[sub]
-					if exists {
-						cfn()
-						delete(mt.readersCfn, sub)
-					}
-					mt.mu.Unlock()
 					tm.logger.Info("stopping target subscription", "name", sub, "target", name)
-					mt.T.StopSubscription(sub)
+					tm.stopTargetSubscription(mt, sub)
 					mt.T.DeleteSubscriptionConfig(sub)
 					mt.appliedSubscriptions = slices.DeleteFunc(mt.appliedSubscriptions, func(s string) bool {
 						return s == sub
@@ -635,54 +634,85 @@ func (tm *TargetsManager) applySubscription(name string, cfg types.SubscriptionC
 	tm.mu.Lock()
 	tm.subscriptions[name] = &cfg
 	tm.logger.Info("subscriptions", "subscriptions", tm.subscriptions)
+	targets := make([]*ManagedTarget, 0, len(tm.targets))
 	for _, mt := range tm.targets {
 		tm.logger.Info("target", "target", mt.Name, "subscriptions", mt.T.Config.Subscriptions)
-		if len(mt.T.Config.Subscriptions) > 0 {
-			if !slices.Contains(mt.T.Config.Subscriptions, name) {
-				tm.logger.Info("subscription not in target's explicit list", "subscription", name, "target", mt.Name)
-				continue
-			}
+		if len(mt.T.Config.Subscriptions) > 0 && !slices.Contains(mt.T.Config.Subscriptions, name) {
+			tm.logger.Info("subscription not in target's explicit list", "subscription", name, "target", mt.Name)
+			continue
 		}
-		tm.logger.Info("(re)starting target subscription", "name", name, "target", mt.Name)
-		// Stop and WAIT for the old subscription to fully terminate
-		mt.mu.Lock()
-		cfn, exists := mt.readersCfn[name]
-		if exists {
-			tm.logger.Info("canceling subscription context", "name", name, "target", mt.Name)
-			cfn() // Cancel the context
-			tm.logger.Info("deleted subscription context", "name", name, "target", mt.Name)
-			delete(mt.readersCfn, name) // Remove from map
+		targets = append(targets, mt)
+	}
+	tm.mu.Unlock()
+
+	for _, mt := range targets {
+		// Hold a read lock so remove() cannot nil mt.T until this restart
+		// finishes. The old reader may call setTargetState (also RLock) while
+		// we wait for it; a write lock here would deadlock.
+		tm.mu.RLock()
+		if tm.targets[mt.Name] != mt || mt.T == nil {
+			tm.mu.RUnlock()
+			continue
 		}
-		mt.mu.Unlock()
-		tm.logger.Info("stopping target subscription", "name", name, "target", mt.Name)
-		mt.T.StopSubscription(name)
-		tm.logger.Info("stopped target subscription", "name", name, "target", mt.Name)
-		// Wait for the reader goroutine to finish
+		tm.logger.Info("(re)starting target subscription", "name", name, "target", mt.Name)
+		tm.stopTargetSubscription(mt, name)
 		mt.T.SetSubscriptionConfig(&cfg)
 		err := tm.startTargetSubscription(mt, &cfg)
+		tm.mu.RUnlock()
 		if err != nil {
 			tm.logger.Error("failed to start target subscription", "subscription", name, "target", mt.Name, "error", err)
 		}
 	}
-	tm.mu.Unlock()
+}
+
+// stopTargetSubscription cancels the per-subscription reader and waits for it
+// to exit so a replacement Subscribe RPC cannot start while the old
+// SubscribeChan goroutine is still unwinding.
+func (tm *TargetsManager) stopTargetSubscription(mt *ManagedTarget, name string) {
+	mt.mu.Lock()
+	cfn := mt.readersCfn[name]
+	done := mt.readersDone[name]
+	delete(mt.readersCfn, name)
+	delete(mt.readersDone, name)
+	mt.mu.Unlock()
+
+	if cfn != nil {
+		tm.logger.Info("canceling subscription context", "name", name, "target", mt.Name)
+		cfn()
+	}
+	if mt.T != nil {
+		mt.T.StopSubscription(name)
+	}
+	if done == nil {
+		return
+	}
+	timer := time.NewTimer(subscriptionReaderStopTimeout)
+	defer timer.Stop()
+	select {
+	case <-done:
+		tm.logger.Info("stopped target subscription", "name", name, "target", mt.Name)
+	case <-timer.C:
+		tm.logger.Warn("timed out waiting for subscription reader to exit", "subscription", name, "target", mt.Name)
+	}
 }
 
 // remove subscription from targets that already reference it and have it running
 func (tm *TargetsManager) removeSubscription(name string) {
 	tm.mu.Lock()
 	delete(tm.subscriptions, name)
+	targets := make([]*ManagedTarget, 0, len(tm.targets))
 	for _, mt := range tm.targets {
-		mt.mu.Lock()
-		cfn, exists := mt.readersCfn[name]
-		if exists {
-			cfn()
-			delete(mt.readersCfn, name)
-		}
-		mt.mu.Unlock()
-		mt.T.StopSubscription(name)
-		mt.T.DeleteSubscriptionConfig(name)
+		targets = append(targets, mt)
 	}
 	tm.mu.Unlock()
+	for _, mt := range targets {
+		tm.mu.RLock()
+		tm.stopTargetSubscription(mt, name)
+		if mt.T != nil {
+			mt.T.DeleteSubscriptionConfig(name)
+		}
+		tm.mu.RUnlock()
+	}
 }
 
 func (tm *TargetsManager) reconcileAssignment(name string) {
@@ -821,8 +851,10 @@ func (tm *TargetsManager) startTargetSubscription(mt *ManagedTarget, cfg *types.
 	mt.T.SetSubscriptionConfig(cfg)
 	mt.readerWG.Add(1)
 	sctx, cfn := context.WithCancel(tm.ctx)
+	done := make(chan struct{})
 	mt.mu.Lock()
 	mt.readersCfn[cfg.Name] = cfn
+	mt.readersDone[cfg.Name] = done
 	mt.mu.Unlock()
 
 	subscriptionOutputs := make(map[string]struct{}, len(cfg.Outputs))
@@ -832,6 +864,7 @@ func (tm *TargetsManager) startTargetSubscription(mt *ManagedTarget, cfg *types.
 	respCh, errCh := mt.T.SubscribeChan(sctx, subreq, cfg.Name)
 	go func() {
 		defer mt.readerWG.Done()
+		defer close(done)
 		// When the goroutine exits (subscription stopped/cancelled), refresh
 		// the target state so the subscriptions map is up-to-date.
 		defer func() {
@@ -841,13 +874,18 @@ func (tm *TargetsManager) startTargetSubscription(mt *ManagedTarget, cfg *types.
 			}
 		}()
 		initialResponse := true
-		for {
+		// Stay in the loop until SubscribeChan closes both channels. Returning
+		// on sctx.Done() alone would let a replacement Subscribe RPC start
+		// while attemptSubscription's defer StopSubscription is still running.
+		for respCh != nil || errCh != nil {
 			select {
-			case <-sctx.Done():
-				return
 			case resp, ok := <-respCh:
 				if !ok {
-					return
+					respCh = nil
+					continue
+				}
+				if sctx.Err() != nil {
+					continue
 				}
 				// The first response confirms the subscription is connected.
 				// Refresh target state so the subscriptions map shows "running".
@@ -887,7 +925,11 @@ func (tm *TargetsManager) startTargetSubscription(mt *ManagedTarget, cfg *types.
 				}
 			case err, ok := <-errCh:
 				if !ok {
-					return
+					errCh = nil
+					continue
+				}
+				if sctx.Err() != nil {
+					continue
 				}
 				// Reset so the next successful response after retry
 				// triggers a state update back to "running".

pkg/collector/managers/targets/targets_manager_test.go

@@ -1,12 +1,14 @@
 package targets_manager
 
 import (
+	"context"
 	"log/slog"
 	"testing"
+	"time"
 
 	"github.com/openconfig/gnmic/pkg/api/types"
-	"github.com/openconfig/gnmic/pkg/config"
 	collstore "github.com/openconfig/gnmic/pkg/collector/store"
+	"github.com/openconfig/gnmic/pkg/config"
 	"github.com/openconfig/gnmic/pkg/pipeline"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/zestor-dev/zestor/store"
@@ -202,3 +204,76 @@ func TestKeys_helper(t *testing.T) {
 		t.Fatalf("keys len = %d", len(got))
 	}
 }
+
+func TestStopTargetSubscription_waitsForReader(t *testing.T) {
+	tm := newTargetsTestManager(t)
+	mt := newManagedTarget("t1", &types.TargetConfig{Name: "t1", Address: "10.0.0.1:57400"}, nil)
+
+	sctx, cfn := context.WithCancel(t.Context())
+	done := make(chan struct{})
+	mt.mu.Lock()
+	mt.readersCfn["sub1"] = cfn
+	mt.readersDone["sub1"] = done
+	mt.mu.Unlock()
+
+	go func() {
+		<-sctx.Done()
+		// The reader also refreshes target state on exit; this must not
+		// deadlock with stopTargetSubscription.
+		tm.setTargetState("t1", collstore.StateRunning)
+		time.Sleep(80 * time.Millisecond)
+		close(done)
+	}()
+
+	started := time.Now()
+	tm.stopTargetSubscription(mt, "sub1")
+	if elapsed := time.Since(started); elapsed < 80*time.Millisecond {
+		t.Fatalf("stopTargetSubscription returned after %s, want to wait for reader", elapsed)
+	}
+	select {
+	case <-done:
+	default:
+		t.Fatal("reader still running after stopTargetSubscription")
+	}
+}
+
+func TestStopTargetSubscription_noReader(t *testing.T) {
+	tm := newTargetsTestManager(t)
+	mt := newManagedTarget("t1", &types.TargetConfig{Name: "t1", Address: "10.0.0.1:57400"}, nil)
+	tm.stopTargetSubscription(mt, "missing")
+}
+
+func TestApplySubscription_doesNotHoldWriteLockWhileWaiting(t *testing.T) {
+	tm := newTargetsTestManager(t)
+	mt := newManagedTarget("t1", &types.TargetConfig{Name: "t1", Address: "10.0.0.1:57400", Subscriptions: []string{"sub1"}}, nil)
+	tm.mu.Lock()
+	tm.targets["t1"] = mt
+	tm.mu.Unlock()
+
+	sctx, cfn := context.WithCancel(t.Context())
+	done := make(chan struct{})
+	mt.mu.Lock()
+	mt.readersCfn["sub1"] = cfn
+	mt.readersDone["sub1"] = done
+	mt.mu.Unlock()
+
+	go func() {
+		<-sctx.Done()
+		tm.setTargetState("t1", collstore.StateRunning)
+		close(done)
+	}()
+
+	finished := make(chan struct{})
+	go func() {
+		defer close(finished)
+		// CreateSubscribeRequest fails (no paths), so we never reach SubscribeChan.
+		// The important part is that waiting for the old reader does not deadlock
+		// with setTargetState (which takes tm.mu.RLock).
+		tm.applySubscription("sub1", types.SubscriptionConfig{Name: "sub1"})
+	}()
+	select {
+	case <-finished:
+	case <-time.After(2 * time.Second):
+		t.Fatal("applySubscription deadlocked waiting for the old reader")
+	}
+}