mesh-sdk/go/kvlimit/kvlimit.go

// Package kvlimit refuses a key-value write whose value is larger than its bucket takes, before it is put, naming the
// bucket, the key, the size and the limit (issue 774, sweep item 5). The bus refuses such a write too, but only after
// the put, as "message size exceeds maximum" (JetStream error 10054); a writer that retries on errors then loops on a
// value that can never land. A refusal from here, or the bus's own one passed through Refused, is a *TooLarge that
// errors.Is ErrTooLarge: a writer stops on it and says it, it never writes the same value again.
//
// A value that may outgrow its bucket is shaped one of two ways, never retried:
//   - split over keys: one key per part that changes on its own (per node, per module), each well under the limit;
//   - or kept in the object store, with only its name and digest in the key.
package kvlimit

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

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

// ModuleStateValueBytes is the most a value in a module's declared state may hold: every module bucket is made with
// it (buscompose.StateMaxValueBytes).
const ModuleStateValueBytes = 256 * 1024

// HeaderRoom is what a write's headers may take of the limit: the bus counts headers and value together, and a create,
// an update or a key's time to live adds one header each. A value closer to the limit than this is refused.
const HeaderRoom = 128

// codeMessageExceedsMaximum is JetStream's error code for a message over the stream's largest (JSStreamMessageExceedsMaximumErr).
const codeMessageExceedsMaximum jetstream.ErrorCode = 10054

// ErrTooLarge is what every *TooLarge is: a value its bucket does not take, which no retry will change.
var ErrTooLarge = errors.New("a value larger than its bucket takes")

// TooLarge is one refused write. Limit is zero when the bus refused it and its limit was not known here.
type TooLarge struct {
	Bucket, Key string
	Size, Limit int64
}

func (e *TooLarge) Error() string {
	limit := "the most its bucket takes"
	if e.Limit > 0 {
		limit = fmt.Sprintf("the %d bytes the bucket takes (headers included)", e.Limit)
	}
	return fmt.Sprintf("the value for %s in %s is %d bytes, over %s: it is refused and written again it is refused again; "+
		"split it over keys or keep it in the object store with its name in the key", e.Key, e.Bucket, e.Size, limit)
}

func (e *TooLarge) Unwrap() error { return ErrTooLarge }

// Check refuses a value of size bytes for key in bucket when it does not fit limit with HeaderRoom to spare; a limit of
// zero or less is no limit.
func Check(bucket, key string, size int, limit int64) error {
	if limit > 0 && int64(size)+HeaderRoom > limit {
		return &TooLarge{Bucket: bucket, Key: key, Size: int64(size), Limit: limit}
	}
	return nil
}

// Refused is err as a *TooLarge when it is the bus's refusal of a value too large (the stream's 10054, the client's
// maximum payload, or either said in words by a runtime that passed it on), and err unchanged otherwise.
func Refused(bucket, key string, size int, limit int64, err error) error {
	if err == nil || errors.Is(err, ErrTooLarge) {
		return err
	}
	var api *jetstream.APIError
	tooLarge := errors.As(err, &api) && api.ErrorCode == codeMessageExceedsMaximum
	tooLarge = tooLarge || errors.Is(err, nats.ErrMaxPayload)
	words := err.Error()
	tooLarge = tooLarge || strings.Contains(words, "message size exceeds maximum") || strings.Contains(words, "maximum payload exceeded")
	if !tooLarge {
		return err
	}
	forget(bucket)
	return fmt.Errorf("%w (the bus: %v)", &TooLarge{Bucket: bucket, Key: key, Size: int64(size), Limit: limit}, err)
}

// limits are the buckets' largest messages as read, by bucket name: a limit is kept once read, and only a limit (an
// unlimited bucket is read again each time). A kept limit that refuses a value is read again before the refusal is
// given, so a bucket made again with a larger limit is never refused by an old one; one made with a smaller limit is
// refused by the bus, and Refused forgets the kept one.
var limits sync.Map // bucket name → int64

func forget(bucket string) { limits.Delete(bucket) }

// LimitOf is the largest message kv's bucket takes, read from its stream and kept; zero when it has no limit or it
// cannot be read (the write then goes, and Refused names the bus's refusal).
func LimitOf(ctx context.Context, kv jetstream.KeyValue) int64 {
	if l, ok := limits.Load(nameOf(kv)); ok {
		return l.(int64)
	}
	return readLimit(ctx, kv)
}

func readLimit(ctx context.Context, kv jetstream.KeyValue) (limit int64) {
	name := nameOf(kv)
	if name == "" {
		return 0
	}
	// A test's fake that implements only the writes it needs panics on Status; its limit is not known.
	defer func() {
		if recover() != nil {
			limit = 0
		}
	}()
	st, err := kv.Status(ctx)
	if err != nil {
		return 0
	}
	s, ok := st.(*jetstream.KeyValueBucketStatus)
	if !ok || s.StreamInfo() == nil || s.StreamInfo().Config.MaxMsgSize <= 0 {
		forget(name)
		return 0
	}
	limit = int64(s.StreamInfo().Config.MaxMsgSize)
	limits.Store(name, limit)
	return limit
}

// fits is the limit of kv's bucket and nil when a value of size bytes fits it, or the *TooLarge refusing it, checked
// against a freshly read limit.
func fits(ctx context.Context, kv jetstream.KeyValue, key string, size int) (int64, error) {
	limit := LimitOf(ctx, kv)
	if Check(nameOf(kv), key, size, limit) == nil {
		return limit, nil
	}
	limit = readLimit(ctx, kv)
	return limit, Check(nameOf(kv), key, size, limit)
}

// nameOf is kv's bucket name, or "" for none or a test's fake that does not say it.
func nameOf(kv jetstream.KeyValue) (name string) {
	if kv == nil {
		return ""
	}
	defer func() {
		if recover() != nil {
			name = ""
		}
	}()
	return kv.Bucket()
}

// Put is kv.Put, refused before it is put when value does not fit the bucket.
func Put(ctx context.Context, kv jetstream.KeyValue, key string, value []byte) (uint64, error) {
	limit, err := fits(ctx, kv, key, len(value))
	if err != nil {
		return 0, err
	}
	rev, err := kv.Put(ctx, key, value)
	return rev, Refused(nameOf(kv), key, len(value), limit, err)
}

// Create is kv.Create, refused before it is put when value does not fit the bucket.
func Create(ctx context.Context, kv jetstream.KeyValue, key string, value []byte, opts ...jetstream.KVCreateOpt) (uint64, error) {
	limit, err := fits(ctx, kv, key, len(value))
	if err != nil {
		return 0, err
	}
	rev, err := kv.Create(ctx, key, value, opts...)
	return rev, Refused(nameOf(kv), key, len(value), limit, err)
}

// Update is kv.Update, refused before it is put when value does not fit the bucket.
func Update(ctx context.Context, kv jetstream.KeyValue, key string, value []byte, revision uint64) (uint64, error) {
	limit, err := fits(ctx, kv, key, len(value))
	if err != nil {
		return 0, err
	}
	rev, err := kv.Update(ctx, key, value, revision)
	return rev, Refused(nameOf(kv), key, len(value), limit, err)
}