Single context for shared background runners in multi-service

micromdm · Jul 3, 2021 · 1406f5c · 1406f5c
1 parent 513d43e
commit 1406f5c
Showing 1 changed file with 15 additions and 10 deletions.
diff --git a/service/multi/multi.go b/service/multi/multi.go
@@ -17,13 +17,18 @@ import (
 type MultiService struct {
 	logger log.Logger
 	svcs   []service.CheckinAndCommandService
+	ctx    context.Context
 }
 
 func New(logger log.Logger, svcs ...service.CheckinAndCommandService) *MultiService {
 	if len(svcs) < 1 {
 		panic("must supply at least one service")
 	}
-	return &MultiService{logger: logger, svcs: svcs}
+	return &MultiService{
+		logger: logger,
+		svcs:   svcs,
+		ctx:    context.Background(),
+	}
 }
 
 type errorRunner func(service.CheckinAndCommandService) error
@@ -43,15 +48,15 @@ func (ms *MultiService) runOthers(r errorRunner) {
 }
 
 // RequestWithContext returns a clone of r and sets its context to ctx.
-func RequestWithContext(r *mdm.Request, ctx context.Context) *mdm.Request {
+func (ms *MultiService) RequestWithContext(r *mdm.Request) *mdm.Request {
 	r2 := r.Clone()
-	r2.Context = ctx
+	r2.Context = ms.ctx
 	return r2
 }
 
 func (ms *MultiService) Authenticate(r *mdm.Request, m *mdm.Authenticate) error {
 	err := ms.svcs[0].Authenticate(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		return svc.Authenticate(rc, m)
 	})
@@ -60,7 +65,7 @@ func (ms *MultiService) Authenticate(r *mdm.Request, m *mdm.Authenticate) error
 
 func (ms *MultiService) TokenUpdate(r *mdm.Request, m *mdm.TokenUpdate) error {
 	err := ms.svcs[0].TokenUpdate(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		return svc.TokenUpdate(rc, m)
 	})
@@ -69,7 +74,7 @@ func (ms *MultiService) TokenUpdate(r *mdm.Request, m *mdm.TokenUpdate) error {
 
 func (ms *MultiService) CheckOut(r *mdm.Request, m *mdm.CheckOut) error {
 	err := ms.svcs[0].CheckOut(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		return svc.CheckOut(rc, m)
 	})
@@ -78,7 +83,7 @@ func (ms *MultiService) CheckOut(r *mdm.Request, m *mdm.CheckOut) error {
 
 func (ms *MultiService) UserAuthenticate(r *mdm.Request, m *mdm.UserAuthenticate) ([]byte, error) {
 	respBytes, err := ms.svcs[0].UserAuthenticate(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		_, err := svc.UserAuthenticate(rc, m)
 		return err
@@ -88,7 +93,7 @@ func (ms *MultiService) UserAuthenticate(r *mdm.Request, m *mdm.UserAuthenticate
 
 func (ms *MultiService) SetBootstrapToken(r *mdm.Request, m *mdm.SetBootstrapToken) error {
 	err := ms.svcs[0].SetBootstrapToken(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		return svc.SetBootstrapToken(rc, m)
 	})
@@ -97,7 +102,7 @@ func (ms *MultiService) SetBootstrapToken(r *mdm.Request, m *mdm.SetBootstrapTok
 
 func (ms *MultiService) GetBootstrapToken(r *mdm.Request, m *mdm.GetBootstrapToken) (*mdm.BootstrapToken, error) {
 	bsToken, err := ms.svcs[0].GetBootstrapToken(r, m)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		_, err := svc.GetBootstrapToken(rc, m)
 		return err
@@ -107,7 +112,7 @@ func (ms *MultiService) GetBootstrapToken(r *mdm.Request, m *mdm.GetBootstrapTok
 
 func (ms *MultiService) CommandAndReportResults(r *mdm.Request, results *mdm.CommandResults) (*mdm.Command, error) {
 	cmd, err := ms.svcs[0].CommandAndReportResults(r, results)
-	rc := RequestWithContext(r, context.Background())
+	rc := ms.RequestWithContext(r)
 	ms.runOthers(func(svc service.CheckinAndCommandService) error {
 		_, err := svc.CommandAndReportResults(rc, results)
 		return err