diff --git a/api/agent.go b/api/agent.go index b707265962..17283fef9d 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 b0b7701a27..cb35a27196 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)