Merge pull request #8609 from hashicorp/dnephin/add-query-options-to-ServiceRegister

api: Add a context to ServiceRegisterOpts

api/agent.go

@@ -3,6 +3,7 @@ package api
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -266,12 +267,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
@@ -688,6 +700,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")
 	}

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)