[SPPSP-5975] optimize "GetOutboundAddress"

projecteru2 · Jul 25, 2022 · afd6e65 · afd6e65
1 parent 688edfe
commit afd6e65
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 16 deletions.
diff --git a/cluster/calcium/calcium_test.go b/cluster/calcium/calcium_test.go
@@ -56,6 +56,7 @@ func NewTestCluster() *Calcium {
 		WALFile:             filepath.Join(walDir, "core.wal.log"),
 		MaxConcurrency:      10,
 		HAKeepaliveInterval: 16 * time.Second,
+		DialTarget:          "8.8.8.8:80",
 	}
 	c.store = &storemocks.Store{}
 	c.scheduler = &schedulermocks.Scheduler{}

diff --git a/cluster/calcium/service.go b/cluster/calcium/service.go
@@ -26,7 +26,7 @@ func (c *Calcium) WatchServiceStatus(ctx context.Context) (<-chan types.ServiceS
 
 // RegisterService writes self service address in store
 func (c *Calcium) RegisterService(ctx context.Context) (unregister func(), err error) {
-	serviceAddress, err := utils.GetOutboundAddress(c.config.Bind)
+	serviceAddress, err := utils.GetOutboundAddress(c.config.Bind, c.config.DialTarget)
 	if err != nil {
 		log.Errorf(ctx, "[RegisterService] failed to get outbound address: %v", err)
 		return

diff --git a/types/config.go b/types/config.go
@@ -18,16 +18,17 @@ const (
 // Config holds eru-core config
 type Config struct {
 	LogLevel            string        `yaml:"log_level" required:"true" default:"INFO"`
-	Bind                string        `yaml:"bind" required:"true" default:"5001"`           // HTTP API address
-	LockTimeout         time.Duration `yaml:"lock_timeout" required:"true" default:"30s"`    // timeout for lock (ttl)
-	GlobalTimeout       time.Duration `yaml:"global_timeout" required:"true" default:"300s"` // timeout for remove, run_and_wait and build, in second
-	ConnectionTimeout   time.Duration `yaml:"connection_timeout" default:"10s"`              // timeout for connections
-	HAKeepaliveInterval time.Duration `yaml:"ha_keepalive_interval" default:"16s"`           // interval for node status watcher
-	Statsd              string        `yaml:"statsd"`                                        // statsd host and port
-	Profile             string        `yaml:"profile"`                                       // profile ip:port
-	CertPath            string        `yaml:"cert_path"`                                     // docker cert files path
-	MaxConcurrency      int64         `yaml:"max_concurrency" default:"20"`                  // concurrently call single runtime in the same time
-	Store               string        `yaml:"store" default:"etcd"`                          // store type
+	Bind                string        `yaml:"bind" required:"true" default:"5001"`               // HTTP API address
+	DialTarget          string        `yaml:"dial_target" required:"false" default:"8.8.8.8:80"` // for getting outbound address
+	LockTimeout         time.Duration `yaml:"lock_timeout" required:"true" default:"30s"`        // timeout for lock (ttl)
+	GlobalTimeout       time.Duration `yaml:"global_timeout" required:"true" default:"300s"`     // timeout for remove, run_and_wait and build, in second
+	ConnectionTimeout   time.Duration `yaml:"connection_timeout" default:"10s"`                  // timeout for connections
+	HAKeepaliveInterval time.Duration `yaml:"ha_keepalive_interval" default:"16s"`               // interval for node status watcher
+	Statsd              string        `yaml:"statsd"`                                            // statsd host and port
+	Profile             string        `yaml:"profile"`                                           // profile ip:port
+	CertPath            string        `yaml:"cert_path"`                                         // docker cert files path
+	MaxConcurrency      int64         `yaml:"max_concurrency" default:"20"`                      // concurrently call single runtime in the same time
+	Store               string        `yaml:"store" default:"etcd"`                              // store type
 
 	Auth       AuthConfig `yaml:"auth"` // grpc auth
 	GRPCConfig GRPCConfig `yaml:"grpc"` // grpc config

diff --git a/utils/service.go b/utils/service.go
@@ -6,15 +6,30 @@ import (
 	"strings"
 )
 
-// GetOutboundAddress finds out self service address
-func GetOutboundAddress(bind string) (string, error) {
-	conn, err := net.Dial("udp", "8.8.8.8:80")
+// GetOutboundAddress finds out self-service address
+func GetOutboundAddress(bind string, dialTarget string) (string, error) {
+	parts := strings.Split(bind, ":")
+	if len(parts) != 2 {
+		return "", fmt.Errorf("invalid bind address %s", bind)
+	}
+	ip := parts[0]
+	port := parts[1]
+
+	address := net.ParseIP(ip)
+	if ip == "" || address == nil || address.IsUnspecified() {
+		return getOutboundAddress(port, dialTarget)
+	}
+
+	return bind, nil
+}
+
+func getOutboundAddress(port string, dialTarget string) (string, error) {
+	conn, err := net.Dial("udp", dialTarget)
 	if err != nil {
 		return "", err
 	}
 	defer conn.Close()
 
 	localAddr := conn.LocalAddr().(*net.UDPAddr)
-	port := strings.Split(bind, ":")[1]
 	return fmt.Sprintf("%s:%s", localAddr.IP, port), nil
 }
diff --git a/utils/service_test.go b/utils/service_test.go
@@ -8,7 +8,7 @@ import (
 
 func TestGetOutboundAddress(t *testing.T) {
 	bind := "1.1.1.1:1234"
-	addr, err := GetOutboundAddress(bind)
+	addr, err := GetOutboundAddress(bind, "8.8.8.8:80")
 	assert.NoError(t, err)
 	assert.Contains(t, addr, "1234")
 }