Repository: nats-io/nats-server
Pull request #1953 (2021-03-03 01:39:34)

JetStream API Changes
In clustered JetStream we need to move API calls out of routes/gateways/leafnodes path.

This moves from explicit imports and subscriptions to one wildcard subscription and a single wildcard export.

Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/accounts.go b/server/accounts.go
index bce86fe7..fb5e6778 100644
--- a/server/accounts.go
+++ b/server/accounts.go
@@ -1707,14 +1707,6 @@ func (a *Account) serviceImportShadowed(from string) bool {
 	return false
 }
 
-// Internal check to see if a service import exists.
-func (a *Account) serviceImportExists(from string) bool {
-	a.mu.RLock()
-	dup := a.imports.services[from]
-	a.mu.RUnlock()
-	return dup != nil
-}
-
 // Add a service import.
 // This does no checks and should only be called by the msg processing code.
 // Use AddServiceImport from above if responding to user input or config changes, etc.
diff --git a/server/consumer.go b/server/consumer.go
index af852fe6..8107d132 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -1677,35 +1677,24 @@ func (o *consumer) processNextMsgReq(_ *subscription, c *client, _, reply string
 	// In case we have to queue up this request. This is all on stack pre-allocated.
 	wr := waitingRequest{client: c, reply: reply, n: batchSize, noWait: noWait, expires: expires}
 
-	// If we are in replay mode, defer to processReplay for delivery.
-	if o.replay {
-		o.waiting.add(&wr)
+	if wr.noWait {
+		if o.maxp > 0 && len(o.pending) >= o.maxp {
+			sendErr(409, "Exceeded MaxAckPending")
+			return
+		}
 		o.mu.Unlock()
-		o.signalNewMessages()
-		return
-	}
-
-	for i := 0; i < batchSize; i++ {
-		// See if we have more messages available.
-		if subj, hdr, msg, seq, dc, ts, err := o.getNextMsg(); err == nil {
-			o.deliverMsg(reply, subj, hdr, msg, seq, dc, ts)
-			// Need to discount this from the total n for the request.
-			wr.n--
-		} else {
-			if wr.noWait {
-				switch err {
-				case errMaxAckPending:
-					sendErr(409, "Exceeded MaxAckPending")
-				default:
-					sendErr(404, "No Messages")
-				}
-				return
-			}
-			o.waiting.add(&wr)
-			break
+		empty := mset.state().Msgs == 0
+		o.mu.Lock()
+		if empty {
+			sendErr(404, "No Messages")
+			return
 		}
 	}
+
+	o.waiting.add(&wr)
 	o.mu.Unlock()
+
+	o.signalNewMessages()
 }
 
 // Increase the delivery count for this message.
diff --git a/server/jetstream.go b/server/jetstream.go
index 987aa216..e5e182cc 100644
--- a/server/jetstream.go
+++ b/server/jetstream.go
@@ -80,6 +80,8 @@ type jetStream struct {
 	accounts      map[*Account]*jsAccount
 	memReserved   int64
 	storeReserved int64
+	apiCalls      int64
+	apiSubs       *Sublist
 	disabled      bool
 }
 
@@ -149,7 +151,7 @@ func (s *Server) EnableJetStream(config *JetStreamConfig) error {
 		cfg.StoreDir = filepath.Join(os.TempDir(), JetStreamStoreDir)
 	}
 
-	s.js = &jetStream{srv: s, config: cfg, accounts: make(map[*Account]*jsAccount)}
+	s.js = &jetStream{srv: s, config: cfg, accounts: make(map[*Account]*jsAccount), apiSubs: NewSublistNoCache()}
 	s.mu.Unlock()
 
 	// FIXME(dlc) - Allow memory only operation?
@@ -229,18 +231,15 @@ func (s *Server) EnableJetStream(config *JetStreamConfig) error {
 // on the system account, and if not go ahead and set them up.
 func (s *Server) checkJetStreamExports() {
 	sacc := s.SystemAccount()
-	if sacc != nil && sacc.getServiceExport(allJsExports[0]) == nil {
+	if sacc != nil && sacc.getServiceExport(jsAllApi) == nil {
 		s.setupJetStreamExports()
 	}
 }
 
 func (s *Server) setupJetStreamExports() {
-	// Setup our internal system exports.
-	sacc := s.SystemAccount()
-	for _, export := range allJsExports {
-		if err := sacc.AddServiceExport(export, nil); err != nil {
-			s.Warnf("Error setting up jetstream service exports: %v", err)
-		}
+	// Setup our internal system export.
+	if err := s.SystemAccount().AddServiceExport(jsAllApi, nil); err != nil {
+		s.Warnf("Error setting up jetstream service exports: %v", err)
 	}
 }
 
@@ -353,14 +352,10 @@ func (a *Account) enableAllJetStreamServiceImports() error {
 	// In case the enabled import exists here.
 	a.removeServiceImport(JSApiAccountInfo)
 
-	sys := s.SystemAccount()
-	for _, export := range allJsExports {
-		if !a.serviceImportExists(export) {
-			if err := a.AddServiceImport(sys, export, _EMPTY_); err != nil {
-				return fmt.Errorf("Error setting up jetstream service imports for account: %v", err)
-			}
-		}
+	if err := a.AddServiceImport(s.SystemAccount(), jsAllApi, _EMPTY_); err != nil {
+		return fmt.Errorf("Error setting up jetstream service imports for account: %v", err)
 	}
+
 	return nil
 }
 
@@ -384,6 +379,7 @@ func (a *Account) enableJetStreamInfoServiceImportOnly() error {
 	if err := a.AddServiceImport(s.SystemAccount(), JSApiAccountInfo, _EMPTY_); err != nil {
 		return fmt.Errorf("Error setting up jetstream service imports for account: %v", err)
 	}
+
 	return nil
 }
 
@@ -516,7 +512,7 @@ func (s *Server) migrateEphemerals() {
 		}
 	}
 
-	// Gove time for migration information to make it out of our server.
+	// Give time for migration information to make it out of our server.
 	if len(consumers) > 0 {
 		time.Sleep(50 * time.Millisecond)
 	}
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index dfe82f25..440db174 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -32,6 +32,9 @@ import (
 
 // Request API subjects for JetStream.
 const (
+	// All API endpoints.
+	jsAllApi = "$JS.API.>"
+
 	// JSApiPrefix
 	JSApiPrefix = "$JS.API"
 
@@ -260,6 +263,8 @@ type ApiResponse struct {
 	Error *ApiError `json:"error,omitempty"`
 }
 
+const JSApiOverloadedType = "io.nats.jetstream.api.v1.system_overloaded"
+
 // ApiPaged includes variables used to create paged responses from the JSON API
 type ApiPaged struct {
 	Total  int `json:"total"`
@@ -620,7 +625,81 @@ var allJsExports = []string{
 	JSApiConsumerDelete,
 }
 
+func (js *jetStream) apiDispatch(sub *subscription, c *client, subject, reply string, rmsg []byte) {
+	js.mu.RLock()
+	s := js.srv
+	rr := js.apiSubs.Match(subject)
+	js.mu.RUnlock()
+
+	// Shortcircuit.
+	if len(rr.psubs)+len(rr.qsubs) == 0 {
+		return
+	}
+
+	// We should only have psubs and only 1 per result.
+	// FIXME(dlc) - Should we respond here with NoResponders or error?
+	if len(rr.psubs) != 1 {
+		s.Warnf("Malformed JetStream API Request: [%s] %q", subject, rmsg)
+		return
+	}
+	jsub := rr.psubs[0]
+
+	// If this is directly from a client connection ok to do in place.
+	if c.kind == CLIENT {
+		jsub.icb(sub, c, subject, reply, rmsg)
+		return
+	}
+
+	// If we are here we have received this request over a non client connection.
+	// We need to make sure not to block. We will spin a Go routine per but also make
+	// sure we do not have too many outstanding.
+	if apiOut := atomic.AddInt64(&js.apiCalls, 1); apiOut > 1024 {
+		atomic.AddInt64(&js.apiCalls, -1)
+		ci, acc, _, msg, err := s.getRequestInfo(c, rmsg)
+		if err == nil {
+			resp := &ApiResponse{Type: JSApiOverloadedType, Error: jsInsufficientErr}
+			s.sendAPIErrResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(&resp))
+		} else {
+			s.Warnf(badAPIRequestT, rmsg)
+		}
+		s.Warnf("JetStream API Overloaded: %d calls outstanding", apiOut)
+		return
+	}
+
+	// If we are here we can properly dispatch this API call.
+	// Copy the message and the client. Client for the pubArgs
+	// but note the JSAPI only uses the hdr index to piece apart
+	// the header from the msg body. No other references are needed.
+	// FIXME(dlc) - Should cleanup eventually and make sending
+	// and receiving internal messages more formal.
+	rmsg = append(rmsg[:0:0], rmsg...)
+	client := &client{srv: s, kind: JETSTREAM}
+	client.pa = c.pa
+
+	// Dispatch the API call to its own Go routine.
+	go func() {
+		defer atomic.AddInt64(&js.apiCalls, -1)
+		jsub.icb(sub, client, subject, reply, rmsg)
+	}()
+}
+
 func (s *Server) setJetStreamExportSubs() error {
+	js := s.getJetStream()
+	if js == nil {
+		return ErrJetStreamNotEnabled
+	}
+
+	// This is the catch all now for all JetStream API calls.
+	if _, err := s.sysSubscribe(jsAllApi, js.apiDispatch); err != nil {
+		return err
+	}
+
+	if err := s.SystemAccount().AddServiceExport(jsAllApi, nil); err != nil {
+		s.Warnf("Error setting up jetstream service exports: %v", err)
+		return err
+	}
+
+	// API handles themselves.
 	pairs := []struct {
 		subject string
 		handler msgHandler
@@ -652,11 +731,16 @@ func (s *Server) setJetStreamExportSubs() error {
 		{JSApiConsumerDelete, s.jsConsumerDeleteRequest},
 	}
 
+	js.mu.Lock()
+	defer js.mu.Unlock()
+
 	for _, p := range pairs {
-		if _, err := s.sysSubscribe(p.subject, p.handler); err != nil {
+		sub := &subscription{subject: []byte(p.subject), icb: p.handler}
+		if err := js.apiSubs.Insert(sub); err != nil {
 			return err
 		}
 	}
+
 	return nil
 }
 
diff --git a/server/monitor.go b/server/monitor.go
index fb298254..abb63758 100644
--- a/server/monitor.go
+++ b/server/monitor.go
@@ -2259,11 +2259,12 @@ type JSInfo struct {
 	Disabled bool            `json:"disabled,omitempty"`
 	Config   JetStreamConfig `json:"config,omitempty"`
 	JetStreamStats
-	StreamCnt    int          `json:"total_streams,omitempty"`
-	ConsumerCnt  int          `json:"total_consumers,omitempty"`
-	MessageCnt   uint64       `json:"total_messages,omitempty"`
-	MessageBytes uint64       `json:"total_message_bytes,omitempty"`
-	Meta         *ClusterInfo `json:"meta_cluster,omitempty"`
+	APICalls  int64        `json:"current_api_calls"`
+	Streams   int          `json:"total_streams,omitempty"`
+	Consumers int          `json:"total_consumers,omitempty"`
+	Messages  uint64       `json:"total_messages,omitempty"`
+	Bytes     uint64       `json:"total_message_bytes,omitempty"`
+	Meta      *ClusterInfo `json:"meta_cluster,omitempty"`
 	// aggregate raft info
 	AccountDetails []*AccountDetail `json:"account_details,omitempty"`
 }
@@ -2396,11 +2397,13 @@ func (s *Server) Jsz(opts *JSzOptions) (*JSInfo, error) {
 		return jsi, nil
 	}
 	accounts := []*jsAccount{}
+
 	s.js.mu.RLock()
 	jsi.Config = s.js.config
 	for _, info := range s.js.accounts {
 		accounts = append(accounts, info)
 	}
+	jsi.APICalls = atomic.LoadInt64(&s.js.apiCalls)
 	s.js.mu.RUnlock()
 
 	jsi.Meta = toClusterInfo(s.js.getMetaGroup())
@@ -2410,19 +2413,20 @@ func (s *Server) Jsz(opts *JSzOptions) (*JSInfo, error) {
 		if jsa.acc().GetName() == opts.Account {
 			filterIdx = i
 		}
-		jsi.StreamCnt += len(jsa.streams)
+		jsi.Streams += len(jsa.streams)
 		jsi.Memory += uint64(jsa.usage.mem)
 		jsi.Store += uint64(jsa.usage.store)
 		jsi.API.Total += jsa.usage.api
 		jsi.API.Errors += jsa.usage.err
 		for _, stream := range jsa.streams {
 			streamState := stream.state()
-			jsi.MessageCnt += streamState.Msgs
-			jsi.MessageBytes += streamState.Bytes
-			jsi.ConsumerCnt += streamState.Consumers
+			jsi.Messages += streamState.Msgs
+			jsi.Bytes += streamState.Bytes
+			jsi.Consumers += streamState.Consumers
 		}
 		jsa.mu.RUnlock()
 	}
+
 	// filter logic
 	if filterIdx != -1 {
 		accounts = []*jsAccount{accounts[filterIdx]}
diff --git a/server/stream.go b/server/stream.go
index c76af3ce..1b091673 100644
--- a/server/stream.go
+++ b/server/stream.go
@@ -2550,8 +2550,6 @@ func (a *Account) RestoreStream(ncfg *StreamConfig, r io.Reader) (*stream, error
 	}
 
 	sd := path.Join(jsa.storeDir, snapsDir)
-	defer os.RemoveAll(sd)
-
 	if _, err := os.Stat(sd); os.IsNotExist(err) {
 		if err := os.MkdirAll(sd, 0755); err != nil {
 			return nil, fmt.Errorf("could not create snapshots directory - %v", err)
@@ -2566,6 +2564,7 @@ func (a *Account) RestoreStream(ncfg *StreamConfig, r io.Reader) (*stream, error
 			return nil, fmt.Errorf("could not create snapshots directory - %v", err)
 		}
 	}
+	defer os.RemoveAll(sdir)
 
 	tr := tar.NewReader(s2.NewReader(r))
 	for {
@@ -2608,6 +2607,11 @@ func (a *Account) RestoreStream(ncfg *StreamConfig, r io.Reader) (*stream, error
 	}
 	// Move into the correct place here.
 	ndir := path.Join(jsa.storeDir, streamsDir, cfg.Name)
+	// Remove old one if for some reason is here.
+	if _, err := os.Stat(ndir); !os.IsNotExist(err) {
+		os.RemoveAll(ndir)
+	}
+	// Move into new location.
 	if err := os.Rename(sdir, ndir); err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/server/events_test.go b/server/events_test.go
index 1bd0a4f1..4aab3c26 100644
--- a/server/events_test.go
+++ b/server/events_test.go
@@ -1257,8 +1257,8 @@ func TestAccountReqMonitoring(t *testing.T) {
 	// query SUBSZ for account
 	if resp, err := ncSys.Request(subsz, nil, time.Second); err != nil {
 		t.Fatalf("Error on request: %v", err)
-	} else if !strings.Contains(string(resp.Data), `"num_subscriptions":25,`) {
-		t.Fatalf("unexpected subs count (expected 0): %v", string(resp.Data))
+	} else if !strings.Contains(string(resp.Data), `"num_subscriptions":1,`) {
+		t.Fatalf("unexpected subs count (expected 1): %v", string(resp.Data))
 	}
 	// create a subscription
 	if sub, err := nc.Subscribe("foo", func(msg *nats.Msg) {}); err != nil {
@@ -1270,8 +1270,8 @@ func TestAccountReqMonitoring(t *testing.T) {
 	// query SUBSZ for account
 	if resp, err := ncSys.Request(subsz, nil, time.Second); err != nil {
 		t.Fatalf("Error on request: %v", err)
-	} else if !strings.Contains(string(resp.Data), `"num_subscriptions":26,`) {
-		t.Fatalf("unexpected subs count (expected 26): %v", string(resp.Data))
+	} else if !strings.Contains(string(resp.Data), `"num_subscriptions":2,`) {
+		t.Fatalf("unexpected subs count (expected 2): %v", string(resp.Data))
 	} else if !strings.Contains(string(resp.Data), `"subject":"foo"`) {
 		t.Fatalf("expected subscription foo: %v", string(resp.Data))
 	}
diff --git a/server/jetstream_cluster_test.go b/server/jetstream_cluster_test.go
index 1371deaf..e03c64f5 100644
--- a/server/jetstream_cluster_test.go
+++ b/server/jetstream_cluster_test.go
@@ -2887,12 +2887,15 @@ func TestJetStreamClusterExtendedAccountInfo(t *testing.T) {
 		return info.JetStreamAccountStats
 	}
 
+	// Wait to accumulate.
+	time.Sleep(250 * time.Millisecond)
+
 	ai := getAccountInfo()
 	if ai.Streams != 3 || ai.Consumers != 3 {
 		t.Fatalf("AccountInfo not correct: %+v", ai)
 	}
-	if ai.API.Total < 10 {
-		t.Fatalf("Expected at least 10 total API calls, got %d", ai.API.Total)
+	if ai.API.Total < 8 {
+		t.Fatalf("Expected at least 8 total API calls, got %d", ai.API.Total)
 	}
 
 	// Now do a failure to make sure we track API errors.
diff --git a/server/monitor_test.go b/server/monitor_test.go
index b32f008c..6f2ba2fa 100644
--- a/server/monitor_test.go
+++ b/server/monitor_test.go
@@ -4006,14 +4006,14 @@ func TestMonitorJsz(t *testing.T) {
 			if len(info.AccountDetails) != 0 {
 				t.Fatalf("expected no account to be returned by %s but got %v", url, info)
 			}
-			if info.StreamCnt == 0 {
-				t.Fatalf("expected stream count to be 2 but got %d", info.StreamCnt)
+			if info.Streams == 0 {
+				t.Fatalf("expected stream count to be 2 but got %d", info.Streams)
 			}
-			if info.ConsumerCnt == 0 {
-				t.Fatalf("expected consumer count to be 2 but got %d", info.ConsumerCnt)
+			if info.Consumers == 0 {
+				t.Fatalf("expected consumer count to be 2 but got %d", info.Consumers)
 			}
-			if info.MessageCnt != 1 {
-				t.Fatalf("expected one message but got %d", info.MessageCnt)
+			if info.Messages != 1 {
+				t.Fatalf("expected one message but got %d", info.Messages)
 			}
 		}
 	})