Refactor some JetStream helper code, add support for specifying JetStream domain

go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/matrix-org/pinecone v0.11.1-0.20230810010612-ea4c33717fd7
 	github.com/matrix-org/util v0.0.0-20221111132719-399730281e66
 	github.com/mattn/go-sqlite3 v1.14.24
-	github.com/nats-io/nats-server/v2 v2.10.23
+	github.com/nats-io/nats-server/v2 v2.10.24
 	github.com/nats-io/nats.go v1.37.0
 	github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646
 	github.com/opentracing/opentracing-go v1.2.0
@@ -115,8 +115,8 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/mschoch/smat v0.2.0 // indirect
-	github.com/nats-io/jwt/v2 v2.5.8 // indirect
-	github.com/nats-io/nkeys v0.4.8 // indirect
+	github.com/nats-io/jwt/v2 v2.7.3 // indirect
+	github.com/nats-io/nkeys v0.4.9 // indirect
 	github.com/nats-io/nuid v1.0.1 // indirect
 	github.com/ncruces/go-strftime v0.1.9 // indirect
 	github.com/onsi/ginkgo/v2 v2.11.0 // indirect

go.sum

@@ -266,14 +266,14 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
 github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/mschoch/smat v0.2.0 h1:8imxQsjDm8yFEAVBe7azKmKSgzSkZXDuKkSq9374khM=
 github.com/mschoch/smat v0.2.0/go.mod h1:kc9mz7DoBKqDyiRL7VZN8KvXQMWeTaVnttLRXOlotKw=
-github.com/nats-io/jwt/v2 v2.5.8 h1:uvdSzwWiEGWGXf+0Q+70qv6AQdvcvxrv9hPM0RiPamE=
-github.com/nats-io/jwt/v2 v2.5.8/go.mod h1:ZdWS1nZa6WMZfFwwgpEaqBV8EPGVgOTDHN/wTbz0Y5A=
-github.com/nats-io/nats-server/v2 v2.10.23 h1:jvfb9cEi5h8UG6HkZgJGdn9f1UPaX3Dohk0PohEekJI=
-github.com/nats-io/nats-server/v2 v2.10.23/go.mod h1:hMFnpDT2XUXsvHglABlFl/uroQCCOcW6X/0esW6GpBk=
+github.com/nats-io/jwt/v2 v2.7.3 h1:6bNPK+FXgBeAqdj4cYQ0F8ViHRbi7woQLq4W29nUAzE=
+github.com/nats-io/jwt/v2 v2.7.3/go.mod h1:GvkcbHhKquj3pkioy5put1wvPxs78UlZ7D/pY+BgZk4=
+github.com/nats-io/nats-server/v2 v2.10.24 h1:KcqqQAD0ZZcG4yLxtvSFJY7CYKVYlnlWoAiVZ6i/IY4=
+github.com/nats-io/nats-server/v2 v2.10.24/go.mod h1:olvKt8E5ZlnjyqBGbAXtxvSQKsPodISK5Eo/euIta4s=
 github.com/nats-io/nats.go v1.37.0 h1:07rauXbVnnJvv1gfIyghFEo6lUcYRY0WXc3x7x0vUxE=
 github.com/nats-io/nats.go v1.37.0/go.mod h1:Ubdu4Nh9exXdSz0RVWRFBbRfrbSxOYd26oF0wkWclB8=
-github.com/nats-io/nkeys v0.4.8 h1:+wee30071y3vCZAYRsnrmIPaOe47A/SkK/UBDPdIV70=
-github.com/nats-io/nkeys v0.4.8/go.mod h1:kqXRgRDPlGy7nGaEDMuYzmiJCIAAWDK0IMBtDmGD0nc=
+github.com/nats-io/nkeys v0.4.9 h1:qe9Faq2Gxwi6RZnZMXfmGMZkg3afLLOtrU+gDZJ35b0=
+github.com/nats-io/nkeys v0.4.9/go.mod h1:jcMqs+FLG+W5YO36OX6wFIFcmpdAns+w1Wm6D3I/evE=
 github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
 github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
 github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=

setup/config/config_jetstream.go

@@ -15,6 +15,8 @@ type JetStream struct {
 	// The prefix to use for stream names for this homeserver - really only
 	// useful if running more than one Dendrite on the same NATS deployment.
 	TopicPrefix string `yaml:"topic_prefix"`
+	// The JetStream domain, if needed.
+	JetStreamDomain string `yaml:"js_domain"`
 	// Keep all storage in memory. This is mostly useful for unit tests.
 	InMemory bool `yaml:"in_memory"`
 	// Disable logging. This is mostly useful for unit tests.

setup/jetstream/helpers.go

@@ -22,7 +22,7 @@ func JetStreamConsumer(
 	f func(ctx context.Context, msgs []*nats.Msg) bool,
 	opts ...nats.SubOpt,
 ) error {
-	defer func() {
+	defer func(durable string) {
 		// If there are existing consumers from before they were pull
 		// consumers, we need to clean up the old push consumers. However,
 		// in order to not affect the interest-based policies, we need to
@@ -33,86 +33,93 @@ func JetStreamConsumer(
 				logrus.WithContext(ctx).Warnf("Failed to clean up old consumer %q", durable)
 			}
 		}
-	}()
+	}(durable)
 
-	name := durable + "Pull"
-	sub, err := js.PullSubscribe(subj, name, opts...)
+	durable = durable + "Pull"
+	sub, err := js.PullSubscribe(subj, durable, opts...)
 	if err != nil {
 		sentry.CaptureException(err)
-		return fmt.Errorf("nats.SubscribeSync: %w", err)
+		logrus.WithContext(ctx).WithError(err).Warnf("Failed to configure durable %q", durable)
+		return err
 	}
-	go func() {
-		for {
-			// If the parent context has given up then there's no point in
-			// carrying on doing anything, so stop the listener.
-			select {
-			case <-ctx.Done():
-				if err := sub.Unsubscribe(); err != nil {
-					logrus.WithContext(ctx).Warnf("Failed to unsubscribe %q", durable)
-				}
-				return
-			default:
-			}
-			// The context behaviour here is surprising — we supply a context
-			// so that we can interrupt the fetch if we want, but NATS will still
-			// enforce its own deadline (roughly 5 seconds by default). Therefore
-			// it is our responsibility to check whether our context expired or
-			// not when a context error is returned. Footguns. Footguns everywhere.
-			msgs, err := sub.Fetch(batch, nats.Context(ctx))
-			if err != nil {
-				if err == context.Canceled || err == context.DeadlineExceeded {
-					// Work out whether it was the JetStream context that expired
-					// or whether it was our supplied context.
-					select {
-					case <-ctx.Done():
-						// The supplied context expired, so we want to stop the
-						// consumer altogether.
-						return
-					default:
-						// The JetStream context expired, so the fetch probably
-						// just timed out and we should try again.
-						continue
-					}
-				} else if errors.Is(err, nats.ErrConsumerDeleted) {
-					// The consumer was deleted so stop.
+	go jetStreamConsumerWorker(ctx, sub, subj, batch, f)
+	return nil
+}
+
+func jetStreamConsumerWorker(
+	ctx context.Context, sub *nats.Subscription, subj string, batch int,
+	f func(ctx context.Context, msgs []*nats.Msg) bool,
+) {
+	for {
+		// If the parent context has given up then there's no point in
+		// carrying on doing anything, so stop the listener.
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+		// The context behaviour here is surprising — we supply a context
+		// so that we can interrupt the fetch if we want, but NATS will still
+		// enforce its own deadline (roughly 5 seconds by default). Therefore
+		// it is our responsibility to check whether our context expired or
+		// not when a context error is returned. Footguns. Footguns everywhere.
+		msgs, err := sub.Fetch(batch, nats.Context(ctx))
+		if err != nil {
+			if err == context.Canceled || err == context.DeadlineExceeded {
+				// Work out whether it was the JetStream context that expired
+				// or whether it was our supplied context.
+				select {
+				case <-ctx.Done():
+					// The supplied context expired, so we want to stop the
+					// consumer altogether.
 					return
-				} else {
-					// Unfortunately, there's no ErrServerShutdown or similar, so we need to compare the string
-					if err.Error() == "nats: Server Shutdown" {
-						logrus.WithContext(ctx).Warn("nats server shutting down")
-						return
-					}
-					// Something else went wrong, so we'll panic.
-					sentry.CaptureException(err)
-					logrus.WithContext(ctx).WithField("subject", subj).Fatal(err)
+				default:
+					// The JetStream context expired, so the fetch probably
+					// just timed out and we should try again.
+					continue
 				}
+			} else if errors.Is(err, nats.ErrTimeout) {
+				// Pull request was invalidated, try again.
+				continue
+			} else if errors.Is(err, nats.ErrConsumerLeadershipChanged) {
+				// Leadership changed so pending pull requests became invalidated,
+				// just try again.
+				continue
+			} else if err.Error() == "nats: Server Shutdown" {
+				// The server is shutting down, but we'll rely on reconnect
+				// behaviour to try and either connect us to another node (if
+				// clustered) or to reconnect when the server comes back up.
+				continue
+			} else {
+				// Something else went wrong.
+				logrus.WithContext(ctx).WithField("subject", subj).WithError(err).Warn("Error on pull subscriber fetch")
+				return
 			}
-			if len(msgs) < 1 {
+		}
+		if len(msgs) < 1 {
+			continue
+		}
+		for _, msg := range msgs {
+			if err = msg.InProgress(nats.Context(ctx)); err != nil {
+				logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.InProgress: %w", err))
+				sentry.CaptureException(err)
 				continue
 			}
+		}
+		if f(ctx, msgs) {
 			for _, msg := range msgs {
-				if err = msg.InProgress(nats.Context(ctx)); err != nil {
-					logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.InProgress: %w", err))
+				if err = msg.AckSync(nats.Context(ctx)); err != nil {
+					logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.AckSync: %w", err))
 					sentry.CaptureException(err)
-					continue
 				}
 			}
-			if f(ctx, msgs) {
-				for _, msg := range msgs {
-					if err = msg.AckSync(nats.Context(ctx)); err != nil {
-						logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.AckSync: %w", err))
-						sentry.CaptureException(err)
-					}
-				}
-			} else {
-				for _, msg := range msgs {
-					if err = msg.Nak(nats.Context(ctx)); err != nil {
-						logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.Nak: %w", err))
-						sentry.CaptureException(err)
-					}
+		} else {
+			for _, msg := range msgs {
+				if err = msg.Nak(nats.Context(ctx)); err != nil {
+					logrus.WithContext(ctx).WithField("subject", subj).Warn(fmt.Errorf("msg.Nak: %w", err))
+					sentry.CaptureException(err)
 				}
 			}
 		}
-	}()
-	return nil
+	}
 }