mesh-catalog/modules/nats/internal/keeper/kept_jetstream.go

package keeper

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"

	"nats-tools/internal/busarchive"
	"nats-tools/internal/bussnap"
)

// KeptJetStream is KeptBus on the keeper's own connection.
type KeptJetStream struct {
	NC   *nats.Conn
	JS   jetstream.JetStream
	Logf func(string, ...any)
}

// Existing answers which of the names are streams the bus holds.
func (b KeptJetStream) Existing(ctx context.Context, names []string) ([]string, error) {
	var out []string
	for _, n := range names {
		_, err := b.JS.Stream(ctx, n)
		switch {
		case err == nil:
			out = append(out, n)
		case errors.Is(err, jetstream.ErrStreamNotFound):
		default:
			return nil, err
		}
	}
	return out, nil
}

// ConsumerOf is whether a durable consumer is the module's: named for it, alone or after the machine that runs it
// (`<node>_<module>`, process.Account.QueueConsumerName).
func ConsumerOf(consumer, module string) bool {
	return consumer == module || strings.HasSuffix(consumer, "_"+module)
}

// Positions are the module's durable consumers on every stream, each with the stream sequence it has acknowledged up to.
func (b KeptJetStream) Positions(ctx context.Context, module string) ([]Position, error) {
	var out []Position
	streams := b.JS.ListStreams(ctx)
	var names []string
	for s := range streams.Info() {
		names = append(names, s.Config.Name)
	}
	if err := streams.Err(); err != nil {
		return nil, err
	}
	for _, name := range names {
		st, err := b.JS.Stream(ctx, name)
		if err != nil {
			return nil, err
		}
		list := st.ListConsumers(ctx)
		for ci := range list.Info() {
			if !ConsumerOf(ci.Name, module) {
				continue
			}
			cfg, err := json.Marshal(ci.Config)
			if err != nil {
				return nil, err
			}
			out = append(out, Position{Stream: name, Consumer: ci.Name, AckFloor: ci.AckFloor.Stream, Config: cfg})
		}
		if err := list.Err(); err != nil {
			return nil, err
		}
	}
	return out, nil
}

// SetPositions makes each consumer again, from the sequence after the one it had acknowledged. A consumer's place cannot
// be edited, so it is removed and made with its own configuration and a new start.
func (b KeptJetStream) SetPositions(ctx context.Context, ps []Position) error {
	for _, p := range ps {
		st, err := b.JS.Stream(ctx, p.Stream)
		if err != nil {
			return fmt.Errorf("the stream %s of %s's position cannot be read: %w", p.Stream, p.Consumer, err)
		}
		var cfg jetstream.ConsumerConfig
		if err := json.Unmarshal(p.Config, &cfg); err != nil {
			return fmt.Errorf("the configuration of %s on %s does not read: %w", p.Consumer, p.Stream, err)
		}
		info, err := st.Info(ctx)
		if err != nil {
			return err
		}
		start := p.AckFloor + 1
		if start > info.State.LastSeq+1 {
			start = info.State.LastSeq + 1
		}
		cfg.DeliverPolicy, cfg.OptStartSeq, cfg.OptStartTime = jetstream.DeliverByStartSequencePolicy, start, nil
		if err := st.DeleteConsumer(ctx, p.Consumer); err != nil && !errors.Is(err, jetstream.ErrConsumerNotFound) {
			return fmt.Errorf("%s on %s cannot be removed to be made again: %w", p.Consumer, p.Stream, err)
		}
		if _, err := st.CreateConsumer(ctx, cfg); err != nil {
			return fmt.Errorf("%s on %s cannot be made again at %d: %w", p.Consumer, p.Stream, start, err)
		}
	}
	return nil
}

// RestoreArchive puts every stream of the archive into the bus, which must hold none of them.
func (b KeptJetStream) RestoreArchive(ctx context.Context, a *busarchive.Archive) error {
	log := b.Logf
	if log == nil {
		log = func(string, ...any) {}
	}
	return bussnap.RestoreAll(ctx, b.NC, a, log)
}

// DeleteStreams removes streams by name.
func (b KeptJetStream) DeleteStreams(ctx context.Context, names []string) error {
	for _, n := range names {
		if err := b.JS.DeleteStream(ctx, n); err != nil && !errors.Is(err, jetstream.ErrStreamNotFound) {
			return fmt.Errorf("%s: %w", n, err)
		}
	}
	return nil
}