use ctx for chooseProtocol

use LookupIPAddr to support context timeout Signed-off-by: Thor <[email protected]>
prometheus · Apr 11, 2019 · 0e6aaa2 · 0e6aaa2
1 parent 80a1218
commit 0e6aaa2
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 38 deletions.
diff --git a/prober/dns.go b/prober/dns.go
@@ -129,7 +129,7 @@ func ProbeDNS(ctx context.Context, target string, module config.Module, registry
 			port = "53"
 			targetAddr = target
 		}
-		ip, _, err = chooseProtocol(module.DNS.IPProtocol, module.DNS.IPProtocolFallback, targetAddr, registry, logger)
+		ip, _, err = chooseProtocol(ctx, module.DNS.IPProtocol, module.DNS.IPProtocolFallback, targetAddr, registry, logger)
 		if err != nil {
 			level.Error(logger).Log("msg", "Error resolving address", "err", err)
 			return false

diff --git a/prober/http.go b/prober/http.go
@@ -271,7 +271,7 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
 		targetHost = targetURL.Host
 	}
 
-	ip, lookupTime, err := chooseProtocol(module.HTTP.IPProtocol, module.HTTP.IPProtocolFallback, targetHost, registry, logger)
+	ip, lookupTime, err := chooseProtocol(ctx, module.HTTP.IPProtocol, module.HTTP.IPProtocolFallback, targetHost, registry, logger)
 	if err != nil {
 		level.Error(logger).Log("msg", "Error resolving address", "err", err)
 		return false

diff --git a/prober/icmp.go b/prober/icmp.go
@@ -79,7 +79,7 @@ func ProbeICMP(ctx context.Context, target string, module config.Module, registr
 
 	registry.MustRegister(durationGaugeVec)
 
-	ip, lookupTime, err := chooseProtocol(module.ICMP.IPProtocol, module.ICMP.IPProtocolFallback, target, registry, logger)
+	ip, lookupTime, err := chooseProtocol(ctx, module.ICMP.IPProtocol, module.ICMP.IPProtocolFallback, target, registry, logger)
 	if err != nil {
 		level.Warn(logger).Log("msg", "Error resolving address", "err", err)
 		return false

diff --git a/prober/tcp.go b/prober/tcp.go
@@ -38,7 +38,7 @@ func dialTCP(ctx context.Context, target string, module config.Module, registry
 		return nil, err
 	}
 
-	ip, _, err := chooseProtocol(module.TCP.IPProtocol, module.TCP.IPProtocolFallback, targetAddress, registry, logger)
+	ip, _, err := chooseProtocol(ctx, module.TCP.IPProtocol, module.TCP.IPProtocolFallback, targetAddress, registry, logger)
 	if err != nil {
 		level.Error(logger).Log("msg", "Error resolving address", "err", err)
 		return nil, err

diff --git a/prober/utils.go b/prober/utils.go
@@ -61,45 +61,45 @@ func chooseProtocol(ctx context.Context, IPProtocol string, fallbackIPProtocol b
 		probeDNSLookupTimeSeconds.Add(lookupTime)
 	}()
 
-	ipC := make(chan *net.IPAddr, 1)
-	errors := make(chan error, 1)
-	go func() {
-		defer close(ipC)
-		defer close(errors)
-		ip, err := net.ResolveIPAddr(IPProtocol, target)
-		if err != nil {
-			if !fallbackIPProtocol {
-				level.Error(logger).Log("msg", "Resolution with IP protocol failed (fallback_ip_protocol is false):", "err", err)
-			} else {
-				level.Warn(logger).Log("msg", "Resolution with IP protocol failed, attempting fallback protocol", "fallback_protocol", fallbackProtocol, "err", err)
-				ip, err = net.ResolveIPAddr(fallbackProtocol, target)
+	resolver := &net.Resolver{}
+	ips, err := resolver.LookupIPAddr(ctx, target)
+	if err != nil {
+		level.Error(logger).Log("msg", "Resolution with IP protocol failed", "err", err)
+		return nil, 0.0, err
+	}
+
+	// Return the IP in the requested protocol format
+	var fallback net.IPAddr
+	for _, ip := range ips {
+		switch IPProtocol {
+		case "ip4":
+			if ip.IP.To4() != nil {
+				level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
+				probeIPProtocolGauge.Set(4)
+				return &ip, lookupTime, nil
 			}
 
-			if err != nil {
-				if IPProtocol == "ip6" {
-					probeIPProtocolGauge.Set(6)
-				} else {
-					probeIPProtocolGauge.Set(4)
-				}
-				errors <- err
-				return
+			// ip4 as fallback
+			fallback = ip
+
+		case "ip6":
+
+			if ip.IP.To4() == nil {
+				level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
+				probeIPProtocolGauge.Set(6)
+				return &ip, lookupTime, nil
 			}
-		}
 
-		if ip.IP.To4() == nil {
-			probeIPProtocolGauge.Set(6)
-		} else {
-			probeIPProtocolGauge.Set(4)
+			// ip6 as fallback
+			fallback = ip
 		}
+	}
 
-		ipC <- ip
-	}()
-
-	select {
-	case <-ctx.Done():
-		return nil, lookupTime, ctx.Err()
-	case ip := <-ipC:
-		level.Info(logger).Log("msg", "Resolved target address", "ip", ip)
-		return ip, lookupTime, nil
+	// Use fallback ip protocol
+	if fallbackProtocol == "ip4" {
+		probeIPProtocolGauge.Set(4)
+	} else {
+		probeIPProtocolGauge.Set(6)
 	}
+	return &fallback, lookupTime, nil
 }