reset consumer state to stream state

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.
@@ -1228,6 +1228,40 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 		o.mu.Lock()
 		o.readStoredState()
 		o.mu.Unlock()
+
+		// only for streams with r1; if we are recovering and the consumer sequence is greater than the last sequence,
+		// we need to reconcile the consumer state with the stream
+		// run only for single replica streams
+		if mset.cfg.Replicas == 1 && o.sseq > 2 && lseq < o.sseq && isRecovering {
+			o.srv.Warnf("Consumer %q recovering with sseq %d greater than stream last seq %d, reconciling state", o.name, o.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,
+			}
+			o.mu.Unlock()
+
+			err := o.store.ForceUpdate(state)
+			if err != nil {
+				s.Warnf("Consumer %q error updating state: %v", o.name, err)
+				return nil, NewJSConsumerStoreFailedError(err)
+			}
+
+			o.srv.Noticef("Consumer %q reconciled state: sseq=%d dseq=%d asflr=%d adflr=%d pending=%d",
+				o.name, o.sseq, o.dseq, o.asflr, o.adflr, len(o.pending))
+		}
 	} else {
 		// Select starting sequence number
 		o.selectStartingSeqNo()
@@ -5712,6 +5746,80 @@ 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) {
+	o.srv.Noticef("Consumer %q reconciling state with stream (lseq=%d): before sseq=%d dseq=%d asflr=%d adflr=%d pending=%d rdc=%d",
+		o.name, streamLastSeq, o.sseq, o.dseq, o.asflr, o.adflr, len(o.pending), len(o.rdc))
+
+	// If an ack floor is higher than stream last sequence, reset everything
+	if o.asflr > streamLastSeq {
+		o.srv.Warnf("Consumer %q ack floor (%d) > stream last seq (%d), resetting all state", o.name, o.asflr, streamLastSeq)
+
+		o.sseq = streamLastSeq + 1
+		o.dseq = 1
+		o.adflr = 0
+		o.asflr = streamLastSeq
+		o.pending = nil
+		o.rdc = nil
+		return
+	}
+
+	// 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
+	}
+
+	// Ensure we don't go beyond the stream
+	if o.sseq > streamLastSeq+1 {
+		o.sseq = streamLastSeq + 1
+	}
+
+	o.srv.Warnf("Consumer %q reconciled state: after sseq=%d dseq=%d asflr=%d adflr=%d pending=%d rdc=%d",
+		o.name, o.sseq, o.dseq, o.asflr, o.adflr, len(o.pending), len(o.rdc))
+}
+
+
 // Will select the starting sequence.
 func (o *consumer) selectStartingSeqNo() {
 	if o.mset == nil || o.mset.store == nil {

server/filestore.go

@@ -11593,6 +11593,51 @@ 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()
+	defer o.mu.Unlock()
+
+	o.state.Delivered = state.Delivered
+	o.state.AckFloor = state.AckFloor
+	o.state.Pending = pending
+	o.state.Redelivered = redelivered
+
+	o.kickFlusher()
+
+	return nil
+}
+
 // 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)