From fe7a1bcfc69147c0b7ff07017740f41dd9bfe4b0 Mon Sep 17 00:00:00 2001 From: Simon Pasquier Date: Tue, 26 Feb 2019 14:47:04 +0100 Subject: [PATCH] discovery/triton: pass context to the client (#5235) Signed-off-by: Simon Pasquier --- discovery/triton/triton.go | 19 +++++++---- discovery/triton/triton_test.go | 60 ++++++++++++++++++--------------- 2 files changed, 45 insertions(+), 34 deletions(-) diff --git a/discovery/triton/triton.go b/discovery/triton/triton.go index 242ac146d..b530a8e41 100644 --- a/discovery/triton/triton.go +++ b/discovery/triton/triton.go @@ -159,7 +159,7 @@ func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) { defer ticker.Stop() // Get an initial set right away. - tg, err := d.refresh() + tg, err := d.refresh(ctx) if err != nil { level.Error(d.logger).Log("msg", "Refreshing targets failed", "err", err) } else { @@ -169,7 +169,7 @@ func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) { for { select { case <-ticker.C: - tg, err := d.refresh() + tg, err := d.refresh(ctx) if err != nil { level.Error(d.logger).Log("msg", "Refreshing targets failed", "err", err) } else { @@ -181,7 +181,7 @@ func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) { } } -func (d *Discovery) refresh() (tg *targetgroup.Group, err error) { +func (d *Discovery) refresh(ctx context.Context) (tg *targetgroup.Group, err error) { t0 := time.Now() defer func() { refreshDuration.Observe(time.Since(t0).Seconds()) @@ -200,22 +200,27 @@ func (d *Discovery) refresh() (tg *targetgroup.Group, err error) { Source: endpoint, } - resp, err := d.client.Get(endpoint) + req, err := http.NewRequest("GET", endpoint, nil) if err != nil { - return tg, fmt.Errorf("an error occurred when requesting targets from the discovery endpoint. %s", err) + return nil, err + } + req = req.WithContext(ctx) + resp, err := d.client.Do(req) + if err != nil { + return nil, fmt.Errorf("an error occurred when requesting targets from the discovery endpoint: %s", err) } defer resp.Body.Close() data, err := ioutil.ReadAll(resp.Body) if err != nil { - return tg, fmt.Errorf("an error occurred when reading the response body. %s", err) + return nil, fmt.Errorf("an error occurred when reading the response body: %s", err) } dr := DiscoveryResponse{} err = json.Unmarshal(data, &dr) if err != nil { - return tg, fmt.Errorf("an error occurred unmarshaling the discovery response json. %s", err) + return nil, fmt.Errorf("an error occurred unmarshaling the discovery response json: %s", err) } for _, container := range dr.Containers { diff --git a/discovery/triton/triton_test.go b/discovery/triton/triton_test.go index 01c4b10a0..7b78c451c 100644 --- a/discovery/triton/triton_test.go +++ b/discovery/triton/triton_test.go @@ -67,8 +67,12 @@ var ( } ) +func newTritonDiscovery(c SDConfig) (*Discovery, error) { + return New(nil, &c) +} + func TestTritonSDNew(t *testing.T) { - td, err := New(nil, &conf) + td, err := newTritonDiscovery(conf) testutil.Ok(t, err) testutil.Assert(t, td != nil, "") testutil.Assert(t, td.client != nil, "") @@ -81,13 +85,13 @@ func TestTritonSDNew(t *testing.T) { } func TestTritonSDNewBadConfig(t *testing.T) { - td, err := New(nil, &badconf) + td, err := newTritonDiscovery(badconf) testutil.NotOk(t, err, "") testutil.Assert(t, td == nil, "") } func TestTritonSDNewGroupsConfig(t *testing.T) { - td, err := New(nil, &groupsconf) + td, err := newTritonDiscovery(groupsconf) testutil.Ok(t, err) testutil.Assert(t, td != nil, "") testutil.Assert(t, td.client != nil, "") @@ -102,14 +106,11 @@ func TestTritonSDNewGroupsConfig(t *testing.T) { func TestTritonSDRun(t *testing.T) { var ( - td, err = New(nil, &conf) + td, _ = newTritonDiscovery(conf) ch = make(chan []*targetgroup.Group) ctx, cancel = context.WithCancel(context.Background()) ) - testutil.Ok(t, err) - testutil.Assert(t, td != nil, "") - wait := make(chan struct{}) go func() { td.Run(ctx, ch) @@ -160,47 +161,52 @@ func TestTritonSDRefreshMultipleTargets(t *testing.T) { func TestTritonSDRefreshNoServer(t *testing.T) { var ( - td, err = New(nil, &conf) + td, _ = newTritonDiscovery(conf) ) - testutil.Ok(t, err) - testutil.Assert(t, td != nil, "") - tg, rerr := td.refresh() - testutil.NotOk(t, rerr, "") - testutil.Equals(t, strings.Contains(rerr.Error(), "an error occurred when requesting targets from the discovery endpoint."), true) - testutil.Assert(t, tg != nil, "") - testutil.Assert(t, tg.Targets == nil, "") + _, err := td.refresh(context.Background()) + testutil.NotOk(t, err, "") + testutil.Equals(t, strings.Contains(err.Error(), "an error occurred when requesting targets from the discovery endpoint"), true) +} + +func TestTritonSDRefreshCancelled(t *testing.T) { + var ( + td, _ = newTritonDiscovery(conf) + ) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := td.refresh(ctx) + testutil.NotOk(t, err, "") + testutil.Equals(t, strings.Contains(err.Error(), context.Canceled.Error()), true) } func testTritonSDRefresh(t *testing.T, dstr string) []model.LabelSet { var ( - td, err = New(nil, &conf) - s = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + td, _ = newTritonDiscovery(conf) + s = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, dstr) })) ) defer s.Close() - u, uperr := url.Parse(s.URL) - testutil.Ok(t, uperr) + u, err := url.Parse(s.URL) + testutil.Ok(t, err) testutil.Assert(t, u != nil, "") - host, strport, sherr := net.SplitHostPort(u.Host) - testutil.Ok(t, sherr) + host, strport, err := net.SplitHostPort(u.Host) + testutil.Ok(t, err) testutil.Assert(t, host != "", "") testutil.Assert(t, strport != "", "") - port, atoierr := strconv.Atoi(strport) - testutil.Ok(t, atoierr) + port, err := strconv.Atoi(strport) + testutil.Ok(t, err) testutil.Assert(t, port != 0, "") td.sdConfig.Port = port - testutil.Ok(t, err) - testutil.Assert(t, td != nil, "") - - tg, err := td.refresh() + tg, err := td.refresh(context.Background()) testutil.Ok(t, err) testutil.Assert(t, tg != nil, "")