vshard.go

package vshard_router //nolint:revive

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"github.com/snksoft/crc"
	"github.com/vmihailenco/msgpack/v5"

	tarantool "github.com/tarantool/go-tarantool/v2"
)

var (
	// ErrInvalidConfig is returned when the configuration is invalid.
	ErrInvalidConfig = fmt.Errorf("config invalid")
	// ErrInvalidInstanceInfo is returned when the instance information is invalid.
	ErrInvalidInstanceInfo = fmt.Errorf("invalid instance info")
	// ErrInvalidReplicasetInfo is returned when the replicaset information is invalid.
	ErrInvalidReplicasetInfo = fmt.Errorf("invalid replicaset info")
	// ErrTopologyProvider is returned when there is an error from the topology provider.
	ErrTopologyProvider = fmt.Errorf("got error from topology provider")
)

// This data struct is instroduced by https://github.com/tarantool/go-vshard-router/issues/39.
// We use an array of atomics to lock-free handling elements of routeMap.
// knownBucketCount reflects a statistic over routeMap.
// knownBucketCount might be inconsistent for a few mksecs, because at first we change routeMap[bucketID],
// only after that we change knownBucketCount: this is not an atomic change of complex state.
// It it is not a problem at all.
//
// While changing `knownBucketCount` we heavily rely on commutative property of algebraic sum operation ("+"),
// due to this property we don't afraid any amount of concurrent modifications.
// See: https://en.wikipedia.org/wiki/Commutative_property
//
// Since RouteMapClean creates a new routeMap, we have to assign knownBucketCount := 0.
// But assign is not a commutative operation, therefore we have to create a completely new atomic variable,
// that reflects a statistic over newly created routeMap.
type consistentView struct {
	routeMap         []atomic.Pointer[Replicaset]
	knownBucketCount atomic.Int32
}

type Router struct {
	cfg Config

	// nameToReplicasetMutex guards not the map itself, but the variable idToReplicaset.
	// nameToReplicaset is an immutable object by our convention.
	// Whenever we add or remove a replicaset, we create a new map object.
	// nameToReplicaset can be modified only by TopologyController methods.
	// Assuming that we rarely add or remove some replicaset,
	// it should be the simplest and most efficient way of handling concurrent access.
	// Additionally, we can safely iterate over a map because it never changes.
	nameToReplicasetMutex sync.RWMutex
	nameToReplicaset      map[string]*Replicaset

	viewMutex sync.RWMutex
	view      *consistentView

	// ----------------------- Map-Reduce -----------------------
	// Storage Ref ID. It must be unique for each ref request
	// and therefore is global and monotonically growing.
	refID atomic.Int64

	cancelDiscovery func()
}

func (r *Router) metrics() MetricsProvider {
	return r.cfg.Metrics
}

func (r *Router) log() LogfProvider {
	return r.cfg.Loggerf
}

func (r *Router) getConsistentView() *consistentView {
	r.viewMutex.RLock()
	view := r.view
	r.viewMutex.RUnlock()

	return view
}

func (r *Router) setConsistentView(view *consistentView) {
	r.viewMutex.Lock()
	r.view = view
	r.viewMutex.Unlock()
}

type Config struct {
	// Providers
	// Loggerf injects a custom logger. By default there is no logger is used.
	Loggerf          LogfProvider     // Loggerf is not required
	Metrics          MetricsProvider  // Metrics is not required
	TopologyProvider TopologyProvider // TopologyProvider is required provider

	// Discovery
	// DiscoveryTimeout is timeout between cron discovery job; by default there is no timeout.
	DiscoveryTimeout time.Duration
	DiscoveryMode    DiscoveryMode
	// DiscoveryWorkStep is a pause between calling buckets_discovery on storage
	// in buckets discovering logic. Default is 10ms.
	DiscoveryWorkStep time.Duration

	// BucketsSearchMode defines policy for Router.Route method.
	// Default value is BucketsSearchLegacy.
	// See BucketsSearchMode constants for more detail.
	BucketsSearchMode BucketsSearchMode

	TotalBucketCount uint64
	User             string
	Password         string
	PoolOpts         tarantool.Opts

	// BucketGetter is an optional argument.
	// You can specify a function that will receive the bucket id from the context.
	// This is useful if you use middleware that inserts the calculated bucket id into the request context.
	BucketGetter func(ctx context.Context) uint64
	// RequestTimeout timeout for requests to Tarantool.
	// Don't rely on using this timeout.
	// This is the difference between the timeout of the library itself
	// that is, our retry timeout if the buckets, for example, move.
	// Currently, it only works for sugar implementations .
	RequestTimeout time.Duration
}

type BucketStatInfo struct {
	BucketID uint64 `msgpack:"id"`
	Status   string `msgpack:"status"`
}

// tnt vshard storage returns map with 'int' keys for bucketStatInfo,
// example: map[id:48 status:active 1:48 2:active].
// But msgpackv5 supports only string keys when decoding maps into structs,
// see issue: https://github.com/vmihailenco/msgpack/issues/372
// To workaround this we decode BucketStatInfo manually.
// When the issue above will be resolved, this code can be (and should be) deleted.
func (bsi *BucketStatInfo) DecodeMsgpack(d *msgpack.Decoder) error {
	nKeys, err := d.DecodeMapLen()
	if err != nil {
		return err
	}

	for i := 0; i < nKeys; i++ {
		key, err := d.DecodeInterface()
		if err != nil {
			return err
		}

		keyName, _ := key.(string)
		switch keyName {
		case "id":
			if err := d.Decode(&bsi.BucketID); err != nil {
				return err
			}
		case "status":
			if err := d.Decode(&bsi.Status); err != nil {
				return err
			}
		default:
			// skip unused value
			if err := d.Skip(); err != nil {
				return err
			}
		}

	}

	return nil
}

// InstanceInfo represents the information about an instance.
// This struct holds the necessary details such as the name, address, and UUID of the instance.
type InstanceInfo struct {
	// Name is a required field for the instance.
	// Starting with Tarantool 3.0, this definition is made into a human-readable name,
	// and it is now mandatory for use in the library.
	// The Name should be a unique identifier for the instance.
	Name string

	// Addr specifies the address of the instance.
	// This can be an IP address or a domain name, depending on how the instance is accessed.
	// It is necessary for connecting to the instance or identifying its location.
	Addr string

	// UUID is an optional field that provides a globally unique identifier (UUID) for the instance.
	// While this information is not mandatory, it can be useful for internal management or tracking purposes.
	// The UUID ensures that each instance can be identified uniquely, but it is not required for basic operations.
	UUID uuid.UUID
}

func (ii InstanceInfo) String() string {
	if ii.UUID == uuid.Nil {
		return fmt.Sprintf("{name: %s, addr: %s}", ii.Name, ii.Addr)
	}

	return fmt.Sprintf("{name: %s, uuid: %s, addr: %s}", ii.Name, ii.UUID, ii.Addr)
}

func (ii InstanceInfo) Validate() error {
	if ii.Name == "" {
		return fmt.Errorf("%w: empty name", ErrInvalidInstanceInfo)
	}

	if ii.Addr == "" {
		return fmt.Errorf("%w: empty addr", ErrInvalidInstanceInfo)
	}

	return nil
}

// --------------------------------------------------------------------------------
// -- Configuration
// --------------------------------------------------------------------------------

func NewRouter(ctx context.Context, cfg Config) (*Router, error) {
	var err error

	cfg, err = prepareCfg(ctx, cfg)
	if err != nil {
		return nil, err
	}

	router := &Router{
		cfg:              cfg,
		nameToReplicaset: make(map[string]*Replicaset),
		view: &consistentView{
			routeMap: make([]atomic.Pointer[Replicaset], cfg.TotalBucketCount+1),
		},
	}

	err = cfg.TopologyProvider.Init(router.Topology())
	if err != nil {
		router.log().Errorf(ctx, "Can't create new topology provider with err: %s", err)

		return nil, fmt.Errorf("%w; cant init topology with err: %w", ErrTopologyProvider, err)
	}

	err = router.DiscoveryAllBuckets(ctx)
	if err != nil {
		router.log().Errorf(ctx, "router.DiscoveryAllBuckets failed: %v", err)
	}

	if cfg.DiscoveryMode == DiscoveryModeOn {
		discoveryCronCtx, cancelFunc := context.WithCancel(ctx)

		// run background cron discovery loop
		// suppress linter warning: Non-inherited new context, use function like `context.WithXXX` instead (contextcheck)
		//nolint:contextcheck
		go router.cronDiscovery(discoveryCronCtx)

		router.cancelDiscovery = cancelFunc
	}

	return router, nil
}

// BucketSet Set a bucket to a replicaset.
func (r *Router) BucketSet(bucketID uint64, rsName string) (*Replicaset, error) {
	nameToReplicasetRef := r.getNameToReplicaset()

	rs := nameToReplicasetRef[rsName]
	if rs == nil {
		return nil, newVShardErrorNoRouteToBucket(bucketID)
	}

	view := r.getConsistentView()

	if oldRs := view.routeMap[bucketID].Swap(rs); oldRs == nil {
		view.knownBucketCount.Add(1)
	}

	return rs, nil
}

func (r *Router) BucketReset(bucketID uint64) {
	view := r.getConsistentView()

	if bucketID > r.cfg.TotalBucketCount {
		return
	}

	if old := view.routeMap[bucketID].Swap(nil); old != nil {
		view.knownBucketCount.Add(-1)
	}
}

func (r *Router) RouteMapClean() {
	newView := &consistentView{
		routeMap: make([]atomic.Pointer[Replicaset], r.cfg.TotalBucketCount+1),
	}

	r.setConsistentView(newView)
}

func prepareCfg(ctx context.Context, cfg Config) (Config, error) {
	const discoveryTimeoutDefault = 1 * time.Minute
	const discoveryWorkStepDefault = 10 * time.Millisecond

	err := validateCfg(cfg)
	if err != nil {
		return Config{}, fmt.Errorf("%v: %v", ErrInvalidConfig, err)
	}

	if cfg.DiscoveryTimeout == 0 {
		cfg.DiscoveryTimeout = discoveryTimeoutDefault
	}

	if cfg.Loggerf == nil {
		cfg.Loggerf = emptyLogfProvider
	}

	// Log tarantool internal events using the same logger as router uses.
	cfg.PoolOpts.Logger = tarantoolOptsLogger{
		loggerf: cfg.Loggerf,
		ctx:     ctx,
	}

	if cfg.Metrics == nil {
		cfg.Metrics = emptyMetricsProvider
	}

	if cfg.DiscoveryWorkStep == 0 {
		cfg.DiscoveryWorkStep = discoveryWorkStepDefault
	}

	return cfg, nil
}

func validateCfg(cfg Config) error {
	if cfg.TopologyProvider == nil {
		return fmt.Errorf("topology provider is nil")
	}

	if cfg.TotalBucketCount == 0 {
		return fmt.Errorf("bucket count must be greater than 0")
	}

	return nil
}

// --------------------------------------------------------------------------------
// -- Other
// --------------------------------------------------------------------------------

func BucketIDStrCRC32(shardKey string, totalBucketCount uint64) uint64 {
	return crc.CalculateCRC(&crc.Parameters{
		Width:      32,
		Polynomial: 0x1EDC6F41,
		FinalXor:   0x0,
		ReflectIn:  true,
		ReflectOut: true,
		Init:       0xFFFFFFFF,
	}, []byte(shardKey))%totalBucketCount + 1
}

// BucketIDStrCRC32 return the bucket identifier from the parameter used for sharding.
func (r *Router) BucketIDStrCRC32(shardKey string) uint64 {
	return BucketIDStrCRC32(shardKey, r.cfg.TotalBucketCount)
}

// BucketCount returns the total number of buckets specified in cfg.
func (r *Router) BucketCount() uint64 {
	return r.cfg.TotalBucketCount
}

// -------------------------------------------------------------------------------_
// -- Bootstrap
// --------------------------------------------------------------------------------

// ClusterBootstrap initializes the cluster by bootstrapping the necessary buckets
// across the available replicasets. It checks the current state of each replicaset
// and creates buckets if required. The function takes a context for managing
// cancellation and deadlines, and a boolean parameter ifNotBootstrapped to control
// error handling. If ifNotBootstrapped is true, the function will log any errors
// encountered during the bootstrapping process but will not halt execution; instead,
// it will return the last error encountered. If ifNotBootstrapped is false, any
// error will result in an immediate return, ensuring that the operation either
// succeeds fully or fails fast.
func (r *Router) ClusterBootstrap(ctx context.Context, ifNotBootstrapped bool) error {
	nameToReplicasetRef := r.getNameToReplicaset()

	rssToBootstrap := make([]Replicaset, 0, len(nameToReplicasetRef))
	var lastErr error

	for _, rs := range nameToReplicasetRef {
		rssToBootstrap = append(rssToBootstrap, *rs)
	}

	err := CalculateEtalonBalance(rssToBootstrap, r.cfg.TotalBucketCount)
	if err != nil {
		return err
	}

	bucketID := uint64(1)
	for id, rs := range rssToBootstrap {
		if rs.EtalonBucketCount > 0 {
			err = rs.BucketForceCreate(ctx, bucketID, rs.EtalonBucketCount)
			if err != nil {
				if ifNotBootstrapped {
					lastErr = err
				} else {
					return err
				}
			} else {
				nextBucketID := bucketID + rs.EtalonBucketCount
				r.log().Infof(ctx, "Buckets from %d to %d are bootstrapped on \"%s\"", bucketID, nextBucketID-1, id)
				bucketID = nextBucketID
			}
		}
	}

	if lastErr != nil {
		return lastErr
	}

	return nil
}