Repository: nats-io/nats-server
Pull request #3128 (2022-05-16 22:00:47)

Allow explicit consumer configuration of replica count and memory storage
Allowing more control over consumers for large number of consumer setups as we have been seeing lately.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 76f12f61..7ac787e3 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -79,6 +79,11 @@ type ConsumerConfig struct {
 	// Ephemeral inactivity threshold.
 	InactiveThreshold time.Duration `json:"inactive_threshold,omitempty"`
 
+	// Generally inherited by parent stream and other markers, now can be configured directly.
+	Replicas int `json:"num_replicas"`
+	// Force memory storage.
+	MemoryStorage bool `json:"mem_storage,omitempty"`
+
 	// Don't add to general clients.
 	Direct bool `json:"direct,omitempty"`
 }
@@ -351,6 +356,11 @@ func checkConsumerCfg(
 	isRecovering bool,
 ) *ApiError {
 
+	// Check if replicas is defined but exceeds parent stream.
+	if config.Replicas > 0 && config.Replicas > cfg.Replicas {
+		return NewJSConsumerReplicasExceedsStreamError()
+	}
+
 	// Check if we have a BackOff defined that MaxDeliver is within range etc.
 	if lbo := len(config.BackOff); lbo > 0 && config.MaxDeliver <= lbo {
 		return NewJSConsumerMaxDeliverBackoffError()
diff --git a/server/errors.json b/server/errors.json
index bf2ff9b6..1192377b 100644
--- a/server/errors.json
+++ b/server/errors.json
@@ -1238,5 +1238,15 @@
     "help": "",
     "url": "",
     "deprecates": ""
+  },
+  {
+    "constant": "JSConsumerReplicasExceedsStream",
+    "code": 400,
+    "error_code": 10126,
+    "description": "consumer config replica count exceeds parent stream",
+    "comment": "",
+    "help": "",
+    "url": "",
+    "deprecates": ""
   }
 ]
\ No newline at end of file
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 64586298..f3ed160c 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -1212,6 +1212,7 @@ func (s *Server) jsStreamCreateRequest(sub *subscription, c *client, _ *Account,
 		}
 		return
 	}
+
 	var cfg StreamConfig
 	if err := json.Unmarshal(msg, &cfg); err != nil {
 		resp.Error = NewJSInvalidJSONError()
@@ -3250,6 +3251,13 @@ func (s *Server) jsConsumerCreate(sub *subscription, c *client, a *Account, subj
 		return
 	}
 
+	// If we are here we are single server mode.
+	if req.Config.Replicas > 1 {
+		resp.Error = NewJSStreamReplicasNotSupportedError()
+		s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+		return
+	}
+
 	stream, err := acc.lookupStream(req.Stream)
 	if err != nil {
 		resp.Error = NewJSStreamNotFoundError(Unless(err))
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index 841e2c5c..689a3b9f 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -1875,10 +1875,11 @@ func (js *jetStream) monitorStream(mset *stream, sa *streamAssignment, sendSnaps
 			// These are not currently assigned so we will need to do so here.
 			if consumers := mset.getPublicConsumers(); len(consumers) > 0 {
 				for _, o := range consumers {
-					rg := cc.createGroupForConsumer(sa)
+					name, cfg := o.String(), o.config()
+					rg := cc.createGroupForConsumer(&cfg, sa)
 					// Pick a preferred leader.
 					rg.setPreferred()
-					name, cfg := o.String(), o.config()
+
 					// Place our initial state here as well for assignment distribution.
 					state, _ := o.store.State()
 					ca := &consumerAssignment{
@@ -2808,8 +2809,9 @@ func (js *jetStream) processClusterCreateStream(acc *Account, sa *streamAssignme
 						js.mu.RUnlock()
 
 						for _, o := range consumers {
-							rg := cc.createGroupForConsumer(sa)
 							name, cfg := o.String(), o.config()
+							rg := cc.createGroupForConsumer(&cfg, sa)
+
 							// Place our initial state here as well for assignment distribution.
 							ca := &consumerAssignment{
 								Group:   rg,
@@ -3139,7 +3141,11 @@ func (js *jetStream) processClusterCreateConsumer(ca *consumerAssignment, state
 
 	if !alreadyRunning {
 		// Process the raft group and make sure its running if needed.
-		js.createRaftGroup(acc.GetName(), rg, mset.config().Storage)
+		storage := mset.config().Storage
+		if ca.Config.MemoryStorage {
+			storage = MemoryStorage
+		}
+		js.createRaftGroup(acc.GetName(), rg, storage)
 	} else {
 		// If we are clustered update the known peers.
 		js.mu.RLock()
@@ -5265,13 +5271,21 @@ func decodeStreamAssignment(buf []byte) (*streamAssignment, error) {
 	return &sa, err
 }
 
-// createGroupForConsumer will create a new group with same peer set as the stream.
-func (cc *jetStreamCluster) createGroupForConsumer(sa *streamAssignment) *raftGroup {
-	peers := sa.Group.Peers
+// createGroupForConsumer will create a new group from same peer set as the stream.
+func (cc *jetStreamCluster) createGroupForConsumer(cfg *ConsumerConfig, sa *streamAssignment) *raftGroup {
+	peers := copyStrings(sa.Group.Peers)
 	if len(peers) == 0 {
 		return nil
 	}
-	return &raftGroup{Name: groupNameForConsumer(peers, sa.Config.Storage), Storage: sa.Config.Storage, Peers: peers}
+	if cfg.Replicas > 0 && cfg.Replicas != len(peers) {
+		rand.Shuffle(len(peers), func(i, j int) { peers[i], peers[j] = peers[j], peers[i] })
+		peers = peers[:cfg.Replicas]
+	}
+	storage := sa.Config.Storage
+	if cfg.MemoryStorage {
+		storage = MemoryStorage
+	}
+	return &raftGroup{Name: groupNameForConsumer(peers, storage), Storage: storage, Peers: peers}
 }
 
 // jsClusteredConsumerRequest is first point of entry to create a consumer with R > 1.
@@ -5372,7 +5386,7 @@ func (s *Server) jsClusteredConsumerRequest(ci *ClientInfo, acc *Account, subjec
 
 	// If this is new consumer.
 	if ca == nil {
-		rg := cc.createGroupForConsumer(sa)
+		rg := cc.createGroupForConsumer(cfg, sa)
 		if rg == nil {
 			resp.Error = NewJSInsufficientResourcesError()
 			s.sendAPIErrResponse(ci, acc, subject, reply, string(rmsg), s.jsonResponse(&resp))
diff --git a/server/jetstream_errors_generated.go b/server/jetstream_errors_generated.go
index 26114931..be507d1a 100644
--- a/server/jetstream_errors_generated.go
+++ b/server/jetstream_errors_generated.go
@@ -155,6 +155,9 @@ const (
 	// JSConsumerReplacementWithDifferentNameErr consumer replacement durable config not the same
 	JSConsumerReplacementWithDifferentNameErr ErrorIdentifier = 10106
 
+	// JSConsumerReplicasExceedsStream consumer config replica count exceeds parent stream
+	JSConsumerReplicasExceedsStream ErrorIdentifier = 10126
+
 	// JSConsumerSmallHeartbeatErr consumer idle heartbeat needs to be >= 100ms
 	JSConsumerSmallHeartbeatErr ErrorIdentifier = 10083
 
@@ -430,6 +433,7 @@ var (
 		JSConsumerPullWithRateLimitErr:             {Code: 400, ErrCode: 10086, Description: "consumer in pull mode can not have rate limit set"},
 		JSConsumerPushMaxWaitingErr:                {Code: 400, ErrCode: 10080, Description: "consumer in push mode can not set max waiting"},
 		JSConsumerReplacementWithDifferentNameErr:  {Code: 400, ErrCode: 10106, Description: "consumer replacement durable config not the same"},
+		JSConsumerReplicasExceedsStream:            {Code: 400, ErrCode: 10126, Description: "consumer config replica count exceeds parent stream"},
 		JSConsumerSmallHeartbeatErr:                {Code: 400, ErrCode: 10083, Description: "consumer idle heartbeat needs to be >= 100ms"},
 		JSConsumerStoreFailedErrF:                  {Code: 500, ErrCode: 10104, Description: "error creating store for consumer: {err}"},
 		JSConsumerWQConsumerNotDeliverAllErr:       {Code: 400, ErrCode: 10101, Description: "consumer must be deliver all on workqueue stream"},
@@ -1071,6 +1075,16 @@ func NewJSConsumerReplacementWithDifferentNameError(opts ...ErrorOption) *ApiErr
 	return ApiErrors[JSConsumerReplacementWithDifferentNameErr]
 }
 
+// NewJSConsumerReplicasExceedsStreamError creates a new JSConsumerReplicasExceedsStream error: "consumer config replica count exceeds parent stream"
+func NewJSConsumerReplicasExceedsStreamError(opts ...ErrorOption) *ApiError {
+	eopts := parseOpts(opts)
+	if ae, ok := eopts.err.(*ApiError); ok {
+		return ae
+	}
+
+	return ApiErrors[JSConsumerReplicasExceedsStream]
+}
+
 // NewJSConsumerSmallHeartbeatError creates a new JSConsumerSmallHeartbeatErr error: "consumer idle heartbeat needs to be >= 100ms"
 func NewJSConsumerSmallHeartbeatError(opts ...ErrorOption) *ApiError {
 	eopts := parseOpts(opts)


Test patch:
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index 521b85d2..78464438 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -10528,3 +10528,89 @@ func TestJetStreamClusterNewHealthz(t *testing.T) {
 	sl = c.restartServer(sl)
 	c.waitOnServerHealthz(sl)
 }
+
+func TestJetStreamClusterConsumerOverrides(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "JSC", 3)
+	defer c.shutdown()
+
+	nc, js := jsClientConnect(t, c.randomServer())
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo"},
+		Replicas: 3,
+	})
+	require_NoError(t, err)
+
+	// Test replica override.
+	// Make sure we can not go "wider" than the parent stream.
+	ccReq := CreateConsumerRequest{
+		Stream: "TEST",
+		Config: ConsumerConfig{
+			Durable:   "d",
+			AckPolicy: AckExplicit,
+			Replicas:  5,
+		},
+	}
+	req, err := json.Marshal(ccReq)
+	require_NoError(t, err)
+
+	ci, err := nc.Request(fmt.Sprintf(JSApiDurableCreateT, "TEST", "d"), req, time.Second)
+	require_NoError(t, err)
+
+	var resp JSApiConsumerCreateResponse
+	err = json.Unmarshal(ci.Data, &resp)
+	require_NoError(t, err)
+
+	if resp.Error == nil || !IsNatsErr(resp.Error, JSConsumerReplicasExceedsStream) {
+		t.Fatalf("Expected an error when replicas > parent stream, got %+v", resp.Error)
+	}
+
+	// Durables inherit the replica count from the stream, so make sure we can override that.
+	ccReq.Config.Replicas = 1
+	req, err = json.Marshal(ccReq)
+	require_NoError(t, err)
+
+	ci, err = nc.Request(fmt.Sprintf(JSApiDurableCreateT, "TEST", "d"), req, time.Second)
+	require_NoError(t, err)
+
+	resp.Error = nil
+	err = json.Unmarshal(ci.Data, &resp)
+	require_NoError(t, err)
+	require_True(t, resp.Error == nil)
+
+	checkCount := func(durable string, expected int) {
+		t.Helper()
+		count := 0
+		for _, s := range c.servers {
+			if mset, err := s.GlobalAccount().lookupStream("TEST"); err == nil {
+				if o := mset.lookupConsumer(durable); o != nil {
+					count++
+				}
+			}
+		}
+		if count != expected {
+			t.Fatalf("Expected %d consumers in cluster, got %d", expected, count)
+		}
+	}
+	checkCount("d", 1)
+
+	// Now override storage and force storage to memory based.
+	ccReq.Config.MemoryStorage = true
+	ccReq.Config.Durable = "m"
+	ccReq.Config.Replicas = 2
+
+	req, err = json.Marshal(ccReq)
+	require_NoError(t, err)
+
+	ci, err = nc.Request(fmt.Sprintf(JSApiDurableCreateT, "TEST", "m"), req, time.Second)
+	require_NoError(t, err)
+
+	resp.Error = nil
+	err = json.Unmarshal(ci.Data, &resp)
+	require_NoError(t, err)
+	require_True(t, resp.Error == nil)
+
+	checkCount("m", 2)
+}