snapshot of https://github.com/nats-io/nats-server.git (main, 2026-09-15; blobs > 131072B omitted)

server/stream_backup.go

@@ -0,0 +1,574 @@
+// Copyright 2026 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package server
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"math"
+	"net"
+	"path"
+	"slices"
+	"strings"
+	"time"
+
+	"github.com/klauspost/compress/s2"
+	"github.com/nats-io/nats-server/v2/server/archive"
+)
+
+type SnapshotConsumerState struct {
+	*ConsumerConfig `json:"config"`
+	*ConsumerState  `json:"state"`
+}
+
+// Create a snapshot of this stream and its consumer's state along with messages.
+// sa is passed in when the stream is clustered, so we can find child consumer assignments.
+func (js *jetStream) CreateStreamSnapshotV2(store StreamStore, deadline time.Duration, includeConsumers bool, sa *streamAssignment) (*SnapshotResult, error) {
+	pr, pw := net.Pipe()
+
+	// Set a write deadline here to protect ourselves.
+	if deadline > 0 {
+		pw.SetWriteDeadline(time.Now().Add(deadline))
+	}
+
+	// We can add to our stream while snapshotting but not "user" delete anything.
+	var state StreamState
+	store.FastState(&state)
+
+	// Stream in separate Go routine.
+	errCh := make(chan error, 1)
+	go js.streamSnapshotV2(store, &state, pw, includeConsumers, sa, errCh)
+
+	return &SnapshotResult{pr, state, errCh}, nil
+}
+
+// Stream our snapshot through S2 compression and the custom archive format.
+func (js *jetStream) streamSnapshotV2(store StreamStore, state *StreamState, w io.WriteCloser, includeConsumers bool, sa *streamAssignment, errCh chan error) {
+	defer close(errCh)
+	defer w.Close()
+
+	enc := s2.NewWriter(w)
+	defer enc.Close()
+
+	tw := archive.NewWriter(enc)
+	defer tw.Close()
+
+	now := time.Now()
+	clustered := js.isClustered()
+
+	writeGeneric := func(name string, mod int64, seq uint64, headerSize, payloadSize int64, buf []byte) error {
+		hdr := &archive.Header{
+			Name:        name,
+			Timestamp:   mod,
+			Sequence:    seq,
+			HeaderSize:  headerSize,
+			PayloadSize: payloadSize,
+		}
+		if err := tw.WriteHeader(hdr); err != nil {
+			return err
+		}
+		if _, err := tw.Write(buf); err != nil {
+			return err
+		}
+		// Need to wait for flush here as the tar/s2 writer is handing off to a
+		// flow-controlled publisher, it's important that we handle backpressure.
+		return tw.Flush()
+	}
+
+	writeStoreMsg := func(msg *StoreMsg) error {
+		return writeGeneric(
+			msg.subj,
+			msg.ts,
+			msg.seq,
+			int64(len(msg.hdr)),
+			int64(len(msg.msg)),
+			msg.buf,
+		)
+	}
+
+	writeConsumerMsg := func(scs SnapshotConsumerState) error {
+		// Bound the consumer state to the stream snapshot. Consumer sequence
+		// numbers are left intact since filtered consumers do not have a
+		// one-to-one mapping between consumer and stream sequences.
+		scs.Delivered.Stream = min(scs.Delivered.Stream, state.LastSeq)
+		scs.AckFloor.Stream = min(scs.AckFloor.Stream, state.LastSeq)
+		for seq := range scs.Pending {
+			if seq > state.LastSeq {
+				delete(scs.Pending, seq)
+			}
+		}
+		for seq := range scs.Redelivered {
+			if seq > state.LastSeq {
+				delete(scs.Redelivered, seq)
+			}
+		}
+		ssj, err := json.Marshal(scs)
+		if err != nil {
+			return err
+		}
+		return writeGeneric(
+			path.Join("consumers", scs.Name),
+			now.UnixNano(),
+			0,
+			0,
+			int64(len(ssj)),
+			ssj,
+		)
+	}
+
+	// If we aren't including consumers here then make sure the consumer count
+	// is set accordingly, this helps on the restore path.
+	var consumerAssignments map[string]*consumerAssignment
+	var consumerStores []ConsumerStore
+	var streamState = *state
+	if !includeConsumers {
+		streamState.Consumers = 0
+	} else if clustered {
+		if sa == nil {
+			errCh <- errors.New("stream assignment not present in clustered mode")
+			return
+		}
+		js.mu.RLock()
+		consumerAssignments = make(map[string]*consumerAssignment, len(sa.consumers))
+		for name, ca := range sa.consumers {
+			consumerAssignments[name] = ca.copyGroup()
+		}
+		streamState.Consumers = len(consumerAssignments)
+		js.mu.RUnlock()
+	} else {
+		consumerStores = slices.Collect(store.Consumers())
+		streamState.Consumers = len(consumerStores)
+	}
+
+	ssj, err := json.Marshal(streamState)
+	if err != nil {
+		errCh <- err
+		return
+	}
+	if err := writeGeneric("state.json", now.UnixNano(), 0, 0, int64(len(ssj)), ssj); err != nil {
+		errCh <- err
+		return
+	}
+
+	// Do consumers first, if the stream is interest/WQ then this may be
+	// important for message retention.
+	if includeConsumers {
+		consumerStateFromInfo := func(ci *ConsumerInfo) *ConsumerState {
+			state := &ConsumerState{
+				Delivered: SequencePair{
+					Consumer: ci.Delivered.Consumer,
+					Stream:   ci.Delivered.Stream,
+				},
+				AckFloor: SequencePair{
+					Consumer: ci.AckFloor.Consumer,
+					Stream:   ci.AckFloor.Stream,
+				},
+			}
+			if ci.NumAckPending > 0 {
+				// Cluster consumer info does not include the sparse pending state.
+				// Roll back to the ack floor so all possibly unacknowledged messages
+				// are delivered again instead of guessing which ones are pending.
+				state.Delivered = state.AckFloor
+			}
+			return state
+		}
+
+		if clustered {
+			for _, ca := range consumerAssignments {
+				ci, err := sysRequest[ConsumerInfo](js.srv, clusterConsumerInfoT, sa.Client.serviceAccount(), sa.Config.Name, ca.Name)
+				if err != nil || ci == nil {
+					errCh <- fmt.Errorf("failed to get consumer state for '%s > %s'", sa.Config.Name, ca.Name)
+					return
+				}
+				if err := writeConsumerMsg(SnapshotConsumerState{
+					ConsumerConfig: ca.Config,
+					ConsumerState:  consumerStateFromInfo(ci),
+				}); err != nil {
+					errCh <- err
+					return
+				}
+			}
+		} else {
+			for _, o := range consumerStores {
+				config := o.GetConfig()
+				state, err := o.State()
+				if err != nil {
+					errCh <- fmt.Errorf("couldn't load consumer '%s' state: %s", config.Name, err)
+					return
+				}
+				if err := writeConsumerMsg(SnapshotConsumerState{
+					ConsumerConfig: config,
+					ConsumerState:  state,
+				}); err != nil {
+					errCh <- err
+					return
+				}
+			}
+		}
+	}
+
+	var sm StoreMsg
+	for seq := state.FirstSeq - 1; seq < state.LastSeq; {
+		if _, seq, err = store.LoadNextMsg(fwcs, true, seq+1, &sm); err != nil {
+			if err == ErrStoreEOF {
+				break
+			}
+			errCh <- fmt.Errorf("couldn't load next message after seq %d: %s", seq+1, err)
+			return
+		}
+		// Concurrent deletions and publishes can move the next available
+		// message beyond the snapshot's original sequence range.
+		if seq > state.LastSeq {
+			break
+		}
+		if err = writeStoreMsg(&sm); err != nil {
+			errCh <- err
+			return
+		}
+	}
+
+	// End of backup sentinel. A clear marker makes it obvious when
+	// a backup has been truncated or not without having to count
+	// messages or from first/last sequence, which may not be possible
+	// during the rewrite of a large stream backup.
+	if err = writeGeneric(_EMPTY_, 0, 0, 0, 0, nil); err != nil {
+		errCh <- err
+	}
+}
+
+// RestoreStreamSnapshotV2 will restore a stream from a snapshot.
+func (a *Account) RestoreStreamV2(ncfg *StreamConfig, r io.Reader) (retMset *stream, retErr error) {
+	dec := s2.NewReader(r)
+	tr := archive.NewReader(dec)
+
+	var nstate StreamState
+
+	// Load the stream state.
+	hdr, err := tr.Next()
+	if err != nil {
+		return nil, err
+	}
+	if hdr.Name != "state.json" {
+		return nil, fmt.Errorf("expected state.json first")
+	}
+	if err := json.NewDecoder(tr).Decode(&nstate); err != nil {
+		return nil, fmt.Errorf("error in state.json: %w", err)
+	}
+	if nstate.Consumers < 0 {
+		return nil, fmt.Errorf("invalid stream state: negative consumer count")
+	}
+	if nstate.FirstSeq == 0 {
+		if nstate.LastSeq != 0 || nstate.Msgs != 0 {
+			return nil, fmt.Errorf("invalid stream state: inconsistent message count or sequence range")
+		}
+	} else if nstate.FirstSeq > nstate.LastSeq {
+		// An empty, advanced stream has FirstSeq exactly one past LastSeq.
+		// Express this by subtracting from FirstSeq so a corrupt LastSeq of
+		// MaxUint64 can not overflow here.
+		if nstate.FirstSeq-1 != nstate.LastSeq || nstate.Msgs != 0 {
+			return nil, fmt.Errorf("invalid stream state: inconsistent message count or sequence range")
+		}
+	} else if nstate.Msgs > nstate.LastSeq-nstate.FirstSeq+1 {
+		return nil, fmt.Errorf("invalid stream state: message count exceeds sequence range")
+	}
+
+	s, jsa, err := a.checkForJetStream()
+	if err != nil {
+		return nil, err
+	}
+	js := jsa.js
+	if js == nil {
+		return nil, NewJSNotEnabledForAccountError()
+	}
+	if _, err := a.lookupStream(ncfg.Name); err == nil {
+		return nil, NewJSStreamNameExistRestoreFailedError()
+	}
+
+	cfg, apiErr := s.checkStreamCfg(ncfg, a, false)
+	if apiErr != nil {
+		return nil, apiErr
+	}
+
+	_, isClustered := jsa.jetStreamAndClustered()
+	jsa.usageMu.RLock()
+	selected, tier, hasTier := jsa.selectLimits(cfg.Replicas)
+	jsa.usageMu.RUnlock()
+	reserved := int64(0)
+	if hasTier {
+		if isClustered {
+			js.mu.RLock()
+			_, reserved = js.tieredStreamAndReservationCount(a.Name, tier, &cfg)
+			js.mu.RUnlock()
+		} else {
+			reserved = jsa.tieredReservation(tier, &cfg)
+		}
+	}
+
+	// Hold the full restore footprint as usage before creating the stream. As
+	// messages are stored, this temporary reservation is converted into actual
+	// usage so other streams cannot consume the capacity needed by the restore.
+	restoreBytes := nstate.Bytes
+	if restoreBytes > math.MaxInt64 {
+		if cfg.Storage == MemoryStorage {
+			return nil, NewJSMemoryResourcesExceededError()
+		}
+		return nil, NewJSStorageResourcesExceededError()
+	}
+	restoreRemaining := int64(restoreBytes)
+	jsa.updateUsage(tier, cfg.Storage, restoreRemaining)
+	defer func() {
+		// Inside a defer func() so that restoreRemaining is captured later.
+		jsa.updateUsage(tier, cfg.Storage, -restoreRemaining)
+	}()
+	releaseRestoreBytes := func(bytes int64) {
+		bytes = min(bytes, restoreRemaining)
+		jsa.updateUsage(tier, cfg.Storage, -bytes)
+		restoreRemaining -= bytes
+	}
+	checkUsageLimits := func() error {
+		lr := int64(max(1, cfg.Replicas))
+		if tier == _EMPTY_ {
+			lr = 1
+		}
+		jsa.usageMu.RLock()
+		usage := jsa.usage[tier]
+		var used, limit int64
+		if cfg.Storage == MemoryStorage {
+			limit = selected.MaxMemory
+			if usage != nil {
+				used = usage.total.mem
+			}
+		} else {
+			limit = selected.MaxStore
+			if usage != nil {
+				used = usage.total.store
+			}
+		}
+		jsa.usageMu.RUnlock()
+		if limit >= 0 && used > mulSaturate(limit, lr) {
+			if cfg.Storage == MemoryStorage {
+				return NewJSMemoryResourcesExceededError()
+			}
+			return NewJSStorageResourcesExceededError()
+		}
+		if js.limitsExceeded(cfg.Storage) {
+			if cfg.Storage == MemoryStorage {
+				return NewJSMemoryResourcesExceededError()
+			}
+			return NewJSStorageResourcesExceededError()
+		}
+		return nil
+	}
+
+	reserveCfg := cfg
+	reserveCfg.MaxBytes = max(reserveCfg.MaxBytes, restoreRemaining)
+	js.mu.RLock()
+	err = js.checkAllLimits(&selected, tier, &reserveCfg, reserved, 0)
+	js.mu.RUnlock()
+	if err != nil {
+		return nil, err
+	}
+	if err := checkUsageLimits(); err != nil {
+		return nil, err
+	}
+
+	mset, err := a.addStreamForRestore(&cfg)
+	if err != nil {
+		return nil, fmt.Errorf("error adding stream: %w", err)
+	}
+	defer func() {
+		var state StreamState
+		mset.store.FastState(&state)
+		mset.mu.Lock()
+		mset.lseq = state.LastSeq
+		mset.mu.Unlock()
+		if err := mset.completeRestore(); err != nil {
+			if err = fmt.Errorf("failed to activate stream %q: %w", cfg.Name, err); retErr == nil {
+				retErr = err
+			}
+			s.Warnf("JetStream stream restore for '%s > %s' failed to activate stream: %v", a.Name, cfg.Name, err)
+		}
+	}()
+
+	// Start off at the right sequence number. This is important in particular
+	// when the backup contains no messages or would restore to no interest.
+	if _, err = mset.store.Compact(nstate.FirstSeq); err != nil {
+		return nil, fmt.Errorf("error purging stream: %w", err)
+	}
+
+	var restoredConsumers, ephemerals []*consumer
+	defer func() {
+		// Consumers must be unconditionally converted and completed, otherwise
+		// a partial restore that fails midway through can leave assets that are
+		// unusable.
+		for _, o := range ephemerals {
+			o.switchToEphemeral()
+		}
+		for _, o := range restoredConsumers {
+			if err := o.completeRestore(); err != nil {
+				if err = fmt.Errorf("failed to activate consumer %q: %w", o.name, err); retErr == nil {
+					retErr = err
+				}
+				s.Warnf("JetStream stream restore for '%s > %s' failed to activate consumers: %v", a.Name, cfg.Name, err)
+			}
+		}
+	}()
+	for range nstate.Consumers {
+		hdr, err := tr.Next()
+		if err != nil {
+			return nil, err
+		}
+		name, found := strings.CutPrefix(hdr.Name, "consumers/")
+		if !found {
+			return nil, fmt.Errorf("expected consumer, found %q", hdr.Name)
+		}
+		var consumer SnapshotConsumerState
+		if err := json.NewDecoder(tr).Decode(&consumer); err != nil {
+			return nil, fmt.Errorf("failed to decode consumer %q state: %w", name, err)
+		}
+		if consumer.ConsumerConfig == nil {
+			return nil, fmt.Errorf("consumer %q is missing config", name)
+		}
+		if consumer.ConsumerState == nil {
+			return nil, fmt.Errorf("consumer %q is missing state", name)
+		}
+		isEphemeral := !isDurableConsumer(consumer.ConsumerConfig)
+		if isEphemeral {
+			// Keep ephemerals alive and interested until all messages have
+			// been restored, then start their normal inactivity lifecycle.
+			consumer.Durable = name
+		}
+		o, err := mset.addConsumerForRestore(consumer.ConsumerConfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to add consumer %q: %w", name, err)
+		}
+		if isEphemeral {
+			ephemerals = append(ephemerals, o)
+		}
+		restoredConsumers = append(restoredConsumers, o)
+		o.mu.Lock()
+		err = o.setStoreState(consumer.ConsumerState)
+		o.mu.Unlock()
+		if err != nil {
+			return nil, fmt.Errorf("failed to set consumer %q state: %w", name, err)
+		}
+	}
+
+	store := mset.store
+	lseq := nstate.FirstSeq - 1
+	eob := false
+	mp := int64(s.getOpts().MaxPayload)
+	for {
+		hdr, err := tr.Next()
+		if err != nil {
+			return nil, err
+		}
+		seq := hdr.Sequence
+		if seq == 0 {
+			// Sentinel "end of backup" if all fields are zero.
+			if hdr.Name == _EMPTY_ && hdr.Timestamp == 0 && hdr.HeaderSize == 0 && hdr.PayloadSize == 0 {
+				eob = true
+				break
+			}
+			return nil, fmt.Errorf("expected message sequence")
+		}
+		if hdr.HeaderSize < 0 || hdr.PayloadSize < 0 || hdr.PayloadSize > math.MaxInt64-hdr.HeaderSize {
+			return nil, fmt.Errorf("invalid message lengths for sequence %d", seq)
+		}
+		declaredSize := hdr.HeaderSize + hdr.PayloadSize
+		if hdr.HeaderSize > mp || hdr.PayloadSize > mp-hdr.HeaderSize {
+			return nil, fmt.Errorf("message sequence %d exceeds maximum payload size", seq)
+		}
+		if mms := int64(cfg.MaxMsgSize); mms >= 0 && (hdr.HeaderSize > mms || hdr.PayloadSize > mms-hdr.HeaderSize) {
+			return nil, fmt.Errorf("message sequence %d exceeds maximum message size", seq)
+		}
+		var storedSizeRaw uint64
+		switch cfg.Storage {
+		case MemoryStorage:
+			if storedSizeRaw = memStoreMsgSizeRaw(len(hdr.Name), int(hdr.HeaderSize), int(hdr.PayloadSize)); storedSizeRaw > math.MaxInt64 {
+				return nil, fmt.Errorf("snapshot message bytes exceed maximum store message size")
+			}
+		default:
+			if storedSizeRaw = fileStoreMsgSizeRaw(len(hdr.Name), int(hdr.HeaderSize), int(hdr.PayloadSize)); storedSizeRaw > rlBadThresh {
+				return nil, fmt.Errorf("snapshot message bytes exceed maximum store message size")
+			}
+		}
+		storedSize := int64(storedSizeRaw)
+		if additional := storedSize - restoreRemaining; additional > 0 {
+			jsa.updateUsage(tier, cfg.Storage, additional)
+			restoreRemaining += additional
+			if err := checkUsageLimits(); err != nil {
+				return nil, err
+			}
+		}
+		buf, err := io.ReadAll(io.LimitReader(tr, declaredSize))
+		if err != nil {
+			return nil, fmt.Errorf("failed to read message sequence %d: %w", seq, err)
+		}
+		if hdr.HeaderSize > int64(len(buf)) {
+			return nil, fmt.Errorf("failed to parse message sequence %d: invalid header length", seq)
+		}
+		if int64(len(buf)) != declaredSize {
+			return nil, fmt.Errorf("failed to read message sequence %d: unexpected payload size", seq)
+		}
+		subj := hdr.Name
+		mhdr := buf[:hdr.HeaderSize]
+		msg := buf[hdr.HeaderSize : hdr.HeaderSize+hdr.PayloadSize]
+		if seq <= lseq {
+			return nil, fmt.Errorf("message sequence %d out of order", seq)
+		}
+		// We could have deleted messages since the last message we stored, if so
+		// we should work out what the gap is and skip those sequences.
+		if gap := seq - lseq - 1; gap > 0 {
+			if err := store.SkipMsgs(lseq+1, gap); err != nil {
+				return nil, fmt.Errorf("failed to process gap: %w", err)
+			}
+		}
+		lseq = seq
+		ttl, err := getMessageTTL(mhdr)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse message TTL: %w", err)
+		}
+		hdrTime := time.Unix(0, hdr.Timestamp)
+		if ttl > 0 && time.Now().After(hdrTime.Add(time.Duration(ttl)*time.Second)) {
+			// If the TTL has exceeded then there isn't much point in storing the message,
+			// but we still need to preserve the sequence.
+			if err := store.SkipMsgs(seq, 1); err != nil {
+				return nil, fmt.Errorf("failed to process expired message sequence %d: %w", seq, err)
+			}
+			releaseRestoreBytes(storedSize)
+			continue
+		}
+		if err = store.StoreRawMsg(subj, mhdr, msg, seq, hdr.Timestamp, ttl, false); err != nil {
+			return nil, fmt.Errorf("failed to store message sequence %d: %w", seq, err)
+		}
+		releaseRestoreBytes(storedSize)
+	}
+
+	if !eob {
+		return mset, fmt.Errorf("backup was truncated")
+	}
+
+	// Need to make sure that we pad out with skip msgs to preserve the last
+	// sequence, otherwise trailing deleted messages could reuse sequence numbers.
+	if lseq < nstate.LastSeq {
+		gap := nstate.LastSeq - lseq
+		if err := store.SkipMsgs(lseq+1, gap); err != nil {
+			return nil, fmt.Errorf("failed to process trailing gap: %w", err)
+		}
+	}
+	return mset, nil
+}