Repository: nats-io/nats-server
Pull request #3129 (2022-05-17 20:20:03)

Enable republishing of messages once stored in a stream.
This enables lightweight distribution of messages to very large number of NATS subscribers.
We add in metadata as headers that allows for gap detection which enables initial value (via JetStream, maybe KV) and real-time NATS core sub updates but all being globally ordered and gap detection allowed, even for a subset of messages.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index f3ed160c..a60f61b2 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -1234,6 +1234,41 @@ func (s *Server) jsStreamCreateRequest(sub *subscription, c *client, _ *Account,
 		return
 	}
 
+	// If we have a republish directive check if we can create a transform here.
+	if cfg.RePublish != nil {
+		// Check to make sure source is a valid subset of the subjects we have.
+		// Also make sure it does not form a cycle.
+		var srcValid bool
+		for _, subj := range cfg.Subjects {
+			if SubjectsCollide(cfg.RePublish.Source, subj) {
+				srcValid = true
+				break
+			}
+		}
+		if !srcValid {
+			resp.Error = NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish source is not valid subset of subjects"))
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
+		var formsCycle bool
+		for _, subj := range cfg.Subjects {
+			if SubjectsCollide(cfg.RePublish.Destination, subj) {
+				formsCycle = true
+				break
+			}
+		}
+		if formsCycle {
+			resp.Error = NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish destination forms a cycle"))
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
+		if _, err := newTransform(cfg.RePublish.Source, cfg.RePublish.Destination); err != nil {
+			resp.Error = NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish not valid"))
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+			return
+		}
+	}
+
 	// Hand off to cluster for processing.
 	if s.JetStreamIsClustered() {
 		s.jsClusteredStreamRequest(ci, acc, subject, reply, rmsg, &cfg)
diff --git a/server/stream.go b/server/stream.go
index 5666c665..7efdb58b 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -58,6 +58,9 @@ type StreamConfig struct {
 	Mirror       *StreamSource   `json:"mirror,omitempty"`
 	Sources      []*StreamSource `json:"sources,omitempty"`
 
+	// Allow republish of the message after being sequenced and stored.
+	RePublish *SubjectMapping `json:"republish,omitempty"`
+
 	// Optional qualifiers. These can not be modified after set to true.
 
 	// Sealed will seal a stream so no messages can get out or in.
@@ -71,6 +74,12 @@ type StreamConfig struct {
 	AllowRollup bool `json:"allow_rollup_hdrs"`
 }
 
+// SubjectMapping allows a source subject to be mapped to a destination subject for republishing.
+type SubjectMapping struct {
+	Source      string `json:"src,omitempty"`
+	Destination string `json:"dest"`
+}
+
 // JSPubAckResponse is a formal response to a publish operation.
 type JSPubAckResponse struct {
 	Error *ApiError `json:"error,omitempty"`
@@ -202,6 +211,9 @@ type stream struct {
 	// Indicates we have direct consumers.
 	directs int
 
+	// For republishing.
+	tr *transform
+
 	// TODO(dlc) - Hide everything below behind two pointers.
 	// Clustered mode.
 	sa       *streamAssignment
@@ -252,6 +264,13 @@ const (
 	JSResponseType        = "Nats-Response-Type"
 )
 
+// Headers for republished messages.
+const (
+	JSStream       = "Nats-Stream"
+	JSSequence     = "Nats-Sequence"
+	JSLastSequence = "Nats-Last-Sequence"
+)
+
 // Rollups, can be subject only or all messages.
 const (
 	JSMsgRollupSubject = "sub"
@@ -396,6 +415,17 @@ func (a *Account) addStreamWithAssignment(config *StreamConfig, fsConfig *FileSt
 		mset.ackq = s.newIPQueue(qpfx + "acks") // of uint64
 	}
 
+	// Check for RePublish.
+	if cfg.RePublish != nil {
+		tr, err := newTransform(cfg.RePublish.Source, cfg.RePublish.Destination)
+		if err != nil {
+			jsa.mu.Unlock()
+			return nil, fmt.Errorf("stream configuration for republish not valid")
+		}
+		// Assign our transform for republishing.
+		mset.tr = tr
+	}
+
 	jsa.streams[cfg.Name] = mset
 	storeDir := filepath.Join(jsa.storeDir, streamsDir, cfg.Name)
 	jsa.mu.Unlock()
@@ -920,7 +950,8 @@ func (s *Server) checkStreamCfg(config *StreamConfig, acc *Account) (StreamConfi
 		return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("duplicates window can not be larger then max age"))
 	}
 	if lim.Duplicates > 0 && cfg.Duplicates > lim.Duplicates {
-		return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("duplicates window can not be larger then server limit of %v", lim.Duplicates.String()))
+		return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("duplicates window can not be larger then server limit of %v",
+			lim.Duplicates.String()))
 	}
 	if cfg.Duplicates > 0 && cfg.Duplicates < 100*time.Millisecond {
 		return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("duplicates window needs to be >= 100ms"))
@@ -1108,6 +1139,35 @@ func (s *Server) checkStreamCfg(config *StreamConfig, acc *Account) (StreamConfi
 		return StreamConfig{}, NewJSStreamMaxStreamBytesExceededError()
 	}
 
+	// If we have a republish directive check if we can create a transform here.
+	if cfg.RePublish != nil {
+		// Check to make sure source is a valid subset of the subjects we have.
+		// Also make sure it does not form a cycle.
+		var srcValid bool
+		for _, subj := range cfg.Subjects {
+			if SubjectsCollide(cfg.RePublish.Source, subj) {
+				srcValid = true
+				break
+			}
+		}
+		if !srcValid {
+			return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish source is not valid subset of subjects"))
+		}
+		var formsCycle bool
+		for _, subj := range cfg.Subjects {
+			if SubjectsCollide(cfg.RePublish.Destination, subj) {
+				formsCycle = true
+				break
+			}
+		}
+		if formsCycle {
+			return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish destination forms a cycle"))
+		}
+		if _, err := newTransform(cfg.RePublish.Source, cfg.RePublish.Destination); err != nil {
+			return StreamConfig{}, NewJSStreamInvalidConfigError(fmt.Errorf("stream configuration for republish not valid"))
+		}
+	}
+
 	return cfg, nil
 }
 
@@ -3406,11 +3466,28 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	clfs := mset.clfs
 	mset.lseq++
 	tierName := mset.tier
+
+	// Republish state if needed.
+	var tsubj string
+	var tlseq uint64
+	if mset.tr != nil {
+		tsubj, _ = mset.tr.transformSubject(subject)
+	}
+	republish := tsubj != _EMPTY_
+
 	// We hold the lock to this point to make sure nothing gets between us since we check for pre-conditions.
 	// Currently can not hold while calling store b/c we have inline storage update calls that may need the lock.
 	// Note that upstream that sets seq/ts should be serialized as much as possible.
 	mset.mu.Unlock()
 
+	// If we are republishing grab last sequence for this exact subject. Aids in gap detection for lightweight clients.
+	if republish {
+		var smv StoreMsg
+		if sm, _ := store.LoadLastMsg(subject, &smv); sm != nil {
+			tlseq = sm.seq
+		}
+	}
+
 	// Store actual msg.
 	if lseq == 0 && ts == 0 {
 		seq, ts, err = store.StoreMsg(subject, hdr, msg)
@@ -3478,6 +3555,13 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 			response = append(pubAck, strconv.FormatUint(seq, 10)...)
 			response = append(response, '}')
 		}
+		// Check for republish.
+		if republish {
+			hdr = genHeader(hdr, JSStream, name)
+			hdr = genHeader(hdr, JSSequence, strconv.FormatUint(seq, 10))
+			hdr = genHeader(hdr, JSLastSequence, strconv.FormatUint(tlseq, 10))
+			mset.outq.send(newJSPubMsg(tsubj, subject, _EMPTY_, copyBytes(hdr), copyBytes(msg), nil, seq))
+		}
 	}
 
 	// Send response here.
diff --git a/server/util.go b/server/util.go
index 24cef1ac..c7f5827e 100644
--- a/server/util.go
+++ b/server/util.go
@@ -314,9 +314,9 @@ func getURLsAsString(urls []*url.URL) []string {
 }
 
 // copyBytes make a new slice of the same size than `src` and copy its content.
-// If `src` is nil, then this returns `nil`
+// If `src` is nil or its length is 0, then this returns `nil`
 func copyBytes(src []byte) []byte {
-	if src == nil {
+	if len(src) == 0 {
 		return nil
 	}
 	dst := make([]byte, len(src))


Test patch:
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index 78464438..4de1e8f9 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -10614,3 +10614,67 @@ func TestJetStreamClusterConsumerOverrides(t *testing.T) {
 
 	checkCount("m", 2)
 }
+
+func TestJetStreamClusterStreamRepublish(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "JSC", 3)
+	defer c.shutdown()
+
+	nc, js := jsClientConnect(t, c.randomServer())
+	defer nc.Close()
+
+	// Do by hand for now.
+	cfg := &StreamConfig{
+		Name:     "RP",
+		Storage:  MemoryStorage,
+		Subjects: []string{"foo", "bar", "baz"},
+		Replicas: 1,
+		RePublish: &SubjectMapping{
+			Source:      ">",
+			Destination: "RP.>",
+		},
+	}
+
+	addStream(t, nc, cfg)
+
+	sub, err := nc.SubscribeSync("RP.>")
+	require_NoError(t, err)
+
+	msg, toSend := []byte("OK TO REPUBLISH?"), 100
+	for i := 0; i < toSend; i++ {
+		_, err = js.PublishAsync("foo", msg)
+		require_NoError(t, err)
+		_, err = js.PublishAsync("bar", msg)
+		require_NoError(t, err)
+		_, err = js.PublishAsync("baz", msg)
+		require_NoError(t, err)
+	}
+	select {
+	case <-js.PublishAsyncComplete():
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Did not receive completion signal")
+	}
+
+	checkSubsPending(t, sub, toSend*3)
+
+	lseq := map[string]int{
+		"foo": 0,
+		"bar": 0,
+		"baz": 0,
+	}
+
+	for i := 1; i <= toSend; i++ {
+		m, err := sub.NextMsg(time.Second)
+		require_NoError(t, err)
+		// Grab info from Header
+		require_True(t, m.Header.Get(JSStream) == "RP")
+		// Make sure sequence is correct.
+		seq, err := strconv.Atoi(m.Header.Get(JSSequence))
+		require_NoError(t, err)
+		require_True(t, seq == i)
+		// Make sure last sequence matches last seq we received on this subject.
+		last, err := strconv.Atoi(m.Header.Get(JSLastSequence))
+		require_NoError(t, err)
+		require_True(t, last == lseq[m.Subject])
+		lseq[m.Subject] = seq
+	}
+}
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 7084a80d..7a3c8023 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -17484,6 +17484,62 @@ func TestJetStreamPullMaxBytes(t *testing.T) {
 	checkSubsPending(t, sub, 0)
 }
 
+func TestJetStreamStreamRepublishCycle(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	if config := s.JetStreamConfig(); config != nil {
+		defer removeDir(t, config.StoreDir)
+	}
+	defer s.Shutdown()
+
+	nc, _ := jsClientConnect(t, s)
+	defer nc.Close()
+
+	// Do by hand for now.
+	cfg := &StreamConfig{
+		Name:     "RPC",
+		Storage:  MemoryStorage,
+		Subjects: []string{"foo.>", "bar.*", "baz"},
+	}
+
+	expectFail := func() {
+		t.Helper()
+		req, err := json.Marshal(cfg)
+		require_NoError(t, err)
+		rmsg, err := nc.Request(fmt.Sprintf(JSApiStreamCreateT, cfg.Name), req, time.Second)
+		require_NoError(t, err)
+		var resp JSApiStreamCreateResponse
+		err = json.Unmarshal(rmsg.Data, &resp)
+		require_NoError(t, err)
+		if resp.Type != JSApiStreamCreateResponseType {
+			t.Fatalf("Invalid response type %s expected %s", resp.Type, JSApiStreamCreateResponseType)
+		}
+		if resp.Error == nil {
+			t.Fatalf("Expected error but got none")
+		}
+		if !strings.Contains(resp.Error.Description, "republish destination forms a cycle") {
+			t.Fatalf("Expected cycle error, got %q", resp.Error.Description)
+		}
+	}
+
+	cfg.RePublish = &SubjectMapping{
+		Source:      "foo.>",
+		Destination: "foo.>",
+	}
+	expectFail()
+
+	cfg.RePublish = &SubjectMapping{
+		Source:      "bar.bar",
+		Destination: "foo.bar",
+	}
+	expectFail()
+
+	cfg.RePublish = &SubjectMapping{
+		Source:      "baz",
+		Destination: "bar.bar",
+	}
+	expectFail()
+}
+
 ///////////////////////////////////////////////////////////////////////////
 // Simple JetStream Benchmarks
 ///////////////////////////////////////////////////////////////////////////