get.pme.sh/pnats@v0.0.0-20240304004023-26bb5a137ed0/server/memstore.go (about)

     1  // Copyright 2019-2024 The NATS Authors
     2  // Licensed under the Apache License, Version 2.0 (the "License");
     3  // you may not use this file except in compliance with the License.
     4  // You may obtain a copy of the License at
     5  //
     6  // http://www.apache.org/licenses/LICENSE-2.0
     7  //
     8  // Unless required by applicable law or agreed to in writing, software
     9  // distributed under the License is distributed on an "AS IS" BASIS,
    10  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package server
    15  
    16  import (
    17  	crand "crypto/rand"
    18  	"encoding/binary"
    19  	"fmt"
    20  	"sort"
    21  	"sync"
    22  	"time"
    23  
    24  	"get.pme.sh/pnats/server/avl"
    25  	"get.pme.sh/pnats/server/stree"
    26  )
    27  
    28  // TODO(dlc) - This is a fairly simplistic approach but should do for now.
    29  type memStore struct {
    30  	mu          sync.RWMutex
    31  	cfg         StreamConfig
    32  	state       StreamState
    33  	msgs        map[uint64]*StoreMsg
    34  	fss         *stree.SubjectTree[SimpleState]
    35  	dmap        avl.SequenceSet
    36  	maxp        int64
    37  	scb         StorageUpdateHandler
    38  	ageChk      *time.Timer
    39  	consumers   int
    40  	receivedAny bool
    41  }
    42  
    43  func newMemStore(cfg *StreamConfig) (*memStore, error) {
    44  	if cfg == nil {
    45  		return nil, fmt.Errorf("config required")
    46  	}
    47  	if cfg.Storage != MemoryStorage {
    48  		return nil, fmt.Errorf("memStore requires memory storage type in config")
    49  	}
    50  	ms := &memStore{
    51  		msgs: make(map[uint64]*StoreMsg),
    52  		fss:  stree.NewSubjectTree[SimpleState](),
    53  		maxp: cfg.MaxMsgsPer,
    54  		cfg:  *cfg,
    55  	}
    56  	if cfg.FirstSeq > 0 {
    57  		if _, err := ms.purge(cfg.FirstSeq); err != nil {
    58  			return nil, err
    59  		}
    60  	}
    61  
    62  	return ms, nil
    63  }
    64  
    65  func (ms *memStore) UpdateConfig(cfg *StreamConfig) error {
    66  	if cfg == nil {
    67  		return fmt.Errorf("config required")
    68  	}
    69  	if cfg.Storage != MemoryStorage {
    70  		return fmt.Errorf("memStore requires memory storage type in config")
    71  	}
    72  
    73  	ms.mu.Lock()
    74  	ms.cfg = *cfg
    75  	// Limits checks and enforcement.
    76  	ms.enforceMsgLimit()
    77  	ms.enforceBytesLimit()
    78  	// Do age timers.
    79  	if ms.ageChk == nil && ms.cfg.MaxAge != 0 {
    80  		ms.startAgeChk()
    81  	}
    82  	if ms.ageChk != nil && ms.cfg.MaxAge == 0 {
    83  		ms.ageChk.Stop()
    84  		ms.ageChk = nil
    85  	}
    86  	// Make sure to update MaxMsgsPer
    87  	maxp := ms.maxp
    88  	ms.maxp = cfg.MaxMsgsPer
    89  	// If the value is smaller we need to enforce that.
    90  	if ms.maxp != 0 && ms.maxp < maxp {
    91  		lm := uint64(ms.maxp)
    92  		ms.fss.Iter(func(subj []byte, ss *SimpleState) bool {
    93  			if ss.Msgs > lm {
    94  				ms.enforcePerSubjectLimit(bytesToString(subj), ss)
    95  			}
    96  			return true
    97  		})
    98  	}
    99  	ms.mu.Unlock()
   100  
   101  	if cfg.MaxAge != 0 {
   102  		ms.expireMsgs()
   103  	}
   104  	return nil
   105  }
   106  
   107  // Stores a raw message with expected sequence number and timestamp.
   108  // Lock should be held.
   109  func (ms *memStore) storeRawMsg(subj string, hdr, msg []byte, seq uint64, ts int64) error {
   110  	if ms.msgs == nil {
   111  		return ErrStoreClosed
   112  	}
   113  
   114  	// Tracking by subject.
   115  	var ss *SimpleState
   116  	var asl bool
   117  	if len(subj) > 0 {
   118  		var ok bool
   119  		if ss, ok = ms.fss.Find(stringToBytes(subj)); ok {
   120  			asl = ms.maxp > 0 && ss.Msgs >= uint64(ms.maxp)
   121  		}
   122  	}
   123  
   124  	// Check if we are discarding new messages when we reach the limit.
   125  	if ms.cfg.Discard == DiscardNew {
   126  		if asl && ms.cfg.DiscardNewPer {
   127  			return ErrMaxMsgsPerSubject
   128  		}
   129  		if ms.cfg.MaxMsgs > 0 && ms.state.Msgs >= uint64(ms.cfg.MaxMsgs) {
   130  			// If we are tracking max messages per subject and are at the limit we will replace, so this is ok.
   131  			if !asl {
   132  				return ErrMaxMsgs
   133  			}
   134  		}
   135  		if ms.cfg.MaxBytes > 0 && ms.state.Bytes+memStoreMsgSize(subj, hdr, msg) >= uint64(ms.cfg.MaxBytes) {
   136  			if !asl {
   137  				return ErrMaxBytes
   138  			}
   139  			// If we are here we are at a subject maximum, need to determine if dropping last message gives us enough room.
   140  			if ss.firstNeedsUpdate {
   141  				ms.recalculateFirstForSubj(subj, ss.First, ss)
   142  			}
   143  			sm, ok := ms.msgs[ss.First]
   144  			if !ok || memStoreMsgSize(sm.subj, sm.hdr, sm.msg) < memStoreMsgSize(subj, hdr, msg) {
   145  				return ErrMaxBytes
   146  			}
   147  		}
   148  	}
   149  
   150  	if seq != ms.state.LastSeq+1 {
   151  		if seq > 0 {
   152  			return ErrSequenceMismatch
   153  		}
   154  		seq = ms.state.LastSeq + 1
   155  	}
   156  
   157  	// Adjust first if needed.
   158  	now := time.Unix(0, ts).UTC()
   159  	if ms.state.Msgs == 0 {
   160  		ms.state.FirstSeq = seq
   161  		ms.state.FirstTime = now
   162  	}
   163  
   164  	// Make copies
   165  	// TODO(dlc) - Maybe be smarter here.
   166  	if len(msg) > 0 {
   167  		msg = copyBytes(msg)
   168  	}
   169  	if len(hdr) > 0 {
   170  		hdr = copyBytes(hdr)
   171  	}
   172  
   173  	// FIXME(dlc) - Could pool at this level?
   174  	sm := &StoreMsg{subj, nil, nil, make([]byte, 0, len(hdr)+len(msg)), seq, ts}
   175  	sm.buf = append(sm.buf, hdr...)
   176  	sm.buf = append(sm.buf, msg...)
   177  	if len(hdr) > 0 {
   178  		sm.hdr = sm.buf[:len(hdr)]
   179  	}
   180  	sm.msg = sm.buf[len(hdr):]
   181  	ms.msgs[seq] = sm
   182  	ms.state.Msgs++
   183  	ms.state.Bytes += memStoreMsgSize(subj, hdr, msg)
   184  	ms.state.LastSeq = seq
   185  	ms.state.LastTime = now
   186  
   187  	// Track per subject.
   188  	if len(subj) > 0 {
   189  		if ss != nil {
   190  			ss.Msgs++
   191  			ss.Last = seq
   192  			// Check per subject limits.
   193  			if ms.maxp > 0 && ss.Msgs > uint64(ms.maxp) {
   194  				ms.enforcePerSubjectLimit(subj, ss)
   195  			}
   196  		} else {
   197  			ms.fss.Insert([]byte(subj), SimpleState{Msgs: 1, First: seq, Last: seq})
   198  		}
   199  	}
   200  
   201  	// Limits checks and enforcement.
   202  	ms.enforceMsgLimit()
   203  	ms.enforceBytesLimit()
   204  
   205  	// Check if we have and need the age expiration timer running.
   206  	if ms.ageChk == nil && ms.cfg.MaxAge != 0 {
   207  		ms.startAgeChk()
   208  	}
   209  	return nil
   210  }
   211  
   212  // StoreRawMsg stores a raw message with expected sequence number and timestamp.
   213  func (ms *memStore) StoreRawMsg(subj string, hdr, msg []byte, seq uint64, ts int64) error {
   214  	ms.mu.Lock()
   215  	err := ms.storeRawMsg(subj, hdr, msg, seq, ts)
   216  	cb := ms.scb
   217  	// Check if first message timestamp requires expiry
   218  	// sooner than initial replica expiry timer set to MaxAge when initializing.
   219  	if !ms.receivedAny && ms.cfg.MaxAge != 0 && ts > 0 {
   220  		ms.receivedAny = true
   221  		// Calculate duration when the next expireMsgs should be called.
   222  		ms.resetAgeChk(int64(time.Millisecond) * 50)
   223  	}
   224  	ms.mu.Unlock()
   225  
   226  	if err == nil && cb != nil {
   227  		cb(1, int64(memStoreMsgSize(subj, hdr, msg)), seq, subj)
   228  	}
   229  
   230  	return err
   231  }
   232  
   233  // Store stores a message.
   234  func (ms *memStore) StoreMsg(subj string, hdr, msg []byte) (uint64, int64, error) {
   235  	ms.mu.Lock()
   236  	seq, ts := ms.state.LastSeq+1, time.Now().UnixNano()
   237  	err := ms.storeRawMsg(subj, hdr, msg, seq, ts)
   238  	cb := ms.scb
   239  	ms.mu.Unlock()
   240  
   241  	if err != nil {
   242  		seq, ts = 0, 0
   243  	} else if cb != nil {
   244  		cb(1, int64(memStoreMsgSize(subj, hdr, msg)), seq, subj)
   245  	}
   246  
   247  	return seq, ts, err
   248  }
   249  
   250  // SkipMsg will use the next sequence number but not store anything.
   251  func (ms *memStore) SkipMsg() uint64 {
   252  	// Grab time.
   253  	now := time.Now().UTC()
   254  
   255  	ms.mu.Lock()
   256  	seq := ms.state.LastSeq + 1
   257  	ms.state.LastSeq = seq
   258  	ms.state.LastTime = now
   259  	if ms.state.Msgs == 0 {
   260  		ms.state.FirstSeq = seq
   261  		ms.state.FirstTime = now
   262  	} else {
   263  		ms.dmap.Insert(seq)
   264  	}
   265  	ms.mu.Unlock()
   266  	return seq
   267  }
   268  
   269  // Skip multiple msgs.
   270  func (ms *memStore) SkipMsgs(seq uint64, num uint64) error {
   271  	// Grab time.
   272  	now := time.Now().UTC()
   273  
   274  	ms.mu.Lock()
   275  	defer ms.mu.Unlock()
   276  
   277  	// Check sequence matches our last sequence.
   278  	if seq != ms.state.LastSeq+1 {
   279  		if seq > 0 {
   280  			return ErrSequenceMismatch
   281  		}
   282  		seq = ms.state.LastSeq + 1
   283  	}
   284  	lseq := seq + num - 1
   285  
   286  	ms.state.LastSeq = lseq
   287  	ms.state.LastTime = now
   288  	if ms.state.Msgs == 0 {
   289  		ms.state.FirstSeq, ms.state.FirstTime = lseq+1, now
   290  	} else {
   291  		for ; seq <= lseq; seq++ {
   292  			ms.dmap.Insert(seq)
   293  		}
   294  	}
   295  	return nil
   296  }
   297  
   298  // RegisterStorageUpdates registers a callback for updates to storage changes.
   299  // It will present number of messages and bytes as a signed integer and an
   300  // optional sequence number of the message if a single.
   301  func (ms *memStore) RegisterStorageUpdates(cb StorageUpdateHandler) {
   302  	ms.mu.Lock()
   303  	ms.scb = cb
   304  	ms.mu.Unlock()
   305  }
   306  
   307  // GetSeqFromTime looks for the first sequence number that has the message
   308  // with >= timestamp.
   309  // FIXME(dlc) - inefficient.
   310  func (ms *memStore) GetSeqFromTime(t time.Time) uint64 {
   311  	ts := t.UnixNano()
   312  	ms.mu.RLock()
   313  	defer ms.mu.RUnlock()
   314  	if len(ms.msgs) == 0 {
   315  		return ms.state.LastSeq + 1
   316  	}
   317  	if ts <= ms.msgs[ms.state.FirstSeq].ts {
   318  		return ms.state.FirstSeq
   319  	}
   320  	// LastSeq is not guaranteed to be present since last does not go backwards.
   321  	var lmsg *StoreMsg
   322  	for lseq := ms.state.LastSeq; lseq > ms.state.FirstSeq; lseq-- {
   323  		if lmsg = ms.msgs[lseq]; lmsg != nil {
   324  			break
   325  		}
   326  	}
   327  	if lmsg == nil {
   328  		return ms.state.FirstSeq
   329  	}
   330  
   331  	last := lmsg.ts
   332  	if ts == last {
   333  		return ms.state.LastSeq
   334  	}
   335  	if ts > last {
   336  		return ms.state.LastSeq + 1
   337  	}
   338  	index := sort.Search(len(ms.msgs), func(i int) bool {
   339  		if msg := ms.msgs[ms.state.FirstSeq+uint64(i)]; msg != nil {
   340  			return msg.ts >= ts
   341  		}
   342  		return false
   343  	})
   344  	return uint64(index) + ms.state.FirstSeq
   345  }
   346  
   347  // FilteredState will return the SimpleState associated with the filtered subject and a proposed starting sequence.
   348  func (ms *memStore) FilteredState(sseq uint64, subj string) SimpleState {
   349  	// This needs to be a write lock, as filteredStateLocked can
   350  	// mutate the per-subject state.
   351  	ms.mu.Lock()
   352  	defer ms.mu.Unlock()
   353  
   354  	return ms.filteredStateLocked(sseq, subj, false)
   355  }
   356  
   357  func (ms *memStore) filteredStateLocked(sseq uint64, filter string, lastPerSubject bool) SimpleState {
   358  	var ss SimpleState
   359  
   360  	if sseq < ms.state.FirstSeq {
   361  		sseq = ms.state.FirstSeq
   362  	}
   363  
   364  	// If past the end no results.
   365  	if sseq > ms.state.LastSeq {
   366  		return ss
   367  	}
   368  
   369  	if filter == _EMPTY_ {
   370  		filter = fwcs
   371  	}
   372  	isAll := filter == fwcs
   373  
   374  	// First check if we can optimize this part.
   375  	// This means we want all and the starting sequence was before this block.
   376  	if isAll && sseq <= ms.state.FirstSeq {
   377  		total := ms.state.Msgs
   378  		if lastPerSubject {
   379  			total = uint64(ms.fss.Size())
   380  		}
   381  		return SimpleState{
   382  			Msgs:  total,
   383  			First: ms.state.FirstSeq,
   384  			Last:  ms.state.LastSeq,
   385  		}
   386  	}
   387  
   388  	tsa := [32]string{}
   389  	fsa := [32]string{}
   390  	fts := tokenizeSubjectIntoSlice(fsa[:0], filter)
   391  	wc := subjectHasWildcard(filter)
   392  
   393  	// 1. See if we match any subs from fss.
   394  	// 2. If we match and the sseq is past ss.Last then we can use meta only.
   395  	// 3. If we match we need to do a partial, break and clear any totals and do a full scan like num pending.
   396  
   397  	isMatch := func(subj string) bool {
   398  		if isAll {
   399  			return true
   400  		}
   401  		if !wc {
   402  			return subj == filter
   403  		}
   404  		tts := tokenizeSubjectIntoSlice(tsa[:0], subj)
   405  		return isSubsetMatchTokenized(tts, fts)
   406  	}
   407  
   408  	update := func(fss *SimpleState) {
   409  		msgs, first, last := fss.Msgs, fss.First, fss.Last
   410  		if lastPerSubject {
   411  			msgs, first = 1, last
   412  		}
   413  		ss.Msgs += msgs
   414  		if ss.First == 0 || first < ss.First {
   415  			ss.First = first
   416  		}
   417  		if last > ss.Last {
   418  			ss.Last = last
   419  		}
   420  	}
   421  
   422  	var havePartial bool
   423  	// We will track start and end sequences as we go.
   424  	ms.fss.Match(stringToBytes(filter), func(subj []byte, fss *SimpleState) {
   425  		subjs := bytesToString(subj)
   426  		if fss.firstNeedsUpdate {
   427  			ms.recalculateFirstForSubj(subjs, fss.First, fss)
   428  		}
   429  		if sseq <= fss.First {
   430  			update(fss)
   431  		} else if sseq <= fss.Last {
   432  			// We matched but it is a partial.
   433  			havePartial = true
   434  			// Don't break here, we will update to keep tracking last.
   435  			update(fss)
   436  		}
   437  	})
   438  
   439  	// If we did not encounter any partials we can return here.
   440  	if !havePartial {
   441  		return ss
   442  	}
   443  
   444  	// If we are here we need to scan the msgs.
   445  	// Capture first and last sequences for scan and then clear what we had.
   446  	first, last := ss.First, ss.Last
   447  	// To track if we decide to exclude we need to calculate first.
   448  	var needScanFirst bool
   449  	if first < sseq {
   450  		first = sseq
   451  		needScanFirst = true
   452  	}
   453  
   454  	// Now we want to check if it is better to scan inclusive and recalculate that way
   455  	// or leave and scan exclusive and adjust our totals.
   456  	// ss.Last is always correct here.
   457  	toScan, toExclude := last-first, first-ms.state.FirstSeq+ms.state.LastSeq-ss.Last
   458  	var seen map[string]bool
   459  	if lastPerSubject {
   460  		seen = make(map[string]bool)
   461  	}
   462  	if toScan < toExclude {
   463  		ss.Msgs, ss.First = 0, 0
   464  		for seq := first; seq <= last; seq++ {
   465  			if sm, ok := ms.msgs[seq]; ok && !seen[sm.subj] && isMatch(sm.subj) {
   466  				ss.Msgs++
   467  				if ss.First == 0 {
   468  					ss.First = seq
   469  				}
   470  				if seen != nil {
   471  					seen[sm.subj] = true
   472  				}
   473  			}
   474  		}
   475  	} else {
   476  		// We will adjust from the totals above by scanning what we need to exclude.
   477  		ss.First = first
   478  		var adjust uint64
   479  		var tss *SimpleState
   480  
   481  		for seq := ms.state.FirstSeq; seq < first; seq++ {
   482  			if sm, ok := ms.msgs[seq]; ok && !seen[sm.subj] && isMatch(sm.subj) {
   483  				if lastPerSubject {
   484  					tss, _ = ms.fss.Find(stringToBytes(sm.subj))
   485  				}
   486  				// If we are last per subject, make sure to only adjust if all messages are before our first.
   487  				if tss == nil || tss.Last < first {
   488  					adjust++
   489  				}
   490  				if seen != nil {
   491  					seen[sm.subj] = true
   492  				}
   493  			}
   494  		}
   495  		// Now do range at end.
   496  		for seq := last + 1; seq < ms.state.LastSeq; seq++ {
   497  			if sm, ok := ms.msgs[seq]; ok && !seen[sm.subj] && isMatch(sm.subj) {
   498  				adjust++
   499  				if seen != nil {
   500  					seen[sm.subj] = true
   501  				}
   502  			}
   503  		}
   504  		ss.Msgs -= adjust
   505  		if needScanFirst {
   506  			for seq := first; seq < last; seq++ {
   507  				if sm, ok := ms.msgs[seq]; ok && isMatch(sm.subj) {
   508  					ss.First = seq
   509  					break
   510  				}
   511  			}
   512  		}
   513  	}
   514  
   515  	return ss
   516  }
   517  
   518  // SubjectsState returns a map of SimpleState for all matching subjects.
   519  func (ms *memStore) SubjectsState(subject string) map[string]SimpleState {
   520  	ms.mu.RLock()
   521  	defer ms.mu.RUnlock()
   522  
   523  	if ms.fss.Size() == 0 {
   524  		return nil
   525  	}
   526  
   527  	if subject == _EMPTY_ {
   528  		subject = fwcs
   529  	}
   530  
   531  	fss := make(map[string]SimpleState)
   532  	ms.fss.Match(stringToBytes(subject), func(subj []byte, ss *SimpleState) {
   533  		subjs := string(subj)
   534  		if ss.firstNeedsUpdate {
   535  			ms.recalculateFirstForSubj(subjs, ss.First, ss)
   536  		}
   537  		oss := fss[subjs]
   538  		if oss.First == 0 { // New
   539  			fss[subjs] = *ss
   540  		} else {
   541  			// Merge here.
   542  			oss.Last, oss.Msgs = ss.Last, oss.Msgs+ss.Msgs
   543  			fss[subjs] = oss
   544  		}
   545  	})
   546  	return fss
   547  }
   548  
   549  func (ms *memStore) MultiLastSeqs(filters []string, maxSeq uint64, maxAllowed int) ([]uint64, error) {
   550  	ms.mu.RLock()
   551  	defer ms.mu.RUnlock()
   552  
   553  	if len(ms.msgs) == 0 {
   554  		return nil, nil
   555  	}
   556  
   557  	// Implied last sequence.
   558  	if maxSeq == 0 {
   559  		maxSeq = ms.state.LastSeq
   560  	}
   561  
   562  	//subs := make(map[string]*SimpleState)
   563  	seqs := make([]uint64, 0, 64)
   564  	seen := make(map[uint64]struct{})
   565  
   566  	addIfNotDupe := func(seq uint64) {
   567  		if _, ok := seen[seq]; !ok {
   568  			seqs = append(seqs, seq)
   569  			seen[seq] = struct{}{}
   570  		}
   571  	}
   572  
   573  	for _, filter := range filters {
   574  		ms.fss.Match(stringToBytes(filter), func(subj []byte, ss *SimpleState) {
   575  			if ss.Last <= maxSeq {
   576  				addIfNotDupe(ss.Last)
   577  			} else if ss.Msgs > 1 {
   578  				// The last is greater than maxSeq.
   579  				s := bytesToString(subj)
   580  				for seq := maxSeq; seq > 0; seq-- {
   581  					if sm, ok := ms.msgs[seq]; ok && sm.subj == s {
   582  						addIfNotDupe(seq)
   583  						break
   584  					}
   585  				}
   586  			}
   587  		})
   588  		// If maxAllowed was sepcified check that we will not exceed that.
   589  		if maxAllowed > 0 && len(seqs) > maxAllowed {
   590  			return nil, ErrTooManyResults
   591  		}
   592  	}
   593  	sort.Slice(seqs, func(i, j int) bool { return seqs[i] < seqs[j] })
   594  	return seqs, nil
   595  }
   596  
   597  // SubjectsTotal return message totals per subject.
   598  func (ms *memStore) SubjectsTotals(filterSubject string) map[string]uint64 {
   599  	ms.mu.RLock()
   600  	defer ms.mu.RUnlock()
   601  
   602  	if ms.fss.Size() == 0 {
   603  		return nil
   604  	}
   605  
   606  	tsa := [32]string{}
   607  	fsa := [32]string{}
   608  	fts := tokenizeSubjectIntoSlice(fsa[:0], filterSubject)
   609  	isAll := filterSubject == _EMPTY_ || filterSubject == fwcs
   610  
   611  	fst := make(map[string]uint64)
   612  	ms.fss.Match(stringToBytes(filterSubject), func(subj []byte, ss *SimpleState) {
   613  		subjs := string(subj)
   614  		if isAll {
   615  			fst[subjs] = ss.Msgs
   616  		} else {
   617  			if tts := tokenizeSubjectIntoSlice(tsa[:0], subjs); isSubsetMatchTokenized(tts, fts) {
   618  				fst[subjs] = ss.Msgs
   619  			}
   620  		}
   621  	})
   622  	return fst
   623  }
   624  
   625  // NumPending will return the number of pending messages matching the filter subject starting at sequence.
   626  func (ms *memStore) NumPending(sseq uint64, filter string, lastPerSubject bool) (total, validThrough uint64) {
   627  	// This needs to be a write lock, as filteredStateLocked can mutate the per-subject state.
   628  	ms.mu.Lock()
   629  	defer ms.mu.Unlock()
   630  
   631  	ss := ms.filteredStateLocked(sseq, filter, lastPerSubject)
   632  	return ss.Msgs, ms.state.LastSeq
   633  }
   634  
   635  // Will check the msg limit for this tracked subject.
   636  // Lock should be held.
   637  func (ms *memStore) enforcePerSubjectLimit(subj string, ss *SimpleState) {
   638  	if ms.maxp <= 0 {
   639  		return
   640  	}
   641  	for nmsgs := ss.Msgs; nmsgs > uint64(ms.maxp); nmsgs = ss.Msgs {
   642  		if ss.firstNeedsUpdate {
   643  			ms.recalculateFirstForSubj(subj, ss.First, ss)
   644  		}
   645  		if !ms.removeMsg(ss.First, false) {
   646  			break
   647  		}
   648  	}
   649  }
   650  
   651  // Will check the msg limit and drop firstSeq msg if needed.
   652  // Lock should be held.
   653  func (ms *memStore) enforceMsgLimit() {
   654  	if ms.cfg.Discard != DiscardOld {
   655  		return
   656  	}
   657  	if ms.cfg.MaxMsgs <= 0 || ms.state.Msgs <= uint64(ms.cfg.MaxMsgs) {
   658  		return
   659  	}
   660  	for nmsgs := ms.state.Msgs; nmsgs > uint64(ms.cfg.MaxMsgs); nmsgs = ms.state.Msgs {
   661  		ms.deleteFirstMsgOrPanic()
   662  	}
   663  }
   664  
   665  // Will check the bytes limit and drop msgs if needed.
   666  // Lock should be held.
   667  func (ms *memStore) enforceBytesLimit() {
   668  	if ms.cfg.Discard != DiscardOld {
   669  		return
   670  	}
   671  	if ms.cfg.MaxBytes <= 0 || ms.state.Bytes <= uint64(ms.cfg.MaxBytes) {
   672  		return
   673  	}
   674  	for bs := ms.state.Bytes; bs > uint64(ms.cfg.MaxBytes); bs = ms.state.Bytes {
   675  		ms.deleteFirstMsgOrPanic()
   676  	}
   677  }
   678  
   679  // Will start the age check timer.
   680  // Lock should be held.
   681  func (ms *memStore) startAgeChk() {
   682  	if ms.ageChk == nil && ms.cfg.MaxAge != 0 {
   683  		ms.ageChk = time.AfterFunc(ms.cfg.MaxAge, ms.expireMsgs)
   684  	}
   685  }
   686  
   687  // Lock should be held.
   688  func (ms *memStore) resetAgeChk(delta int64) {
   689  	if ms.cfg.MaxAge == 0 {
   690  		return
   691  	}
   692  
   693  	fireIn := ms.cfg.MaxAge
   694  	if delta > 0 && time.Duration(delta) < fireIn {
   695  		if fireIn = time.Duration(delta); fireIn < time.Second {
   696  			// Only fire at most once a second.
   697  			// Excessive firing can effect ingest performance.
   698  			fireIn = time.Second
   699  		}
   700  	}
   701  	if ms.ageChk != nil {
   702  		ms.ageChk.Reset(fireIn)
   703  	} else {
   704  		ms.ageChk = time.AfterFunc(fireIn, ms.expireMsgs)
   705  	}
   706  }
   707  
   708  // Will expire msgs that are too old.
   709  func (ms *memStore) expireMsgs() {
   710  	ms.mu.RLock()
   711  	now := time.Now().UnixNano()
   712  	minAge := now - int64(ms.cfg.MaxAge)
   713  	ms.mu.RUnlock()
   714  
   715  	for {
   716  		ms.mu.Lock()
   717  		if sm, ok := ms.msgs[ms.state.FirstSeq]; ok && sm.ts <= minAge {
   718  			ms.deleteFirstMsgOrPanic()
   719  			// Recalculate in case we are expiring a bunch.
   720  			now = time.Now().UnixNano()
   721  			minAge = now - int64(ms.cfg.MaxAge)
   722  			ms.mu.Unlock()
   723  		} else {
   724  			// We will exit here
   725  			if len(ms.msgs) == 0 {
   726  				if ms.ageChk != nil {
   727  					ms.ageChk.Stop()
   728  					ms.ageChk = nil
   729  				}
   730  			} else {
   731  				var fireIn time.Duration
   732  				if sm == nil {
   733  					fireIn = ms.cfg.MaxAge
   734  				} else {
   735  					fireIn = time.Duration(sm.ts - minAge)
   736  				}
   737  				if ms.ageChk != nil {
   738  					ms.ageChk.Reset(fireIn)
   739  				} else {
   740  					ms.ageChk = time.AfterFunc(fireIn, ms.expireMsgs)
   741  				}
   742  			}
   743  			ms.mu.Unlock()
   744  			break
   745  		}
   746  	}
   747  }
   748  
   749  // PurgeEx will remove messages based on subject filters, sequence and number of messages to keep.
   750  // Will return the number of purged messages.
   751  func (ms *memStore) PurgeEx(subject string, sequence, keep uint64) (purged uint64, err error) {
   752  	if subject == _EMPTY_ || subject == fwcs {
   753  		if keep == 0 && sequence == 0 {
   754  			return ms.Purge()
   755  		}
   756  		if sequence > 1 {
   757  			return ms.Compact(sequence)
   758  		} else if keep > 0 {
   759  			ms.mu.RLock()
   760  			msgs, lseq := ms.state.Msgs, ms.state.LastSeq
   761  			ms.mu.RUnlock()
   762  			if keep >= msgs {
   763  				return 0, nil
   764  			}
   765  			return ms.Compact(lseq - keep + 1)
   766  		}
   767  		return 0, nil
   768  
   769  	}
   770  	eq := compareFn(subject)
   771  	if ss := ms.FilteredState(1, subject); ss.Msgs > 0 {
   772  		if keep > 0 {
   773  			if keep >= ss.Msgs {
   774  				return 0, nil
   775  			}
   776  			ss.Msgs -= keep
   777  		}
   778  		last := ss.Last
   779  		if sequence > 1 {
   780  			last = sequence - 1
   781  		}
   782  		ms.mu.Lock()
   783  		for seq := ss.First; seq <= last; seq++ {
   784  			if sm, ok := ms.msgs[seq]; ok && eq(sm.subj, subject) {
   785  				if ok := ms.removeMsg(sm.seq, false); ok {
   786  					purged++
   787  					if purged >= ss.Msgs {
   788  						break
   789  					}
   790  				}
   791  			}
   792  		}
   793  		ms.mu.Unlock()
   794  	}
   795  	return purged, nil
   796  }
   797  
   798  // Purge will remove all messages from this store.
   799  // Will return the number of purged messages.
   800  func (ms *memStore) Purge() (uint64, error) {
   801  	ms.mu.RLock()
   802  	first := ms.state.LastSeq + 1
   803  	ms.mu.RUnlock()
   804  	return ms.purge(first)
   805  }
   806  
   807  func (ms *memStore) purge(fseq uint64) (uint64, error) {
   808  	ms.mu.Lock()
   809  	purged := uint64(len(ms.msgs))
   810  	cb := ms.scb
   811  	bytes := int64(ms.state.Bytes)
   812  	if fseq < ms.state.LastSeq {
   813  		ms.mu.Unlock()
   814  		return 0, fmt.Errorf("partial purges not supported on memory store")
   815  	}
   816  	ms.state.FirstSeq = fseq
   817  	ms.state.LastSeq = fseq - 1
   818  	ms.state.FirstTime = time.Time{}
   819  	ms.state.Bytes = 0
   820  	ms.state.Msgs = 0
   821  	ms.msgs = make(map[uint64]*StoreMsg)
   822  	ms.fss = stree.NewSubjectTree[SimpleState]()
   823  	ms.mu.Unlock()
   824  
   825  	if cb != nil {
   826  		cb(-int64(purged), -bytes, 0, _EMPTY_)
   827  	}
   828  
   829  	return purged, nil
   830  }
   831  
   832  // Compact will remove all messages from this store up to
   833  // but not including the seq parameter.
   834  // Will return the number of purged messages.
   835  func (ms *memStore) Compact(seq uint64) (uint64, error) {
   836  	if seq == 0 {
   837  		return ms.Purge()
   838  	}
   839  
   840  	var purged, bytes uint64
   841  
   842  	ms.mu.Lock()
   843  	cb := ms.scb
   844  	if seq <= ms.state.LastSeq {
   845  		sm, ok := ms.msgs[seq]
   846  		if !ok {
   847  			ms.mu.Unlock()
   848  			return 0, ErrStoreMsgNotFound
   849  		}
   850  		fseq := ms.state.FirstSeq
   851  		ms.state.FirstSeq = seq
   852  		ms.state.FirstTime = time.Unix(0, sm.ts).UTC()
   853  
   854  		for seq := seq - 1; seq >= fseq; seq-- {
   855  			if sm := ms.msgs[seq]; sm != nil {
   856  				bytes += memStoreMsgSize(sm.subj, sm.hdr, sm.msg)
   857  				purged++
   858  				delete(ms.msgs, seq)
   859  				ms.removeSeqPerSubject(sm.subj, seq)
   860  			}
   861  		}
   862  		if purged > ms.state.Msgs {
   863  			purged = ms.state.Msgs
   864  		}
   865  		ms.state.Msgs -= purged
   866  		if bytes > ms.state.Bytes {
   867  			bytes = ms.state.Bytes
   868  		}
   869  		ms.state.Bytes -= bytes
   870  	} else {
   871  		// We are compacting past the end of our range. Do purge and set sequences correctly
   872  		// such that the next message placed will have seq.
   873  		purged = uint64(len(ms.msgs))
   874  		bytes = ms.state.Bytes
   875  		ms.state.Bytes = 0
   876  		ms.state.Msgs = 0
   877  		ms.state.FirstSeq = seq
   878  		ms.state.FirstTime = time.Time{}
   879  		ms.state.LastSeq = seq - 1
   880  		ms.msgs = make(map[uint64]*StoreMsg)
   881  	}
   882  	ms.mu.Unlock()
   883  
   884  	if cb != nil {
   885  		cb(-int64(purged), -int64(bytes), 0, _EMPTY_)
   886  	}
   887  
   888  	return purged, nil
   889  }
   890  
   891  // Will completely reset our store.
   892  func (ms *memStore) reset() error {
   893  
   894  	ms.mu.Lock()
   895  	var purged, bytes uint64
   896  	cb := ms.scb
   897  	if cb != nil {
   898  		for _, sm := range ms.msgs {
   899  			purged++
   900  			bytes += memStoreMsgSize(sm.subj, sm.hdr, sm.msg)
   901  		}
   902  	}
   903  
   904  	// Reset
   905  	ms.state.FirstSeq = 0
   906  	ms.state.FirstTime = time.Time{}
   907  	ms.state.LastSeq = 0
   908  	ms.state.LastTime = time.Now().UTC()
   909  	// Update msgs and bytes.
   910  	ms.state.Msgs = 0
   911  	ms.state.Bytes = 0
   912  	// Reset msgs and fss.
   913  	ms.msgs = make(map[uint64]*StoreMsg)
   914  	ms.fss = stree.NewSubjectTree[SimpleState]()
   915  
   916  	ms.mu.Unlock()
   917  
   918  	if cb != nil {
   919  		cb(-int64(purged), -int64(bytes), 0, _EMPTY_)
   920  	}
   921  
   922  	return nil
   923  }
   924  
   925  // Truncate will truncate a stream store up to seq. Sequence needs to be valid.
   926  func (ms *memStore) Truncate(seq uint64) error {
   927  	// Check for request to reset.
   928  	if seq == 0 {
   929  		return ms.reset()
   930  	}
   931  
   932  	var purged, bytes uint64
   933  
   934  	ms.mu.Lock()
   935  	lsm, ok := ms.msgs[seq]
   936  	if !ok {
   937  		ms.mu.Unlock()
   938  		return ErrInvalidSequence
   939  	}
   940  
   941  	for i := ms.state.LastSeq; i > seq; i-- {
   942  		if sm := ms.msgs[i]; sm != nil {
   943  			purged++
   944  			bytes += memStoreMsgSize(sm.subj, sm.hdr, sm.msg)
   945  			delete(ms.msgs, i)
   946  			ms.removeSeqPerSubject(sm.subj, i)
   947  		}
   948  	}
   949  	// Reset last.
   950  	ms.state.LastSeq = lsm.seq
   951  	ms.state.LastTime = time.Unix(0, lsm.ts).UTC()
   952  	// Update msgs and bytes.
   953  	if purged > ms.state.Msgs {
   954  		purged = ms.state.Msgs
   955  	}
   956  	ms.state.Msgs -= purged
   957  	if bytes > ms.state.Bytes {
   958  		bytes = ms.state.Bytes
   959  	}
   960  	ms.state.Bytes -= bytes
   961  
   962  	cb := ms.scb
   963  	ms.mu.Unlock()
   964  
   965  	if cb != nil {
   966  		cb(-int64(purged), -int64(bytes), 0, _EMPTY_)
   967  	}
   968  
   969  	return nil
   970  }
   971  
   972  func (ms *memStore) deleteFirstMsgOrPanic() {
   973  	if !ms.deleteFirstMsg() {
   974  		panic("jetstream memstore has inconsistent state, can't find first seq msg")
   975  	}
   976  }
   977  
   978  func (ms *memStore) deleteFirstMsg() bool {
   979  	return ms.removeMsg(ms.state.FirstSeq, false)
   980  }
   981  
   982  // LoadMsg will lookup the message by sequence number and return it if found.
   983  func (ms *memStore) LoadMsg(seq uint64, smp *StoreMsg) (*StoreMsg, error) {
   984  	ms.mu.RLock()
   985  	sm, ok := ms.msgs[seq]
   986  	last := ms.state.LastSeq
   987  	ms.mu.RUnlock()
   988  
   989  	if !ok || sm == nil {
   990  		var err = ErrStoreEOF
   991  		if seq <= last {
   992  			err = ErrStoreMsgNotFound
   993  		}
   994  		return nil, err
   995  	}
   996  
   997  	if smp == nil {
   998  		smp = new(StoreMsg)
   999  	}
  1000  	sm.copy(smp)
  1001  	return smp, nil
  1002  }
  1003  
  1004  // LoadLastMsg will return the last message we have that matches a given subject.
  1005  // The subject can be a wildcard.
  1006  func (ms *memStore) LoadLastMsg(subject string, smp *StoreMsg) (*StoreMsg, error) {
  1007  	var sm *StoreMsg
  1008  	var ok bool
  1009  
  1010  	// This needs to be a write lock, as filteredStateLocked can
  1011  	// mutate the per-subject state.
  1012  	ms.mu.Lock()
  1013  	defer ms.mu.Unlock()
  1014  
  1015  	if subject == _EMPTY_ || subject == fwcs {
  1016  		sm, ok = ms.msgs[ms.state.LastSeq]
  1017  	} else if subjectIsLiteral(subject) {
  1018  		var ss *SimpleState
  1019  		if ss, ok = ms.fss.Find(stringToBytes(subject)); ok && ss.Msgs > 0 {
  1020  			sm, ok = ms.msgs[ss.Last]
  1021  		}
  1022  	} else if ss := ms.filteredStateLocked(1, subject, true); ss.Msgs > 0 {
  1023  		sm, ok = ms.msgs[ss.Last]
  1024  	}
  1025  	if !ok || sm == nil {
  1026  		return nil, ErrStoreMsgNotFound
  1027  	}
  1028  
  1029  	if smp == nil {
  1030  		smp = new(StoreMsg)
  1031  	}
  1032  	sm.copy(smp)
  1033  	return smp, nil
  1034  }
  1035  
  1036  // LoadNextMsg will find the next message matching the filter subject starting at the start sequence.
  1037  // The filter subject can be a wildcard.
  1038  func (ms *memStore) LoadNextMsg(filter string, wc bool, start uint64, smp *StoreMsg) (*StoreMsg, uint64, error) {
  1039  	ms.mu.Lock()
  1040  	defer ms.mu.Unlock()
  1041  
  1042  	if start < ms.state.FirstSeq {
  1043  		start = ms.state.FirstSeq
  1044  	}
  1045  
  1046  	// If past the end no results.
  1047  	if start > ms.state.LastSeq {
  1048  		return nil, ms.state.LastSeq, ErrStoreEOF
  1049  	}
  1050  
  1051  	if filter == _EMPTY_ {
  1052  		filter = fwcs
  1053  	}
  1054  	isAll := filter == fwcs
  1055  
  1056  	// Skip scan of ms.fss is number of messages in the block are less than
  1057  	// 1/2 the number of subjects in ms.fss. Or we have a wc and lots of fss entries.
  1058  	const linearScanMaxFSS = 256
  1059  	doLinearScan := isAll || 2*int(ms.state.LastSeq-start) < ms.fss.Size() || (wc && ms.fss.Size() > linearScanMaxFSS)
  1060  
  1061  	// Initial setup.
  1062  	fseq, lseq := start, ms.state.LastSeq
  1063  
  1064  	if !doLinearScan {
  1065  		subs := []string{filter}
  1066  		if wc || isAll {
  1067  			subs = subs[:0]
  1068  			ms.fss.Match(stringToBytes(filter), func(subj []byte, val *SimpleState) {
  1069  				subs = append(subs, string(subj))
  1070  			})
  1071  		}
  1072  		fseq, lseq = ms.state.LastSeq, uint64(0)
  1073  		for _, subj := range subs {
  1074  			ss, ok := ms.fss.Find(stringToBytes(subj))
  1075  			if !ok {
  1076  				continue
  1077  			}
  1078  			if ss.firstNeedsUpdate {
  1079  				ms.recalculateFirstForSubj(subj, ss.First, ss)
  1080  			}
  1081  			if ss.First < fseq {
  1082  				fseq = ss.First
  1083  			}
  1084  			if ss.Last > lseq {
  1085  				lseq = ss.Last
  1086  			}
  1087  		}
  1088  		if fseq < start {
  1089  			fseq = start
  1090  		}
  1091  	}
  1092  
  1093  	eq := subjectsEqual
  1094  	if wc {
  1095  		eq = subjectIsSubsetMatch
  1096  	}
  1097  
  1098  	for nseq := fseq; nseq <= lseq; nseq++ {
  1099  		if sm, ok := ms.msgs[nseq]; ok && (isAll || eq(sm.subj, filter)) {
  1100  			if smp == nil {
  1101  				smp = new(StoreMsg)
  1102  			}
  1103  			sm.copy(smp)
  1104  			return smp, nseq, nil
  1105  		}
  1106  	}
  1107  	return nil, ms.state.LastSeq, ErrStoreEOF
  1108  }
  1109  
  1110  // RemoveMsg will remove the message from this store.
  1111  // Will return the number of bytes removed.
  1112  func (ms *memStore) RemoveMsg(seq uint64) (bool, error) {
  1113  	ms.mu.Lock()
  1114  	removed := ms.removeMsg(seq, false)
  1115  	ms.mu.Unlock()
  1116  	return removed, nil
  1117  }
  1118  
  1119  // EraseMsg will remove the message and rewrite its contents.
  1120  func (ms *memStore) EraseMsg(seq uint64) (bool, error) {
  1121  	ms.mu.Lock()
  1122  	removed := ms.removeMsg(seq, true)
  1123  	ms.mu.Unlock()
  1124  	return removed, nil
  1125  }
  1126  
  1127  // Performs logic to update first sequence number.
  1128  // Lock should be held.
  1129  func (ms *memStore) updateFirstSeq(seq uint64) {
  1130  	if seq != ms.state.FirstSeq {
  1131  		// Interior delete.
  1132  		return
  1133  	}
  1134  	var nsm *StoreMsg
  1135  	var ok bool
  1136  	for nseq := ms.state.FirstSeq + 1; nseq <= ms.state.LastSeq; nseq++ {
  1137  		if nsm, ok = ms.msgs[nseq]; ok {
  1138  			break
  1139  		}
  1140  	}
  1141  	oldFirst := ms.state.FirstSeq
  1142  	if nsm != nil {
  1143  		ms.state.FirstSeq = nsm.seq
  1144  		ms.state.FirstTime = time.Unix(0, nsm.ts).UTC()
  1145  	} else {
  1146  		// Like purge.
  1147  		ms.state.FirstSeq = ms.state.LastSeq + 1
  1148  		ms.state.FirstTime = time.Time{}
  1149  	}
  1150  
  1151  	if oldFirst == ms.state.FirstSeq-1 {
  1152  		ms.dmap.Delete(oldFirst)
  1153  	} else {
  1154  		for seq := oldFirst; seq < ms.state.FirstSeq; seq++ {
  1155  			ms.dmap.Delete(seq)
  1156  		}
  1157  	}
  1158  }
  1159  
  1160  // Remove a seq from the fss and select new first.
  1161  // Lock should be held.
  1162  func (ms *memStore) removeSeqPerSubject(subj string, seq uint64) {
  1163  	ss, ok := ms.fss.Find(stringToBytes(subj))
  1164  	if !ok {
  1165  		return
  1166  	}
  1167  	if ss.Msgs == 1 {
  1168  		ms.fss.Delete(stringToBytes(subj))
  1169  		return
  1170  	}
  1171  	ss.Msgs--
  1172  
  1173  	// If we know we only have 1 msg left don't need to search for next first.
  1174  	if ss.Msgs == 1 {
  1175  		if seq == ss.Last {
  1176  			ss.Last = ss.First
  1177  		} else {
  1178  			ss.First = ss.Last
  1179  		}
  1180  		ss.firstNeedsUpdate = false
  1181  	} else {
  1182  		ss.firstNeedsUpdate = seq == ss.First || ss.firstNeedsUpdate
  1183  	}
  1184  }
  1185  
  1186  // Will recalculate the first sequence for this subject in this block.
  1187  // Lock should be held.
  1188  func (ms *memStore) recalculateFirstForSubj(subj string, startSeq uint64, ss *SimpleState) {
  1189  	for tseq := startSeq + 1; tseq <= ss.Last; tseq++ {
  1190  		if sm := ms.msgs[tseq]; sm != nil && sm.subj == subj {
  1191  			ss.First = tseq
  1192  			ss.firstNeedsUpdate = false
  1193  			return
  1194  		}
  1195  	}
  1196  }
  1197  
  1198  // Removes the message referenced by seq.
  1199  // Lock should be held.
  1200  func (ms *memStore) removeMsg(seq uint64, secure bool) bool {
  1201  	var ss uint64
  1202  	sm, ok := ms.msgs[seq]
  1203  	if !ok {
  1204  		return false
  1205  	}
  1206  
  1207  	ss = memStoreMsgSize(sm.subj, sm.hdr, sm.msg)
  1208  
  1209  	delete(ms.msgs, seq)
  1210  	if ms.state.Msgs > 0 {
  1211  		ms.state.Msgs--
  1212  		if ss > ms.state.Bytes {
  1213  			ss = ms.state.Bytes
  1214  		}
  1215  		ms.state.Bytes -= ss
  1216  	}
  1217  	ms.dmap.Insert(seq)
  1218  	ms.updateFirstSeq(seq)
  1219  
  1220  	if secure {
  1221  		if len(sm.hdr) > 0 {
  1222  			sm.hdr = make([]byte, len(sm.hdr))
  1223  			crand.Read(sm.hdr)
  1224  		}
  1225  		if len(sm.msg) > 0 {
  1226  			sm.msg = make([]byte, len(sm.msg))
  1227  			crand.Read(sm.msg)
  1228  		}
  1229  		sm.seq, sm.ts = 0, 0
  1230  	}
  1231  
  1232  	// Remove any per subject tracking.
  1233  	ms.removeSeqPerSubject(sm.subj, seq)
  1234  
  1235  	if ms.scb != nil {
  1236  		// We do not want to hold any locks here.
  1237  		ms.mu.Unlock()
  1238  		delta := int64(ss)
  1239  		ms.scb(-1, -delta, seq, sm.subj)
  1240  		ms.mu.Lock()
  1241  	}
  1242  
  1243  	return ok
  1244  }
  1245  
  1246  // Type returns the type of the underlying store.
  1247  func (ms *memStore) Type() StorageType {
  1248  	return MemoryStorage
  1249  }
  1250  
  1251  // FastState will fill in state with only the following.
  1252  // Msgs, Bytes, First and Last Sequence and Time and NumDeleted.
  1253  func (ms *memStore) FastState(state *StreamState) {
  1254  	ms.mu.RLock()
  1255  	state.Msgs = ms.state.Msgs
  1256  	state.Bytes = ms.state.Bytes
  1257  	state.FirstSeq = ms.state.FirstSeq
  1258  	state.FirstTime = ms.state.FirstTime
  1259  	state.LastSeq = ms.state.LastSeq
  1260  	state.LastTime = ms.state.LastTime
  1261  	if state.LastSeq > state.FirstSeq {
  1262  		state.NumDeleted = int((state.LastSeq - state.FirstSeq + 1) - state.Msgs)
  1263  		if state.NumDeleted < 0 {
  1264  			state.NumDeleted = 0
  1265  		}
  1266  	}
  1267  	state.Consumers = ms.consumers
  1268  	state.NumSubjects = ms.fss.Size()
  1269  	ms.mu.RUnlock()
  1270  }
  1271  
  1272  func (ms *memStore) State() StreamState {
  1273  	ms.mu.RLock()
  1274  	defer ms.mu.RUnlock()
  1275  
  1276  	state := ms.state
  1277  	state.Consumers = ms.consumers
  1278  	state.NumSubjects = ms.fss.Size()
  1279  	state.Deleted = nil
  1280  
  1281  	// Calculate interior delete details.
  1282  	if numDeleted := int((state.LastSeq - state.FirstSeq + 1) - state.Msgs); numDeleted > 0 {
  1283  		state.Deleted = make([]uint64, 0, numDeleted)
  1284  		fseq, lseq := state.FirstSeq, state.LastSeq
  1285  		ms.dmap.Range(func(seq uint64) bool {
  1286  			if seq < fseq || seq > lseq {
  1287  				ms.dmap.Delete(seq)
  1288  			} else {
  1289  				state.Deleted = append(state.Deleted, seq)
  1290  			}
  1291  			return true
  1292  		})
  1293  	}
  1294  	if len(state.Deleted) > 0 {
  1295  		state.NumDeleted = len(state.Deleted)
  1296  	}
  1297  
  1298  	return state
  1299  }
  1300  
  1301  func (ms *memStore) Utilization() (total, reported uint64, err error) {
  1302  	ms.mu.RLock()
  1303  	defer ms.mu.RUnlock()
  1304  	return ms.state.Bytes, ms.state.Bytes, nil
  1305  }
  1306  
  1307  func memStoreMsgSize(subj string, hdr, msg []byte) uint64 {
  1308  	return uint64(len(subj) + len(hdr) + len(msg) + 16) // 8*2 for seq + age
  1309  }
  1310  
  1311  // Delete is same as Stop for memory store.
  1312  func (ms *memStore) Delete() error {
  1313  	return ms.Stop()
  1314  }
  1315  
  1316  func (ms *memStore) Stop() error {
  1317  	// These can't come back, so stop is same as Delete.
  1318  	ms.Purge()
  1319  	ms.mu.Lock()
  1320  	if ms.ageChk != nil {
  1321  		ms.ageChk.Stop()
  1322  		ms.ageChk = nil
  1323  	}
  1324  	ms.msgs = nil
  1325  	ms.mu.Unlock()
  1326  	return nil
  1327  }
  1328  
  1329  func (ms *memStore) isClosed() bool {
  1330  	ms.mu.RLock()
  1331  	defer ms.mu.RUnlock()
  1332  	return ms.msgs == nil
  1333  }
  1334  
  1335  type consumerMemStore struct {
  1336  	mu     sync.Mutex
  1337  	ms     StreamStore
  1338  	cfg    ConsumerConfig
  1339  	state  ConsumerState
  1340  	closed bool
  1341  }
  1342  
  1343  func (ms *memStore) ConsumerStore(name string, cfg *ConsumerConfig) (ConsumerStore, error) {
  1344  	if ms == nil {
  1345  		return nil, fmt.Errorf("memstore is nil")
  1346  	}
  1347  	if ms.isClosed() {
  1348  		return nil, ErrStoreClosed
  1349  	}
  1350  	if cfg == nil || name == _EMPTY_ {
  1351  		return nil, fmt.Errorf("bad consumer config")
  1352  	}
  1353  	o := &consumerMemStore{ms: ms, cfg: *cfg}
  1354  	ms.AddConsumer(o)
  1355  	return o, nil
  1356  }
  1357  
  1358  func (ms *memStore) AddConsumer(o ConsumerStore) error {
  1359  	ms.mu.Lock()
  1360  	ms.consumers++
  1361  	ms.mu.Unlock()
  1362  	return nil
  1363  }
  1364  
  1365  func (ms *memStore) RemoveConsumer(o ConsumerStore) error {
  1366  	ms.mu.Lock()
  1367  	if ms.consumers > 0 {
  1368  		ms.consumers--
  1369  	}
  1370  	ms.mu.Unlock()
  1371  	return nil
  1372  }
  1373  
  1374  func (ms *memStore) Snapshot(_ time.Duration, _, _ bool) (*SnapshotResult, error) {
  1375  	return nil, fmt.Errorf("no impl")
  1376  }
  1377  
  1378  // Binary encoded state snapshot, >= v2.10 server.
  1379  func (ms *memStore) EncodedStreamState(failed uint64) ([]byte, error) {
  1380  	ms.mu.RLock()
  1381  	defer ms.mu.RUnlock()
  1382  
  1383  	// Quick calculate num deleted.
  1384  	numDeleted := int((ms.state.LastSeq - ms.state.FirstSeq + 1) - ms.state.Msgs)
  1385  	if numDeleted < 0 {
  1386  		numDeleted = 0
  1387  	}
  1388  
  1389  	// Encoded is Msgs, Bytes, FirstSeq, LastSeq, Failed, NumDeleted and optional DeletedBlocks
  1390  	var buf [1024]byte
  1391  	buf[0], buf[1] = streamStateMagic, streamStateVersion
  1392  	n := hdrLen
  1393  	n += binary.PutUvarint(buf[n:], ms.state.Msgs)
  1394  	n += binary.PutUvarint(buf[n:], ms.state.Bytes)
  1395  	n += binary.PutUvarint(buf[n:], ms.state.FirstSeq)
  1396  	n += binary.PutUvarint(buf[n:], ms.state.LastSeq)
  1397  	n += binary.PutUvarint(buf[n:], failed)
  1398  	n += binary.PutUvarint(buf[n:], uint64(numDeleted))
  1399  
  1400  	b := buf[0:n]
  1401  
  1402  	if numDeleted > 0 {
  1403  		buf, err := ms.dmap.Encode(nil)
  1404  		if err != nil {
  1405  			return nil, err
  1406  		}
  1407  		b = append(b, buf...)
  1408  	}
  1409  
  1410  	return b, nil
  1411  }
  1412  
  1413  // SyncDeleted will make sure this stream has same deleted state as dbs.
  1414  func (ms *memStore) SyncDeleted(dbs DeleteBlocks) {
  1415  	// For now we share one dmap, so if we have one entry here check if states are the same.
  1416  	// Note this will work for any DeleteBlock type, but we expect this to be a dmap too.
  1417  	if len(dbs) == 1 {
  1418  		ms.mu.RLock()
  1419  		min, max, num := ms.dmap.State()
  1420  		ms.mu.RUnlock()
  1421  		if pmin, pmax, pnum := dbs[0].State(); pmin == min && pmax == max && pnum == num {
  1422  			return
  1423  		}
  1424  	}
  1425  	for _, db := range dbs {
  1426  		db.Range(func(dseq uint64) bool {
  1427  			ms.RemoveMsg(dseq)
  1428  			return true
  1429  		})
  1430  	}
  1431  }
  1432  
  1433  func (o *consumerMemStore) Update(state *ConsumerState) error {
  1434  	// Sanity checks.
  1435  	if state.AckFloor.Consumer > state.Delivered.Consumer {
  1436  		return fmt.Errorf("bad ack floor for consumer")
  1437  	}
  1438  	if state.AckFloor.Stream > state.Delivered.Stream {
  1439  		return fmt.Errorf("bad ack floor for stream")
  1440  	}
  1441  
  1442  	// Copy to our state.
  1443  	var pending map[uint64]*Pending
  1444  	var redelivered map[uint64]uint64
  1445  	if len(state.Pending) > 0 {
  1446  		pending = make(map[uint64]*Pending, len(state.Pending))
  1447  		for seq, p := range state.Pending {
  1448  			pending[seq] = &Pending{p.Sequence, p.Timestamp}
  1449  		}
  1450  		for seq := range pending {
  1451  			if seq <= state.AckFloor.Stream || seq > state.Delivered.Stream {
  1452  				return fmt.Errorf("bad pending entry, sequence [%d] out of range", seq)
  1453  			}
  1454  		}
  1455  	}
  1456  	if len(state.Redelivered) > 0 {
  1457  		redelivered = make(map[uint64]uint64, len(state.Redelivered))
  1458  		for seq, dc := range state.Redelivered {
  1459  			redelivered[seq] = dc
  1460  		}
  1461  	}
  1462  
  1463  	// Replace our state.
  1464  	o.mu.Lock()
  1465  
  1466  	// Check to see if this is an outdated update.
  1467  	if state.Delivered.Consumer < o.state.Delivered.Consumer {
  1468  		o.mu.Unlock()
  1469  		return fmt.Errorf("old update ignored")
  1470  	}
  1471  
  1472  	o.state.Delivered = state.Delivered
  1473  	o.state.AckFloor = state.AckFloor
  1474  	o.state.Pending = pending
  1475  	o.state.Redelivered = redelivered
  1476  	o.mu.Unlock()
  1477  
  1478  	return nil
  1479  }
  1480  
  1481  // SetStarting sets our starting stream sequence.
  1482  func (o *consumerMemStore) SetStarting(sseq uint64) error {
  1483  	o.mu.Lock()
  1484  	o.state.Delivered.Stream = sseq
  1485  	o.mu.Unlock()
  1486  	return nil
  1487  }
  1488  
  1489  // HasState returns if this store has a recorded state.
  1490  func (o *consumerMemStore) HasState() bool {
  1491  	return false
  1492  }
  1493  
  1494  func (o *consumerMemStore) UpdateDelivered(dseq, sseq, dc uint64, ts int64) error {
  1495  	o.mu.Lock()
  1496  	defer o.mu.Unlock()
  1497  
  1498  	if dc != 1 && o.cfg.AckPolicy == AckNone {
  1499  		return ErrNoAckPolicy
  1500  	}
  1501  
  1502  	if dseq <= o.state.AckFloor.Consumer {
  1503  		return nil
  1504  	}
  1505  
  1506  	// See if we expect an ack for this.
  1507  	if o.cfg.AckPolicy != AckNone {
  1508  		// Need to create pending records here.
  1509  		if o.state.Pending == nil {
  1510  			o.state.Pending = make(map[uint64]*Pending)
  1511  		}
  1512  		var p *Pending
  1513  		// Check for an update to a message already delivered.
  1514  		if sseq <= o.state.Delivered.Stream {
  1515  			if p = o.state.Pending[sseq]; p != nil {
  1516  				p.Sequence, p.Timestamp = dseq, ts
  1517  			}
  1518  		} else {
  1519  			// Add to pending.
  1520  			o.state.Pending[sseq] = &Pending{dseq, ts}
  1521  		}
  1522  		// Update delivered as needed.
  1523  		if dseq > o.state.Delivered.Consumer {
  1524  			o.state.Delivered.Consumer = dseq
  1525  		}
  1526  		if sseq > o.state.Delivered.Stream {
  1527  			o.state.Delivered.Stream = sseq
  1528  		}
  1529  
  1530  		if dc > 1 {
  1531  			if maxdc := uint64(o.cfg.MaxDeliver); maxdc > 0 && dc > maxdc {
  1532  				// Make sure to remove from pending.
  1533  				delete(o.state.Pending, sseq)
  1534  			}
  1535  			if o.state.Redelivered == nil {
  1536  				o.state.Redelivered = make(map[uint64]uint64)
  1537  			}
  1538  			// Only update if greater than what we already have.
  1539  			if o.state.Redelivered[sseq] < dc-1 {
  1540  				o.state.Redelivered[sseq] = dc - 1
  1541  			}
  1542  		}
  1543  	} else {
  1544  		// For AckNone just update delivered and ackfloor at the same time.
  1545  		if dseq > o.state.Delivered.Consumer {
  1546  			o.state.Delivered.Consumer = dseq
  1547  			o.state.AckFloor.Consumer = dseq
  1548  		}
  1549  		if sseq > o.state.Delivered.Stream {
  1550  			o.state.Delivered.Stream = sseq
  1551  			o.state.AckFloor.Stream = sseq
  1552  		}
  1553  	}
  1554  
  1555  	return nil
  1556  }
  1557  
  1558  func (o *consumerMemStore) UpdateAcks(dseq, sseq uint64) error {
  1559  	o.mu.Lock()
  1560  	defer o.mu.Unlock()
  1561  
  1562  	if o.cfg.AckPolicy == AckNone {
  1563  		return ErrNoAckPolicy
  1564  	}
  1565  	if len(o.state.Pending) == 0 || o.state.Pending[sseq] == nil {
  1566  		return ErrStoreMsgNotFound
  1567  	}
  1568  
  1569  	// On restarts the old leader may get a replay from the raft logs that are old.
  1570  	if dseq <= o.state.AckFloor.Consumer {
  1571  		return nil
  1572  	}
  1573  
  1574  	// Check for AckAll here.
  1575  	if o.cfg.AckPolicy == AckAll {
  1576  		sgap := sseq - o.state.AckFloor.Stream
  1577  		o.state.AckFloor.Consumer = dseq
  1578  		o.state.AckFloor.Stream = sseq
  1579  		for seq := sseq; seq > sseq-sgap; seq-- {
  1580  			delete(o.state.Pending, seq)
  1581  			if len(o.state.Redelivered) > 0 {
  1582  				delete(o.state.Redelivered, seq)
  1583  			}
  1584  		}
  1585  		return nil
  1586  	}
  1587  
  1588  	// AckExplicit
  1589  
  1590  	// First delete from our pending state.
  1591  	if p, ok := o.state.Pending[sseq]; ok {
  1592  		delete(o.state.Pending, sseq)
  1593  		dseq = p.Sequence // Use the original.
  1594  	}
  1595  	// Now remove from redelivered.
  1596  	if len(o.state.Redelivered) > 0 {
  1597  		delete(o.state.Redelivered, sseq)
  1598  	}
  1599  
  1600  	if len(o.state.Pending) == 0 {
  1601  		o.state.AckFloor.Consumer = o.state.Delivered.Consumer
  1602  		o.state.AckFloor.Stream = o.state.Delivered.Stream
  1603  	} else if dseq == o.state.AckFloor.Consumer+1 {
  1604  		first := o.state.AckFloor.Consumer == 0
  1605  		o.state.AckFloor.Consumer = dseq
  1606  		o.state.AckFloor.Stream = sseq
  1607  
  1608  		if !first && o.state.Delivered.Consumer > dseq {
  1609  			for ss := sseq + 1; ss < o.state.Delivered.Stream; ss++ {
  1610  				if p, ok := o.state.Pending[ss]; ok {
  1611  					if p.Sequence > 0 {
  1612  						o.state.AckFloor.Consumer = p.Sequence - 1
  1613  						o.state.AckFloor.Stream = ss - 1
  1614  					}
  1615  					break
  1616  				}
  1617  			}
  1618  		}
  1619  	}
  1620  
  1621  	return nil
  1622  }
  1623  
  1624  func (o *consumerMemStore) UpdateConfig(cfg *ConsumerConfig) error {
  1625  	o.mu.Lock()
  1626  	defer o.mu.Unlock()
  1627  
  1628  	// This is mostly unchecked here. We are assuming the upper layers have done sanity checking.
  1629  	o.cfg = *cfg
  1630  	return nil
  1631  }
  1632  
  1633  func (o *consumerMemStore) Stop() error {
  1634  	o.mu.Lock()
  1635  	o.closed = true
  1636  	ms := o.ms
  1637  	o.mu.Unlock()
  1638  	ms.RemoveConsumer(o)
  1639  	return nil
  1640  }
  1641  
  1642  func (o *consumerMemStore) Delete() error {
  1643  	return o.Stop()
  1644  }
  1645  
  1646  func (o *consumerMemStore) StreamDelete() error {
  1647  	return o.Stop()
  1648  }
  1649  
  1650  func (o *consumerMemStore) State() (*ConsumerState, error) {
  1651  	return o.stateWithCopy(true)
  1652  }
  1653  
  1654  // This will not copy pending or redelivered, so should only be done under the
  1655  // consumer owner's lock.
  1656  func (o *consumerMemStore) BorrowState() (*ConsumerState, error) {
  1657  	return o.stateWithCopy(false)
  1658  }
  1659  
  1660  func (o *consumerMemStore) stateWithCopy(doCopy bool) (*ConsumerState, error) {
  1661  	o.mu.Lock()
  1662  	defer o.mu.Unlock()
  1663  
  1664  	if o.closed {
  1665  		return nil, ErrStoreClosed
  1666  	}
  1667  
  1668  	state := &ConsumerState{}
  1669  
  1670  	state.Delivered = o.state.Delivered
  1671  	state.AckFloor = o.state.AckFloor
  1672  	if len(o.state.Pending) > 0 {
  1673  		if doCopy {
  1674  			state.Pending = o.copyPending()
  1675  		} else {
  1676  			state.Pending = o.state.Pending
  1677  		}
  1678  	}
  1679  	if len(o.state.Redelivered) > 0 {
  1680  		if doCopy {
  1681  			state.Redelivered = o.copyRedelivered()
  1682  		} else {
  1683  			state.Redelivered = o.state.Redelivered
  1684  		}
  1685  	}
  1686  	return state, nil
  1687  }
  1688  
  1689  // EncodeState for this consumer store.
  1690  func (o *consumerMemStore) EncodedState() ([]byte, error) {
  1691  	o.mu.Lock()
  1692  	defer o.mu.Unlock()
  1693  
  1694  	if o.closed {
  1695  		return nil, ErrStoreClosed
  1696  	}
  1697  
  1698  	return encodeConsumerState(&o.state), nil
  1699  }
  1700  
  1701  func (o *consumerMemStore) copyPending() map[uint64]*Pending {
  1702  	pending := make(map[uint64]*Pending, len(o.state.Pending))
  1703  	for seq, p := range o.state.Pending {
  1704  		pending[seq] = &Pending{p.Sequence, p.Timestamp}
  1705  	}
  1706  	return pending
  1707  }
  1708  
  1709  func (o *consumerMemStore) copyRedelivered() map[uint64]uint64 {
  1710  	redelivered := make(map[uint64]uint64, len(o.state.Redelivered))
  1711  	for seq, dc := range o.state.Redelivered {
  1712  		redelivered[seq] = dc
  1713  	}
  1714  	return redelivered
  1715  }
  1716  
  1717  // Type returns the type of the underlying store.
  1718  func (o *consumerMemStore) Type() StorageType { return MemoryStorage }
  1719  
  1720  // Templates
  1721  type templateMemStore struct{}
  1722  
  1723  func newTemplateMemStore() *templateMemStore {
  1724  	return &templateMemStore{}
  1725  }
  1726  
  1727  // No-ops for memstore.
  1728  func (ts *templateMemStore) Store(t *streamTemplate) error  { return nil }
  1729  func (ts *templateMemStore) Delete(t *streamTemplate) error { return nil }