github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/lbcache/cache.go

/*
 * Copyright 2021 CloudWeGo 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 lbcache combine balancer with resolver and cache the resolve result
package lbcache

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

	"golang.org/x/sync/singleflight"

	"github.com/cloudwego/kitex/pkg/diagnosis"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/loadbalance"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

const (
	defaultRefreshInterval = 5 * time.Second
	defaultExpireInterval  = 15 * time.Second
)

var (
	balancerFactories    sync.Map // key: resolver name + loadbalance name
	balancerFactoriesSfg singleflight.Group
)

// Options for create builder
type Options struct {
	// refresh discovery result timely
	RefreshInterval time.Duration

	// Balancer expire check interval
	// we need remove idle Balancers for resource saving
	ExpireInterval time.Duration

	// DiagnosisService is used register info for diagnosis
	DiagnosisService diagnosis.Service

	// Cacheable is used to indicate that if the factory could be shared between multi clients
	Cacheable bool
}

func (v *Options) check() {
	if v.RefreshInterval <= 0 {
		v.RefreshInterval = defaultRefreshInterval
	}
	if v.ExpireInterval <= 0 {
		v.ExpireInterval = defaultExpireInterval
	}
}

// Hookable add hook for rebalancer events
type Hookable interface {
	// register loadbalance rebalance hook for Rebalance events
	RegisterRebalanceHook(func(ch *discovery.Change)) (index int)
	DeregisterRebalanceHook(index int)
	// register loadbalance delete hook for Delete events
	RegisterDeleteHook(func(ch *discovery.Change)) (index int)
	DeregisterDeleteHook(index int)
}

// BalancerFactory get or create a balancer with given target
// if it has the same key(reslover.Target(target)), we will cache and reuse the Balance
type BalancerFactory struct {
	Hookable
	opts       Options
	cache      sync.Map // key -> LoadBalancer
	resolver   discovery.Resolver
	balancer   loadbalance.Loadbalancer
	rebalancer loadbalance.Rebalancer
	sfg        singleflight.Group
}

func cacheKey(resolver, balancer string, opts Options) string {
	return fmt.Sprintf("%s|%s|{%s %s}", resolver, balancer, opts.RefreshInterval, opts.ExpireInterval)
}

func newBalancerFactory(resolver discovery.Resolver, balancer loadbalance.Loadbalancer, opts Options) *BalancerFactory {
	b := &BalancerFactory{
		opts:     opts,
		resolver: resolver,
		balancer: balancer,
	}
	if rb, ok := balancer.(loadbalance.Rebalancer); ok {
		hrb := newHookRebalancer(rb)
		b.rebalancer = hrb
		b.Hookable = hrb
	} else {
		b.Hookable = noopHookRebalancer{}
	}
	go b.watcher()
	return b
}

// NewBalancerFactory get or create a balancer factory for balancer instance
// cache key with resolver name, balancer name and options
func NewBalancerFactory(resolver discovery.Resolver, balancer loadbalance.Loadbalancer, opts Options) *BalancerFactory {
	opts.check()
	if !opts.Cacheable {
		return newBalancerFactory(resolver, balancer, opts)
	}
	uniqueKey := cacheKey(resolver.Name(), balancer.Name(), opts)
	val, ok := balancerFactories.Load(uniqueKey)
	if ok {
		return val.(*BalancerFactory)
	}
	val, _, _ = balancerFactoriesSfg.Do(uniqueKey, func() (interface{}, error) {
		b := newBalancerFactory(resolver, balancer, opts)
		balancerFactories.Store(uniqueKey, b)
		return b, nil
	})
	return val.(*BalancerFactory)
}

// watch expired balancer
func (b *BalancerFactory) watcher() {
	for range time.Tick(b.opts.ExpireInterval) {
		b.cache.Range(func(key, value interface{}) bool {
			bl := value.(*Balancer)
			if atomic.CompareAndSwapInt32(&bl.expire, 0, 1) {
				// 1. set expire flag
				// 2. wait next ticker for collect, maybe the balancer is used again
				// (avoid being immediate delete the balancer which had been created recently)
			} else {
				b.cache.Delete(key)
				bl.close()
			}
			return true
		})
	}
}

// cache key with resolver name prefix avoid conflict for balancer
func renameResultCacheKey(res *discovery.Result, resolverName string) {
	res.CacheKey = resolverName + ":" + res.CacheKey
}

// Get create a new balancer if not exists
func (b *BalancerFactory) Get(ctx context.Context, target rpcinfo.EndpointInfo) (*Balancer, error) {
	desc := b.resolver.Target(ctx, target)
	val, ok := b.cache.Load(desc)
	if ok {
		return val.(*Balancer), nil
	}
	val, err, _ := b.sfg.Do(desc, func() (interface{}, error) {
		res, err := b.resolver.Resolve(ctx, desc)
		if err != nil {
			return nil, err
		}
		renameResultCacheKey(&res, b.resolver.Name())
		bl := &Balancer{
			b:      b,
			target: desc,
		}
		bl.res.Store(res)
		bl.sharedTicker = getSharedTicker(bl, b.opts.RefreshInterval)
		b.cache.Store(desc, bl)
		return bl, nil
	})
	if err != nil {
		return nil, err
	}
	return val.(*Balancer), nil
}

// Balancer same with loadbalance.Loadbalancer but without resolver.Result that
// has been cached
type Balancer struct {
	b            *BalancerFactory
	target       string       // a description returned from the resolver's Target method
	res          atomic.Value // newest and previous discovery result
	expire       int32        // 0 = normal, 1 = expire and collect next ticker
	sharedTicker *utils.SharedTicker
}

func (bl *Balancer) Refresh() {
	res, err := bl.b.resolver.Resolve(context.Background(), bl.target)
	if err != nil {
		klog.Warnf("KITEX: resolver refresh failed, key=%s error=%s", bl.target, err.Error())
		return
	}
	renameResultCacheKey(&res, bl.b.resolver.Name())
	prev := bl.res.Load().(discovery.Result)
	if bl.b.rebalancer != nil {
		if ch, ok := bl.b.resolver.Diff(res.CacheKey, prev, res); ok {
			bl.b.rebalancer.Rebalance(ch)
		}
	}
	// replace previous result
	bl.res.Store(res)
}

// Tick implements the interface utils.TickerTask.
func (bl *Balancer) Tick() {
	bl.Refresh()
}

// GetResult returns the discovery result that the Balancer holds.
func (bl *Balancer) GetResult() (res discovery.Result, ok bool) {
	if v := bl.res.Load(); v != nil {
		return v.(discovery.Result), true
	}
	return
}

// GetPicker equal to loadbalance.Balancer without pass discovery.Result, because we cache the result
func (bl *Balancer) GetPicker() loadbalance.Picker {
	atomic.StoreInt32(&bl.expire, 0)
	res := bl.res.Load().(discovery.Result)
	return bl.b.balancer.GetPicker(res)
}

func (bl *Balancer) close() {
	// notice the under rebalancer
	if rb, ok := bl.b.balancer.(loadbalance.Rebalancer); ok {
		// notice to rebalancing
		rb.Delete(discovery.Change{
			Result: discovery.Result{
				Cacheable: true,
				CacheKey:  bl.res.Load().(discovery.Result).CacheKey,
			},
		})
	}
	// delete from sharedTicker
	bl.sharedTicker.Delete(bl)
}

const unknown = "unknown"

func Dump() interface{} {
	type instInfo struct {
		Address string
		Weight  int
	}
	cacheDump := make(map[string]interface{})
	balancerFactories.Range(func(key, val interface{}) bool {
		cacheKey := key.(string)
		if bf, ok := val.(*BalancerFactory); ok {
			routeMap := make(map[string]interface{})
			cacheDump[cacheKey] = routeMap
			bf.cache.Range(func(k, v interface{}) bool {
				routeKey := k.(string)
				if bl, ok := v.(*Balancer); ok {
					if dr, ok := bl.res.Load().(discovery.Result); ok {
						insts := make([]instInfo, 0, len(dr.Instances))
						for i := range dr.Instances {
							inst := dr.Instances[i]
							addr := fmt.Sprintf("%s://%s", inst.Address().Network(), inst.Address().String())
							insts = append(insts, instInfo{Address: addr, Weight: inst.Weight()})
						}
						routeMap[routeKey] = insts
					} else {
						routeMap[routeKey] = unknown
					}
				} else {
					routeMap[routeKey] = unknown
				}
				return true
			})
		} else {
			cacheDump[cacheKey] = unknown
		}
		return true
	})
	return cacheDump
}