Repository: nats-io/nats-server
Pull request #5400 (2024-05-08 23:57:01)

Cherry picks for v2.10.15 RC.9
Includes:

- https://github.com/nats-io/nats-server/pull/5398
- https://github.com/nats-io/nats-server/pull/5399



Patch:
diff --git a/server/accounts.go b/server/accounts.go
index d1b9a7f8..2f80b606 100644
--- a/server/accounts.go
+++ b/server/accounts.go
@@ -82,7 +82,7 @@ type Account struct {
 	js           *jsAccount
 	jsLimits     map[string]JetStreamAccountLimits
 	limits
-	expired      bool
+	expired      atomic.Bool
 	incomplete   bool
 	signingKeys  map[string]jwt.Scope
 	extAuth      *jwt.ExternalAuthorization
@@ -759,7 +759,7 @@ func (a *Account) selectMappedSubject(dest string) (string, bool) {
 		return dest, false
 	}
 
-	a.mu.Lock()
+	a.mu.RLock()
 	// In case we have to tokenize for subset matching.
 	tsa := [32]string{}
 	tts := tsa[:0]
@@ -790,7 +790,7 @@ func (a *Account) selectMappedSubject(dest string) (string, bool) {
 	}
 
 	if m == nil {
-		a.mu.Unlock()
+		a.mu.RUnlock()
 		return dest, false
 	}
 
@@ -829,7 +829,7 @@ func (a *Account) selectMappedSubject(dest string) (string, bool) {
 		}
 	}
 
-	a.mu.Unlock()
+	a.mu.RUnlock()
 	return ndest, true
 }
 
@@ -2189,7 +2189,7 @@ const (
 // This is where all service export responses are handled.
 func (a *Account) processServiceImportResponse(sub *subscription, c *client, _ *Account, subject, reply string, msg []byte) {
 	a.mu.RLock()
-	if a.expired || len(a.exports.responses) == 0 {
+	if a.expired.Load() || len(a.exports.responses) == 0 {
 		a.mu.RUnlock()
 		return
 	}
@@ -2664,7 +2664,7 @@ func (a *Account) getWildcardServiceExport(from string) *serviceExport {
 // These are import stream specific versions for when an activation expires.
 func (a *Account) streamActivationExpired(exportAcc *Account, subject string) {
 	a.mu.RLock()
-	if a.expired || a.imports.streams == nil {
+	if a.expired.Load() || a.imports.streams == nil {
 		a.mu.RUnlock()
 		return
 	}
@@ -2699,7 +2699,7 @@ func (a *Account) streamActivationExpired(exportAcc *Account, subject string) {
 // These are import service specific versions for when an activation expires.
 func (a *Account) serviceActivationExpired(subject string) {
 	a.mu.RLock()
-	if a.expired || a.imports.services == nil {
+	if a.expired.Load() || a.imports.services == nil {
 		a.mu.RUnlock()
 		return
 	}
@@ -2960,18 +2960,13 @@ func (a *Account) checkServiceImportAuthorizedNoLock(account *Account, subject s
 
 // IsExpired returns expiration status.
 func (a *Account) IsExpired() bool {
-	a.mu.RLock()
-	exp := a.expired
-	a.mu.RUnlock()
-	return exp
+	return a.expired.Load()
 }
 
 // Called when an account has expired.
 func (a *Account) expiredTimeout() {
 	// Mark expired first.
-	a.mu.Lock()
-	a.expired = true
-	a.mu.Unlock()
+	a.expired.Store(true)
 
 	// Collect the clients and expire them.
 	cs := a.getClients()
@@ -3016,17 +3011,17 @@ func (a *Account) checkExpiration(claims *jwt.ClaimsData) {
 
 	a.clearExpirationTimer()
 	if claims.Expires == 0 {
-		a.expired = false
+		a.expired.Store(false)
 		return
 	}
 	tn := time.Now().Unix()
 	if claims.Expires <= tn {
-		a.expired = true
+		a.expired.Store(true)
 		return
 	}
 	expiresAt := time.Duration(claims.Expires - tn)
 	a.setExpirationTimer(expiresAt * time.Second)
-	a.expired = false
+	a.expired.Store(false)
 }
 
 // hasIssuer returns true if the issuer matches the account
diff --git a/server/client.go b/server/client.go
index 67b85207..66c8d1e7 100644
--- a/server/client.go
+++ b/server/client.go
@@ -474,8 +474,8 @@ func (rcf readCacheFlag) isSet(c readCacheFlag) bool {
 }
 
 const (
-	defaultMaxPerAccountCacheSize   = 4096
-	defaultPrunePerAccountCacheSize = 256
+	defaultMaxPerAccountCacheSize   = 8192
+	defaultPrunePerAccountCacheSize = 1024
 	defaultClosedSubsCheckInterval  = 5 * time.Minute
 )
 
@@ -5453,13 +5453,13 @@ func (c *client) getAccAndResultFromCache() (*Account, *SublistResult) {
 		// Match against the account sublist.
 		r = sl.Match(string(c.pa.subject))
 
-		// Store in our cache
-		c.in.pacache[string(c.pa.pacache)] = &perAccountCache{acc, r, atomic.LoadUint64(&sl.genid)}
-
 		// Check if we need to prune.
-		if len(c.in.pacache) > maxPerAccountCacheSize {
+		if len(c.in.pacache) >= maxPerAccountCacheSize {
 			c.prunePerAccountCache()
 		}
+
+		// Store in our cache,make sure to do so after we prune.
+		c.in.pacache[string(c.pa.pacache)] = &perAccountCache{acc, r, atomic.LoadUint64(&sl.genid)}
 	}
 	return acc, r
 }
diff --git a/server/monitor.go b/server/monitor.go
index 0467d821..11cd8645 100644
--- a/server/monitor.go
+++ b/server/monitor.go
@@ -1,4 +1,4 @@
-// Copyright 2013-2023 The NATS Authors
+// Copyright 2013-2024 The NATS Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -2697,7 +2697,7 @@ func (s *Server) accountInfo(accName string) (*AccountInfo, error) {
 		accName,
 		a.updated.UTC(),
 		isSys,
-		a.expired,
+		a.expired.Load(),
 		!a.incomplete,
 		a.js != nil,
 		a.numLocalLeafNodes(),
diff --git a/server/route.go b/server/route.go
index 1b2e3f75..0341f798 100644
--- a/server/route.go
+++ b/server/route.go
@@ -1398,7 +1398,7 @@ func (c *client) processRemoteSub(argo []byte, hasOrigin bool) (err error) {
 		var isNew bool
 		if acc, isNew = srv.LookupOrRegisterAccount(accountName); isNew {
 			acc.mu.Lock()
-			acc.expired = true
+			acc.expired.Store(true)
 			acc.incomplete = true
 			acc.mu.Unlock()
 		}
diff --git a/server/server.go b/server/server.go
index be5cfc31..89e9c7de 100644
--- a/server/server.go
+++ b/server/server.go
@@ -2478,12 +2478,7 @@ func (s *Server) Shutdown() {
 	}
 
 	// Kick websocket server
-	if s.websocket.server != nil {
-		doneExpected++
-		s.websocket.server.Close()
-		s.websocket.server = nil
-		s.websocket.listener = nil
-	}
+	doneExpected += s.closeWebsocketServer()
 
 	// Kick MQTT accept loop
 	if s.mqtt.listener != nil {
@@ -2571,6 +2566,24 @@ func (s *Server) Shutdown() {
 	close(s.shutdownComplete)
 }
 
+// Close the websocket server if running. If so, returns 1, else 0.
+// Server lock held on entry.
+func (s *Server) closeWebsocketServer() int {
+	ws := &s.websocket
+	ws.mu.Lock()
+	hs := ws.server
+	if hs != nil {
+		ws.server = nil
+		ws.listener = nil
+	}
+	ws.mu.Unlock()
+	if hs != nil {
+		hs.Close()
+		return 1
+	}
+	return 0
+}
+
 // WaitForShutdown will block until the server has been fully shutdown.
 func (s *Server) WaitForShutdown() {
 	<-s.shutdownComplete
@@ -4090,12 +4103,7 @@ func (s *Server) lameDuckMode() {
 	expected := 1
 	s.listener.Close()
 	s.listener = nil
-	if s.websocket.server != nil {
-		expected++
-		s.websocket.server.Close()
-		s.websocket.server = nil
-		s.websocket.listener = nil
-	}
+	expected += s.closeWebsocketServer()
 	s.ldmCh = make(chan bool, expected)
 	opts := s.getOpts()
 	gp := opts.LameDuckGracePeriod
@@ -4396,8 +4404,10 @@ func (s *Server) changeRateLimitLogInterval(d time.Duration) {
 
 // DisconnectClientByID disconnects a client by connection ID
 func (s *Server) DisconnectClientByID(id uint64) error {
-	client := s.clients[id]
-	if client != nil {
+	if s == nil {
+		return ErrServerNotRunning
+	}
+	if client := s.getClient(id); client != nil {
 		client.closeConnection(Kicked)
 		return nil
 	}
@@ -4406,23 +4416,27 @@ func (s *Server) DisconnectClientByID(id uint64) error {
 
 // LDMClientByID sends a Lame Duck Mode info message to a client by connection ID
 func (s *Server) LDMClientByID(id uint64) error {
+	if s == nil {
+		return ErrServerNotRunning
+	}
+	s.mu.RLock()
+	c := s.clients[id]
+	if c == nil {
+		s.mu.RUnlock()
+		return errors.New("no such client id")
+	}
 	info := s.copyInfo()
 	info.LameDuckMode = true
-
-	c := s.clients[id]
-	if c != nil {
-		c.mu.Lock()
-		defer c.mu.Unlock()
-		if c.opts.Protocol >= ClientProtoInfo &&
-			c.flags.isSet(firstPongSent) {
-			// sendInfo takes care of checking if the connection is still
-			// valid or not, so don't duplicate tests here.
-			c.Debugf("sending Lame Duck Mode info to client")
-			c.enqueueProto(c.generateClientInfoJSON(info))
-			return nil
-		} else {
-			return errors.New("ClientProtoInfo < ClientOps.Protocol or first pong not sent")
-		}
+	s.mu.RUnlock()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.opts.Protocol >= ClientProtoInfo && c.flags.isSet(firstPongSent) {
+		// sendInfo takes care of checking if the connection is still
+		// valid or not, so don't duplicate tests here.
+		c.Debugf("Sending Lame Duck Mode info to client")
+		c.enqueueProto(c.generateClientInfoJSON(info))
+		return nil
+	} else {
+		return errors.New("client does not support Lame Duck Mode or is not ready to receive the notification")
 	}
-	return errors.New("no such client id")
 }


Test patch:
diff --git a/server/events_test.go b/server/events_test.go
index 7e6e7b9e..296a49a8 100644
--- a/server/events_test.go
+++ b/server/events_test.go
@@ -3459,10 +3459,28 @@ func TestServerEventsLDMKick(t *testing.T) {
 	reqkick := KickClientReq{CID: cid}
 	reqkickpayload, _ := json.Marshal(reqkick)
 
-	_, err = ncs.Request(fmt.Sprintf("$SYS.REQ.SERVER.%s.LDM", s.ID()), reqldmpayload, time.Second)
-	if err != nil {
-		t.Fatalf("Error trying to publish the LDM request: %v", err)
+	// Test for data races when getting the client by ID
+	uc := createUserCreds(t, s, akp2)
+	totalClients := 100
+	someClients := make([]*nats.Conn, 0, totalClients)
+	for i := 0; i < totalClients; i++ {
+		nc, err := nats.Connect(s.ClientURL(), uc)
+		require_NoError(t, err)
+		defer nc.Close()
+		someClients = append(someClients, nc)
 	}
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < totalClients; i++ {
+			someClients[i].Close()
+		}
+	}()
+	defer wg.Wait()
+
+	_, err = ncs.Request(fmt.Sprintf("$SYS.REQ.SERVER.%s.LDM", s.ID()), reqldmpayload, time.Second)
+	require_NoError(t, err)
 
 	select {
 	case <-ldmed:
@@ -3471,9 +3489,7 @@ func TestServerEventsLDMKick(t *testing.T) {
 	}
 
 	_, err = ncs.Request(fmt.Sprintf("$SYS.REQ.SERVER.%s.KICK", s.ID()), reqkickpayload, time.Second)
-	if err != nil {
-		t.Fatalf("Error trying to publish the KICK request: %v", err)
-	}
+	require_NoError(t, err)
 
 	select {
 	case <-disconnected:
diff --git a/server/jwt_test.go b/server/jwt_test.go
index 1dc228cc..ab7de041 100644
--- a/server/jwt_test.go
+++ b/server/jwt_test.go
@@ -2964,7 +2964,7 @@ func TestJWTAccountLimitsMaxConnsAfterExpired(t *testing.T) {
 	// is now lower, some clients should have been removed.
 	acc, _ := s.LookupAccount(fooPub)
 	acc.mu.Lock()
-	acc.expired = true
+	acc.expired.Store(true)
 	acc.updated = time.Now().UTC().Add(-2 * time.Second) // work around updating to quickly
 	acc.mu.Unlock()