Repository: nats-io/nats-server
Pull request #3126 (2022-05-15 20:45:15)

Support for MaxBytes for pull requests.
Getting us prepped for move to all pull consumers for clients and for the internal consumers for mirrors/sources.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 357e1e97..76f12f61 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -68,8 +68,9 @@ type ConsumerConfig struct {
 	HeadersOnly     bool            `json:"headers_only,omitempty"`
 
 	// Pull based options.
-	MaxRequestBatch   int           `json:"max_batch,omitempty"`
-	MaxRequestExpires time.Duration `json:"max_expires,omitempty"`
+	MaxRequestBatch    int           `json:"max_batch,omitempty"`
+	MaxRequestExpires  time.Duration `json:"max_expires,omitempty"`
+	MaxRequestMaxBytes int           `json:"max_bytes,omitempty"`
 
 	// Push based consumers.
 	DeliverSubject string `json:"deliver_subject,omitempty"`
@@ -2296,36 +2297,36 @@ func (o *consumer) needAck(sseq uint64) bool {
 }
 
 // Helper for the next message requests.
-func nextReqFromMsg(msg []byte) (time.Time, int, bool, time.Duration, time.Time, error) {
+func nextReqFromMsg(msg []byte) (time.Time, int, int, bool, time.Duration, time.Time, error) {
 	req := bytes.TrimSpace(msg)
 
 	switch {
 	case len(req) == 0:
-		return time.Time{}, 1, false, 0, time.Time{}, nil
+		return time.Time{}, 1, 0, false, 0, time.Time{}, nil
 
 	case req[0] == '{':
 		var cr JSApiConsumerGetNextRequest
 		if err := json.Unmarshal(req, &cr); err != nil {
-			return time.Time{}, -1, false, 0, time.Time{}, err
+			return time.Time{}, -1, 0, false, 0, time.Time{}, err
 		}
 		var hbt time.Time
 		if cr.Heartbeat > 0 {
 			if cr.Heartbeat*2 > cr.Expires {
-				return time.Time{}, 1, false, 0, time.Time{}, errors.New("heartbeat value too large")
+				return time.Time{}, 1, 0, false, 0, time.Time{}, errors.New("heartbeat value too large")
 			}
 			hbt = time.Now().Add(cr.Heartbeat)
 		}
 		if cr.Expires == time.Duration(0) {
-			return time.Time{}, cr.Batch, cr.NoWait, cr.Heartbeat, hbt, nil
+			return time.Time{}, cr.Batch, cr.MaxBytes, cr.NoWait, cr.Heartbeat, hbt, nil
 		}
-		return time.Now().Add(cr.Expires), cr.Batch, cr.NoWait, cr.Heartbeat, hbt, nil
+		return time.Now().Add(cr.Expires), cr.Batch, cr.MaxBytes, cr.NoWait, cr.Heartbeat, hbt, nil
 	default:
 		if n, err := strconv.Atoi(string(req)); err == nil {
-			return time.Time{}, n, false, 0, time.Time{}, nil
+			return time.Time{}, n, 0, false, 0, time.Time{}, nil
 		}
 	}
 
-	return time.Time{}, 1, false, 0, time.Time{}, nil
+	return time.Time{}, 1, 0, false, 0, time.Time{}, nil
 }
 
 // Represents a request that is on the internal waiting queue
@@ -2335,6 +2336,7 @@ type waitingRequest struct {
 	reply    string
 	n        int // For batching
 	d        int
+	b        int // For max bytes tracking.
 	expires  time.Time
 	received time.Time
 	hb       time.Duration
@@ -2501,12 +2503,39 @@ func (o *consumer) pendingRequestReplies() []string {
 // Return next waiting request. This will check for expirations but not noWait or interest.
 // That will be handled by processWaiting.
 // Lock should be held.
-func (o *consumer) nextWaiting() *waitingRequest {
+func (o *consumer) nextWaiting(sz int) *waitingRequest {
 	if o.waiting == nil || o.waiting.isEmpty() {
 		return nil
 	}
 	for wr := o.waiting.peek(); !o.waiting.isEmpty(); wr = o.waiting.peek() {
-		if wr == nil || wr.expires.IsZero() || time.Now().Before(wr.expires) {
+		if wr == nil {
+			break
+		}
+		// Check if we have max bytes set.
+		if wr.b > 0 {
+			if sz <= wr.b {
+				wr.b -= sz
+				// If we are right now at zero, set batch to 1 to deliver this one but stop after.
+				if wr.b == 0 {
+					wr.n = 1
+				}
+			} else {
+				// If we have not delivered anything to the requestor let them know.
+				if wr.d == 0 {
+					hdr := []byte("NATS/1.0 408 Message Size Exceeds MaxBytes\r\n\r\n")
+					o.outq.send(newJSPubMsg(wr.reply, _EMPTY_, _EMPTY_, hdr, nil, nil, 0))
+				}
+				// Remove the current one, no longer valid due to max bytes limit.
+				o.waiting.removeCurrent()
+				if o.node != nil {
+					o.removeClusterPendingRequest(wr.reply)
+				}
+				wr.recycle()
+				continue
+			}
+		}
+
+		if wr.expires.IsZero() || time.Now().Before(wr.expires) {
 			rr := wr.acc.sl.Match(wr.interest)
 			if len(rr.psubs)+len(rr.qsubs) > 0 {
 				return o.waiting.pop()
@@ -2561,7 +2590,7 @@ func (o *consumer) processNextMsgRequest(reply string, msg []byte) {
 	}
 
 	// Check payload here to see if they sent in batch size or a formal request.
-	expires, batchSize, noWait, hb, hbt, err := nextReqFromMsg(msg)
+	expires, batchSize, maxBytes, noWait, hb, hbt, err := nextReqFromMsg(msg)
 	if err != nil {
 		sendErr(400, fmt.Sprintf("Bad Request - %v", err))
 		return
@@ -2578,6 +2607,11 @@ func (o *consumer) processNextMsgRequest(reply string, msg []byte) {
 		return
 	}
 
+	if maxBytes > 0 && o.cfg.MaxRequestMaxBytes > 0 && maxBytes > o.cfg.MaxRequestMaxBytes {
+		sendErr(409, fmt.Sprintf("Exceeded MaxRequestMaxBytes of %v", o.cfg.MaxRequestMaxBytes))
+		return
+	}
+
 	// If we have the max number of requests already pending try to expire.
 	if o.waiting.isFull() {
 		// Try to expire some of the requests.
@@ -2607,9 +2641,10 @@ func (o *consumer) processNextMsgRequest(reply string, msg []byte) {
 	// If we receive this request though an account export, we need to track that interest subject and account.
 	acc, interest := trackDownAccountAndInterest(o.acc, reply)
 
-	// In case we have to queue up this request.
+	// Create a waiting request.
 	wr := wrPool.Get().(*waitingRequest)
 	wr.acc, wr.interest, wr.reply, wr.n, wr.d, wr.noWait, wr.expires, wr.hb, wr.hbt = acc, interest, reply, batchSize, 0, noWait, expires, hb, hbt
+	wr.b = maxBytes
 	wr.received = time.Now()
 
 	if err := o.waiting.add(wr); err != nil {
@@ -2946,7 +2981,7 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 
 		if o.isPushMode() {
 			dsubj = o.dsubj
-		} else if wr := o.nextWaiting(); wr != nil {
+		} else if wr := o.nextWaiting(len(pmsg.hdr) + len(pmsg.msg)); wr != nil {
 			dsubj = wr.reply
 			if done := wr.recycleIfDone(); done && o.node != nil {
 				o.removeClusterPendingRequest(dsubj)
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 40daa8dd..64586298 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -608,6 +608,7 @@ const JSApiConsumerListResponseType = "io.nats.jetstream.api.v1.consumer_list_re
 type JSApiConsumerGetNextRequest struct {
 	Expires   time.Duration `json:"expires,omitempty"`
 	Batch     int           `json:"batch,omitempty"`
+	MaxBytes  int           `json:"max_bytes,omitempty"`
 	NoWait    bool          `json:"no_wait,omitempty"`
 	Heartbeat time.Duration `json:"idle_heartbeat,omitempty"`
 }


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 150b12f5..7084a80d 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -650,7 +650,7 @@ func TestJetStreamConsumerMaxDeliveries(t *testing.T) {
 
 func TestJetStreamNextReqFromMsg(t *testing.T) {
 	bef := time.Now()
-	expires, _, _, _, _, err := nextReqFromMsg([]byte(`{"expires":5000000000}`)) // nanoseconds
+	expires, _, _, _, _, _, err := nextReqFromMsg([]byte(`{"expires":5000000000}`)) // nanoseconds
 	require_NoError(t, err)
 	now := time.Now()
 	if expires.Before(bef.Add(5*time.Second)) || expires.After(now.Add(5*time.Second)) {
@@ -14435,15 +14435,16 @@ func TestJetStreamPullConsumerRequestMaximums(t *testing.T) {
 	require_NoError(t, err)
 
 	_, err = mset.addConsumer(&ConsumerConfig{
-		Durable:           "dlc",
-		MaxRequestBatch:   10,
-		MaxRequestExpires: time.Second,
-		AckPolicy:         AckExplicit,
+		Durable:            "dlc",
+		MaxRequestBatch:    10,
+		MaxRequestMaxBytes: 10_000,
+		MaxRequestExpires:  time.Second,
+		AckPolicy:          AckExplicit,
 	})
 	require_NoError(t, err)
 
-	genReq := func(b int, e time.Duration) []byte {
-		req := &JSApiConsumerGetNextRequest{Batch: b, Expires: e}
+	genReq := func(b, mb int, e time.Duration) []byte {
+		req := &JSApiConsumerGetNextRequest{Batch: b, Expires: e, MaxBytes: mb}
 		jreq, err := json.Marshal(req)
 		require_NoError(t, err)
 		return jreq
@@ -14452,14 +14453,21 @@ func TestJetStreamPullConsumerRequestMaximums(t *testing.T) {
 	rsubj := fmt.Sprintf(JSApiRequestNextT, "TEST", "dlc")
 
 	// Exceeds max batch size.
-	resp, err := nc.Request(rsubj, genReq(11, 100*time.Millisecond), time.Second)
+	resp, err := nc.Request(rsubj, genReq(11, 0, 100*time.Millisecond), time.Second)
 	require_NoError(t, err)
 	if status := resp.Header.Get("Status"); status != "409" {
 		t.Fatalf("Expected a 409 status code, got %q", status)
 	}
 
 	// Exceeds max expires.
-	resp, err = nc.Request(rsubj, genReq(1, 10*time.Minute), time.Second)
+	resp, err = nc.Request(rsubj, genReq(1, 0, 10*time.Minute), time.Second)
+	require_NoError(t, err)
+	if status := resp.Header.Get("Status"); status != "409" {
+		t.Fatalf("Expected a 409 status code, got %q", status)
+	}
+
+	// Exceeds max bytes.
+	resp, err = nc.Request(rsubj, genReq(10, 10_000*2, 10*time.Minute), time.Second)
 	require_NoError(t, err)
 	if status := resp.Header.Get("Status"); status != "409" {
 		t.Fatalf("Expected a 409 status code, got %q", status)
@@ -17369,6 +17377,113 @@ func TestJetStreamWorkQueueSourceNamingRestart(t *testing.T) {
 	}
 }
 
+func TestJetStreamPullMaxBytes(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	if config := s.JetStreamConfig(); config != nil {
+		defer removeDir(t, config.StoreDir)
+	}
+	defer s.Shutdown()
+
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name: "TEST",
+	})
+	require_NoError(t, err)
+
+	// Put in ~2MB, each ~100k
+	msz := 99_980
+	total, msg := 20, []byte(strings.Repeat("Z", msz))
+
+	for i := 0; i < total; i++ {
+		if _, err := js.Publish("TEST", msg); err != nil {
+			t.Fatalf("Unexpected publish error: %v", err)
+		}
+	}
+
+	_, err = js.AddConsumer("TEST", &nats.ConsumerConfig{
+		Durable:   "pr",
+		AckPolicy: nats.AckExplicitPolicy,
+	})
+	require_NoError(t, err)
+
+	req := &JSApiConsumerGetNextRequest{MaxBytes: 100, NoWait: true}
+	jreq, _ := json.Marshal(req)
+
+	subj := fmt.Sprintf(JSApiRequestNextT, "TEST", "pr")
+	reply := "_pr_"
+	sub, _ := nc.SubscribeSync(reply)
+	defer sub.Unsubscribe()
+
+	checkHeader := func(m *nats.Msg, expected *nats.Header) {
+		t.Helper()
+		if len(m.Data) != 0 {
+			t.Fatalf("Did not expect data, got %d bytes", len(m.Data))
+		}
+		if !reflect.DeepEqual(&m.Header, expected) {
+			t.Fatalf("Expected %+v hdr, got %+v", expected, m.Header)
+		}
+	}
+
+	// If we ask for less MaxBytes then a single message make sure we get an error.
+	badReq := &nats.Header{"Status": []string{"408"}, "Description": []string{"Message Size Exceeds MaxBytes"}}
+
+	nc.PublishRequest(subj, reply, jreq)
+	m, err := sub.NextMsg(time.Second)
+	require_NoError(t, err)
+	checkSubsPending(t, sub, 0)
+	checkHeader(m, badReq)
+
+	// If we request a ton of max bytes make sure batch size overrides.
+	req = &JSApiConsumerGetNextRequest{Batch: 1, MaxBytes: 10_000_000, NoWait: true}
+	jreq, _ = json.Marshal(req)
+	nc.PublishRequest(subj, reply, jreq)
+	checkSubsPending(t, sub, 1)
+
+	m, err = sub.NextMsg(time.Second)
+	require_NoError(t, err)
+	require_True(t, len(m.Data) == msz)
+	require_True(t, len(m.Header) == 0)
+	checkSubsPending(t, sub, 0)
+
+	// Same but with batch > 1
+	req = &JSApiConsumerGetNextRequest{Batch: 5, MaxBytes: 10_000_000, NoWait: true}
+	jreq, _ = json.Marshal(req)
+	nc.PublishRequest(subj, reply, jreq)
+	checkSubsPending(t, sub, 5)
+	for i := 0; i < 5; i++ {
+		m, err = sub.NextMsg(time.Second)
+		require_NoError(t, err)
+		require_True(t, len(m.Data) == msz)
+		require_True(t, len(m.Header) == 0)
+	}
+	checkSubsPending(t, sub, 0)
+
+	// Now ask for large batch but make sure we are limited by batch size.
+	req = &JSApiConsumerGetNextRequest{Batch: 1_000, MaxBytes: msz * 5, NoWait: true}
+	jreq, _ = json.Marshal(req)
+	nc.PublishRequest(subj, reply, jreq)
+	checkSubsPending(t, sub, 5)
+	for i := 0; i < 5; i++ {
+		m, err = sub.NextMsg(time.Second)
+		require_NoError(t, err)
+		require_True(t, len(m.Data) == msz)
+		require_True(t, len(m.Header) == 0)
+	}
+	checkSubsPending(t, sub, 0)
+
+	req = &JSApiConsumerGetNextRequest{Batch: 1_000, MaxBytes: msz + 20, NoWait: true}
+	jreq, _ = json.Marshal(req)
+	nc.PublishRequest(subj, reply, jreq)
+	checkSubsPending(t, sub, 1)
+	m, err = sub.NextMsg(time.Second)
+	require_NoError(t, err)
+	require_True(t, len(m.Data) == msz)
+	require_True(t, len(m.Header) == 0)
+	checkSubsPending(t, sub, 0)
+}
+
 ///////////////////////////////////////////////////////////////////////////
 // Simple JetStream Benchmarks
 ///////////////////////////////////////////////////////////////////////////