Repository: nats-io/nats-server
Pull request #1966 (2021-03-05 16:49:33)

Enable opt-in flow control for push based consumers.
Enable slow start flow control for push based consumers using flow control messages.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index f54329e5..33ea9be3 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -62,6 +62,7 @@ type ConsumerConfig struct {
 	MaxWaiting      int           `json:"max_waiting,omitempty"`
 	MaxAckPending   int           `json:"max_ack_pending,omitempty"`
 	Heartbeat       time.Duration `json:"idle_heartbeat,omitempty"`
+	FlowControl     bool          `json:"flow_control,omitempty"`
 }
 
 type CreateConsumerRequest struct {
@@ -188,6 +189,11 @@ type consumer struct {
 	ackSubj           string
 	nextMsgSubj       string
 	maxp              int
+	pblimit           int
+	maxpb             int
+	pbytes            int
+	pfcs              int
+	fcSub             *subscription
 	outq              *jsOutQ
 	pending           map[uint64]*Pending
 	ptmr              *time.Timer
@@ -225,6 +231,9 @@ const (
 	// JsDeleteWaitTimeDefault is the default amount of time we will wait for non-durable
 	// observables to be in an inactive state before deleting them.
 	JsDeleteWaitTimeDefault = 5 * time.Second
+	// JsFlowControlMaxPending specifies default pending bytes during flow control that can be
+	// outstanding.
+	JsFlowControlMaxPending = 64 * 1024 * 1024
 )
 
 func (mset *stream) addConsumer(config *ConsumerConfig) (*consumer, error) {
@@ -288,6 +297,9 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 		if config.Heartbeat > 0 {
 			return nil, fmt.Errorf("consumer idle heartbeat requires a push based consumer")
 		}
+		if config.FlowControl {
+			return nil, fmt.Errorf("consumer flow control requires a push based consumer")
+		}
 	}
 
 	// Setup proper default for ack wait if we are in explicit ack mode.
@@ -660,6 +672,16 @@ func (o *consumer) setLeader(isLeader bool) {
 			return
 		}
 
+		// Check on flow control settings.
+		if o.cfg.FlowControl {
+			o.setMaxPendingBytes(JsFlowControlMaxPending)
+			if o.fcSub, err = o.subscribeInternal(jsFlowControl, o.processFlowControl); err != nil {
+				o.mu.Unlock()
+				o.deleteWithoutAdvisory()
+				return
+			}
+		}
+
 		// Setup initial pending.
 		o.setInitialPending()
 
@@ -691,8 +713,10 @@ func (o *consumer) setLeader(isLeader bool) {
 		o.mu.Lock()
 		o.unsubscribe(o.ackSub)
 		o.unsubscribe(o.reqSub)
+		o.unsubscribe(o.fcSub)
 		o.ackSub = nil
 		o.reqSub = nil
+		o.fcSub = nil
 		if o.infoSub != nil {
 			o.srv.sysUnsubscribe(o.infoSub)
 			o.infoSub = nil
@@ -1862,9 +1886,14 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 			return
 		}
 
-		// If we are in push mode and not active let's stop sending.
-		if o.isPushMode() && !o.active {
-			goto waitForMsgs
+		// If we are in push mode and not active or under flowcontrol let's stop sending.
+		if o.isPushMode() {
+			if !o.active {
+				goto waitForMsgs
+			}
+			if o.maxpb > 0 && o.pbytes > o.maxpb {
+				goto waitForMsgs
+			}
 		}
 
 		// If we are in pull mode and no one is waiting already break and wait.
@@ -1922,6 +1951,7 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 			}
 		}
 
+		// Do actual delivery.
 		o.deliverMsg(dsubj, subj, hdr, msg, seq, dc, ts)
 
 		o.mu.Unlock()
@@ -1968,6 +1998,15 @@ func (o *consumer) ackReply(sseq, dseq, dc uint64, ts int64, pending uint64) str
 	return fmt.Sprintf(o.ackReplyT, dc, sseq, dseq, ts, pending)
 }
 
+// Used mostly for testing. Sets max pending bytes for flow control setups.
+func (o *consumer) setMaxPendingBytes(limit int) {
+	o.pblimit = limit
+	o.maxpb = limit / 8
+	if o.maxpb == 0 {
+		o.maxpb = 1
+	}
+}
+
 // Deliver a msg to the consumer.
 // Lock should be held and o.mset validated to be non-nil.
 func (o *consumer) deliverMsg(dsubj, subj string, hdr, msg []byte, seq, dc uint64, ts int64) {
@@ -1979,13 +2018,6 @@ func (o *consumer) deliverMsg(dsubj, subj string, hdr, msg []byte, seq, dc uint6
 		o.sgap--
 	}
 
-	if len(hdr) > 0 {
-		hdr = append(hdr[:0:0], hdr...)
-	}
-	if len(msg) > 0 {
-		msg = append(msg[:0:0], msg...)
-	}
-
 	dseq := o.dseq
 	o.dseq++
 
@@ -2008,10 +2040,85 @@ func (o *consumer) deliverMsg(dsubj, subj string, hdr, msg []byte, seq, dc uint6
 		o.asflr = seq
 	}
 
+	// Flow control.
+	if o.maxpb > 0 {
+		o.pbytes += pmsg.size()
+		if o.needFlowControl() {
+			o.sendFlowControl()
+		}
+	}
+
 	// FIXME(dlc) - Capture errors?
 	o.updateDelivered(dseq, seq, dc, ts)
 }
 
+func (o *consumer) needFlowControl() bool {
+	if o.maxpb == 0 {
+		return false
+	}
+	// Decide whether to send a flow control message which we will need the user to respond.
+	// We send if we are at the limit or over, and at 25%, 50% and 75%.
+	if o.pbytes >= o.maxpb {
+		return true
+	} else if o.pfcs == 0 && o.pbytes > o.maxpb/4 {
+		return true
+	} else if o.pfcs == 1 && o.pbytes > o.maxpb/2 {
+		return true
+	} else if o.pfcs == 2 && o.pbytes > o.maxpb*3/4 {
+		return true
+	}
+	return false
+}
+
+func (o *consumer) processFlowControl(_ *subscription, c *client, subj, _ string, _ []byte) {
+	sz, err := strconv.Atoi(tokenAt(subj, 5))
+	if err != nil {
+		o.srv.Warnf("Bad flow control response subject: %q", subj)
+		return
+	}
+
+	o.mu.Lock()
+	defer o.mu.Unlock()
+	wasBlocked := o.pbytes >= o.maxpb
+
+	// For slow starts and ramping up.
+	if o.maxpb < o.pblimit {
+		o.maxpb *= 2
+		if o.maxpb > o.pblimit {
+			o.maxpb = o.pblimit
+		}
+	}
+
+	// Update accounting.
+	o.pbytes -= sz
+	o.pfcs--
+
+	// In case they are sent out of order or we get duplicates etc.
+	if o.pbytes < 0 {
+		o.pbytes = 0
+	}
+	if o.pfcs < 0 {
+		o.pfcs = 0
+	}
+
+	if wasBlocked {
+		o.signalNewMessages()
+	}
+}
+
+// sendFlowControl will send a flow control packet to the consumer.
+// Lock should be held.
+func (o *consumer) sendFlowControl() {
+	if !o.isPushMode() {
+		return
+	}
+	subj := o.cfg.DeliverSubject
+	o.pfcs++
+	reply := fmt.Sprintf(jsFlowControlT, o.stream, o.name, o.pbytes)
+	hdr := []byte("NATS/1.0 100 FlowControl Request\r\n\r\n")
+	o.outq.send(&jsPubMsg{subj, subj, reply, hdr, nil, nil, 0, nil})
+}
+
 // Tracks our outstanding pending acks. Only applicable to AckExplicit mode.
 // Lock should be held.
 func (o *consumer) trackPending(sseq, dseq uint64) {
@@ -2622,7 +2729,7 @@ func (o *consumer) setInitialPending() {
 func (o *consumer) decStreamPending(sseq uint64, subj string) {
 	o.mu.Lock()
 	// Ignore if we have already seen this one.
-	if sseq >= o.sseq && o.sgap > 0 && o.isFilteredMatch(subj) {
+	if sseq >= o.sseq && o.sgap > 0 && o.isFilteredMatch(subj) && o.sgap > 0 {
 		o.sgap--
 	}
 	o.mu.Unlock()
diff --git a/server/filestore.go b/server/filestore.go
index dc54ed8b..b4bda1c5 100644
--- a/server/filestore.go
+++ b/server/filestore.go
@@ -579,6 +579,7 @@ func (mb *msgBlock) rebuildState() (*LostStreamData, error) {
 					truncate(index)
 					return gatherLost(lbuf - index), errBadMsg
 				}
+				copy(mb.lchk[0:], checksum)
 			}
 
 			if mb.first.seq == 0 {
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index c66357e0..b9ae1899 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -180,6 +180,11 @@ const (
 	jsAckT   = "$JS.ACK.%s.%s"
 	jsAckPre = "$JS.ACK."
 
+	// jsFlowControlT is the template for flow control messages.
+	jsFlowControlT = "$JS.FC.%s.%s.%d"
+	// jsFlowControl is for FC responses.
+	jsFlowControl = "$JS.FC.*.*.*"
+
 	// JSAdvisoryPrefix is a prefix for all JetStream advisories.
 	JSAdvisoryPrefix = "$JS.EVENT.ADVISORY"
 
diff --git a/server/signal.go b/server/signal.go
index 254e3ecb..ee6bba52 100644
--- a/server/signal.go
+++ b/server/signal.go
@@ -40,7 +40,7 @@ func (s *Server) handleSignals() {
 	}
 	c := make(chan os.Signal, 1)
 
-	signal.Notify(c, syscall.SIGINT, syscall.SIGUSR1, syscall.SIGUSR2, syscall.SIGHUP)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2, syscall.SIGHUP)
 
 	go func() {
 		for {
@@ -48,7 +48,7 @@ func (s *Server) handleSignals() {
 			case sig := <-c:
 				s.Debugf("Trapped %q signal", sig)
 				switch sig {
-				case syscall.SIGINT:
+				case syscall.SIGINT, syscall.SIGTERM:
 					s.Shutdown()
 					os.Exit(0)
 				case syscall.SIGUSR1:
diff --git a/server/stream.go b/server/stream.go
index 60f63f01..49e42df3 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -2268,6 +2268,13 @@ type jsPubMsg struct {
 	next  *jsPubMsg
 }
 
+func (pm *jsPubMsg) size() int {
+	if pm == nil {
+		return 0
+	}
+	return len(pm.subj) + len(pm.reply) + len(pm.hdr) + len(pm.msg)
+}
+
 // Forms a linked list for sending internal system messages.
 type jsOutQ struct {
 	mu   sync.Mutex


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index a017d443..854f79c3 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -7031,6 +7031,108 @@ func TestJetStreamSimpleFileRecovery(t *testing.T) {
 	}
 }
 
+func TestJetStreamPushConsumerFlowControl(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer s.Shutdown()
+
+	// Forced cleanup of all persisted state.
+	if config := s.JetStreamConfig(); config != nil {
+		defer os.RemoveAll(config.StoreDir)
+	}
+
+	// Client for API requests.
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	if _, err := js.AddStream(&nats.StreamConfig{Name: "TEST"}); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	sub, err := nc.SubscribeSync(nats.NewInbox())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer sub.Unsubscribe()
+
+	obsReq := CreateConsumerRequest{
+		Stream: "TEST",
+		Config: ConsumerConfig{
+			Durable:        "dlc",
+			DeliverSubject: sub.Subject,
+			FlowControl:    true,
+		},
+	}
+	req, err := json.Marshal(obsReq)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	resp, err := nc.Request(fmt.Sprintf(JSApiDurableCreateT, "TEST", "dlc"), req, time.Second)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	var ccResp JSApiConsumerCreateResponse
+	if err = json.Unmarshal(resp.Data, &ccResp); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if ccResp.Error != nil {
+		t.Fatalf("Unexpected error: %+v", ccResp.Error)
+	}
+
+	// Grab the low level consumer so we can manually set the fc max.
+	if mset, err := s.GlobalAccount().lookupStream("TEST"); err != nil {
+		t.Fatalf("Error looking up stream: %v", err)
+	} else if obs := mset.lookupConsumer("dlc"); obs == nil {
+		t.Fatalf("Error looking up stream: %v", err)
+	} else {
+		obs.setMaxPendingBytes(16 * 1024)
+	}
+
+	msgSize := 1024
+	msg := make([]byte, msgSize)
+	rand.Read(msg)
+
+	sendBatch := func(n int) {
+		for i := 0; i < n; i++ {
+			if _, err := js.Publish("TEST", msg); err != nil {
+				t.Fatalf("Unexpected publish error: %v", err)
+			}
+		}
+	}
+
+	checkSubPending := func(numExpected int) {
+		t.Helper()
+		checkFor(t, time.Second, 100*time.Millisecond, func() error {
+			if nmsgs, _, err := sub.Pending(); err != nil || nmsgs != numExpected {
+				return fmt.Errorf("Did not receive correct number of messages: %d vs %d", nmsgs, numExpected)
+			}
+			return nil
+		})
+	}
+
+	sendBatch(100)
+	checkSubPending(4) // First two and two flowcontrol from slow start pause.
+
+	var n int
+	for m, err := sub.NextMsg(time.Second); err == nil; m, err = sub.NextMsg(time.Second) {
+		if m.Subject == "TEST" {
+			n++
+		} else {
+			// This should be a FC control message.
+			if m.Header.Get("Status") != "100" {
+				t.Fatalf("Expected a 100 status code, got %q", m.Header.Get("Status"))
+			}
+			if m.Header.Get("Description") != "FlowControl Request" {
+				t.Fatalf("Wrong description, got %q", m.Header.Get("Description"))
+			}
+			m.Respond(nil)
+		}
+	}
+
+	if n != 100 {
+		t.Fatalf("Expected to receive all 100 messages but got %d", n)
+	}
+}
+
 func TestJetStreamPushConsumerIdleHeartbeats(t *testing.T) {
 	s := RunBasicJetStreamServer()
 	defer s.Shutdown()
@@ -7102,7 +7204,7 @@ func TestJetStreamPushConsumerIdleHeartbeats(t *testing.T) {
 		t.Fatalf("Expected a 100 status code, got %q", m.Header.Get("Status"))
 	}
 	if m.Header.Get("Description") != "Idle Heartbeat" {
-		t.Fatalf("Wrong description , got %q", m.Header.Get("Description"))
+		t.Fatalf("Wrong description, got %q", m.Header.Get("Description"))
 	}
 }