83 lines
1.4 KiB
Go
83 lines
1.4 KiB
Go
|
package invoke
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"google.golang.org/grpc/resolver"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
scheme = "bifrost"
|
||
|
|
)
|
||
|
|
|
||
|
|
type CustomBuilder struct{}
|
||
|
|
|
||
|
|
func (cb *CustomBuilder) Scheme() string {
|
||
|
|
return scheme
|
||
|
|
}
|
||
|
|
|
||
|
|
func (cb *CustomBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
|
||
|
|
cr := &customResolver{
|
||
|
|
cc: cc,
|
||
|
|
target: target,
|
||
|
|
}
|
||
|
|
|
||
|
|
cr.ResolveNow(resolver.ResolveNowOptions{})
|
||
|
|
|
||
|
|
return cr, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type customResolver struct {
|
||
|
|
sync.Mutex
|
||
|
|
target resolver.Target
|
||
|
|
cc resolver.ClientConn
|
||
|
|
ips map[string]string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (cr *customResolver) ResolveNow(o resolver.ResolveNowOptions) {
|
||
|
|
var (
|
||
|
|
addrs = make([]resolver.Address, 0)
|
||
|
|
hp []string
|
||
|
|
)
|
||
|
|
|
||
|
|
cr.Lock()
|
||
|
|
defer cr.Unlock()
|
||
|
|
|
||
|
|
if hp = strings.Split(cr.target.URL.Host, ":"); len(hp) >= 2 {
|
||
|
|
if ip, ok := pool[hp[0]]; ok {
|
||
|
|
addr := fmt.Sprintf("%s:%s", ip, hp[1])
|
||
|
|
addrs = append(addrs, resolver.Address{Addr: addr})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
_ = cr.cc.UpdateState(resolver.State{Addresses: addrs})
|
||
|
|
}
|
||
|
|
|
||
|
|
func (cr *customResolver) Close() {}
|
||
|
|
|
||
|
|
var (
|
||
|
|
cb = &CustomBuilder{}
|
||
|
|
pool = make(map[string]string)
|
||
|
|
)
|
||
|
|
|
||
|
|
func init() {
|
||
|
|
resolver.Register(cb)
|
||
|
|
}
|
||
|
|
|
||
|
|
type CustomDomain struct {
|
||
|
|
Domain string
|
||
|
|
IP string
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewCustomBuilder(cds ...CustomDomain) resolver.Builder {
|
||
|
|
locker.Lock()
|
||
|
|
defer locker.Unlock()
|
||
|
|
|
||
|
|
for _, cd := range cds {
|
||
|
|
pool[cd.Domain] = cd.IP
|
||
|
|
}
|
||
|
|
|
||
|
|
return cb
|
||
|
|
}
|