Merge pull request juju#17488 from SimonRichardson/api-context-common…

…-packages

juju#17488

Adds context.Context to the first argument of the api/common type methods. This is the ongoing process of removing context.TODO from the api packages. There are a lot of these to do, so I'm working them through piecemeal.

The rationale for this is to allow tracing from the api clients to the apiserver. Thus we'll be able to trace all the calls to and from the agents. This with parca should give us insights to Juju when running in production.

## Checklist

<!-- If an item is not applicable, use `~strikethrough~`. -->

- [x] Code style: imports ordered, good names, simple structure, etc
- [x] Comments saying why design decisions were made
- [x] Go unit tests, with comments saying what you're testing

## QA steps

```sh
$ juju bootstrap lxd test
$ juju add-model default
$ juju deploy ubuntu
```

api/agent/machiner/machine.go

@@ -89,8 +89,8 @@ func (m *Machine) EnsureDead() error {
 }
 
 // Watch returns a watcher for observing changes to the machine.
-func (m *Machine) Watch() (watcher.NotifyWatcher, error) {
-	return common.Watch(m.client.facade, "Watch", m.tag)
+func (m *Machine) Watch(ctx context.Context) (watcher.NotifyWatcher, error) {
+	return common.Watch(ctx, m.client.facade, "Watch", m.tag)
 }
 
 // Jobs returns a list of jobs for the machine.

api/agent/machiner/machiner_test.go

@@ -242,7 +242,7 @@ func (s *machinerSuite) TestWatch(c *gc.C) {
 	client := machiner.NewClient(apiCaller)
 	m, err := client.Machine(context.Background(), tag)
 	c.Assert(err, jc.ErrorIsNil)
-	_, err = m.Watch()
+	_, err = m.Watch(context.Background())
 	c.Assert(err, gc.ErrorMatches, "FAIL")
 	c.Assert(calls, gc.Equals, 2)
 }

api/agent/uniter/application.go

@@ -43,8 +43,8 @@ func (s *Application) String() string {
 }
 
 // Watch returns a watcher for observing changes to an application.
-func (s *Application) Watch() (watcher.NotifyWatcher, error) {
-	return common.Watch(s.client.facade, "Watch", s.tag)
+func (s *Application) Watch(ctx context.Context) (watcher.NotifyWatcher, error) {
+	return common.Watch(ctx, s.client.facade, "Watch", s.tag)
 }
 
 // Life returns the application's current life state.

api/agent/uniter/application_test.go

@@ -128,7 +128,7 @@ func (s *applicationSuite) TestWatch(c *gc.C) {
 	app, err := client.Application(context.Background(), names.NewApplicationTag("mysql"))
 	c.Assert(err, jc.ErrorIsNil)
 
-	w, err := app.Watch()
+	w, err := app.Watch(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	wc := watchertest.NewNotifyWatcherC(c, w)
 	defer wc.AssertStops()

api/agent/uniter/unit.go

@@ -160,8 +160,8 @@ func (u *Unit) EnsureDead() error {
 }
 
 // Watch returns a watcher for observing changes to the unit.
-func (u *Unit) Watch() (watcher.NotifyWatcher, error) {
-	return common.Watch(u.client.facade, "Watch", u.tag)
+func (u *Unit) Watch(ctx context.Context) (watcher.NotifyWatcher, error) {
+	return common.Watch(ctx, u.client.facade, "Watch", u.tag)
 }
 
 // WatchRelations returns a StringsWatcher that notifies of changes to
@@ -547,8 +547,8 @@ func (u *Unit) WatchActionNotifications() (watcher.StringsWatcher, error) {
 
 // WatchUpgradeSeriesNotifications returns a NotifyWatcher for observing the
 // state of a series upgrade.
-func (u *Unit) WatchUpgradeSeriesNotifications() (watcher.NotifyWatcher, error) {
-	return u.client.WatchUpgradeSeriesNotifications()
+func (u *Unit) WatchUpgradeSeriesNotifications(ctx context.Context) (watcher.NotifyWatcher, error) {
+	return u.client.WatchUpgradeSeriesNotifications(ctx)
 }
 
 // LogActionMessage logs a progress message for the specified action.
@@ -565,13 +565,13 @@ func (u *Unit) LogActionMessage(tag names.ActionTag, message string) error {
 }
 
 // UpgradeSeriesStatus returns the upgrade series status of a unit from remote state
-func (u *Unit) UpgradeSeriesStatus() (model.UpgradeSeriesStatus, string, error) {
-	return u.client.UpgradeSeriesUnitStatus()
+func (u *Unit) UpgradeSeriesStatus(ctx context.Context) (model.UpgradeSeriesStatus, string, error) {
+	return u.client.UpgradeSeriesUnitStatus(ctx)
 }
 
 // SetUpgradeSeriesStatus sets the upgrade series status of the unit in the remote state
-func (u *Unit) SetUpgradeSeriesStatus(status model.UpgradeSeriesStatus, reason string) error {
-	return u.client.SetUpgradeSeriesUnitStatus(status, reason)
+func (u *Unit) SetUpgradeSeriesStatus(ctx context.Context, status model.UpgradeSeriesStatus, reason string) error {
+	return u.client.SetUpgradeSeriesUnitStatus(ctx, status, reason)
 }
 
 // RequestReboot sets the reboot flag for its machine agent
@@ -726,14 +726,14 @@ func (u *Unit) NetworkInfo(bindings []string, relationId *int) (map[string]param
 
 // State returns the state persisted by the charm running in this unit
 // and the state internal to the uniter for this unit.
-func (u *Unit) State() (params.UnitStateResult, error) {
-	return u.client.State()
+func (u *Unit) State(ctx context.Context) (params.UnitStateResult, error) {
+	return u.client.State(ctx)
 }
 
 // SetState sets the state persisted by the charm running in this unit
 // and the state internal to the uniter for this unit.
-func (u *Unit) SetState(unitState params.SetUnitStateArg) error {
-	return u.client.SetState(unitState)
+func (u *Unit) SetState(ctx context.Context, unitState params.SetUnitStateArg) error {
+	return u.client.SetState(ctx, unitState)
 }
 
 // CommitHookChanges batches together all required API calls for applying

api/agent/uniter/unit_test.go

@@ -250,7 +250,7 @@ func (s *unitSuite) TestWatch(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	w, err := unit.Watch()
+	w, err := unit.Watch(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	wc := watchertest.NewNotifyWatcherC(c, w)
 	defer wc.AssertStops()
@@ -639,7 +639,7 @@ func (s *unitSuite) TestWatchUpgradeSeriesNotifications(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	w, err := unit.WatchUpgradeSeriesNotifications()
+	w, err := unit.WatchUpgradeSeriesNotifications(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	wc := watchertest.NewNotifyWatcherC(c, w)
 	defer wc.AssertStops()
@@ -673,7 +673,7 @@ func (s *unitSuite) TestUpgradeSeriesStatus(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	err := unit.SetUpgradeSeriesStatus(model.UpgradeSeriesCompleted, "done")
+	err := unit.SetUpgradeSeriesStatus(context.Background(), model.UpgradeSeriesCompleted, "done")
 	c.Assert(err, gc.ErrorMatches, "biff")
 }
 
@@ -694,7 +694,7 @@ func (s *unitSuite) TestSetUpgradeSeriesStatus(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	seriesStatus, target, err := unit.UpgradeSeriesStatus()
+	seriesStatus, target, err := unit.UpgradeSeriesStatus(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Check(seriesStatus, gc.Equals, model.UpgradeSeriesCompleted)
 	c.Check(target, gc.Equals, "focal")
@@ -750,7 +750,7 @@ func (s *unitSuite) TestUnitState(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	result, err := unit.State()
+	result, err := unit.State(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(result, jc.DeepEquals, unitState)
 }
@@ -776,7 +776,7 @@ func (s *unitSuite) TestSetState(c *gc.C) {
 	client := uniter.NewClient(apiCaller, names.NewUnitTag("mysql/0"))
 
 	unit := uniter.CreateUnit(client, names.NewUnitTag("mysql/0"))
-	err := unit.SetState(unitState)
+	err := unit.SetState(context.Background(), unitState)
 	c.Assert(err, gc.ErrorMatches, "biff")
 }
 

api/client/modelmanager/modelmanager_test.go

@@ -4,6 +4,7 @@
 package modelmanager_test
 
 import (
+	"context"
 	"regexp"
 	"time"
 
@@ -325,7 +326,7 @@ func (s *modelmanagerSuite) TestModelStatus(c *gc.C) {
 	mockFacadeCaller.EXPECT().FacadeCall(gomock.Any(), "ModelStatus", args, res).SetArg(3, ress).Return(nil)
 	client := common.NewModelStatusAPI(mockFacadeCaller)
 
-	results, err := client.ModelStatus(coretesting.ModelTag, coretesting.ModelTag)
+	results, err := client.ModelStatus(context.Background(), coretesting.ModelTag, coretesting.ModelTag)
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(results[0], jc.DeepEquals, base.ModelStatus{
 		UUID:               coretesting.ModelTag.Id(),
@@ -354,7 +355,7 @@ func (s *modelmanagerSuite) TestModelStatusEmpty(c *gc.C) {
 	mockFacadeCaller.EXPECT().FacadeCall(gomock.Any(), "ModelStatus", args, res).SetArg(3, ress).Return(nil)
 	client := common.NewModelStatusAPI(mockFacadeCaller)
 
-	results, err := client.ModelStatus()
+	results, err := client.ModelStatus(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(results, jc.DeepEquals, []base.ModelStatus{})
 }
@@ -375,7 +376,7 @@ func (s *modelmanagerSuite) TestModelStatusError(c *gc.C) {
 	mockFacadeCaller := basemocks.NewMockFacadeCaller(ctrl)
 	mockFacadeCaller.EXPECT().FacadeCall(gomock.Any(), "ModelStatus", args, res).Return(errors.New("model error"))
 	client := common.NewModelStatusAPI(mockFacadeCaller)
-	out, err := client.ModelStatus(coretesting.ModelTag, coretesting.ModelTag)
+	out, err := client.ModelStatus(context.Background(), coretesting.ModelTag, coretesting.ModelTag)
 	c.Assert(err, gc.ErrorMatches, "model error")
 	c.Assert(out, gc.IsNil)
 }

api/common/apiaddresser.go

@@ -28,9 +28,9 @@ func NewAPIAddresser(facade base.FacadeCaller) *APIAddresser {
 }
 
 // APIAddresses returns the list of addresses used to connect to the API.
-func (a *APIAddresser) APIAddresses() ([]string, error) {
+func (a *APIAddresser) APIAddresses(ctx context.Context) ([]string, error) {
 	var result params.StringsResult
-	err := a.facade.FacadeCall(context.TODO(), "APIAddresses", nil, &result)
+	err := a.facade.FacadeCall(ctx, "APIAddresses", nil, &result)
 	if err != nil {
 		return nil, err
 	}
@@ -42,19 +42,19 @@ func (a *APIAddresser) APIAddresses() ([]string, error) {
 }
 
 // APIHostPorts returns the host/port addresses of the API servers.
-func (a *APIAddresser) APIHostPorts() ([]network.ProviderHostPorts, error) {
+func (a *APIAddresser) APIHostPorts(ctx context.Context) ([]network.ProviderHostPorts, error) {
 	var result params.APIHostPortsResult
-	err := a.facade.FacadeCall(context.TODO(), "APIHostPorts", nil, &result)
+	err := a.facade.FacadeCall(ctx, "APIHostPorts", nil, &result)
 	if err != nil {
 		return nil, err
 	}
 	return params.ToProviderHostsPorts(result.Servers), nil
 }
 
 // WatchAPIHostPorts watches the host/port addresses of the API servers.
-func (a *APIAddresser) WatchAPIHostPorts() (watcher.NotifyWatcher, error) {
+func (a *APIAddresser) WatchAPIHostPorts(ctx context.Context) (watcher.NotifyWatcher, error) {
 	var result params.NotifyWatchResult
-	err := a.facade.FacadeCall(context.TODO(), "WatchAPIHostPorts", nil, &result)
+	err := a.facade.FacadeCall(ctx, "WatchAPIHostPorts", nil, &result)
 	if err != nil {
 		return nil, err
 	}

api/common/apiaddresser_test.go

@@ -4,6 +4,7 @@
 package common_test
 
 import (
+	"context"
 	"time"
 
 	jujutesting "github.com/juju/testing"
@@ -38,7 +39,7 @@ func (s *apiaddresserSuite) TestAPIAddresses(c *gc.C) {
 	facade.EXPECT().FacadeCall(gomock.Any(), "APIAddresses", nil, gomock.Any()).SetArg(3, result).Return(nil)
 
 	client := common.NewAPIAddresser(facade)
-	addresses, err := client.APIAddresses()
+	addresses, err := client.APIAddresses(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(addresses, gc.DeepEquals, []string{"0.1.2.3:1234"})
 }
@@ -81,7 +82,7 @@ func (s *apiaddresserSuite) TestAPIHostPorts(c *gc.C) {
 		{corenetwork.ProviderHostPort{ProviderAddress: corenetwork.NewMachineAddress("3.4.5.6").AsProviderAddress(), NetPort: 3456}},
 	}
 
-	serverAddrs, err := client.APIHostPorts()
+	serverAddrs, err := client.APIHostPorts(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Assert(serverAddrs, gc.DeepEquals, expectServerAddrs)
 }
@@ -100,7 +101,7 @@ func (s *apiaddresserSuite) TestWatchAPIHostPorts(c *gc.C) {
 	facade.EXPECT().RawAPICaller().Return(caller)
 
 	client := common.NewAPIAddresser(facade)
-	w, err := client.WatchAPIHostPorts()
+	w, err := client.WatchAPIHostPorts(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 
 	// watch for the changes

api/common/controllerconfig.go

@@ -24,9 +24,9 @@ func NewControllerConfig(facade base.FacadeCaller) *ControllerConfigAPI {
 }
 
 // ControllerConfig returns the current controller configuration.
-func (e *ControllerConfigAPI) ControllerConfig() (controller.Config, error) {
+func (e *ControllerConfigAPI) ControllerConfig(ctx context.Context) (controller.Config, error) {
 	var result params.ControllerConfigResult
-	err := e.facade.FacadeCall(context.TODO(), "ControllerConfig", nil, &result)
+	err := e.facade.FacadeCall(ctx, "ControllerConfig", nil, &result)
 	if err != nil {
 		return nil, err
 	}

api/common/leadership.go

@@ -41,9 +41,9 @@ func NewLeadershipPinningAPIFromFacade(facade base.FacadeCaller) *LeadershipPinn
 // PinnedLeadership returns a collection of application names for which
 // leadership is currently pinned, with the entities requiring each
 // application's pinned behaviour.
-func (a *LeadershipPinningAPI) PinnedLeadership() (map[string][]names.Tag, error) {
+func (a *LeadershipPinningAPI) PinnedLeadership(ctx context.Context) (map[string][]names.Tag, error) {
 	var callResult params.PinnedLeadershipResult
-	err := a.facade.FacadeCall(context.TODO(), "PinnedLeadership", nil, &callResult)
+	err := a.facade.FacadeCall(ctx, "PinnedLeadership", nil, &callResult)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -73,8 +73,8 @@ func (a *LeadershipPinningAPI) PinnedLeadership() (map[string][]names.Tag, error
 // If the caller is not a machine agent, an error will be returned.
 // The return is a collection of applications determined to be running on the
 // machine, with the result of each individual pin operation.
-func (a *LeadershipPinningAPI) PinMachineApplications() (map[string]error, error) {
-	res, err := a.pinMachineAppsOps("PinMachineApplications")
+func (a *LeadershipPinningAPI) PinMachineApplications(ctx context.Context) (map[string]error, error) {
+	res, err := a.pinMachineAppsOps(ctx, "PinMachineApplications")
 	return res, errors.Trace(err)
 }
 
@@ -83,16 +83,16 @@ func (a *LeadershipPinningAPI) PinMachineApplications() (map[string]error, error
 // If the caller is not a machine agent, an error will be returned.
 // The return is a collection of applications determined to be running on the
 // machine, with the result of each individual unpin operation.
-func (a *LeadershipPinningAPI) UnpinMachineApplications() (map[string]error, error) {
-	res, err := a.pinMachineAppsOps("UnpinMachineApplications")
+func (a *LeadershipPinningAPI) UnpinMachineApplications(ctx context.Context) (map[string]error, error) {
+	res, err := a.pinMachineAppsOps(ctx, "UnpinMachineApplications")
 	return res, errors.Trace(err)
 }
 
 // pinMachineAppsOps makes a facade call to the input method name and
 // transforms the response into map.
-func (a *LeadershipPinningAPI) pinMachineAppsOps(callName string) (map[string]error, error) {
+func (a *LeadershipPinningAPI) pinMachineAppsOps(ctx context.Context, callName string) (map[string]error, error) {
 	var callResult params.PinApplicationsResults
-	err := a.facade.FacadeCall(context.TODO(), callName, nil, &callResult)
+	err := a.facade.FacadeCall(ctx, callName, nil, &callResult)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

api/common/leadership_test.go

@@ -4,6 +4,7 @@
 package common_test
 
 import (
+	"context"
 	"errors"
 
 	"github.com/juju/names/v5"
@@ -41,7 +42,7 @@ func (s *LeadershipSuite) TestPinnedLeadership(c *gc.C) {
 	resultSource := params.PinnedLeadershipResult{Result: pinned}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "PinnedLeadership", nil, gomock.Any()).SetArg(3, resultSource)
 
-	res, err := s.client.PinnedLeadership()
+	res, err := s.client.PinnedLeadership(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Check(res, gc.DeepEquals, map[string][]names.Tag{"redis": {names.NewMachineTag("0"), names.NewMachineTag("1")}})
 }
@@ -52,7 +53,7 @@ func (s *LeadershipSuite) TestPinnedLeadershipError(c *gc.C) {
 	resultSource := params.PinnedLeadershipResult{Error: apiservererrors.ServerError(errors.New("splat"))}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "PinnedLeadership", nil, gomock.Any()).SetArg(3, resultSource)
 
-	_, err := s.client.PinnedLeadership()
+	_, err := s.client.PinnedLeadership(context.Background())
 	c.Assert(err, gc.ErrorMatches, "splat")
 }
 
@@ -62,7 +63,7 @@ func (s *LeadershipSuite) TestPinMachineApplicationsSuccess(c *gc.C) {
 	resultSource := params.PinApplicationsResults{Results: s.pinApplicationsServerSuccessResults()}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "PinMachineApplications", nil, gomock.Any()).SetArg(3, resultSource)
 
-	res, err := s.client.PinMachineApplications()
+	res, err := s.client.PinMachineApplications(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Check(res, gc.DeepEquals, s.pinApplicationsClientSuccessResults())
 }
@@ -76,7 +77,7 @@ func (s *LeadershipSuite) TestPinMachineApplicationsPartialError(c *gc.C) {
 	resultSource := params.PinApplicationsResults{Results: results}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "PinMachineApplications", nil, gomock.Any()).SetArg(3, resultSource)
 
-	res, err := s.client.PinMachineApplications()
+	res, err := s.client.PinMachineApplications(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 
 	exp := s.pinApplicationsClientSuccessResults()
@@ -90,7 +91,7 @@ func (s *LeadershipSuite) TestUnpinMachineApplicationsSuccess(c *gc.C) {
 	resultSource := params.PinApplicationsResults{Results: s.pinApplicationsServerSuccessResults()}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "UnpinMachineApplications", nil, gomock.Any()).SetArg(3, resultSource)
 
-	res, err := s.client.UnpinMachineApplications()
+	res, err := s.client.UnpinMachineApplications(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 	c.Check(res, gc.DeepEquals, s.pinApplicationsClientSuccessResults())
 }
@@ -113,7 +114,7 @@ func (s *LeadershipSuite) TestUnpinMachineApplicationsPartialError(c *gc.C) {
 	resultSource := params.PinApplicationsResults{Results: results}
 	s.facade.EXPECT().FacadeCall(gomock.Any(), "UnpinMachineApplications", nil, gomock.Any()).SetArg(3, resultSource)
 
-	res, err := s.client.UnpinMachineApplications()
+	res, err := s.client.UnpinMachineApplications(context.Background())
 	c.Assert(err, jc.ErrorIsNil)
 
 	exp := s.pinApplicationsClientSuccessResults()

api/common/life.go

@@ -23,7 +23,7 @@ func Life(ctx context.Context, caller base.FacadeCaller, tags []names.Tag) ([]pa
 	var result params.LifeResults
 	entities := make([]params.Entity, len(tags))
 	for i, t := range tags {
-		entities[i] = params.Entity{t.String()}
+		entities[i] = params.Entity{Tag: t.String()}
 	}
 	args := params.Entities{Entities: entities}
 	if err := caller.FacadeCall(ctx, "Life", args, &result); err != nil {