Repository: nats-io/nats-server (go)

Issue:
Allow publishers to specify a different subject for "Nats-Expected-Last-Subject-Sequence"
### Proposed change

I'd like to request a change to allow publishers to provide a specific subject for use when nats-server enforced "Nats-Expected-Last-Subject-Sequence". Adding an additional header field that allows publishers to optionally change the subject used for "Nats-Expected-Last-Subject-Sequence" from the subject of the message being published to to a subject provided in the header.

The proposal is for a "Nats-Expected-Last-Subject-Sequence-Subject" header that when set with a string value becomes the subject used during enforcement of the last sequence. The header does not necessarily need to be that string exactly, the functionality is the actual proposal.

### Use case

The specific use case I am trying to solve here is to allow publishing to complex and dynamic subjects while still enabling optimistic concurrency control.

For example in an event sourcing application an aggregate might be defined against the subject:

- **orders.[order_id]** 
  - orders.1
  - orders.2
  - orders.3
  - ...

Today we can use "Nats-Expected-Last-Subject-Sequence" to ensure that when we publish to `orders.2` we are publishing against the known last sequence for messages on the `orders.2` subject.

However, it is often desirable to have more complex subject patterns. For example image a pattern that looks like this:

- **orders.[order_id].[event_type].[visibility]**
  - orders.1.created.public
  - orders.1.approved.private
  - orders.2.created.public
  - ...

Today we cannot use "Nats-Expected-Last-Subject-Sequence" with this pattern and still have OCC for the aggregate `orders.2` when publishing. If a publisher attempts to set the expected sequence and then publish to `order.2.rejected.private` the sequence for that subject is likely 0 as opposed to sequence for the last event published to `order.2.>`

By allowing publishers to override and customize the subject used in the expect last subject sequence check OCC can be enforced for the aggregate itself versus the subject being published to. For example by setting the proposed "Nats-Expected-Last-Subject-Sequence-Subject" to "order.2.>"

Relevant Slack Conversations:

- https://natsio.slack.com/archives/C055AB5M1KM/p1710872807434689
- https://natsio.slack.com/archives/CM3T6T7JQ/p1687313617160579

### Contribution

Possibly

Pull request description:
[ADDED] Publish header "Nats-Expected-Last-Subject-Sequence-Subject"
This change adds a new header
"Nats-Expected-Last-Subject-Sequence-Subject" when when paired with "Nats-Expected-Last-Subject-Sequence" allows publishers to customize the subject used when the server enforces
"Nats-Expected-Last-Subject-Sequence". Publishers can specify a alternative subject to be used that includes wildcards.

Resolves #5280

Signed-off-by: Caleb Champlin <caleb.champlin@gmail.com>


Patch:
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index a59ae021..4eaa09e1 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -7853,9 +7853,15 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 		// Expected last sequence per subject.
 		// We can check for last sequence per subject but only if the expected seq <= lseq.
 		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists && store != nil && seq <= lseq {
+			// Allow override of the subject used for the check.
+			seqSubj := subject
+			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
+				seqSubj = optSubj
+			}
+
 			var smv StoreMsg
 			var fseq uint64
-			sm, err := store.LoadLastMsg(subject, &smv)
+			sm, err := store.LoadLastMsg(seqSubj, &smv)
 			if sm != nil {
 				fseq = sm.seq
 			}
diff --git a/server/stream.go b/server/stream.go
index 36296993..1230203a 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -354,18 +354,19 @@ const (
 
 // Headers for published messages.
 const (
-	JSMsgId               = "Nats-Msg-Id"
-	JSExpectedStream      = "Nats-Expected-Stream"
-	JSExpectedLastSeq     = "Nats-Expected-Last-Sequence"
-	JSExpectedLastSubjSeq = "Nats-Expected-Last-Subject-Sequence"
-	JSExpectedLastMsgId   = "Nats-Expected-Last-Msg-Id"
-	JSStreamSource        = "Nats-Stream-Source"
-	JSLastConsumerSeq     = "Nats-Last-Consumer"
-	JSLastStreamSeq       = "Nats-Last-Stream"
-	JSConsumerStalled     = "Nats-Consumer-Stalled"
-	JSMsgRollup           = "Nats-Rollup"
-	JSMsgSize             = "Nats-Msg-Size"
-	JSResponseType        = "Nats-Response-Type"
+	JSMsgId                   = "Nats-Msg-Id"
+	JSExpectedStream          = "Nats-Expected-Stream"
+	JSExpectedLastSeq         = "Nats-Expected-Last-Sequence"
+	JSExpectedLastSubjSeq     = "Nats-Expected-Last-Subject-Sequence"
+	JSExpectedLastSubjSeqSubj = "Nats-Expected-Last-Subject-Sequence-Subject"
+	JSExpectedLastMsgId       = "Nats-Expected-Last-Msg-Id"
+	JSStreamSource            = "Nats-Stream-Source"
+	JSLastConsumerSeq         = "Nats-Last-Consumer"
+	JSLastStreamSeq           = "Nats-Last-Stream"
+	JSConsumerStalled         = "Nats-Consumer-Stalled"
+	JSMsgRollup               = "Nats-Rollup"
+	JSMsgSize                 = "Nats-Msg-Size"
+	JSResponseType            = "Nats-Response-Type"
 )
 
 // Headers for republished messages and direct gets.
@@ -4114,6 +4115,11 @@ func getExpectedLastSeqPerSubject(hdr []byte) (uint64, bool) {
 	return uint64(parseInt64(bseq)), true
 }
 
+// Fast lookup of expected subject for the expected stream sequence per subject.
+func getExpectedLastSeqPerSubjectForSubject(hdr []byte) string {
+	return string(getHeader(JSExpectedLastSubjSeqSubj, hdr))
+}
+
 // Signal if we are clustered. Will acquire rlock.
 func (mset *stream) IsClustered() bool {
 	mset.mu.RLock()
@@ -4680,10 +4686,16 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 
 		// Expected last sequence per subject.
 		if seq, exists := getExpectedLastSeqPerSubject(hdr); exists {
+			// Allow override of the subject used for the check.
+			seqSubj := subject
+			if optSubj := getExpectedLastSeqPerSubjectForSubject(hdr); optSubj != _EMPTY_ {
+				seqSubj = optSubj
+			}
+
 			// TODO(dlc) - We could make a new store func that does this all in one.
 			var smv StoreMsg
 			var fseq uint64
-			sm, err := store.LoadLastMsg(subject, &smv)
+			sm, err := store.LoadLastMsg(seqSubj, &smv)
 			if sm != nil {
 				fseq = sm.seq
 			}


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index 8bb40f65..875ed38b 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -11270,6 +11270,100 @@ func TestJetStreamLastSequenceBySubject(t *testing.T) {
 	}
 }
 
+func TestJetStreamLastSequenceBySubjectWithSubject(t *testing.T) {
+	for _, st := range []StorageType{FileStorage, MemoryStorage} {
+		t.Run(st.String(), func(t *testing.T) {
+			c := createJetStreamClusterExplicit(t, "JSC", 3)
+			defer c.shutdown()
+
+			nc, js := jsClientConnect(t, c.randomServer())
+			defer nc.Close()
+
+			cfg := StreamConfig{
+				Name:       "KV",
+				Subjects:   []string{"kv.>"},
+				Storage:    st,
+				Replicas:   3,
+				MaxMsgsPer: 1,
+			}
+
+			req, err := json.Marshal(cfg)
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			// Do manually for now.
+			m, err := nc.Request(fmt.Sprintf(JSApiStreamCreateT, cfg.Name), req, time.Second)
+			require_NoError(t, err)
+			si, err := js.StreamInfo("KV")
+			if err != nil {
+				t.Fatalf("Unexpected error: %v, respmsg: %q", err, string(m.Data))
+			}
+			if si == nil || si.Config.Name != "KV" {
+				t.Fatalf("StreamInfo is not correct %+v", si)
+			}
+
+			js.PublishAsync("kv.1.foo", []byte("1:1")) // Last is 1 for kv.1.foo; 1 for kv.1.*;
+			js.PublishAsync("kv.1.bar", []byte("1:2")) // Last is 2 for kv.1.bar; 2 for kv.1.*;
+			js.PublishAsync("kv.2.foo", []byte("2:1")) // Last is 3 for kv.2.foo; 3 for kv.2.*;
+			js.PublishAsync("kv.3.bar", []byte("3:1")) // Last is 4 for kv.3.bar; 4 for kv.3.*;
+			js.PublishAsync("kv.1.baz", []byte("1:3")) // Last is 5 for kv.1.baz; 5 for kv.1.*;
+			js.PublishAsync("kv.1.bar", []byte("1:4")) // Last is 6 for kv.1.baz; 6 for kv.1.*;
+			js.PublishAsync("kv.2.baz", []byte("2:2")) // Last is 7 for kv.2.baz; 7 for kv.2.*;
+
+			select {
+			case <-js.PublishAsyncComplete():
+			case <-time.After(time.Second):
+				t.Fatalf("Did not receive completion signal")
+			}
+
+			// Now make sure we get an error if the last sequence is not correct per subject.
+			pubAndCheck := func(subj, filterSubject, seq string, ok bool) {
+				t.Helper()
+				m := nats.NewMsg(subj)
+				m.Data = []byte("HELLO")
+
+				// Expect last to be seq.
+				m.Header.Set(JSExpectedLastSubjSeq, seq)
+
+				// Constrain the sequence restriction to a specific subject
+				// e.g. "kv.1.*" for kv.1.foo, kv.1.bar, kv.1.baz; kv.2.* for kv.2.foo, kv.2.baz; kv.3.* for kv.3.bar
+				m.Header.Set(JSExpectedLastSubjSeqSubj, filterSubject)
+				_, err := js.PublishMsg(m)
+				if ok && err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				if !ok && err == nil {
+					t.Fatalf("Expected to get an error and got none")
+				}
+			}
+
+			pubAndCheck("kv.1.foo", "kv.1.*", "0", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "0", false)
+			pubAndCheck("kv.1.xxx", "kv.1.*", "0", false)
+			pubAndCheck("kv.1.foo", "kv.1.*", "1", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "1", false)
+			pubAndCheck("kv.1.xxx", "kv.1.*", "1", false)
+			pubAndCheck("kv.2.foo", "kv.2.*", "1", false)
+			pubAndCheck("kv.2.bar", "kv.2.*", "1", false)
+			pubAndCheck("kv.2.xxx", "kv.2.*", "1", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "2", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "3", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "4", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "5", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "6", true) // Last is 8 for kv.1.bar; 8 for kv.1.*;
+			pubAndCheck("kv.1.baz", "kv.1.*", "2", false)
+			pubAndCheck("kv.1.bar", "kv.1.*", "7", false)
+			pubAndCheck("kv.1.xxx", "kv.1.*", "8", true) // Last is 9 for kv.1.xxx; 9 for kv.1.*;
+			pubAndCheck("kv.2.foo", "kv.2.*", "2", false)
+			pubAndCheck("kv.2.foo", "kv.2.*", "7", true)  // Last is 10 for kv.2.foo; 10 for kv.2.*;
+			pubAndCheck("kv.xxx", "kv.*", "0", true)      // Last is 0 for kv.xxx; 0 for kv.*;
+			pubAndCheck("kv.xxx", "kv.*.*", "0", false)   // Last is 11 for kv.xxx; 11 for kv.*.*;
+			pubAndCheck("kv.3.xxx", "kv.3.*", "4", true)  // Last is 12 for kv.3.xxx; 12 for kv.3.*;
+			pubAndCheck("kv.3.xyz", "kv.3.*", "12", true) // Last is 13 for kv.3.xyz; 13 for kv.3.*;
+		})
+	}
+}
+
 func TestJetStreamFilteredConsumersWithWiderFilter(t *testing.T) {
 	s := RunBasicJetStreamServer(t)
 	defer s.Shutdown()