Repository: nats-io/nats-server
Pull request #4013 (2023-04-02 21:30:47)

[IMPROVED] Protect against concurrent creation of streams and consumers.
Also make sure we have exited monitoring routines when doing resets for both streams and consumers.

Signed-off-by: Derek Collison <derek@nats.io>


Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 23067214..3f4f9103 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -299,6 +299,8 @@ type consumer struct {
 	prOk      bool
 	uch       chan struct{}
 	retention RetentionPolicy
+
+	monitorWg sync.WaitGroup
 	inMonitor bool
 
 	// R>1 proposals
@@ -4559,8 +4561,8 @@ func (o *consumer) clearMonitorRunning() {
 
 // Test whether we are in the monitor routine.
 func (o *consumer) isMonitorRunning() bool {
-	o.mu.Lock()
-	defer o.mu.Unlock()
+	o.mu.RLock()
+	defer o.mu.RUnlock()
 	return o.inMonitor
 }
 
diff --git a/server/jetstream.go b/server/jetstream.go
index e5452e99..89313447 100644
--- a/server/jetstream.go
+++ b/server/jetstream.go
@@ -137,6 +137,7 @@ type jsAccount struct {
 	js        *jetStream
 	account   *Account
 	storeDir  string
+	inflight  sync.Map
 	streams   map[string]*stream
 	templates map[string]*streamTemplate
 	store     TemplateStore
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index b54c2c10..9c794c80 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -436,8 +436,11 @@ func (cc *jetStreamCluster) isStreamCurrent(account, stream string) bool {
 
 // isStreamHealthy will determine if the stream is up to date or very close.
 // For R1 it will make sure the stream is present on this server.
-// Read lock should be held.
-func (cc *jetStreamCluster) isStreamHealthy(account, stream string) bool {
+func (js *jetStream) isStreamHealthy(account, stream string) bool {
+	js.mu.RLock()
+	defer js.mu.RUnlock()
+	cc := js.cluster
+
 	if cc == nil {
 		// Non-clustered mode
 		return true
@@ -477,8 +480,11 @@ func (cc *jetStreamCluster) isStreamHealthy(account, stream string) bool {
 
 // isConsumerCurrent will determine if the consumer is up to date.
 // For R1 it will make sure the consunmer is present on this server.
-// Read lock should be held.
-func (cc *jetStreamCluster) isConsumerCurrent(account, stream, consumer string) bool {
+func (js *jetStream) isConsumerCurrent(account, stream, consumer string) bool {
+	js.mu.RLock()
+	defer js.mu.RUnlock()
+	cc := js.cluster
+
 	if cc == nil {
 		// Non-clustered mode
 		return true
@@ -2486,41 +2492,42 @@ func (mset *stream) resetClusteredState(err error) bool {
 
 	// Preserve our current state and messages unless we have a first sequence mismatch.
 	shouldDelete := err == errFirstSequenceMismatch
+
+	mset.monitorWg.Wait()
+	mset.resetAndWaitOnConsumers()
+	// Stop our stream.
 	mset.stop(shouldDelete, false)
 
 	if sa != nil {
-		s.Warnf("Resetting stream cluster state for '%s > %s'", sa.Client.serviceAccount(), sa.Config.Name)
 		js.mu.Lock()
+		s.Warnf("Resetting stream cluster state for '%s > %s'", sa.Client.serviceAccount(), sa.Config.Name)
+		// Now wipe groups from assignments.
 		sa.Group.node = nil
+		var consumers []*consumerAssignment
+		if cc := js.cluster; cc != nil && cc.meta != nil {
+			ourID := cc.meta.ID()
+			for _, ca := range sa.consumers {
+				if rg := ca.Group; rg != nil && rg.isMember(ourID) {
+					rg.node = nil // Erase group raft/node state.
+					consumers = append(consumers, ca)
+				}
+			}
+		}
 		js.mu.Unlock()
-		go js.restartClustered(acc, sa)
-	}
-	return true
-}
 
-// This will reset the stream and consumers.
-// Should be done in separate go routine.
-func (js *jetStream) restartClustered(acc *Account, sa *streamAssignment) {
-	// Check and collect consumers first.
-	js.mu.RLock()
-	var consumers []*consumerAssignment
-	if cc := js.cluster; cc != nil && cc.meta != nil {
-		ourID := cc.meta.ID()
-		for _, ca := range sa.consumers {
-			if rg := ca.Group; rg != nil && rg.isMember(ourID) {
-				rg.node = nil // Erase group raft/node state.
-				consumers = append(consumers, ca)
+		// restart in a separate Go routine.
+		// This will reset the stream and consumers.
+		go func() {
+			// Reset stream.
+			js.processClusterCreateStream(acc, sa)
+			// Reset consumers.
+			for _, ca := range consumers {
+				js.processClusterCreateConsumer(ca, nil, false)
 			}
-		}
+		}()
 	}
-	js.mu.RUnlock()
 
-	// Reset stream.
-	js.processClusterCreateStream(acc, sa)
-	// Reset consumers.
-	for _, ca := range consumers {
-		js.processClusterCreateConsumer(ca, nil, false)
-	}
+	return true
 }
 
 func isControlHdr(hdr []byte) bool {
@@ -3974,6 +3981,7 @@ func (js *jetStream) processClusterCreateConsumer(ca *consumerAssignment, state
 			// Clustered consumer.
 			// Start our monitoring routine if needed.
 			if !alreadyRunning && !o.isMonitorRunning() {
+				o.monitorWg.Add(1)
 				s.startGoRoutine(func() { js.monitorConsumer(o, ca) })
 			}
 			// For existing consumer, only send response if not recovering.
@@ -4172,6 +4180,8 @@ func (js *jetStream) monitorConsumer(o *consumer, ca *consumerAssignment) {
 	s, n, cc := js.server(), o.raftNode(), js.cluster
 	defer s.grWG.Done()
 
+	defer o.monitorWg.Done()
+
 	if n == nil {
 		s.Warnf("No RAFT group for '%s > %s > %s'", o.acc.Name, ca.Stream, ca.Name)
 		return
diff --git a/server/monitor.go b/server/monitor.go
index 3b3488b0..624091ce 100644
--- a/server/monitor.go
+++ b/server/monitor.go
@@ -3129,7 +3129,7 @@ func (s *Server) healthz(opts *HealthzOptions) *HealthStatus {
 		for stream, sa := range asa {
 			if sa.Group.isMember(ourID) {
 				// Make sure we can look up
-				if !cc.isStreamHealthy(acc, stream) {
+				if !js.isStreamHealthy(acc, stream) {
 					health.Status = na
 					health.Error = fmt.Sprintf("JetStream stream '%s > %s' is not current", acc, stream)
 					return health
@@ -3137,7 +3137,7 @@ func (s *Server) healthz(opts *HealthzOptions) *HealthStatus {
 				// Now check consumers.
 				for consumer, ca := range sa.consumers {
 					if ca.Group.isMember(ourID) {
-						if !cc.isConsumerCurrent(acc, stream, consumer) {
+						if !js.isConsumerCurrent(acc, stream, consumer) {
 							health.Status = na
 							health.Error = fmt.Sprintf("JetStream consumer '%s > %s > %s' is not current", acc, stream, consumer)
 							return health
diff --git a/server/stream.go b/server/stream.go
index 89cbb807..f9db14d1 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -366,10 +366,23 @@ func (a *Account) addStreamWithAssignment(config *StreamConfig, fsConfig *FileSt
 		return nil, ApiErrors[JSStreamReplicasNotSupportedErr]
 	}
 
+	// Make sure we are ok when these are done in parallel.
+	v, loaded := jsa.inflight.LoadOrStore(cfg.Name, &sync.WaitGroup{})
+	wg := v.(*sync.WaitGroup)
+	if loaded {
+		wg.Wait()
+	} else {
+		wg.Add(1)
+		defer func() {
+			jsa.inflight.Delete(cfg.Name)
+			wg.Done()
+		}()
+	}
+
 	js, isClustered := jsa.jetStreamAndClustered()
-	jsa.mu.RLock()
+	jsa.mu.Lock()
 	if mset, ok := jsa.streams[cfg.Name]; ok {
-		jsa.mu.RUnlock()
+		jsa.mu.Unlock()
 		// Check to see if configs are same.
 		ocfg := mset.config()
 		if reflect.DeepEqual(ocfg, cfg) {
@@ -388,7 +401,8 @@ func (a *Account) addStreamWithAssignment(config *StreamConfig, fsConfig *FileSt
 	if !isClustered {
 		reserved = jsa.tieredReservation(tier, &cfg)
 	}
-	jsa.mu.RUnlock()
+	jsa.mu.Unlock()
+
 	if !hasTier {
 		return nil, NewJSNoLimitsError()
 	}
@@ -4427,6 +4441,26 @@ func (mset *stream) internalLoop() {
 	}
 }
 
+// Used to break consumers out of their
+func (mset *stream) resetAndWaitOnConsumers() {
+	mset.mu.RLock()
+	consumers := make([]*consumer, 0, len(mset.consumers))
+	for _, o := range mset.consumers {
+		consumers = append(consumers, o)
+	}
+	mset.mu.RUnlock()
+
+	for _, o := range consumers {
+		if node := o.raftNode(); node != nil {
+			if o.IsLeader() {
+				node.StepDown()
+			}
+			node.Delete()
+		}
+		o.monitorWg.Wait()
+	}
+}
+
 // Internal function to delete a stream.
 func (mset *stream) delete() error {
 	if mset == nil {
@@ -4445,7 +4479,7 @@ func (mset *stream) stop(deleteFlag, advisory bool) error {
 		return NewJSNotEnabledForAccountError()
 	}
 
-	// Remove from our account map.
+	// Remove from our account map first.
 	jsa.mu.Lock()
 	delete(jsa.streams, mset.cfg.Name)
 	accName := jsa.account.Name
@@ -4474,6 +4508,19 @@ func (mset *stream) stop(deleteFlag, advisory bool) error {
 			mset.cancelSourceConsumer(si.iname)
 		}
 	}
+
+	// Cluster cleanup
+	var sa *streamAssignment
+	if n := mset.node; n != nil {
+		if deleteFlag {
+			n.Delete()
+			sa = mset.sa
+		} else if n.NeedSnapshot() {
+			// Attempt snapshot on clean exit.
+			n.InstallSnapshot(mset.stateSnapshotLocked())
+			n.Stop()
+		}
+	}
 	mset.mu.Unlock()
 
 	for _, o := range obs {
@@ -4481,9 +4528,10 @@ func (mset *stream) stop(deleteFlag, advisory bool) error {
 		// TODO(dlc) - If we have an err here we don't want to stop
 		// but should we log?
 		o.stopWithFlags(deleteFlag, deleteFlag, false, advisory)
+		o.monitorWg.Wait()
 	}
-	mset.mu.Lock()
 
+	mset.mu.Lock()
 	// Stop responding to sync requests.
 	mset.stopClusterSubs()
 	// Unsubscribe from direct stream.
@@ -4495,19 +4543,6 @@ func (mset *stream) stop(deleteFlag, advisory bool) error {
 		mset.infoSub = nil
 	}
 
-	// Cluster cleanup
-	var sa *streamAssignment
-	if n := mset.node; n != nil {
-		if deleteFlag {
-			n.Delete()
-			sa = mset.sa
-		} else if n.NeedSnapshot() {
-			// Attempt snapshot on clean exit.
-			n.InstallSnapshot(mset.stateSnapshotLocked())
-			n.Stop()
-		}
-	}
-
 	// Send stream delete advisory after the consumers.
 	if deleteFlag && advisory {
 		mset.sendDeleteAdvisoryLocked()


Test patch:
diff --git a/server/norace_test.go b/server/norace_test.go
index 17b20f36..2e9381e4 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -3607,7 +3607,7 @@ func TestNoRaceJetStreamClusterCorruptWAL(t *testing.T) {
 	fs = o.raftNode().(*raft).wal.(*fileStore)
 	state = fs.State()
 	err = fs.Truncate(state.FirstSeq)
-	require_NoError(t, err)
+	require_True(t, err == nil || err == ErrInvalidSequence)
 	state = fs.State()
 
 	sub, err = js.PullSubscribe("foo", "dlc")
@@ -7687,3 +7687,124 @@ func TestNoRaceJetStreamClusterUnbalancedInterestMultipleFilteredConsumers(t *te
 		require_True(t, numPreAcks == 0)
 	}
 }
+
+func TestNoRaceParallelStreamAndConsumerCreation(t *testing.T) {
+	s := RunBasicJetStreamServer(t)
+	defer s.Shutdown()
+
+	// stream config.
+	scfg := &StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo", "bar"},
+		MaxMsgs:  10,
+		Storage:  FileStorage,
+		Replicas: 1,
+	}
+
+	// Will do these direct against the low level API to really make
+	// sure parallel creation ok.
+	np := 1000
+	startCh := make(chan bool)
+	errCh := make(chan error, np)
+	wg := sync.WaitGroup{}
+	wg.Add(np)
+
+	var streams sync.Map
+
+	for i := 0; i < np; i++ {
+		go func() {
+			defer wg.Done()
+
+			// Make them all fire at once.
+			<-startCh
+
+			if mset, err := s.GlobalAccount().addStream(scfg); err != nil {
+				t.Logf("Stream create got an error: %v", err)
+				errCh <- err
+			} else {
+				streams.Store(mset, true)
+			}
+		}()
+	}
+	time.Sleep(100 * time.Millisecond)
+	close(startCh)
+	wg.Wait()
+
+	// Check for no errors.
+	if len(errCh) > 0 {
+		t.Fatalf("Expected no errors, got %d", len(errCh))
+	}
+
+	// Now make sure we really only created one stream.
+	var numStreams int
+	streams.Range(func(k, v any) bool {
+		numStreams++
+		return true
+	})
+	if numStreams > 1 {
+		t.Fatalf("Expected only one stream to be really created, got %d out of %d attempts", numStreams, np)
+	}
+
+	// Also make sure we cleanup the inflight entries for streams.
+	gacc := s.GlobalAccount()
+	_, jsa, err := gacc.checkForJetStream()
+	require_NoError(t, err)
+	var numEntries int
+	jsa.inflight.Range(func(k, v any) bool {
+		numEntries++
+		return true
+	})
+	if numEntries > 0 {
+		t.Fatalf("Expected no inflight entries to be left over, got %d", numEntries)
+	}
+
+	// Now do consumers.
+	mset, err := gacc.lookupStream("TEST")
+	require_NoError(t, err)
+
+	cfg := &ConsumerConfig{
+		DeliverSubject: "to",
+		Name:           "DLC",
+		AckPolicy:      AckExplicit,
+	}
+
+	startCh = make(chan bool)
+	errCh = make(chan error, np)
+	wg.Add(np)
+
+	var consumers sync.Map
+
+	for i := 0; i < np; i++ {
+		go func() {
+			defer wg.Done()
+
+			// Make them all fire at once.
+			<-startCh
+
+			if _, err = mset.addConsumer(cfg); err != nil {
+				t.Logf("Consumer create got an error: %v", err)
+				errCh <- err
+			} else {
+				consumers.Store(mset, true)
+			}
+		}()
+	}
+	time.Sleep(100 * time.Millisecond)
+	close(startCh)
+	wg.Wait()
+
+	// Check for no errors.
+	if len(errCh) > 0 {
+		t.Fatalf("Expected no errors, got %d", len(errCh))
+	}
+
+	// Now make sure we really only created one stream.
+	var numConsumers int
+	consumers.Range(func(k, v any) bool {
+		numConsumers++
+		return true
+	})
+	if numConsumers > 1 {
+		t.Fatalf("Expected only one consumer to be really created, got %d out of %d attempts", numConsumers, np)
+	}
+}