Added timeout to blackbox exporter (#181)

diff --git a/README.md b/README.md
index 8e072744ff2ac74e2f57dac30040e78f14c93813..1accdfce2cca93659196a1a71b08943caf5a5e2c 100644 (file)
--- a/README.md
+++ b/README.md
@@ -39,6 +39,8 @@ Additionally, an [example configuration](https://github.com/prometheus/blackbox_
 HTTP, HTTPS (via the `http` prober), DNS, TCP socket and ICMP (see permissions section) are currently supported.
 Additional modules can be defined to meet your needs.
 
+The timeout of each probe is automatically determined from the `scrape_timeout` in the [Prometheus config](https://prometheus.io/docs/operating/configuration/#configuration-file), slightly reduced to allow for network delays. 
+This can be further limited by the `timeout` in the Blackbox exporter config file. If neither is specified, it defaults to 10 seconds.
 
 ## Prometheus Configuration
 

diff --git a/dns.go b/dns.go
index 727c1dc2682b9689326a8ccc5814f7f491779da4..6718faf44c0473e5bd24c664258fca70b1376468 100644 (file)
--- a/dns.go
+++ b/dns.go
@@ -14,8 +14,10 @@
 package main
 
 import (
+       "context"
        "net"
        "regexp"
+       "time"
 
        "github.com/miekg/dns"
        "github.com/prometheus/client_golang/prometheus"
@@ -80,7 +82,7 @@ func validRcode(rcode int, valid []string) bool {
        return false
 }
 
-func probeDNS(target string, module Module, registry *prometheus.Registry) bool {
+func probeDNS(ctx context.Context, target string, module Module, registry *prometheus.Registry) bool {
        var numAnswer, numAuthority, numAdditional int
        var dialProtocol string
        probeDNSAnswerRRSGauge := prometheus.NewGauge(prometheus.GaugeOpts{
@@ -137,8 +139,6 @@ func probeDNS(target string, module Module, registry *prometheus.Registry) bool
 
        client := new(dns.Client)
        client.Net = dialProtocol
-       client.Timeout = module.Timeout
-
        qt := dns.TypeANY
        if module.DNS.QueryType != "" {
                var ok bool
@@ -150,6 +150,10 @@ func probeDNS(target string, module Module, registry *prometheus.Registry) bool
        }
        msg := new(dns.Msg)
        msg.SetQuestion(dns.Fqdn(module.DNS.QueryName), qt)
+
+       timeoutDeadline, _ := ctx.Deadline()
+       client.Timeout = timeoutDeadline.Sub(time.Now())
+
        response, _, err := client.Exchange(msg, target)
        if err != nil {
                log.Warnf("Error while sending a DNS query: %s", err)

diff --git a/dns_test.go b/dns_test.go
index ff36a4f367f89cf35555aaa07f5124a31fc1a4c5..93113b2294889db5902d37619ac259d85d9b772b 100644 (file)
--- a/dns_test.go
+++ b/dns_test.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+       "context"
        "net"
        "runtime"
        "testing"
@@ -118,7 +119,10 @@ func TestRecursiveDNSResponse(t *testing.T) {
                        test.Probe.TransportProtocol = protocol
                        registry := prometheus.NewPedanticRegistry()
                        registry.Gather()
-                       result := probeDNS(addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
+
+                       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+                       defer cancel()
+                       result := probeDNS(testCTX, addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
                        if result != test.ShouldSucceed {
                                t.Fatalf("Test %d had unexpected result: %v", i, result)
                        }
@@ -243,7 +247,9 @@ func TestAuthoritativeDNSResponse(t *testing.T) {
                for i, test := range tests {
                        test.Probe.TransportProtocol = protocol
                        registry := prometheus.NewRegistry()
-                       result := probeDNS(addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
+                       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+                       defer cancel()
+                       result := probeDNS(testCTX, addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
                        if result != test.ShouldSucceed {
                                t.Fatalf("Test %d had unexpected result: %v", i, result)
                        }
@@ -299,7 +305,9 @@ func TestServfailDNSResponse(t *testing.T) {
                for i, test := range tests {
                        test.Probe.TransportProtocol = protocol
                        registry := prometheus.NewRegistry()
-                       result := probeDNS(addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
+                       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+                       defer cancel()
+                       result := probeDNS(testCTX, addr.String(), Module{Timeout: time.Second, DNS: test.Probe}, registry)
                        if result != test.ShouldSucceed {
                                t.Fatalf("Test %d had unexpected result: %v", i, result)
                        }
@@ -345,7 +353,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry := prometheus.NewRegistry()
-               result := probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result := probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if !result {
                        t.Fatalf("DNS protocol: \"%v4\" connection test failed, expected success.", protocol)
                }
@@ -369,7 +379,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry = prometheus.NewRegistry()
-               result = probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result = probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if !result {
                        t.Fatalf("DNS protocol: \"%v6\" connection test failed, expected success.", protocol)
                }
@@ -392,7 +404,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry = prometheus.NewRegistry()
-               result = probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result = probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if !result {
                        t.Fatalf("DNS protocol: \"%v\", preferred \"ip6\" connection test failed, expected success.", protocol)
                }
@@ -415,7 +429,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry = prometheus.NewRegistry()
-               result = probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result = probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if !result {
                        t.Fatalf("DNS protocol: \"%v\", preferred \"ip4\" connection test failed, expected success.", protocol)
                }
@@ -438,7 +454,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry = prometheus.NewRegistry()
-               result = probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result = probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if !result {
                        t.Fatalf("DNS protocol: \"%v\" connection test failed, expected success.", protocol)
                }
@@ -460,7 +478,9 @@ func TestDNSProtocol(t *testing.T) {
                        },
                }
                registry = prometheus.NewRegistry()
-               result = probeDNS(net.JoinHostPort("localhost", port), module, registry)
+               testCTX, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result = probeDNS(testCTX, net.JoinHostPort("localhost", port), module, registry)
                if protocol == "udp" {
                        if !result {
                                t.Fatalf("DNS test connection with protocol %s failed, expected success.", protocol)

diff --git a/http.go b/http.go
index 217a2f6a2f1b9c8907515d03ecc7cf646527a18c..8c75d2e5d4249aff6ce1d6e18ab04c11a6587e57 100644 (file)
--- a/http.go
+++ b/http.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+       "context"
        "errors"
        "io"
        "io/ioutil"
@@ -58,9 +59,8 @@ func matchRegularExpressions(reader io.Reader, httpConfig HTTPProbe) bool {
        return true
 }
 
-func probeHTTP(target string, module Module, registry *prometheus.Registry) (success bool) {
+func probeHTTP(ctx context.Context, target string, module Module, registry *prometheus.Registry) (success bool) {
        var redirects int
-
        var (
                contentLengthGauge = prometheus.NewGauge(prometheus.GaugeOpts{
                        Name: "content_length",
@@ -127,7 +127,6 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
                log.Errorf("Error generating HTTP client: %v", err)
                return false
        }
-       client.Timeout = module.Timeout
 
        client.CheckRedirect = func(_ *http.Request, via []*http.Request) error {
                redirects = len(via)
@@ -143,6 +142,7 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
 
        request, err := http.NewRequest(httpConfig.Method, target, nil)
        request.Host = targetURL.Host
+       request = request.WithContext(ctx)
        if targetPort == "" {
                targetURL.Host = ip.String()
        } else {
@@ -162,14 +162,15 @@ func probeHTTP(target string, module Module, registry *prometheus.Registry) (suc
                request.Header.Set(key, value)
        }
 
-       // If a body is configured, add it to the request
+       // If a body is configured, add it to the request.
        if httpConfig.Body != "" {
                request.Body = ioutil.NopCloser(strings.NewReader(httpConfig.Body))
        }
        resp, err := client.Do(request)
        // Err won't be nil if redirects were turned off. See https://github.com/golang/go/issues/3795
        if err != nil && resp == nil {
-               log.Warnf("Error for HTTP request to %s: %s", target, err)
+               log.Errorf("Error for HTTP request to %s: %s", target, err)
+               success = false
        } else {
                defer resp.Body.Close()
                if len(httpConfig.ValidStatusCodes) != 0 {

diff --git a/http_test.go b/http_test.go
index 443d6890b2aeefa1d6f957e48ae2cc0de9b5fa56..f897dfd8f65ba8ef4f77d56a3ad9532043627175 100644 (file)
--- a/http_test.go
+++ b/http_test.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+       "context"
        "fmt"
        "net/http"
        "net/http/httptest"
@@ -49,7 +50,9 @@ func TestHTTPStatusCodes(t *testing.T) {
                defer ts.Close()
                registry := prometheus.NewRegistry()
                recorder := httptest.NewRecorder()
-               result := probeHTTP(ts.URL,
+               testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+               defer cancel()
+               result := probeHTTP(testCTX, ts.URL,
                        Module{Timeout: time.Second, HTTP: HTTPProbe{ValidStatusCodes: test.ValidStatusCodes}}, registry)
                body := recorder.Body.String()
                if result != test.ShouldSucceed {
@@ -74,7 +77,7 @@ func TestValidHTTPVersion(t *testing.T) {
                defer ts.Close()
                recorder := httptest.NewRecorder()
                registry := prometheus.NewRegistry()
-               result := probeHTTP(ts.URL,
+               result := probeHTTP(context.Background(), ts.URL,
                        Module{Timeout: time.Second, HTTP: HTTPProbe{
                                ValidHTTPVersions: test.ValidHTTPVersions,
                        }}, registry)
@@ -96,7 +99,9 @@ func TestRedirectFollowed(t *testing.T) {
        // Follow redirect, should succeed with 200.
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL, Module{Timeout: time.Second, HTTP: HTTPProbe{}}, registry)
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL, Module{Timeout: time.Second, HTTP: HTTPProbe{}}, registry)
        body := recorder.Body.String()
        if !result {
                t.Fatalf("Redirect test failed unexpectedly, got %s", body)
@@ -121,7 +126,9 @@ func TestRedirectNotFollowed(t *testing.T) {
        // Follow redirect, should succeed with 200.
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{NoFollowRedirects: true, ValidStatusCodes: []int{302}}}, registry)
        body := recorder.Body.String()
        if !result {
@@ -140,7 +147,9 @@ func TestPost(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{Method: "POST"}}, registry)
        body := recorder.Body.String()
        if !result {
@@ -155,7 +164,9 @@ func TestBasicAuth(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{
                        HTTPClientConfig: config.HTTPClientConfig{
                                TLSConfig: config.TLSConfig{InsecureSkipVerify: false},
@@ -175,7 +186,9 @@ func TestBearerToken(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{
                        HTTPClientConfig: config.HTTPClientConfig{
                                BearerToken: config.Secret("mysecret"),
@@ -194,7 +207,9 @@ func TestFailIfNotSSL(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfNotSSL: true}}, registry)
        body := recorder.Body.String()
        if result {
@@ -218,7 +233,9 @@ func TestFailIfMatchesRegexp(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfMatchesRegexp: []string{"could not connect to database"}}}, registry)
        body := recorder.Body.String()
        if result {
@@ -232,7 +249,7 @@ func TestFailIfMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfMatchesRegexp: []string{"could not connect to database"}}}, registry)
        body = recorder.Body.String()
        if !result {
@@ -248,7 +265,7 @@ func TestFailIfMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfMatchesRegexp: []string{"could not connect to database", "internal error"}}}, registry)
        body = recorder.Body.String()
        if result {
@@ -262,7 +279,7 @@ func TestFailIfMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfMatchesRegexp: []string{"could not connect to database", "internal error"}}}, registry)
        body = recorder.Body.String()
        if !result {
@@ -278,7 +295,9 @@ func TestFailIfNotMatchesRegexp(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfNotMatchesRegexp: []string{"Download the latest version here"}}}, registry)
        body := recorder.Body.String()
        if result {
@@ -292,7 +311,7 @@ func TestFailIfNotMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfNotMatchesRegexp: []string{"Download the latest version here"}}}, registry)
        body = recorder.Body.String()
        if !result {
@@ -308,7 +327,7 @@ func TestFailIfNotMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfNotMatchesRegexp: []string{"Download the latest version here", "Copyright 2015"}}}, registry)
        body = recorder.Body.String()
        if result {
@@ -322,7 +341,7 @@ func TestFailIfNotMatchesRegexp(t *testing.T) {
 
        recorder = httptest.NewRecorder()
        registry = prometheus.NewRegistry()
-       result = probeHTTP(ts.URL,
+       result = probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{FailIfNotMatchesRegexp: []string{"Download the latest version here", "Copyright 2015"}}}, registry)
        body = recorder.Body.String()
        if !result {
@@ -352,7 +371,9 @@ func TestHTTPHeaders(t *testing.T) {
        }))
        defer ts.Close()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL, Module{Timeout: time.Second, HTTP: HTTPProbe{
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL, Module{Timeout: time.Second, HTTP: HTTPProbe{
                Headers: headers,
        }}, registry)
        if !result {
@@ -367,7 +388,9 @@ func TestFailIfSelfSignedCA(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{
                        HTTPClientConfig: config.HTTPClientConfig{
                                TLSConfig: config.TLSConfig{InsecureSkipVerify: false},
@@ -394,7 +417,9 @@ func TestSucceedIfSelfSignedCA(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{
                        HTTPClientConfig: config.HTTPClientConfig{
                                TLSConfig: config.TLSConfig{InsecureSkipVerify: true},
@@ -421,7 +446,9 @@ func TestTLSConfigIsIgnoredForPlainHTTP(t *testing.T) {
 
        recorder := httptest.NewRecorder()
        registry := prometheus.NewRegistry()
-       result := probeHTTP(ts.URL,
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       result := probeHTTP(testCTX, ts.URL,
                Module{Timeout: time.Second, HTTP: HTTPProbe{
                        HTTPClientConfig: config.HTTPClientConfig{
                                TLSConfig: config.TLSConfig{InsecureSkipVerify: false},

diff --git a/icmp.go b/icmp.go
index 6a61070a28b60bda33687c196cf8aa1f91e05ee4..56754e7ee7256f71d8ebb75147aaf52afdd5148d 100644 (file)
--- a/icmp.go
+++ b/icmp.go
@@ -15,6 +15,7 @@ package main
 
 import (
        "bytes"
+       "context"
        "net"
        "os"
        "sync"
@@ -40,14 +41,14 @@ func getICMPSequence() uint16 {
        return icmpSequence
 }
 
-func probeICMP(target string, module Module, registry *prometheus.Registry) (success bool) {
+func probeICMP(ctx context.Context, target string, module Module, registry *prometheus.Registry) (success bool) {
        var (
                socket      *icmp.PacketConn
                requestType icmp.Type
                replyType   icmp.Type
        )
-
-       deadline := time.Now().Add(module.Timeout)
+       timeoutDeadline, _ := ctx.Deadline()
+       deadline := time.Now().Add(timeoutDeadline.Sub(time.Now()))
 
        ip, err := chooseProtocol(module.ICMP.PreferredIPProtocol, target, registry)
        if err != nil {

diff --git a/main.go b/main.go
index 5eff6bcfc92e27b679494928e38a552870a220fb..6a9bfe944819b17f63b934f6dbb7d591d07a7369 100644 (file)
--- a/main.go
+++ b/main.go
@@ -14,12 +14,14 @@
 package main
 
 import (
+       "context"
        "flag"
        "fmt"
        "io/ioutil"
        "net/http"
        "os"
        "os/signal"
+       "strconv"
        "syscall"
        "time"
 
@@ -31,7 +33,17 @@ import (
        "github.com/prometheus/common/version"
 )
 
-var Probers = map[string]func(string, Module, *prometheus.Registry) bool{
+var (
+       sc = &SafeConfig{
+               C: &Config{},
+       }
+       configFile    = flag.String("config.file", "blackbox.yml", "Blackbox exporter configuration file.")
+       listenAddress = flag.String("web.listen-address", ":9115", "The address to listen on for HTTP requests.")
+       showVersion   = flag.Bool("version", false, "Print version information.")
+       timeoutOffset = flag.Float64("timeout-offset", 0.5, "Offset to subtract from timeout in seconds.")
+)
+
+var Probers = map[string]func(context.Context, string, Module, *prometheus.Registry) bool{
        "http": probeHTTP,
        "tcp":  probeTCP,
        "icmp": probeICMP,
@@ -60,7 +72,40 @@ func (sc *SafeConfig) reloadConfig(confFile string) (err error) {
        return nil
 }
 
-func probeHandler(w http.ResponseWriter, r *http.Request, conf *Config) {
+func probeHandler(w http.ResponseWriter, r *http.Request, c *Config) {
+
+       moduleName := r.URL.Query().Get("module")
+       if moduleName == "" {
+               moduleName = "http_2xx"
+       }
+       module, ok := c.Modules[moduleName]
+       if !ok {
+               http.Error(w, fmt.Sprintf("Unknown module %q", moduleName), 400)
+               return
+       }
+
+       // If a timeout is configured via the Prometheus header, add it to the request.
+       var prometheusTimeout string
+       if r.Header["X-Prometheus-Scrape-Timeout-Seconds"] != nil {
+               prometheusTimeout = r.Header["X-Prometheus-Scrape-Timeout-Seconds"][0]
+       }
+
+       timeoutSeconds, err := strconv.ParseFloat(prometheusTimeout, 64)
+       if err != nil {
+               http.Error(w, fmt.Sprintf("Failed to parse timeout from Prometheus header: %s", err), http.StatusInternalServerError)
+               return
+       }
+       if timeoutSeconds == 0 {
+               timeoutSeconds = 10
+       }
+
+       if module.Timeout.Seconds() < timeoutSeconds && module.Timeout.Seconds() > 0 {
+               timeoutSeconds = module.Timeout.Seconds()
+       }
+       ctx, cancel := context.WithTimeout(context.Background(), time.Duration((timeoutSeconds-*timeoutOffset)*1e9))
+       defer cancel()
+       r = r.WithContext(ctx)
+
        probeSuccessGauge := prometheus.NewGauge(prometheus.GaugeOpts{
                Name: "probe_success",
                Help: "Displays whether or not the probe was a success",
@@ -76,15 +121,6 @@ func probeHandler(w http.ResponseWriter, r *http.Request, conf *Config) {
                return
        }
 
-       moduleName := params.Get("module")
-       if moduleName == "" {
-               moduleName = "http_2xx"
-       }
-       module, ok := conf.Modules[moduleName]
-       if !ok {
-               http.Error(w, fmt.Sprintf("Unknown module %q", moduleName), 400)
-               return
-       }
        prober, ok := Probers[module.Prober]
        if !ok {
                http.Error(w, fmt.Sprintf("Unknown prober %q", module.Prober), 400)
@@ -95,7 +131,7 @@ func probeHandler(w http.ResponseWriter, r *http.Request, conf *Config) {
        registry := prometheus.NewRegistry()
        registry.MustRegister(probeSuccessGauge)
        registry.MustRegister(probeDurationGauge)
-       success := prober(target, module, registry)
+       success := prober(ctx, target, module, registry)
        probeDurationGauge.Set(time.Since(start).Seconds())
        if success {
                probeSuccessGauge.Set(1)
@@ -109,15 +145,6 @@ func init() {
 }
 
 func main() {
-
-       var (
-               configFile    = flag.String("config.file", "blackbox.yml", "Blackbox exporter configuration file.")
-               listenAddress = flag.String("web.listen-address", ":9115", "The address to listen on for HTTP requests.")
-               showVersion   = flag.Bool("version", false, "Print version information.")
-               sc            = &SafeConfig{
-                       C: &Config{},
-               }
-       )
        flag.Parse()
 
        if *showVersion {
@@ -153,15 +180,6 @@ func main() {
                }
        }()
 
-       http.Handle("/metrics", prometheus.Handler())
-       http.HandleFunc("/probe",
-               func(w http.ResponseWriter, r *http.Request) {
-                       sc.RLock()
-                       c := sc.C
-                       sc.RUnlock()
-
-                       probeHandler(w, r, c)
-               })
        http.HandleFunc("/-/reload",
                func(w http.ResponseWriter, r *http.Request) {
                        if r.Method != "POST" {
@@ -176,7 +194,13 @@ func main() {
                                http.Error(w, fmt.Sprintf("failed to reload config: %s", err), http.StatusInternalServerError)
                        }
                })
-
+       http.Handle("/metrics", prometheus.Handler())
+       http.HandleFunc("/probe", func(w http.ResponseWriter, r *http.Request) {
+               sc.Lock()
+               conf := sc.C
+               sc.Unlock()
+               probeHandler(w, r, conf)
+       })
        http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
                w.Write([]byte(`<html>
                        <head><title>Blackbox Exporter</title></head>

diff --git a/main_test.go b/main_test.go
new file mode 100644 (file)
index 0000000..740cb81
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,42 @@
+package main
+
+import (
+       "net/http"
+       "net/http/httptest"
+       "testing"
+       "time"
+)
+
+var c = &Config{
+       Modules: map[string]Module{
+               "http_2xx": Module{
+                       Prober:  "http",
+                       Timeout: 10 * time.Second,
+               },
+       },
+}
+
+func TestPrometheusTimeout(t *testing.T) {
+
+       ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+               time.Sleep(2 * time.Second)
+       }))
+       defer ts.Close()
+
+       req, err := http.NewRequest("GET", "?target="+ts.URL, nil)
+       if err != nil {
+               t.Fatal(err)
+       }
+       req.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", "1")
+
+       rr := httptest.NewRecorder()
+       handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+               probeHandler(w, r, c)
+       })
+
+       handler.ServeHTTP(rr, req)
+
+       if status := rr.Code; status != http.StatusOK {
+               t.Errorf("probe request handler returned wrong status code: %v, want %v", status, http.StatusOK)
+       }
+}

diff --git a/tcp.go b/tcp.go
index d823f42ba053f1ee81977fc8a320c87062ddbc9d..4a3138718a12cb4f73a1a3087aa80c4b908acdd1 100644 (file)
--- a/tcp.go
+++ b/tcp.go
@@ -15,6 +15,7 @@ package main
 
 import (
        "bufio"
+       "context"
        "crypto/tls"
        "fmt"
        "net"
@@ -26,10 +27,9 @@ import (
        "github.com/prometheus/common/log"
 )
 
-func dialTCP(target string, module Module, registry *prometheus.Registry) (net.Conn, error) {
+func dialTCP(ctx context.Context, target string, module Module, registry *prometheus.Registry) (net.Conn, error) {
        var dialProtocol, dialTarget string
-       dialer := &net.Dialer{Timeout: module.Timeout}
-
+       dialer := &net.Dialer{}
        targetAddress, port, err := net.SplitHostPort(target)
        if err != nil {
                log.Errorf("Error splitting target address and port: %v", err)
@@ -50,24 +50,26 @@ func dialTCP(target string, module Module, registry *prometheus.Registry) (net.C
        dialTarget = net.JoinHostPort(ip.String(), port)
 
        if !module.TCP.TLS {
-               return dialer.Dial(dialProtocol, dialTarget)
+               return dialer.DialContext(ctx, dialProtocol, dialTarget)
        }
        tlsConfig, err := config.NewTLSConfig(&module.TCP.TLSConfig)
        if err != nil {
                log.Errorf("Error creating TLS configuration: %v", err)
                return nil, err
        }
+       timeoutDeadline, _ := ctx.Deadline()
+       dialer.Deadline = timeoutDeadline
+
        return tls.DialWithDialer(dialer, dialProtocol, dialTarget, tlsConfig)
 }
 
-func probeTCP(target string, module Module, registry *prometheus.Registry) bool {
+func probeTCP(ctx context.Context, target string, module Module, registry *prometheus.Registry) bool {
        probeSSLEarliestCertExpiry := prometheus.NewGauge(prometheus.GaugeOpts{
                Name: "probe_ssl_earliest_cert_expiry",
                Help: "Returns earliest SSL cert expiry date",
        })
-       registry.MustRegister(probeSSLEarliestCertExpiry)
        deadline := time.Now().Add(module.Timeout)
-       conn, err := dialTCP(target, module, registry)
+       conn, err := dialTCP(ctx, target, module, registry)
        if err != nil {
                log.Errorf("Error dialing TCP: %v", err)
                return false
@@ -83,6 +85,7 @@ func probeTCP(target string, module Module, registry *prometheus.Registry) bool
        }
        if module.TCP.TLS {
                state := conn.(*tls.Conn).ConnectionState()
+               registry.MustRegister(probeSSLEarliestCertExpiry)
                probeSSLEarliestCertExpiry.Set(float64(getEarliestCertExpiry(&state).UnixNano()) / 1e9)
        }
        scanner := bufio.NewScanner(conn)

diff --git a/tcp_test.go b/tcp_test.go
index 7add94c7f4f058008bd1cbd6ef9e936b44c3dae7..222ab06badf1f8f8da43f85eabf2052abcd81a4a 100644 (file)
--- a/tcp_test.go
+++ b/tcp_test.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+       "context"
        "fmt"
        "net"
        "runtime"
@@ -39,8 +40,10 @@ func TestTCPConnection(t *testing.T) {
                conn.Close()
                ch <- struct{}{}
        }()
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
        registry := prometheus.NewRegistry()
-       if !probeTCP(ln.Addr().String(), Module{Timeout: time.Second}, registry) {
+       if !probeTCP(testCTX, ln.Addr().String(), Module{Timeout: time.Second}, registry) {
                t.Fatalf("TCP module failed, expected success.")
        }
        <-ch
@@ -49,7 +52,9 @@ func TestTCPConnection(t *testing.T) {
 func TestTCPConnectionFails(t *testing.T) {
        // Invalid port number.
        registry := prometheus.NewRegistry()
-       if probeTCP(":0", Module{Timeout: time.Second}, registry) {
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+       if probeTCP(testCTX, ":0", Module{Timeout: time.Second}, registry) {
                t.Fatalf("TCP module suceeded, expected failure.")
        }
 }
@@ -61,6 +66,9 @@ func TestTCPConnectionQueryResponseIRC(t *testing.T) {
        }
        defer ln.Close()
 
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+
        module := Module{
                Timeout: time.Second,
                TCP: TCPProbe{
@@ -87,7 +95,7 @@ func TestTCPConnectionQueryResponseIRC(t *testing.T) {
                ch <- struct{}{}
        }()
        registry := prometheus.NewRegistry()
-       if !probeTCP(ln.Addr().String(), module, registry) {
+       if !probeTCP(testCTX, ln.Addr().String(), module, registry) {
                t.Fatalf("TCP module failed, expected success.")
        }
        <-ch
@@ -106,7 +114,7 @@ func TestTCPConnectionQueryResponseIRC(t *testing.T) {
                ch <- struct{}{}
        }()
        registry = prometheus.NewRegistry()
-       if probeTCP(ln.Addr().String(), module, registry) {
+       if probeTCP(testCTX, ln.Addr().String(), module, registry) {
                t.Fatalf("TCP module succeeded, expected failure.")
        }
        <-ch
@@ -119,6 +127,9 @@ func TestTCPConnectionQueryResponseMatching(t *testing.T) {
        }
        defer ln.Close()
 
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+
        module := Module{
                Timeout: time.Second,
                TCP: TCPProbe{
@@ -145,7 +156,7 @@ func TestTCPConnectionQueryResponseMatching(t *testing.T) {
                ch <- version
        }()
        registry := prometheus.NewRegistry()
-       if !probeTCP(ln.Addr().String(), module, registry) {
+       if !probeTCP(testCTX, ln.Addr().String(), module, registry) {
                t.Fatalf("TCP module failed, expected success.")
        }
        if got, want := <-ch, "OpenSSH_6.9p1"; got != want {
@@ -171,6 +182,9 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
        defer ln.Close()
 
+       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+       defer cancel()
+
        _, port, _ := net.SplitHostPort(ln.Addr().String())
 
        // Force IPv4
@@ -182,7 +196,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry := prometheus.NewRegistry()
-       result := probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result := probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP protocol: \"tcp4\" connection test failed, expected success.")
        }
@@ -202,7 +216,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry = prometheus.NewRegistry()
-       result = probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result = probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP protocol: \"tcp6\" connection test failed, expected success.")
        }
@@ -224,7 +238,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry = prometheus.NewRegistry()
-       result = probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result = probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP protocol: \"tcp\", prefer: \"ip4\" connection test failed, expected success.")
        }
@@ -246,7 +260,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry = prometheus.NewRegistry()
-       result = probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result = probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP protocol: \"tcp\", prefer: \"ip6\" connection test failed, expected success.")
        }
@@ -266,7 +280,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry = prometheus.NewRegistry()
-       result = probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result = probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP protocol: \"tcp\" connection test failed, expected success.")
        }
@@ -286,7 +300,7 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
 
        registry = prometheus.NewRegistry()
-       result = probeTCP(net.JoinHostPort("localhost", port), module, registry)
+       result = probeTCP(testCTX, net.JoinHostPort("localhost", port), module, registry)
        if !result {
                t.Fatalf("TCP connection test with protocol unspecified failed, expected success.")
        }