reconcile consumer state to stream seq

Author: tilakraj94

server/consumer.go

@@ -963,7 +963,7 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 	}
 
 	mset.mu.RLock()
-	s, js, jsa, cfg, acc := mset.srv, mset.js, mset.jsa, mset.cfg, mset.acc
+	s, js, jsa, cfg, acc, lseq := mset.srv, mset.js, mset.jsa, mset.cfg, mset.acc, mset.lseq
 	mset.mu.RUnlock()
 
 	// If we do not have the consumer currently assigned to us in cluster mode we will proceed but warn.
@@ -1227,7 +1227,49 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 		// Restore our saved state.
 		o.mu.Lock()
 		o.readStoredState()
+
+		replicas := o.cfg.replicas(&mset.cfg)
+
+		// Starting sequence represents the next sequence to be delivered, so decrement it
+		// since that's the minimum amount the stream should have as its last sequence.
+		sseq := o.sseq
+		if sseq > 0 {
+			sseq--
+		}
+
 		o.mu.Unlock()
+
+		// A stream observing data loss rolls back in its sequence. Check if we need to reconcile the consumer state
+		// to ensure new messages aren't skipped.
+		// Only performed for non-replicated consumers for now.
+		if replicas == 1 && lseq < sseq && isRecovering {
+			s.Warnf("JetStream consumer '%s > %s > %s' delivered sequence %d past last stream sequence of %d",
+				o.acc.Name, o.stream, o.name, sseq, lseq)
+
+			o.mu.Lock()
+			o.reconcileStateWithStream(lseq)
+
+			// Save the reconciled state
+			state := &ConsumerState{
+				Delivered: SequencePair{
+					Stream:   o.sseq - 1,
+					Consumer: o.dseq - 1,
+				},
+				AckFloor: SequencePair{
+					Stream:   o.asflr,
+					Consumer: o.adflr,
+				},
+				Pending:     o.pending,
+				Redelivered: o.rdc,
+			}
+			err := o.store.ForceUpdate(state)
+			o.mu.Unlock()
+			if err != nil {
+				s.Errorf("JetStream consumer '%s > %s > %s' errored while updating state: %v", o.acc.Name, o.stream, o.name, err)
+				mset.mu.Unlock()
+				return nil, NewJSConsumerStoreFailedError(err)
+			}
+		}
 	} else {
 		// Select starting sequence number
 		o.selectStartingSeqNo()
@@ -5710,6 +5752,67 @@ func (o *consumer) hasSkipListPending() bool {
 	return o.lss != nil && len(o.lss.seqs) > 0
 }
 
+// reconcileStateWithStream reconciles consumer state when the stream has reverted
+// due to data loss (e.g., VM crash). This handles the case where consumer state
+// is ahead of the stream's last sequence.
+// Lock should be held.
+func (o *consumer) reconcileStateWithStream(streamLastSeq uint64) {
+	// If an ack floor is higher than stream last sequence,
+	// reset back down but keep the highest known sequences.
+	if o.asflr > streamLastSeq {
+		o.asflr = streamLastSeq
+		// Delivery floor is one below the delivered sequence,
+		// but if it is zero somehow, ensure we don't underflow.
+		o.adflr = o.dseq
+		if o.adflr > 0 {
+			o.adflr--
+		}
+		o.pending = nil
+		o.rdc = nil
+	}
+
+	// Remove pending entries that are beyond the stream's last sequence
+	if len(o.pending) > 0 {
+		for seq := range o.pending {
+			if seq > streamLastSeq {
+				delete(o.pending, seq)
+			}
+		}
+	}
+
+	// Remove redelivered entries that are beyond the stream's last sequence
+	if len(o.rdc) > 0 {
+		for seq := range o.rdc {
+			if seq > streamLastSeq {
+				delete(o.rdc, seq)
+			}
+		}
+	}
+
+	// Update starting sequence and delivery sequence based on pending state
+	if len(o.pending) == 0 {
+		o.sseq = o.asflr + 1
+		o.dseq = o.adflr + 1
+	} else {
+		// Find highest stream sequence in pending
+		var maxStreamSeq uint64
+		var maxConsumerSeq uint64
+
+		for streamSeq, p := range o.pending {
+			if streamSeq > maxStreamSeq {
+				maxStreamSeq = streamSeq
+			}
+			if p.Sequence > maxConsumerSeq {
+				maxConsumerSeq = p.Sequence
+			}
+		}
+
+		// Set next sequences based on highest pending
+		o.sseq = maxStreamSeq + 1
+		o.dseq = maxConsumerSeq + 1
+	}
+}
+
 // Will select the starting sequence.
 func (o *consumer) selectStartingSeqNo() {
 	if o.mset == nil || o.mset.store == nil {

server/filestore.go

@@ -11654,6 +11654,52 @@ func (o *consumerFileStore) Update(state *ConsumerState) error {
 	return nil
 }
 
+// ForceUpdate updates the consumer state without the backwards check.
+// This is used during recovery when we need to reset the consumer to an earlier sequence.
+func (o *consumerFileStore) ForceUpdate(state *ConsumerState) error {
+	o.fs.warn("Consumer %q force updating state: %+v", o.name, state)
+	// Sanity checks.
+	if state.AckFloor.Consumer > state.Delivered.Consumer {
+		return fmt.Errorf("bad ack floor for consumer")
+	}
+	if state.AckFloor.Stream > state.Delivered.Stream {
+		return fmt.Errorf("bad ack floor for stream")
+	}
+
+	// Copy to our state.
+	var pending map[uint64]*Pending
+	var redelivered map[uint64]uint64
+	if len(state.Pending) > 0 {
+		pending = make(map[uint64]*Pending, len(state.Pending))
+		for seq, p := range state.Pending {
+			pending[seq] = &Pending{p.Sequence, p.Timestamp}
+			if seq <= state.AckFloor.Stream || seq > state.Delivered.Stream {
+				return fmt.Errorf("bad pending entry, sequence [%d] out of range", seq)
+			}
+		}
+	}
+	if len(state.Redelivered) > 0 {
+		redelivered = make(map[uint64]uint64, len(state.Redelivered))
+		for seq, dc := range state.Redelivered {
+			redelivered[seq] = dc
+		}
+	}
+
+	// Replace our state.
+	o.mu.Lock()
+
+	o.state.Delivered = state.Delivered
+	o.state.AckFloor = state.AckFloor
+	o.state.Pending = pending
+	o.state.Redelivered = redelivered
+	buf, err := o.encodeState()
+	o.mu.Unlock()
+	if err != nil {
+		return err
+	}
+	return o.writeState(buf)
+}
+
 // Will encrypt the state with our asset key. Will be a no-op if encryption not enabled.
 // Lock should be held.
 func (o *consumerFileStore) encryptState(buf []byte) ([]byte, error) {

server/filestore_test.go

@@ -8518,6 +8518,94 @@ func TestFileStoreRecoverFullStateDetectCorruptState(t *testing.T) {
 	require_Error(t, err, errCorruptState)
 }
 
+func TestFileStoreResetConsumerToStreamState(t *testing.T) {
+	fs, err := newFileStore(
+		FileStoreConfig{StoreDir: t.TempDir()},
+		StreamConfig{Name: "zzz", Subjects: []string{"foo.*"}, Storage: FileStorage})
+	require_NoError(t, err)
+	defer fs.Stop()
+
+	msg := []byte("abc")
+	for i := 1; i <= 30; i++ {
+		_, _, err = fs.StoreMsg(fmt.Sprintf("foo.%d", i), nil, msg, 0)
+		require_NoError(t, err)
+	}
+
+	err = fs.writeFullState()
+	require_NoError(t, err)
+
+	obs, err := fs.ConsumerStore("c1", time.Now(), &ConsumerConfig{
+		Durable:       "c1",
+		FilterSubject: "foo.*",
+		AckPolicy:     AckNone,
+		DeliverPolicy: DeliverAll,
+	})
+
+	require_NoError(t, err)
+	defer obs.Stop()
+
+	state := &ConsumerState{}
+	state.Delivered = SequencePair{
+		Consumer: 5,
+		Stream:   5,
+	}
+
+	state.AckFloor = SequencePair{
+		Consumer: 5,
+		Stream:   5,
+	}
+
+	// set to 5
+	err = obs.Update(state)
+	require_NoError(t, err)
+
+	currState, err := obs.State()
+	require_NoError(t, err)
+
+	require_Equal(t, fs.state.LastSeq, uint64(30))
+	require_Equal(t, fs.state.FirstSeq, uint64(1))
+	require_Equal(t, currState.AckFloor, state.AckFloor)
+	require_Equal(t, currState.Delivered, state.Delivered)
+	require_Equal(t, len(currState.Redelivered), len(state.Redelivered))
+	require_Equal(t, len(currState.Pending), len(state.Pending))
+
+	// reset our state to 0.
+	fs.state.FirstSeq = 0
+	fs.state.LastSeq = 0
+
+	// set back to lower values
+	newState := &ConsumerState{}
+	newState.Delivered = SequencePair{
+		Consumer: 1,
+		Stream:   4,
+	}
+
+	newState.AckFloor = SequencePair{
+		Consumer: 1,
+		Stream:   3,
+	}
+
+	// update should fail but force update should pass
+	err = obs.Update(newState)
+	require_Error(t, err, fmt.Errorf("old update ignored"))
+
+	err = obs.ForceUpdate(newState)
+	require_NoError(t, err)
+
+	fullState, err := obs.State()
+	require_NoError(t, err)
+
+	borrowedState, err := obs.BorrowState()
+	require_NoError(t, err)
+
+	require_Equal(t, fullState.Delivered, borrowedState.Delivered)
+	require_Equal(t, fullState.AckFloor, borrowedState.AckFloor)
+	require_Equal(t, newState.Delivered, borrowedState.Delivered)
+	require_Equal(t, newState.AckFloor, borrowedState.AckFloor)
+	require_Equal(t, len(newState.Redelivered), len(borrowedState.Redelivered))
+	require_Equal(t, len(newState.Pending), len(borrowedState.Pending))
+}
+
 func TestFileStoreNumPendingMulti(t *testing.T) {
 	fs, err := newFileStore(
 		FileStoreConfig{StoreDir: t.TempDir()},

server/memstore.go

@@ -2311,6 +2311,46 @@ func (o *consumerMemStore) Update(state *ConsumerState) error {
 	return nil
 }
 
+func (o *consumerMemStore) ForceUpdate(state *ConsumerState) error {
+	// Sanity checks.
+	if state.AckFloor.Consumer > state.Delivered.Consumer {
+		return fmt.Errorf("bad ack floor for consumer")
+	}
+	if state.AckFloor.Stream > state.Delivered.Stream {
+		return fmt.Errorf("bad ack floor for stream")
+	}
+
+	// Copy to our state.
+	var pending map[uint64]*Pending
+	var redelivered map[uint64]uint64
+	if len(state.Pending) > 0 {
+		pending = make(map[uint64]*Pending, len(state.Pending))
+		for seq, p := range state.Pending {
+			pending[seq] = &Pending{p.Sequence, p.Timestamp}
+			if seq <= state.AckFloor.Stream || seq > state.Delivered.Stream {
+				return fmt.Errorf("bad pending entry, sequence [%d] out of range", seq)
+			}
+		}
+	}
+	if len(state.Redelivered) > 0 {
+		redelivered = make(map[uint64]uint64, len(state.Redelivered))
+		for seq, dc := range state.Redelivered {
+			redelivered[seq] = dc
+		}
+	}
+
+	// Replace our state.
+	o.mu.Lock()
+	defer o.mu.Unlock()
+
+	o.state.Delivered = state.Delivered
+	o.state.AckFloor = state.AckFloor
+	o.state.Pending = pending
+	o.state.Redelivered = redelivered
+
+	return nil
+}
+
 // SetStarting sets our starting stream sequence.
 func (o *consumerMemStore) SetStarting(sseq uint64) error {
 	o.mu.Lock()

server/store.go

@@ -363,6 +363,7 @@ type ConsumerStore interface {
 	UpdateAcks(dseq, sseq uint64) error
 	UpdateConfig(cfg *ConsumerConfig) error
 	Update(*ConsumerState) error
+	ForceUpdate(*ConsumerState) error
 	State() (*ConsumerState, error)
 	BorrowState() (*ConsumerState, error)
 	EncodedState() ([]byte, error)