mesh-sdk/go/pernode/pernode.go

// Package pernode reads a state its owner keeps one bucket per node as one state (novox/hq ADR 0593 §6: each node's
// engine writes its own `composed-<node>`, and the provisioner, the assigner and the view read them all), over any
// connection that satisfies Bus: a process's own (process.Conn) and the runtime's (node-tools) alike, so the rules are
// written once.
//
// A reader names the state as its manifest's `reads` does (`<module>.<state>`, such as `mesh-host.composed`) with the
// prefix of its buckets (`composed-`): each bucket `<prefix><node>` holds keys whose first part is that node, so a key
// names its bucket. Get reads one key from its node's bucket; Keys and Watch go over every bucket of the family the bus
// holds, and Watch takes in a node's bucket created after it began, as the bus says it (its creation advisory), so a
// node added later is heard with no timer. It reads only: the family has no write, and the owner's nodes write as their
// grants allow.
package pernode

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"sync"

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

// What a reader of a per-node state is granted beside each bucket's read (the grants part composes them): the listing
// of stream names (ListSubject) and the creation of any stream (CreatedAdvisory), which says a name only. A subject's
// wildcard is a whole token, so `KV_composed-*` cannot be granted: every creation is heard and the family's picked.
const (
	ListSubject     = "$JS.API.STREAM.NAMES"
	CreatedAdvisory = "$JS.EVENT.ADVISORY.STREAM.CREATED.*"
	// createdPrefix starts the subject of a key-value bucket's creation advisory, the bucket's name after it.
	createdPrefix = "$JS.EVENT.ADVISORY.STREAM.CREATED.KV_"
)

// Bus is what a connection gives the family: the names of the key-value buckets the bus holds, a bucket opened by name
// (jetstream.ErrBucketNotFound when it is not there), and the creation advisories (CreatedAdvisory) heard, each by its
// subject, until stop.
type Bus interface {
	BucketNames(ctx context.Context) ([]string, error)
	OpenBucket(ctx context.Context, name string) (jetstream.KeyValue, error)
	SubscribeCreated(heard func(subject string)) (stop func(), err error)
}

// Family is a per-node state: its read name, its buckets' prefix, and the connection it is read over.
type Family struct {
	Read   string
	Prefix string
	Bus    Bus
}

// Change is one change a watch hands over: the node whose bucket it is, the entry (its operation says put, delete or
// purge), and whether it was there when that bucket's watch began.
type Change struct {
	Node    string
	Entry   jetstream.KeyValueEntry
	Current bool
}

// nodeName is a node a bucket and a subject token accept.
var nodeName = regexp.MustCompile(`^[A-Za-z0-9_-]+$`)

// Check refuses a family with no read name `<module>.<state>`, or a prefix a bucket name cannot start with or that
// does not end in `-` or `_`, since a prefix without its separator would read a look-alike family (`composed` reading
// `composedx-ace`).
func (f Family) Check() error {
	module, name, ok := strings.Cut(f.Read, ".")
	if !ok || module == "" || name == "" || strings.Contains(name, ".") {
		return fmt.Errorf("%q is not a state another module keeps, <module>.<state>", f.Read)
	}
	if !nodeName.MatchString(strings.TrimRight(f.Prefix, "-_")) || !strings.HasSuffix(f.Prefix, "-") && !strings.HasSuffix(f.Prefix, "_") {
		return fmt.Errorf("the state %s: %q is not a prefix of buckets, letters, digits, '-' and '_' ending in '-' or '_'", f.Read, f.Prefix)
	}
	if f.Bus == nil {
		return fmt.Errorf("the state %s is read over no connection", f.Read)
	}
	return nil
}

// Bucket is the bucket of a node.
func (f Family) Bucket(node string) (string, error) {
	if err := f.Check(); err != nil {
		return "", err
	}
	if !nodeName.MatchString(node) {
		return "", fmt.Errorf("the state %s: %q is not a node a bucket is named for", f.Read, node)
	}
	return f.Prefix + node, nil
}

// NodeOf is the node a key of the family names: its first part.
func NodeOf(key string) (string, error) {
	node, _, ok := strings.Cut(key, ".")
	if !ok || !nodeName.MatchString(node) {
		return "", fmt.Errorf("%q names no node: a key of a per-node state starts <node>.", key)
	}
	return node, nil
}

// nodeOfBucket is the node of a bucket name of the family, and whether it is one.
func (f Family) nodeOfBucket(name string) (string, bool) {
	node, ok := strings.CutPrefix(name, f.Prefix)
	return node, ok && nodeName.MatchString(node)
}

// Nodes are the nodes whose bucket of the family the bus holds now, sorted.
func (f Family) Nodes(ctx context.Context) ([]string, error) {
	if err := f.Check(); err != nil {
		return nil, err
	}
	names, err := f.Bus.BucketNames(ctx)
	if err != nil {
		return nil, fmt.Errorf("the state %s: listing the buckets %s*: %w", f.Read, f.Prefix, err)
	}
	var nodes []string
	for _, name := range names {
		if node, ok := f.nodeOfBucket(name); ok {
			nodes = append(nodes, node)
		}
	}
	sort.Strings(nodes)
	return nodes, nil
}

// Open is the bucket of a node, opened; nil, nil when the bus holds none.
func (f Family) Open(ctx context.Context, node string) (jetstream.KeyValue, error) {
	bucket, err := f.Bucket(node)
	if err != nil {
		return nil, err
	}
	kv, err := f.Bus.OpenBucket(ctx, bucket)
	if errors.Is(err, jetstream.ErrBucketNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("the state %s (bucket %s): %w", f.Read, bucket, err)
	}
	return kv, nil
}

// Get is a key's current entry, read from the bucket of the node the key names, or nil when it has none or that node
// has no bucket.
func (f Family) Get(ctx context.Context, key string) (jetstream.KeyValueEntry, error) {
	node, err := NodeOf(key)
	if err != nil {
		return nil, err
	}
	kv, err := f.Open(ctx, node)
	if err != nil || kv == nil {
		return nil, err
	}
	e, err := kv.Get(ctx, key)
	if errors.Is(err, jetstream.ErrKeyNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("the state %s, key %s: %w", f.Read, key, err)
	}
	return e, nil
}

// Keys is every key with a current value in every bucket of the family, sorted.
func (f Family) Keys(ctx context.Context) ([]string, error) {
	nodes, err := f.Nodes(ctx)
	if err != nil {
		return nil, err
	}
	keys := []string{}
	for _, node := range nodes {
		kv, err := f.Open(ctx, node)
		if err != nil {
			return nil, err
		}
		if kv == nil {
			continue
		}
		k, err := kv.Keys(ctx)
		if errors.Is(err, jetstream.ErrNoKeysFound) {
			continue
		}
		if err != nil {
			return nil, fmt.Errorf("the state %s on %s, keys: %w", f.Read, node, err)
		}
		keys = append(keys, k...)
	}
	sort.Strings(keys)
	return keys, nil
}

// Pattern is a key narrowing as the bus names it: `*` stays one name, `**` is the rest, empty is every key.
func Pattern(key string) string {
	if key == "" {
		return jetstream.AllKeys
	}
	tokens := strings.Split(key, ".")
	for i, t := range tokens {
		if t == "**" {
			tokens[i] = ">"
		}
	}
	return strings.Join(tokens, ".")
}

// Watch hands handle, one call at a time, the current values of every bucket of the family the bus holds (each
// Current, deletes left out), then every change, and the current values and changes of a bucket the bus creates later.
// It returns once the current values of the buckets there now are handled; an error before that is returned, and an
// error from handle after it ends that bucket's watch. key narrows it within each bucket (Pattern); a key whose first
// part names a node watches only that node's bucket. ctx ends every watch and the hearing of new buckets.
func (f Family) Watch(ctx context.Context, key string, handle func(Change) error) error {
	if err := f.Check(); err != nil {
		return err
	}
	var mu sync.Mutex
	one := func(c Change) error {
		mu.Lock()
		defer mu.Unlock()
		return handle(c)
	}
	var wmu sync.Mutex
	watched := map[string]bool{}
	watch := func(node string) error {
		wmu.Lock()
		if watched[node] {
			wmu.Unlock()
			return nil
		}
		watched[node] = true
		wmu.Unlock()
		kv, err := f.Open(ctx, node)
		if err != nil || kv == nil {
			return err
		}
		return f.watchBucket(ctx, kv, node, key, one)
	}
	if first, _, _ := strings.Cut(key, "."); key != "" && first != "*" && first != "**" {
		return watch(first)
	}
	// Heard before the buckets are listed, so a bucket created between the two is not missed.
	stop, err := f.Bus.SubscribeCreated(func(subject string) {
		bucket, isKV := strings.CutPrefix(subject, createdPrefix)
		if node, ok := f.nodeOfBucket(bucket); isKV && ok {
			go func() { _ = watch(node) }()
		}
	})
	if err != nil {
		return fmt.Errorf("the state %s: hearing new buckets %s*: %w", f.Read, f.Prefix, err)
	}
	go func() { <-ctx.Done(); stop() }()
	nodes, err := f.Nodes(ctx)
	if err != nil {
		stop()
		return err
	}
	for _, node := range nodes {
		if err := watch(node); err != nil {
			stop()
			return err
		}
	}
	return nil
}

// watchBucket hands one bucket's current values, then its changes from a goroutine, until ctx ends or handle errs.
func (f Family) watchBucket(ctx context.Context, kv jetstream.KeyValue, node, key string, handle func(Change) error) error {
	w, err := kv.Watch(ctx, Pattern(key))
	if err != nil {
		return fmt.Errorf("the state %s on %s, watching %q: %w", f.Read, node, key, err)
	}
	for {
		select {
		case <-ctx.Done():
			_ = w.Stop()
			return ctx.Err()
		case e, open := <-w.Updates():
			if !open {
				return fmt.Errorf("the state %s on %s: the watch ended", f.Read, node)
			}
			if e == nil {
				go func() {
					defer func() { _ = w.Stop() }()
					for {
						select {
						case <-ctx.Done():
							return
						case e, open := <-w.Updates():
							if !open || e == nil {
								return
							}
							if handle(Change{Node: node, Entry: e}) != nil {
								return
							}
						}
					}
				}()
				return nil
			}
			if e.Operation() != jetstream.KeyValuePut {
				continue
			}
			if err := handle(Change{Node: node, Entry: e, Current: true}); err != nil {
				_ = w.Stop()
				return err
			}
		}
	}
}

// OverNATS is the Bus of a plain connection: its JetStream's bucket names and buckets, and the creation advisories
// subscribed on it. A connection that words a refused open its own way (process.Conn) gives its own OpenBucket.
func OverNATS(nc *nats.Conn) Bus { return natsBus{nc: nc} }

type natsBus struct{ nc *nats.Conn }

func (b natsBus) BucketNames(ctx context.Context) ([]string, error) {
	js, err := jetstream.New(b.nc)
	if err != nil {
		return nil, err
	}
	lister := js.KeyValueStoreNames(ctx)
	var names []string
	for name := range lister.Name() {
		names = append(names, name)
	}
	if err := lister.Error(); err != nil && !errors.Is(err, jetstream.ErrEndOfData) {
		return nil, err
	}
	return names, nil
}

func (b natsBus) OpenBucket(ctx context.Context, name string) (jetstream.KeyValue, error) {
	js, err := jetstream.New(b.nc)
	if err != nil {
		return nil, err
	}
	return js.KeyValue(ctx, name)
}

func (b natsBus) SubscribeCreated(heard func(subject string)) (func(), error) {
	sub, err := b.nc.Subscribe(CreatedAdvisory, func(m *nats.Msg) { heard(m.Subject) })
	if err != nil {
		return nil, err
	}
	return func() { _ = sub.Unsubscribe() }, nil
}