github.com/cloudwego/kitex@v0.9.0/pkg/discovery/discovery.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 discovery defines interfaces for service discovery.
// Developers that are willing to customize service discovery
// should implement their own Resolver and supply it with the
// option WithResolver at client's creation.
package discovery

import (
	"context"
	"net"

	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

// DefaultWeight is the default weight for an instance.
const DefaultWeight = 10

// Result contains the result of service discovery process.
// Cacheable tells whether the instance list can/should be cached.
// When Cacheable is true, CacheKey can be used to map the instance list in cache.
type Result struct {
	Cacheable bool
	CacheKey  string
	Instances []Instance
}

// Change contains the difference between the current discovery result and the previous one.
// It is designed for providing detail information when dispatching an event for service
// discovery result change.
// Since the loadbalancer may rely on caching the result of resolver to improve performance,
// the resolver implementation should dispatch an event when result changes.
type Change struct {
	Result  Result
	Added   []Instance
	Updated []Instance
	Removed []Instance
}

// Resolver resolves the target endpoint into a list of Instance.
type Resolver interface {
	// Target should return a description for the given target that is suitable for being a key for cache.
	Target(ctx context.Context, target rpcinfo.EndpointInfo) (description string)

	// Resolve returns a list of instances for the given description of a target.
	Resolve(ctx context.Context, desc string) (Result, error)

	// Diff computes the difference between two results.
	// When `next` is cacheable, the Change should be cacheable, too. And the `Result` field's CacheKey in
	// the return value should be set with the given cacheKey.
	Diff(cacheKey string, prev, next Result) (Change, bool)

	// Name returns the name of the resolver.
	Name() string
}

// DefaultDiff provides a natural implementation for the Diff method of the Resolver interface.
func DefaultDiff(cacheKey string, prev, next Result) (Change, bool) {
	ch := Change{
		Result: Result{
			Cacheable: next.Cacheable,
			CacheKey:  cacheKey,
			Instances: next.Instances,
		},
	}

	prevMap := make(map[string]struct{}, len(prev.Instances))
	for _, ins := range prev.Instances {
		prevMap[ins.Address().String()] = struct{}{}
	}

	nextMap := make(map[string]struct{}, len(next.Instances))
	for _, ins := range next.Instances {
		addr := ins.Address().String()
		nextMap[addr] = struct{}{}
		if _, found := prevMap[addr]; !found {
			ch.Added = append(ch.Added, ins)
		}
	}

	for _, ins := range prev.Instances {
		if _, found := nextMap[ins.Address().String()]; !found {
			ch.Removed = append(ch.Removed, ins)
		}
	}
	return ch, len(ch.Added)+len(ch.Removed) != 0
}

type instance struct {
	addr   net.Addr
	weight int
	tags   map[string]string
}

func (i *instance) Address() net.Addr {
	return i.addr
}

func (i *instance) Weight() int {
	return i.weight
}

func (i *instance) Tag(key string) (value string, exist bool) {
	value, exist = i.tags[key]
	return
}

// NewInstance creates a Instance using the given network, address and tags
func NewInstance(network, address string, weight int, tags map[string]string) Instance {
	return &instance{
		addr:   utils.NewNetAddr(network, address),
		weight: weight,
		tags:   tags,
	}
}

// SynthesizedResolver synthesizes a Resolver using a resolve function.
type SynthesizedResolver struct {
	TargetFunc  func(ctx context.Context, target rpcinfo.EndpointInfo) string
	ResolveFunc func(ctx context.Context, key string) (Result, error)
	DiffFunc    func(key string, prev, next Result) (Change, bool)
	NameFunc    func() string
}

// Target implements the Resolver interface.
func (sr SynthesizedResolver) Target(ctx context.Context, target rpcinfo.EndpointInfo) string {
	if sr.TargetFunc == nil {
		return ""
	}
	return sr.TargetFunc(ctx, target)
}

// Resolve implements the Resolver interface.
func (sr SynthesizedResolver) Resolve(ctx context.Context, key string) (Result, error) {
	return sr.ResolveFunc(ctx, key)
}

// Diff implements the Resolver interface.
func (sr SynthesizedResolver) Diff(key string, prev, next Result) (Change, bool) {
	if sr.DiffFunc == nil {
		return DefaultDiff(key, prev, next)
	}
	return sr.DiffFunc(key, prev, next)
}

// Name implements the Resolver interface
func (sr SynthesizedResolver) Name() string {
	if sr.NameFunc == nil {
		return ""
	}
	return sr.NameFunc()
}

// Instance contains information of an instance from the target service.
type Instance interface {
	Address() net.Addr
	Weight() int
	Tag(key string) (value string, exist bool)
}