Add Server.Ready() to ensure assets are up-to-date

Signed-off-by: Byron Ruth <[email protected]>

pkg/drivers/nats/new.go

@@ -95,12 +95,13 @@ func newBackend(ctx context.Context, connection string, tlsInfo tls.Config, lega
 		// TODO: limit the number of retries?
 		var retries int
 		for {
-			if ns.ReadyForConnections(5 * time.Second) {
+			if ns.Ready() {
 				logrus.Infof("embedded NATS server is ready for client connections")
 				break
 			}
 			retries++
 			logrus.Infof("waiting for embedded NATS server to be ready: %d", retries)
+			time.Sleep(100 * time.Millisecond)
 		}
 
 		// Use the local server's client URL.
@@ -173,13 +174,6 @@ func newBackend(ctx context.Context, connection string, tlsInfo tls.Config, lega
 
 	logrus.Infof("bucket initialized: %s", config.bucket)
 
-	if ns != nil {
-		for !ns.JetStreamIsCurrent() || !ns.JetStreamIsStreamCurrent("$G", fmt.Sprintf("KV_%s", config.bucket)) {
-			logrus.Warnf("waiting for JetStream to be current")
-			time.Sleep(time.Second)
-		}
-	}
-
 	ekv := NewKeyValue(ctx, bucket, js)
 
 	// Reference the global logger, since it appears log levels are

pkg/drivers/nats/server/interface.go

@@ -2,16 +2,13 @@ package server
 
 import (
 	"net"
-	"time"
 )
 
 type Server interface {
 	Start()
+	Ready() bool
 	Shutdown()
 	ClientURL() string
-	ReadyForConnections(wait time.Duration) bool
-	JetStreamIsCurrent() bool
-	JetStreamIsStreamCurrent(account, stream string) bool
 	InProcessConn() (net.Conn, error)
 }
 

pkg/drivers/nats/server/server.go

@@ -4,15 +4,65 @@
 package server
 
 import (
+	"bytes"
+	"encoding/json"
 	"fmt"
+	"net/http"
+	"net/url"
 
 	"github.com/nats-io/nats-server/v2/server"
+	"github.com/sirupsen/logrus"
 )
 
 const (
 	Embedded = true
 )
 
+type responseWriter struct {
+	code   int
+	header http.Header
+	body   *bytes.Buffer
+}
+
+func (w *responseWriter) Header() http.Header {
+	return w.header
+}
+
+func (w *responseWriter) Write(b []byte) (int, error) {
+	return w.body.Write(b)
+}
+
+func (w *responseWriter) WriteHeader(code int) {
+	w.code = code
+}
+
+type embeddedServer struct {
+	*server.Server
+}
+
+func (s *embeddedServer) Ready() bool {
+	rw := responseWriter{
+		header: http.Header{},
+		body:   &bytes.Buffer{},
+	}
+
+	r := http.Request{
+		Method: "GET",
+		URL: &url.URL{
+			Path: "/healthz",
+		},
+		Header: http.Header{},
+	}
+
+	s.Server.HandleHealthz(&rw, &r)
+
+	var hs server.HealthStatus
+	json.NewDecoder(rw.body).Decode(&hs)
+	logrus.Debugf("embedded NATS server health: %#v", hs)
+
+	return hs.Status == "ok"
+}
+
 func New(c *Config) (Server, error) {
 	opts := &server.Options{}
 
@@ -52,5 +102,5 @@ func New(c *Config) (Server, error) {
 		srv.ConfigureLogger()
 	}
 
-	return srv, nil
+	return &embeddedServer{Server: srv}, nil
 }