Repository: nats-io/nats.go
Pull request #1303 (2023-06-08 11:23:53)

[IMPROVED] Lock handling for pull consumers
This PRs simplifies mutex usage in pull consumer JetStream API


Patch:
diff --git a/jetstream/errors.go b/jetstream/errors.go
index 77ddf3d..7125c12 100644
--- a/jetstream/errors.go
+++ b/jetstream/errors.go
@@ -109,6 +109,9 @@ var (
 	// ErrInvalidStreamName is returned when the provided stream name is invalid (contains '.').
 	ErrInvalidStreamName JetStreamError = &jsError{message: "invalid stream name"}
 
+	// ErrInvalidSubject is returned when the provided subject name is invalid.
+	ErrInvalidSubject JetStreamError = &jsError{message: "invalid subject name"}
+
 	// ErrInvalidConsumerName is returned when the provided consumer name is invalid (contains '.').
 	ErrInvalidConsumerName JetStreamError = &jsError{message: "invalid consumer name"}
 
diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go
index ed68d25..563699a 100644
--- a/jetstream/jetstream.go
+++ b/jetstream/jetstream.go
@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"regexp"
 	"strings"
 
 	"github.com/nats-io/nats.go"
@@ -68,6 +69,8 @@ type (
 		UpdateStream(context.Context, StreamConfig) (Stream, error)
 		// Stream returns a [Stream] hook for a given stream name
 		Stream(context.Context, string) (Stream, error)
+		// StreamNameBySubject returns a stream name stream listening on given subject
+		StreamNameBySubject(context.Context, string) (string, error)
 		// DeleteStream removes a stream with given name
 		DeleteStream(context.Context, string) error
 		// ListStreams returns StreamInfoLister enabling iterating over a channel of stream infos
@@ -186,8 +189,14 @@ type (
 		apiPaged
 		Streams []string `json:"streams"`
 	}
+
+	streamsRequest struct {
+		Subject string `json:"subject,omitempty"`
+	}
 )
 
+var subjectRegexp = regexp.MustCompile(`^[^ >]*[>]?$`)
+
 // New returns a enw JetStream instance
 //
 // Available options:
@@ -494,6 +503,16 @@ func validateStreamName(stream string) error {
 	return nil
 }
 
+func validateSubject(subject string) error {
+	if subject == "" {
+		return fmt.Errorf("%w: %s", ErrInvalidSubject, "subject cannot be empty")
+	}
+	if !subjectRegexp.MatchString(subject) {
+		return fmt.Errorf("%w: %s", ErrInvalidSubject, subject)
+	}
+	return nil
+}
+
 func (js *jetStream) AccountInfo(ctx context.Context) (*AccountInfo, error) {
 	var resp accountInfoResponse
 
@@ -591,6 +610,32 @@ func (js *jetStream) StreamNames(ctx context.Context) StreamNameLister {
 	return l
 }
 
+func (js *jetStream) StreamNameBySubject(ctx context.Context, subject string) (string, error) {
+	if err := validateSubject(subject); err != nil {
+		return "", err
+	}
+	streamsSubject := apiSubj(js.apiPrefix, apiStreams)
+
+	r := &streamsRequest{Subject: subject}
+	req, err := json.Marshal(r)
+	if err != nil {
+		return "", err
+	}
+	var resp streamNamesResponse
+	_, err = js.apiRequestJSON(ctx, streamsSubject, &resp, req)
+	if err != nil {
+		return "", err
+	}
+	if resp.Error != nil {
+		return "", resp.Error
+	}
+	if len(resp.Streams) == 0 {
+		return "", ErrStreamNotFound
+	}
+
+	return resp.Streams[0], nil
+}
+
 // Name returns a channel allowing retrieval of stream names returned by [StreamNames]
 func (s *streamLister) Name() <-chan string {
 	return s.names
diff --git a/jetstream/pull.go b/jetstream/pull.go
index b9cef00..16ed9ca 100644
--- a/jetstream/pull.go
+++ b/jetstream/pull.go
@@ -186,22 +186,8 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 		}
 		defer func() {
 			sub.Lock()
-			pending := sub.pending
+			sub.checkPending()
 			sub.Unlock()
-
-			if pending.msgCount < consumeOpts.ThresholdMessages ||
-				(pending.byteCount < consumeOpts.ThresholdBytes && sub.consumeOpts.MaxBytes != 0) &&
-					atomic.LoadUint32(&sub.fetchInProgress) == 1 {
-
-				sub.fetchNext <- &pullRequest{
-					Expires:   sub.consumeOpts.Expires,
-					Batch:     sub.consumeOpts.MaxMessages - pending.msgCount,
-					MaxBytes:  sub.consumeOpts.MaxBytes - pending.byteCount,
-					Heartbeat: sub.consumeOpts.Heartbeat,
-				}
-
-				sub.resetPendingMsgs()
-			}
 		}()
 		if !userMsg {
 			// heartbeat message
@@ -225,7 +211,9 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 			return
 		}
 		handler(p.jetStream.toJSMsg(msg))
+		sub.Lock()
 		sub.decrementPendingMsgs(msg)
+		sub.Unlock()
 	}
 	inbox := nats.NewInbox()
 	sub.subscription, err = p.jetStream.conn.Subscribe(inbox, internalHandler)
@@ -233,6 +221,7 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 		return nil, err
 	}
 
+	sub.Lock()
 	// initial pull
 	sub.resetPendingMsgs()
 	if err := sub.pull(&pullRequest{
@@ -243,6 +232,7 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 	}, subject); err != nil {
 		sub.errs <- err
 	}
+	sub.Unlock()
 
 	go func() {
 		isConnected := true
@@ -292,11 +282,12 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 							MaxBytes:  sub.consumeOpts.MaxBytes,
 							Heartbeat: sub.consumeOpts.Heartbeat,
 						}
-						sub.Unlock()
 						sub.resetPendingMsgs()
 					}
+					sub.Unlock()
 				}
 			case err := <-sub.errs:
+				sub.Lock()
 				if sub.consumeOpts.ErrHandler != nil {
 					sub.consumeOpts.ErrHandler(sub, err)
 				}
@@ -309,6 +300,7 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 					}
 					sub.resetPendingMsgs()
 				}
+				sub.Unlock()
 			}
 		}
 	}()
@@ -318,22 +310,43 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 	return sub, nil
 }
 
+// resetPendingMsgs resets pending message count and byte count
+// to the values set in consumeOpts
+// lock should be held before calling this method
 func (s *pullSubscription) resetPendingMsgs() {
-	s.Lock()
-	defer s.Unlock()
 	s.pending.msgCount = s.consumeOpts.MaxMessages
 	s.pending.byteCount = s.consumeOpts.MaxBytes
 }
 
+// decrementPendingMsgs decrements pending message count and byte count
+// lock should be held before calling this method
 func (s *pullSubscription) decrementPendingMsgs(msg *nats.Msg) {
-	s.Lock()
-	defer s.Unlock()
 	s.pending.msgCount--
 	if s.consumeOpts.MaxBytes != 0 {
 		s.pending.byteCount -= msg.Size()
 	}
 }
 
+// checkPending verifies whether there are enough messages in
+// the buffer to trigger a new pull request.
+// lock should be held before calling this method
+func (s *pullSubscription) checkPending() {
+	if s.pending.msgCount < s.consumeOpts.ThresholdMessages ||
+		(s.pending.byteCount < s.consumeOpts.ThresholdBytes && s.consumeOpts.MaxBytes != 0) &&
+			atomic.LoadUint32(&s.fetchInProgress) == 1 {
+
+		s.fetchNext <- &pullRequest{
+			Expires:   s.consumeOpts.Expires,
+			Batch:     s.consumeOpts.MaxMessages - s.pending.msgCount,
+			MaxBytes:  s.consumeOpts.MaxBytes - s.pending.byteCount,
+			Heartbeat: s.consumeOpts.Heartbeat,
+		}
+
+		s.pending.msgCount = s.consumeOpts.MaxMessages
+		s.pending.byteCount = s.consumeOpts.MaxBytes
+	}
+}
+
 // Messages returns MessagesContext, allowing continuously iterating over messages on a stream.
 //
 // Available options:
@@ -425,21 +438,7 @@ func (s *pullSubscription) Next() (Msg, error) {
 
 	isConnected := true
 	for {
-		if s.pending.msgCount < s.consumeOpts.ThresholdMessages ||
-			(s.pending.byteCount < s.consumeOpts.ThresholdBytes && s.consumeOpts.MaxBytes != 0) &&
-				atomic.LoadUint32(&s.fetchInProgress) == 1 {
-
-			s.fetchNext <- &pullRequest{
-				Expires:   s.consumeOpts.Expires,
-				Batch:     s.consumeOpts.MaxMessages - s.pending.msgCount,
-				MaxBytes:  s.consumeOpts.MaxBytes - s.pending.byteCount,
-				Heartbeat: s.consumeOpts.Heartbeat,
-			}
-			s.pending.msgCount = s.consumeOpts.MaxMessages
-			if s.consumeOpts.MaxBytes > 0 {
-				s.pending.byteCount = s.consumeOpts.MaxBytes
-			}
-		}
+		s.checkPending()
 		select {
 		case msg := <-s.msgs:
 			if hbMonitor != nil {
@@ -457,10 +456,7 @@ func (s *pullSubscription) Next() (Msg, error) {
 				}
 				continue
 			}
-			s.pending.msgCount--
-			if s.consumeOpts.MaxBytes > 0 {
-				s.pending.byteCount -= msg.Size()
-			}
+			s.decrementPendingMsgs(msg)
 			return s.consumer.jetStream.toJSMsg(msg), nil
 		case err := <-s.errs:
 			if errors.Is(err, ErrNoHeartbeat) {
diff --git a/micro/service.go b/micro/service.go
index 445c81f..6c01570 100644
--- a/micro/service.go
+++ b/micro/service.go
@@ -264,7 +264,7 @@ var (
 	// this regular expression is suggested regexp for semver validation: https://semver.org/
 	semVerRegexp  = regexp.MustCompile(`^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$`)
 	nameRegexp    = regexp.MustCompile(`^[A-Za-z0-9\-_]+$`)
-	subjectRegexp = regexp.MustCompile(`^[^ >]+[>]?$`)
+	subjectRegexp = regexp.MustCompile(`^[^ >]*[>]?$`)
 )
 
 // Common errors returned by the Service framework.


Test patch:
diff --git a/jetstream/message_test.go b/jetstream/jetstream_test.go
similarity index 75%
rename from jetstream/message_test.go
rename to jetstream/jetstream_test.go
index 653ee07..e7f2317 100644
--- a/jetstream/message_test.go
+++ b/jetstream/jetstream_test.go
@@ -15,6 +15,7 @@ package jetstream
 
 import (
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
@@ -80,3 +81,38 @@ func TestMessageMetadata(t *testing.T) {
 		})
 	}
 }
+
+func TestValidateSubject(t *testing.T) {
+	tests := []struct {
+		subject   string
+		withError bool
+	}{
+		{"test.A", false},
+		{"test.*", false},
+		{"*", false},
+		{"*.*", false},
+		{"test.*.A", false},
+		{"test.>", false},
+		{">", false},
+		{">.", true},
+		{"test.>.A", true},
+		{"", true},
+		{"test A", true},
+	}
+
+	for _, test := range tests {
+		tName := fmt.Sprintf("subj=%s,err=%t", test.subject, test.withError)
+		t.Run(tName, func(t *testing.T) {
+			err := validateSubject(test.subject)
+			if test.withError {
+				if err == nil {
+					t.Fatal("Expected error; got nil")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+		})
+	}
+}
diff --git a/jetstream/test/jetstream_test.go b/jetstream/test/jetstream_test.go
index a9323ec..3416c14 100644
--- a/jetstream/test/jetstream_test.go
+++ b/jetstream/test/jetstream_test.go
@@ -928,3 +928,81 @@ func TestJetStream_DeleteConsumer(t *testing.T) {
 		})
 	}
 }
+
+func TestStreamNameBySubject(t *testing.T) {
+	tests := []struct {
+		name      string
+		subject   string
+		withError error
+		expected  string
+	}{
+		{
+			name:     "get stream name by subject explicit",
+			subject:  "FOO.123",
+			expected: "foo",
+		},
+		{
+			name:     "get stream name by subject with wildcard",
+			subject:  "BAR.*",
+			expected: "bar",
+		},
+		{
+			name:     "match more than one stream, return the first one",
+			subject:  ">",
+			expected: "",
+		},
+		{
+			name:      "stream not found",
+			subject:   "BAR.XYZ",
+			withError: jetstream.ErrStreamNotFound,
+		},
+		{
+			name:      "invalid subject",
+			subject:   "FOO.>.123",
+			withError: jetstream.ErrInvalidSubject,
+		},
+	}
+
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+	nc, err := nats.Connect(srv.ClientURL())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	js, err := jetstream.New(nc)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "foo", Subjects: []string{"FOO.*"}})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	_, err = js.CreateStream(ctx, jetstream.StreamConfig{Name: "bar", Subjects: []string{"BAR.ABC"}})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			name, err := js.StreamNameBySubject(ctx, test.subject)
+			if test.withError != nil {
+				if err == nil || !errors.Is(err, test.withError) {
+					t.Fatalf("Expected error: %v; got: %v", test.withError, err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if test.expected != "" && name != test.expected {
+				t.Fatalf("Unexpected stream name; want: %s; got: %s", test.expected, name)
+			}
+
+		})
+	}
+}