Repository: nats-io/nats-server
Pull request #5955 (2024-10-02 22:10:05)

Refactor delayed API responses
This is similar to @neilalexander PR (#5948) but has a slightly different approach. We use the struct to have a `next` so that we can use a list since the list will be maintained ordered. The for-loop is in one level, as opposed to two, which means that pushing to the IPQ will be "immediately" popped, even if we have a timer initialized and waiting to fire.
The timer is reset (not recreated everytime) only if the added response ends-up being at the head of the list.

Added a test that verifies that it works properly.

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>


Patch:
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 93b26fc5..bab99b72 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -1012,30 +1012,128 @@ func (s *Server) sendAPIErrResponse(ci *ClientInfo, acc *Account, subject, reply
 
 const errRespDelay = 500 * time.Millisecond
 
-func (s *Server) sendDelayedAPIErrResponse(ci *ClientInfo, acc *Account, subject, reply, request, response string, rg *raftGroup) {
-	js := s.getJetStream()
-	if js == nil {
+type delayedAPIResponse struct {
+	ci       *ClientInfo
+	acc      *Account
+	subject  string
+	reply    string
+	request  string
+	response string
+	rg       *raftGroup
+	deadline time.Time
+	next     *delayedAPIResponse
+}
+
+// Add `r` in the list that is maintained ordered by the `delayedAPIResponse.deadline` time.
+func addDelayedResponse(head, tail **delayedAPIResponse, r *delayedAPIResponse) {
+	// Check if list empty.
+	if *head == nil {
+		*head, *tail = r, r
 		return
 	}
-	var quitCh <-chan struct{}
-	js.mu.RLock()
-	if rg != nil && rg.node != nil {
-		quitCh = rg.node.QuitC()
+	// Check if it should be added at the end, which is if after or equal to the tail.
+	if r.deadline.After((*tail).deadline) || r.deadline.Equal((*tail).deadline) {
+		(*tail).next, *tail = r, r
+		return
 	}
-	js.mu.RUnlock()
+	// Find its spot in the list.
+	var prev *delayedAPIResponse
+	for c := *head; c != nil; c = c.next {
+		// We insert only if we are stricly before the current `c`.
+		if r.deadline.Before(c.deadline) {
+			r.next = c
+			if prev != nil {
+				prev.next = r
+			} else {
+				*head = r
+			}
+			return
+		}
+		prev = c
+	}
+}
 
-	s.startGoRoutine(func() {
-		defer s.grWG.Done()
+func (s *Server) delayedAPIResponder() {
+	defer s.grWG.Done()
+	var (
+		head, tail *delayedAPIResponse // Linked list.
+		r          *delayedAPIResponse // Updated by calling next().
+		rqch       <-chan struct{}     // Quit channel of the Raft group (if present).
+		tm         = time.NewTimer(time.Hour)
+	)
+	next := func() {
+		r, rqch = nil, nil
+		// Check that JetStream is still on. Do not exit the go routine
+		// since JS can be enabled/disabled. The go routine will exit
+		// only if server is shutdown.
+		js := s.getJetStream()
+		if js == nil {
+			// Reset head and tail here. Also drain the ipQueue.
+			head, tail = nil, nil
+			s.delayedAPIResponses.drain()
+			// Fall back into next "if" that resets timer.
+		}
+		// If there are no delayed messages then delay the timer for
+		// a while.
+		if head == nil {
+			tm.Reset(time.Hour)
+			return
+		}
+		// Get the first expected message and then reset the timer.
+		r = head
+		js.mu.RLock()
+		if r.rg != nil && r.rg.node != nil {
+			// If there's an attached Raft group to the delayed response
+			// then pull out the quit channel, so that we don't bother
+			// sending responses for entities which are now no longer
+			// running.
+			rqch = r.rg.node.QuitC()
+		}
+		js.mu.RUnlock()
+		tm.Reset(time.Until(r.deadline))
+	}
+	pop := func() {
+		if head == nil {
+			return
+		}
+		head = head.next
+		if head == nil {
+			tail = nil
+		}
+	}
+	for {
 		select {
-		case <-quitCh:
+		case <-s.delayedAPIResponses.ch:
+			v, ok := s.delayedAPIResponses.popOne()
+			if !ok {
+				continue
+			}
+			// Add it to the list, and if ends up being the head, set things up.
+			addDelayedResponse(&head, &tail, v)
+			if v == head {
+				next()
+			}
 		case <-s.quitCh:
-		case <-time.After(errRespDelay):
-			acc.trackAPIErr()
-			if reply != _EMPTY_ {
-				s.sendInternalAccountMsg(nil, reply, response)
+			return
+		case <-rqch:
+			// If we were the head, drop and setup things for next.
+			if r != nil && r == head {
+				pop()
 			}
-			s.sendJetStreamAPIAuditAdvisory(ci, acc, subject, request, response)
+			next()
+		case <-tm.C:
+			if r != nil {
+				s.sendAPIErrResponse(r.ci, r.acc, r.subject, r.reply, r.request, r.response)
+				pop()
+			}
+			next()
 		}
+	}
+}
+
+func (s *Server) sendDelayedAPIErrResponse(ci *ClientInfo, acc *Account, subject, reply, request, response string, rg *raftGroup, duration time.Duration) {
+	s.delayedAPIResponses.push(&delayedAPIResponse{
+		ci, acc, subject, reply, request, response, rg, time.Now().Add(duration), nil,
 	})
 }
 
@@ -1887,12 +1985,12 @@ func (s *Server) jsStreamInfoRequest(sub *subscription, c *client, a *Account, s
 			if js.isLeaderless() {
 				resp.Error = NewJSClusterNotAvailError()
 				// Delaying an error response gives the leader a chance to respond before us
-				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil)
+				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil, errRespDelay)
 			}
 			return
 		} else if isLeader && offline {
 			resp.Error = NewJSStreamOfflineError()
-			s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil)
+			s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil, errRespDelay)
 			return
 		}
 
@@ -1904,7 +2002,7 @@ func (s *Server) jsStreamInfoRequest(sub *subscription, c *client, a *Account, s
 			if js.isLeaderless() {
 				resp.Error = NewJSClusterNotAvailError()
 				// Delaying an error response gives the leader a chance to respond before us
-				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), sa.Group)
+				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), sa.Group, errRespDelay)
 				return
 			}
 
@@ -4400,12 +4498,12 @@ func (s *Server) jsConsumerInfoRequest(sub *subscription, c *client, _ *Account,
 			if isLeaderLess {
 				resp.Error = NewJSClusterNotAvailError()
 				// Delaying an error response gives the leader a chance to respond before us
-				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil)
+				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil, errRespDelay)
 			}
 			return
 		} else if isLeader && offline {
 			resp.Error = NewJSConsumerOfflineError()
-			s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil)
+			s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil, errRespDelay)
 			return
 		}
 
@@ -4421,7 +4519,7 @@ func (s *Server) jsConsumerInfoRequest(sub *subscription, c *client, _ *Account,
 			if isLeaderLess {
 				resp.Error = NewJSClusterNotAvailError()
 				// Delaying an error response gives the leader a chance to respond before us
-				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), ca.Group)
+				s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), ca.Group, errRespDelay)
 				return
 			}
 
@@ -4464,7 +4562,7 @@ func (s *Server) jsConsumerInfoRequest(sub *subscription, c *client, _ *Account,
 			if node.GroupLeader() != _EMPTY_ || node.HadPreviousLeader() {
 				if leaderNotPartOfGroup {
 					resp.Error = NewJSConsumerOfflineError()
-					s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil)
+					s.sendDelayedAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp), nil, errRespDelay)
 				}
 				return
 			}
diff --git a/server/server.go b/server/server.go
index f8e5a45a..f7e0d2b0 100644
--- a/server/server.go
+++ b/server/server.go
@@ -362,6 +362,9 @@ type Server struct {
 	// Queue to process JS API requests that come from routes (or gateways)
 	jsAPIRoutedReqs *ipQueue[*jsAPIRoutedReq]
 
+	// Delayed API responses.
+	delayedAPIResponses *ipQueue[*delayedAPIResponse]
+
 	// Whether moving NRG traffic into accounts is permitted on this server.
 	// Controls whether or not the account NRG capability is set in statsz.
 	// Currently used by unit tests to simulate nodes not supporting account NRG.
@@ -730,6 +733,11 @@ func NewServer(opts *Options) (*Server, error) {
 		syncOutSem:         make(chan struct{}, maxConcurrentSyncRequests),
 	}
 
+	// Delayed API response queue. Create regardless if JetStream is configured
+	// or not (since it can be enabled/disabled with config reload, we want this
+	// queue to exist at all times).
+	s.delayedAPIResponses = newIPQueue[*delayedAPIResponse](s, "delayed API responses")
+
 	// By default we'll allow account NRG.
 	s.accountNRGAllowed.Store(true)
 
@@ -2396,6 +2404,11 @@ func (s *Server) Start() {
 		}
 	}
 
+	// Delayed API response handling. Start regardless of JetStream being
+	// currently configured or not (since it can be enabled/disabled with
+	// configuration reload).
+	s.startGoRoutine(s.delayedAPIResponder)
+
 	// Start OCSP Stapling monitoring for TLS certificates if enabled. Hook TLS handshake for
 	// OCSP check on peers (LEAF and CLIENT kind) if enabled.
 	s.startOCSPMonitoring()


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 9f8cdeb7..c3723c4a 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -24455,3 +24455,90 @@ func TestJetStreamMirroringClipStartSeq(t *testing.T) {
 		require_Equal(t, o.sseq, 11)
 	}
 }
+
+func TestJetStreamDelayedAPIResponses(t *testing.T) {
+	tdir := t.TempDir()
+	tmpl := `
+		listen: 127.0.0.1:-1
+		%sjetstream: {max_mem_store: 64GB, max_file_store: 10TB, store_dir: %q}
+	`
+	conf := createConfFile(t, []byte(fmt.Sprintf(tmpl, "", tdir)))
+
+	s, _ := RunServerWithConfig(conf)
+	defer s.Shutdown()
+
+	nc := natsConnect(t, s.ClientURL())
+	defer nc.Close()
+
+	sub := natsSubSync(t, nc, JSAuditAdvisory)
+
+	acc := s.GlobalAccount()
+
+	// Send B, A, D, C and exected to receive A, B, C, D
+	s.sendDelayedAPIErrResponse(nil, acc, "B", _EMPTY_, "request2", "response2", nil, 500*time.Millisecond)
+	time.Sleep(50 * time.Millisecond)
+	s.sendDelayedAPIErrResponse(nil, acc, "A", _EMPTY_, "request1", "response1", nil, 200*time.Millisecond)
+	time.Sleep(50 * time.Millisecond)
+	s.sendDelayedAPIErrResponse(nil, acc, "D", _EMPTY_, "request4", "response4", nil, 800*time.Millisecond)
+	time.Sleep(50 * time.Millisecond)
+	s.sendDelayedAPIErrResponse(nil, acc, "C", _EMPTY_, "request3", "response3", nil, 650*time.Millisecond)
+
+	check := func(req, resp string) {
+		t.Helper()
+		msg := natsNexMsg(t, sub, time.Second)
+		var audit JSAPIAudit
+		err := json.Unmarshal(msg.Data, &audit)
+		require_NoError(t, err)
+		require_Equal(t, audit.Request, req)
+		require_Equal(t, audit.Response, resp)
+	}
+	check("request1", "response1")
+	check("request2", "response2")
+	check("request3", "response3")
+	check("request4", "response4")
+
+	// Verify that if a raft group is canceled, the delayed API response is canceled too.
+	node := &raft{quit: make(chan struct{})}
+	g := &raftGroup{node: node}
+	// Send delayed API response with this raft group
+	s.sendDelayedAPIErrResponse(nil, acc, "E", _EMPTY_, "request5", "response5", g, 250*time.Millisecond)
+	time.Sleep(50 * time.Millisecond)
+	// Send that one without a group
+	s.sendDelayedAPIErrResponse(nil, acc, "F", _EMPTY_, "request6", "response6", nil, 400*time.Millisecond)
+	// Close the "request5"'s channel.
+	close(node.quit)
+	// So we should receive request6, not 5.
+	check("request6", "response6")
+
+	// Check config reload.
+	s.sendDelayedAPIErrResponse(nil, acc, "G", _EMPTY_, "request7", "response7", nil, 400*time.Millisecond)
+	time.Sleep(50 * time.Millisecond)
+	// Send a bunch more
+	for i := 0; i < 10; i++ {
+		s.sendDelayedAPIErrResponse(nil, acc, "H", _EMPTY_, "request8", "response8", nil, 500*time.Millisecond)
+	}
+	// Config reload to disable JS.
+	reloadUpdateConfig(t, s, conf, fmt.Sprintf(tmpl, "#", tdir))
+	check("request7", "response7")
+	// We should not receive more.
+	if msg, err := sub.NextMsg(600 * time.Millisecond); err != nats.ErrTimeout {
+		t.Fatalf("Did not expect to receive: %s", msg.Data)
+	}
+	// Check that the queue is empty.
+	s.mu.RLock()
+	q := s.delayedAPIResponses
+	s.mu.RUnlock()
+	require_Equal(t, q.len(), 0)
+
+	// Restore JS and check delayed response can be received.
+	reloadUpdateConfig(t, s, conf, fmt.Sprintf(tmpl, "", tdir))
+	// Wait until js is re-enabled
+	checkFor(t, 10*time.Second, 50*time.Millisecond, func() error {
+		if s.getJetStream() == nil {
+			return ErrJetStreamNotEnabled
+		}
+		return nil
+	})
+	s.sendDelayedAPIErrResponse(nil, acc, "I", _EMPTY_, "request9", "response9", nil, 100*time.Millisecond)
+	check("request9", "response9")
+}