Repository: nats-io/nats-server
Pull request #4981 (2024-01-21 03:49:32)

[ADDED] Batch support for direct gets from streams.
This adds batch support for direct gets. This also works with NextFor as well.

We will not hold any state in the server, so we need to protect against slow consumers. The request also adds in a MaxBytes which can allow the client to control how much data the server will blast to it. If the client request does not contain a non-zero MaxBytes, the server will use its opts.MaxPending.

After a batch has completed, even if cut short due to MaxBytes, the server will send an end of batch (EOB) empty message with a status of 204 in the header, and a header Nats-Num-Pending which reports how many messages remain.

```go
// JSApiMsgGetRequest get a message request.
type JSApiMsgGetRequest struct {
	Seq     uint64 `json:"seq,omitempty"`
	LastFor string `json:"last_by_subj,omitempty"`
	NextFor string `json:"next_by_subj,omitempty"`

	// Batch support. Used to request more then one msg at a time.
	// Can be used with simple starting seq, but also NextFor with wildcards.
	Batch int `json:"batch,omitempty"`
	// This will make sure we limit how much data we blast out. If not set we will
	// inherit the slow consumer default max setting of the server. Default is MAX_PENDING_SIZE.
	MaxBytes int `json:"max_bytes,omitempty"`
}
```


Signed-off-by: Derek Collison <derek@nats.io>


Patch:
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 14369d6c..5c501cd4 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -631,6 +631,13 @@ type JSApiMsgGetRequest struct {
 	Seq     uint64 `json:"seq,omitempty"`
 	LastFor string `json:"last_by_subj,omitempty"`
 	NextFor string `json:"next_by_subj,omitempty"`
+
+	// Batch support. Used to request more then one msg at a time.
+	// Can be used with simple starting seq, but also NextFor with wildcards.
+	Batch int `json:"batch,omitempty"`
+	// This will make sure we limit how much data we blast out. If not set we will
+	// inherit the slow consumer default max setting of the server. Default is MAX_PENDING_SIZE.
+	MaxBytes int `json:"max_bytes,omitempty"`
 }
 
 type JSApiMsgGetResponse struct {
@@ -3094,6 +3101,13 @@ func (s *Server) jsMsgGetRequest(sub *subscription, c *client, _ *Account, subje
 		return
 	}
 
+	// This version does not support batch.
+	if req.Batch > 0 || req.MaxBytes > 0 {
+		resp.Error = NewJSBadRequestError()
+		s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+		return
+	}
+
 	// Check that we do not have both options set.
 	if req.Seq > 0 && req.LastFor != _EMPTY_ || req.Seq == 0 && req.LastFor == _EMPTY_ && req.NextFor == _EMPTY_ {
 		resp.Error = NewJSBadRequestError()
diff --git a/server/stream.go b/server/stream.go
index f032f285..6971dcf5 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -350,6 +350,7 @@ const (
 	JSTimeStamp    = "Nats-Time-Stamp"
 	JSSubject      = "Nats-Subject"
 	JSLastSequence = "Nats-Last-Sequence"
+	JSNumPending   = "Nats-Num-Pending"
 )
 
 // Rollups, can be subject only or all messages.
@@ -4032,12 +4033,10 @@ func (mset *stream) processDirectGetRequest(_ *subscription, c *client, _ *Accou
 		return
 	}
 	// Check that we do not have both options set.
-	if req.Seq > 0 && req.LastFor != _EMPTY_ {
-		hdr := []byte("NATS/1.0 408 Bad Request\r\n\r\n")
-		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
-		return
-	}
-	if req.LastFor != _EMPTY_ && req.NextFor != _EMPTY_ {
+	// We do not allow batch mode for lastFor requests.
+	if (req.Seq > 0 && req.LastFor != _EMPTY_) ||
+		(req.LastFor != _EMPTY_ && req.NextFor != _EMPTY_) ||
+		(req.LastFor != _EMPTY_ && req.Batch > 0) {
 		hdr := []byte("NATS/1.0 408 Bad Request\r\n\r\n")
 		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
 		return
@@ -4099,41 +4098,86 @@ func (mset *stream) processDirectGetLastBySubjectRequest(_ *subscription, c *cli
 // Do actual work on a direct msg request.
 // This could be called in a Go routine if we are inline for a non-client connection.
 func (mset *stream) getDirectRequest(req *JSApiMsgGetRequest, reply string) {
-	var svp StoreMsg
-	var sm *StoreMsg
-	var err error
-
 	mset.mu.RLock()
-	store, name := mset.store, mset.cfg.Name
+	store, name, s := mset.store, mset.cfg.Name, mset.srv
 	mset.mu.RUnlock()
 
-	if req.Seq > 0 && req.NextFor == _EMPTY_ {
-		sm, err = store.LoadMsg(req.Seq, &svp)
-	} else if req.NextFor != _EMPTY_ {
-		sm, _, err = store.LoadNextMsg(req.NextFor, subjectHasWildcard(req.NextFor), req.Seq, &svp)
-	} else {
-		sm, err = store.LoadLastMsg(req.LastFor, &svp)
-	}
-	if err != nil {
-		hdr := []byte("NATS/1.0 404 Message Not Found\r\n\r\n")
-		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
-		return
-	}
+	seq := req.Seq
+	wc := subjectHasWildcard(req.NextFor)
+	batch := req.Batch
+	if batch == 0 {
+		batch = 1
+	}
+	// Grab MaxBytes
+	mb := req.MaxBytes
+	if mb == 0 && s != nil {
+		// Fill in with the server's MaxPending.
+		mb = int(s.opts.MaxPending)
+	}
+	// Track what we have sent.
+	var sentBytes int
+
+	// Loop over batch, which defaults to 1.
+	for i := 0; i < batch; i++ {
+		var (
+			svp StoreMsg
+			sm  *StoreMsg
+			err error
+		)
+		if seq > 0 && req.NextFor == _EMPTY_ {
+			// Only do direct lookup for first in a batch.
+			if i == 0 {
+				sm, err = store.LoadMsg(seq, &svp)
+			} else {
+				// We want to use load next with fwcs to step over deleted msgs.
+				sm, seq, err = store.LoadNextMsg(fwcs, true, seq, &svp)
+			}
+			// Bump for next loop if applicable.
+			seq++
+		} else if req.NextFor != _EMPTY_ {
+			sm, seq, err = store.LoadNextMsg(req.NextFor, wc, seq, &svp)
+			seq++
+		} else {
+			// Batch is not applicable here, this is checked before we get here.
+			sm, err = store.LoadLastMsg(req.LastFor, &svp)
+		}
+		if err != nil {
+			// For batches, if we stop early we want to do EOB logic below.
+			if batch > 1 && i > 0 {
+				break
+			}
+			hdr := []byte("NATS/1.0 404 Message Not Found\r\n\r\n")
+			mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
+			return
+		}
 
-	hdr := sm.hdr
-	ts := time.Unix(0, sm.ts).UTC()
+		hdr := sm.hdr
+		ts := time.Unix(0, sm.ts).UTC()
 
-	if len(hdr) == 0 {
-		const ht = "NATS/1.0\r\nNats-Stream: %s\r\nNats-Subject: %s\r\nNats-Sequence: %d\r\nNats-Time-Stamp: %s\r\n\r\n"
-		hdr = []byte(fmt.Sprintf(ht, name, sm.subj, sm.seq, ts.Format(time.RFC3339Nano)))
-	} else {
-		hdr = copyBytes(hdr)
-		hdr = genHeader(hdr, JSStream, name)
-		hdr = genHeader(hdr, JSSubject, sm.subj)
-		hdr = genHeader(hdr, JSSequence, strconv.FormatUint(sm.seq, 10))
-		hdr = genHeader(hdr, JSTimeStamp, ts.Format(time.RFC3339Nano))
+		if len(hdr) == 0 {
+			const ht = "NATS/1.0\r\nNats-Stream: %s\r\nNats-Subject: %s\r\nNats-Sequence: %d\r\nNats-Time-Stamp: %s\r\n\r\n"
+			hdr = []byte(fmt.Sprintf(ht, name, sm.subj, sm.seq, ts.Format(time.RFC3339Nano)))
+		} else {
+			hdr = copyBytes(hdr)
+			hdr = genHeader(hdr, JSStream, name)
+			hdr = genHeader(hdr, JSSubject, sm.subj)
+			hdr = genHeader(hdr, JSSequence, strconv.FormatUint(sm.seq, 10))
+			hdr = genHeader(hdr, JSTimeStamp, ts.Format(time.RFC3339Nano))
+		}
+		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, sm.msg, nil, 0))
+		// Check if we have exceeded max bytes.
+		sentBytes += len(sm.subj) + len(sm.hdr) + len(sm.msg)
+		if sentBytes >= mb {
+			break
+		}
+	}
+
+	// If batch was requested send EOB.
+	if batch > 1 {
+		np, _ := store.NumPending(seq, req.NextFor, false)
+		hdr := []byte(fmt.Sprintf("NATS/1.0 204 EOB\r\nNats-Pending-Messages: %d\r\n\r\n", np))
+		mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
 	}
-	mset.outq.send(newJSPubMsg(reply, _EMPTY_, _EMPTY_, hdr, sm.msg, nil, 0))
 }
 
 // processInboundJetStreamMsg handles processing messages bound for a stream.


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index bd279ef2..a5c7130c 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -22089,3 +22089,161 @@ func TestJetStreamConsumerPendingForKV(t *testing.T) {
 		})
 	}
 }
+
+func TestJetStreamDirectGetBatch(t *testing.T) {
+	s := RunBasicJetStreamServer(t)
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo.*"},
+	})
+	require_NoError(t, err)
+
+	// Add in messages
+	for i := 0; i < 333; i++ {
+		js.PublishAsync("foo.foo", []byte("HELLO"))
+		js.PublishAsync("foo.bar", []byte("WORLD"))
+		js.PublishAsync("foo.baz", []byte("AGAIN"))
+	}
+	select {
+	case <-js.PublishAsyncComplete():
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Did not receive completion signal")
+	}
+
+	// DirectGet is required for batch. Make sure we error correctly if not enabled.
+	mreq := &JSApiMsgGetRequest{Seq: 1, Batch: 10}
+	req, _ := json.Marshal(mreq)
+	rr, err := nc.Request("$JS.API.STREAM.MSG.GET.TEST", req, time.Second)
+	require_NoError(t, err)
+	var resp JSApiMsgGetResponse
+	json.Unmarshal(rr.Data, &resp)
+	require_True(t, resp.Error != nil)
+	require_Equal(t, resp.Error.Code, NewJSBadRequestError().Code)
+
+	// Update stream to support direct.
+	_, err = js.UpdateStream(&nats.StreamConfig{
+		Name:        "TEST",
+		Subjects:    []string{"foo.*"},
+		AllowDirect: true,
+	})
+	require_NoError(t, err)
+
+	// Direct subjects.
+	sendRequest := func(mreq *JSApiMsgGetRequest) *nats.Subscription {
+		t.Helper()
+		req, _ := json.Marshal(mreq)
+		// We will get multiple responses so can't do normal request.
+		reply := nats.NewInbox()
+		sub, err := nc.SubscribeSync(reply)
+		require_NoError(t, err)
+		err = nc.PublishRequest("$JS.API.DIRECT.GET.TEST", reply, req)
+		require_NoError(t, err)
+		return sub
+	}
+
+	// Batch sizes greater than 1 will have a nil message as the end marker.
+	checkResponses := func(sub *nats.Subscription, numPending int, expected ...string) {
+		t.Helper()
+		defer sub.Unsubscribe()
+		checkSubsPending(t, sub, len(expected))
+		for i := 0; i < len(expected); i++ {
+			msg, err := sub.NextMsg(10 * time.Millisecond)
+			require_NoError(t, err)
+			// If expected is _EMPTY_ that signals we expect a EOB marker.
+			if subj := expected[i]; subj != _EMPTY_ {
+				// Make sure subject is correct.
+				require_Equal(t, expected[i], msg.Header.Get("Nats-Subject"))
+				// Should have Data field non-zero
+				require_True(t, len(msg.Data) > 0)
+			} else {
+				// Check for properly formatted EOB marker.
+				// Should have no body.
+				require_Equal(t, len(msg.Data), 0)
+				// We mark status as 204 - No Content
+				require_Equal(t, msg.Header.Get("Status"), "204")
+				// Check description is EOB
+				require_Equal(t, msg.Header.Get("Description"), "EOB")
+				// Check we have NumPending and its correct.
+				require_Equal(t, strconv.Itoa(numPending), msg.Header.Get("Nats-Pending-Messages"))
+			}
+		}
+	}
+
+	// Run some simple tests.
+	sub := sendRequest(&JSApiMsgGetRequest{Seq: 1, Batch: 2})
+	checkResponses(sub, 997, "foo.foo", "foo.bar", _EMPTY_)
+
+	sub = sendRequest(&JSApiMsgGetRequest{Seq: 1, Batch: 3})
+	checkResponses(sub, 996, "foo.foo", "foo.bar", "foo.baz", _EMPTY_)
+
+	// Test NextFor works
+	sub = sendRequest(&JSApiMsgGetRequest{Seq: 1, Batch: 3, NextFor: "foo.*"})
+	checkResponses(sub, 996, "foo.foo", "foo.bar", "foo.baz", _EMPTY_)
+
+	sub = sendRequest(&JSApiMsgGetRequest{Seq: 1, Batch: 3, NextFor: "foo.baz"})
+	checkResponses(sub, 330, "foo.baz", "foo.baz", "foo.baz", _EMPTY_)
+
+	// Test stopping early by starting at 997 with only 3 messages.
+	sub = sendRequest(&JSApiMsgGetRequest{Seq: 997, Batch: 10, NextFor: "foo.*"})
+	checkResponses(sub, 0, "foo.foo", "foo.bar", "foo.baz", _EMPTY_)
+}
+
+func TestJetStreamDirectGetBatchMaxBytes(t *testing.T) {
+	s := RunBasicJetStreamServer(t)
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:        "TEST",
+		Subjects:    []string{"foo.*"},
+		AllowDirect: true,
+		Compression: nats.S2Compression,
+	})
+	require_NoError(t, err)
+
+	msg := bytes.Repeat([]byte("Z"), 512*1024)
+	// Add in messages
+	for i := 0; i < 333; i++ {
+		js.PublishAsync("foo.foo", msg)
+		js.PublishAsync("foo.bar", msg)
+		js.PublishAsync("foo.baz", msg)
+	}
+	select {
+	case <-js.PublishAsyncComplete():
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Did not receive completion signal")
+	}
+
+	sendRequestAndCheck := func(mreq *JSApiMsgGetRequest, numExpected int) {
+		t.Helper()
+		req, _ := json.Marshal(mreq)
+		// We will get multiple responses so can't do normal request.
+		reply := nats.NewInbox()
+		sub, err := nc.SubscribeSync(reply)
+		require_NoError(t, err)
+		defer sub.Unsubscribe()
+		err = nc.PublishRequest("$JS.API.DIRECT.GET.TEST", reply, req)
+		require_NoError(t, err)
+		// Make sure we get correct number of responses.
+		checkSubsPending(t, sub, numExpected)
+	}
+
+	// Total msg size being sent back to us.
+	msgSize := len(msg) + len("foo.foo")
+	// We should get 1 msg and 1 EOB
+	sendRequestAndCheck(&JSApiMsgGetRequest{Seq: 1, Batch: 3, MaxBytes: msgSize}, 2)
+
+	// Test NextFor tracks as well.
+	sendRequestAndCheck(&JSApiMsgGetRequest{Seq: 1, NextFor: "foo.bar", Batch: 3, MaxBytes: 2 * msgSize}, 3)
+
+	// Now test no MaxBytes to inherit server max_num_pending.
+	expected := (int(s.getOpts().MaxPending) / msgSize) + 1
+	sendRequestAndCheck(&JSApiMsgGetRequest{Seq: 1, Batch: 200}, expected+1)
+}