From b7bfa5e6dfac1525ae200de60fb9b9ab2f967288 Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Thu, 3 Sep 2020 12:22:03 -0400
Subject: [PATCH] api: Add a context to ServiceRegisterOpts

So that operations can timeout instead of hanging forever
---
 api/agent.go      | 15 ++++++++++++++-
 api/agent_test.go | 18 +++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/api/agent.go b/api/agent.go
index b70726596..17283fef9 100644
--- a/api/agent.go
+++ b/api/agent.go
@@ -3,6 +3,7 @@ package api
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -264,12 +265,23 @@ type AgentServiceRegistration struct {
 	Namespace         string                          `json:",omitempty" bexpr:"-" hash:"ignore"`
 }
 
-//ServiceRegisterOpts is used to pass extra options to the service register.
+// ServiceRegisterOpts is used to pass extra options to the service register.
 type ServiceRegisterOpts struct {
 	//Missing healthchecks will be deleted from the agent.
 	//Using this parameter allows to idempotently register a service and its checks without
 	//having to manually deregister checks.
 	ReplaceExistingChecks bool
+
+	// ctx is an optional context pass through to the underlying HTTP
+	// request layer. Use WithContext() to set the context.
+	ctx context.Context
+}
+
+// WithContext sets the context to be used for the request on a new ServiceRegisterOpts,
+// and returns the opts.
+func (o ServiceRegisterOpts) WithContext(ctx context.Context) ServiceRegisterOpts {
+	o.ctx = ctx
+	return o
 }
 
 // AgentCheckRegistration is used to register a new check
@@ -686,6 +698,7 @@ func (a *Agent) ServiceRegisterOpts(service *AgentServiceRegistration, opts Serv
 func (a *Agent) serviceRegister(service *AgentServiceRegistration, opts ServiceRegisterOpts) error {
 	r := a.c.newRequest("PUT", "/v1/agent/service/register")
 	r.obj = service
+	r.ctx = opts.ctx
 	if opts.ReplaceExistingChecks {
 		r.params.Set("replace-existing-checks", "true")
 	}
diff --git a/api/agent_test.go b/api/agent_test.go
index b0b7701a2..cb35a2719 100644
--- a/api/agent_test.go
+++ b/api/agent_test.go
@@ -1,7 +1,9 @@
 package api
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -208,7 +210,9 @@ func TestAPI_AgentServiceAndReplaceChecks(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
-	if err := agent.ServiceRegisterOpts(regupdate, ServiceRegisterOpts{ReplaceExistingChecks: true}); err != nil {
+	ctx := context.Background()
+	opts := ServiceRegisterOpts{ReplaceExistingChecks: true}.WithContext(ctx)
+	if err := agent.ServiceRegisterOpts(regupdate, opts); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
@@ -247,6 +251,18 @@ func TestAPI_AgentServiceAndReplaceChecks(t *testing.T) {
 	}
 }
 
+func TestAgent_ServiceRegisterOpts_WithContextTimeout(t *testing.T) {
+	c, err := NewClient(DefaultConfig())
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
+	t.Cleanup(cancel)
+
+	opts := ServiceRegisterOpts{}.WithContext(ctx)
+	err = c.Agent().ServiceRegisterOpts(&AgentServiceRegistration{}, opts)
+	require.True(t, errors.Is(err, context.DeadlineExceeded), "expected timeout")
+}
+
 func TestAPI_AgentServices(t *testing.T) {
 	t.Parallel()
 	c, s := makeClient(t)