mirror of https://github.com/tikv/client-go.git
61 lines
1.5 KiB
Go
61 lines
1.5 KiB
Go
package util
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type dnsF func(ctx context.Context, target string) (net.Conn, error)
|
|
|
|
func wrapWithDomain(target, domain string) (string, error) {
|
|
if len(domain) == 0 {
|
|
return target, nil
|
|
}
|
|
strlist := strings.Split(target, ":")
|
|
if len(strlist) != 2 {
|
|
return "", fmt.Errorf("target %s is not valid", target)
|
|
}
|
|
address := strlist[0]
|
|
port := strlist[1]
|
|
return fmt.Sprintf("%s.%s:%s", address, domain, port), nil
|
|
}
|
|
|
|
// GetDefaultDNSDialer is a util which built for connecting TiKV on k8s.
|
|
// Here is an example:
|
|
// coreDNSAddr := "8.8.8.8:53"
|
|
// domain := "cluster.local"
|
|
// dialer := grpc.WithContextDialer(util.GetCustomDNSDialer(coreDNSAddr, domain))
|
|
// cli, err := rawkv.NewClientWithOpts(
|
|
//
|
|
// context.TODO(),
|
|
// []string{"pd0.pd:2379"},
|
|
// rawkv.WithPDOptions(opt.WithGRPCDialOptions(dialer)),
|
|
// rawkv.WithGRPCDialOptions(dialer),
|
|
//
|
|
// )
|
|
// cli.Close()
|
|
func GetCustomDNSDialer(dnsServer, dnsDomain string) dnsF {
|
|
return func(ctx context.Context, target string) (net.Conn, error) {
|
|
r := net.Resolver{
|
|
PreferGo: true,
|
|
Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
|
|
d := net.Dialer{
|
|
Timeout: time.Millisecond * time.Duration(10000),
|
|
}
|
|
return d.DialContext(ctx, network, dnsServer)
|
|
},
|
|
}
|
|
dialer := &net.Dialer{
|
|
Resolver: &r,
|
|
}
|
|
addr, err := wrapWithDomain(target, dnsDomain)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return dialer.DialContext(ctx, "tcp", addr)
|
|
}
|
|
}
|