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 }